diff --git a/cosipy/cpkernel/cosipy_core.py b/cosipy/cpkernel/cosipy_core.py index a1a594c..084debf 100644 --- a/cosipy/cpkernel/cosipy_core.py +++ b/cosipy/cpkernel/cosipy_core.py @@ -15,7 +15,43 @@ from cosipy.modules.roughness import updateRoughness from cosipy.modules.surfaceTemperature import update_surface_temperature -def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_data=None, opt_dict=None): + +def init_nan_array_1d(nt: int) -> np.ndarray: + """Initialise and fill an array with NaNs. + + Args: + nt: Array size (time dimension). + + Returns: + NaN array. + """ + if not Config.WRF_X_CSPY: + x = np.full(nt, np.nan) + else: + x = None + + return x + + +def init_nan_array_2d(nt: int, max_layers: int) -> np.ndarray: + """Initialise and fill an array with NaNs. + + Args: + nt: Array's temporal resolution. + max_layers: Array's spatial resolution. + + Returns: + 2D NaN array. + """ + if not Config.WRF_X_CSPY and Config.full_field: + x = np.full((nt, max_layers), np.nan) + else: + x = None + + return x + + +def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_data=None): """Cosipy core function. The calculations are performed on a single core. @@ -26,7 +62,7 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat indX (int): The grid cell's X index. GRID_RESTART (xarray.Dataset): Use a restart dataset instead of creating an initial profile. Default ``None``. - stake_name (list): Stake names. Default ``None``. + stake_names (list): Stake names. Default ``None``. stake_data (pd.Dataframe): Stake measurements. Default ``None``. Returns: @@ -50,9 +86,12 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat spread_snow_transfer_function = Constants.spread_snow_transfer_function constant_density = Constants.constant_density albedo_fresh_snow = Constants.albedo_fresh_snow + albedo_firn = Constants.albedo_firn + WRF_X_CSPY = Config.WRF_X_CSPY # Replace values from constants.py if coupled - if Config.WRF_X_CSPY: + # TODO: This only affects the current module scope instead of global. + if WRF_X_CSPY: dt = int(DATA.DT.values) max_layers = int(DATA.max_layers.values) z = float(DATA.ZLVL.values) @@ -73,62 +112,68 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat roughness_firn = opt_dict[10] aging_factor_roughness = opt_dict[11] - #read_opt(opt_dict, globals()) - # Local variables - nt = len(DATA.time.values) #accessing DATA is expensive - _RRR = np.full(nt, np.nan) - _RAIN = np.full(nt, np.nan) - _SNOWFALL = np.full(nt, np.nan) - _LWin = np.full(nt, np.nan) - _LWout = np.full(nt, np.nan) - _H = np.full(nt, np.nan) - _LE = np.full(nt, np.nan) - _B = np.full(nt, np.nan) - _QRR = np.full(nt, np.nan) - _MB = np.full(nt, np.nan) - _surfMB = np.full(nt, np.nan) - _MB = np.full(nt, np.nan) - _Q = np.full(nt, np.nan) - _SNOWHEIGHT = np.full(nt, np.nan) - _TOTALHEIGHT = np.full(nt, np.nan) - _TS = np.full(nt, np.nan) - _ALBEDO = np.full(nt, np.nan) - _ME = np.full(nt, np.nan) - _intMB = np.full(nt, np.nan) - _EVAPORATION = np.full(nt, np.nan) - _SUBLIMATION = np.full(nt, np.nan) - _CONDENSATION = np.full(nt, np.nan) - _DEPOSITION = np.full(nt, np.nan) - _REFREEZE = np.full(nt, np.nan) - _NLAYERS = np.full(nt, np.nan) - _subM = np.full(nt, np.nan) - _Z0 = np.full(nt, np.nan) - _surfM = np.full(nt, np.nan) - _MOL = np.full(nt, np.nan) - - _LAYER_HEIGHT = np.full((nt,max_layers), np.nan) - _LAYER_RHO = np.full((nt,max_layers), np.nan) - _LAYER_T = np.full((nt,max_layers), np.nan) - _LAYER_LWC = np.full((nt,max_layers), np.nan) - _LAYER_CC = np.full((nt,max_layers), np.nan) - _LAYER_POROSITY = np.full((nt,max_layers), np.nan) - _LAYER_ICE_FRACTION = np.full((nt,max_layers), np.nan) - _LAYER_IRREDUCIBLE_WATER = np.full((nt,max_layers), np.nan) - _LAYER_REFREEZE = np.full((nt,max_layers), np.nan) + nt = len(DATA.time.values) # accessing DATA is expensive + """ + Local variables bypass local array creation for WRF_X_CSPY + TODO: Implement more elegant solution. + """ + if not WRF_X_CSPY: + _RRR = init_nan_array_1d(nt) + _RAIN = init_nan_array_1d(nt) + _SNOWFALL = init_nan_array_1d(nt) + _LWin = init_nan_array_1d(nt) + _LWout = init_nan_array_1d(nt) + _H = init_nan_array_1d(nt) + _LE = init_nan_array_1d(nt) + _B = init_nan_array_1d(nt) + _QRR = init_nan_array_1d(nt) + _MB = init_nan_array_1d(nt) + _surfMB = init_nan_array_1d(nt) + _MB = init_nan_array_1d(nt) + _Q = init_nan_array_1d(nt) + _SNOWHEIGHT = init_nan_array_1d(nt) + _TOTALHEIGHT = init_nan_array_1d(nt) + _TS = init_nan_array_1d(nt) + _ALBEDO = init_nan_array_1d(nt) + _ME = init_nan_array_1d(nt) + _intMB = init_nan_array_1d(nt) + _EVAPORATION = init_nan_array_1d(nt) + _SUBLIMATION = init_nan_array_1d(nt) + _CONDENSATION = init_nan_array_1d(nt) + _DEPOSITION = init_nan_array_1d(nt) + _REFREEZE = init_nan_array_1d(nt) + _NLAYERS = init_nan_array_1d(nt) + _subM = init_nan_array_1d(nt) + _Z0 = init_nan_array_1d(nt) + _surfM = init_nan_array_1d(nt) + _MOL = init_nan_array_1d(nt) + _new_snow_height = init_nan_array_1d(nt) + _new_snow_timestamp = init_nan_array_1d(nt) + _old_snow_timestamp = init_nan_array_1d(nt) + + _LAYER_HEIGHT = init_nan_array_2d(nt, max_layers) + _LAYER_RHO = init_nan_array_2d(nt, max_layers) + _LAYER_T = init_nan_array_2d(nt, max_layers) + _LAYER_LWC = init_nan_array_2d(nt, max_layers) + _LAYER_CC = init_nan_array_2d(nt, max_layers) + _LAYER_POROSITY = init_nan_array_2d(nt, max_layers) + _LAYER_ICE_FRACTION = init_nan_array_2d(nt, max_layers) + _LAYER_IRREDUCIBLE_WATER = init_nan_array_2d(nt, max_layers) + _LAYER_REFREEZE = init_nan_array_2d(nt, max_layers) #-------------------------------------------- # Initialize snowpack or load restart grid #-------------------------------------------- if GRID_RESTART is None: - GRID = init_snowpack(DATA, opt_dict) + GRID = init_snowpack(DATA) else: GRID = load_snowpack(GRID_RESTART) # Create the local output datasets if not coupled RESTART = None - if not Config.WRF_X_CSPY: - IO = IOClass(DATA, opt_dict) + if not WRF_X_CSPY: + IOClass(DATA, opt_dict) RESTART = IO.create_local_restart_dataset() # hours since the last snowfall (albedo module) @@ -161,6 +206,8 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat if Config.force_use_TP: SNOWF = None + LWin = np.array(nt * [None]) + N = np.array(nt * [None]) if ('LWin' in DATA) and ('N' in DATA): LWin = DATA.LWin.values N = DATA.N.values @@ -174,17 +221,18 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat if Config.force_use_N: LWin = None + SLOPE = 0. if 'SLOPE' in DATA: SLOPE = DATA.SLOPE.values - else: - SLOPE = 0.0 # Initial cumulative mass balance variable MB_cum = 0 - + # Initial snow albedo and surface temperature for Bougamont et al. 2005 albedo surface_temperature = 270.0 albedo_snow = albedo_fresh_snow + if WRF_X_CSPY: + albedo_snow = albedo_firn if Config.stake_evaluation: # Create pandas dataframe for stake evaluation @@ -192,7 +240,7 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat #-------------------------------------------- # TIME LOOP - #-------------------------------------------- + #-------------------------------------------- for t in np.arange(nt): # Check grid @@ -212,14 +260,15 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat # Derive snowfall [m] and rain rates [m w.e.] if (SNOWF is not None) and (RRR is not None): SNOWFALL = SNOWF[t] - #ensure rain is not negative - RAIN = np.maximum(RRR[t]-SNOWFALL*(density_fresh_snow/water_density) * 1000.0, 0.0) + RAIN = RRR[t]-SNOWFALL*(density_fresh_snow/water_density) * 1000.0 elif SNOWF is not None: SNOWFALL = SNOWF[t] - else: + elif RRR is not None: # Else convert total precipitation [mm] to snowheight [m]; liquid/solid fraction SNOWFALL = (RRR[t]/1000.0)*(water_density/density_fresh_snow)*(0.5*(-np.tanh(((T2[t]-zero_temperature) - center_snow_transfer_function) * spread_snow_transfer_function) + 1.0)) RAIN = RRR[t]-SNOWFALL*(density_fresh_snow/water_density) * 1000.0 + else: + raise ValueError("No SNOWFALL or RRR data provided.") # if snowfall is smaller than the threshold if SNOWFALL 0.0: # Add a new snow node on top - GRID.add_fresh_snow(SNOWFALL, density_fresh_snow, np.minimum(float(T2[t]),zero_temperature), 0.0) + GRID.add_fresh_snow(SNOWFALL, density_fresh_snow, np.minimum(float(T2[t]),zero_temperature), 0.0, dt) else: GRID.set_fresh_snow_props_update_time(dt) # Guarantee that solar radiation is greater equal zero - if (G[t]<0.0): - G[t] = 0.0 + G[t] = max(G[t],0.0) #-------------------------------------------- # Merge grid layers, if necessary @@ -247,12 +295,12 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat #-------------------------------------------- # Calculate albedo and roughness length changes if first layer is snow #-------------------------------------------- - alpha, albedo_snow = updateAlbedo(GRID,surface_temperature,albedo_snow, opt_dict) + alpha, albedo_snow = updateAlbedo(GRID,surface_temperature,albedo_snow,opt_dict) #-------------------------------------------- # Update roughness length #-------------------------------------------- - z0 = updateRoughness(GRID, opt_dict) + z0 = updateRoughness(GRID,opt_dict) #-------------------------------------------- # Surface Energy Balance @@ -270,32 +318,37 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat # Calculate residual net shortwave radiation (penetrating part removed) sw_radiation_net = SWnet - G_penetrating - if LWin is not None: - # Find new surface temperature (LW is used from the input file) + if (LWin is not None) and (N is not None): + # Find new surface temperature + fun, surface_temperature, lw_radiation_in, lw_radiation_out, sensible_heat_flux, latent_heat_flux, \ + ground_heat_flux, rain_heat_flux, rho, Lv, MOL, Cs_t, Cs_q, q0, q2 \ + = update_surface_temperature(GRID, dt, z, z0, T2[t], RH2[t], PRES[t], sw_radiation_net, \ + U2[t], RAIN, SLOPE, LWin=LWin[t], N=N[t]) + elif LWin is not None: fun, surface_temperature, lw_radiation_in, lw_radiation_out, sensible_heat_flux, latent_heat_flux, \ ground_heat_flux, rain_heat_flux, rho, Lv, MOL, Cs_t, Cs_q, q0, q2 \ - = update_surface_temperature(GRID, dt, z, z0, T2[t], RH2[t], PRES[t], - sw_radiation_net, U2[t], RAIN, SLOPE, LWin=LWin[t], opt_dict=opt_dict) + = update_surface_temperature(GRID, dt, z, z0, T2[t], RH2[t], PRES[t], sw_radiation_net, \ + U2[t], RAIN, SLOPE, LWin=LWin[t], opt_dict=opt_dict) else: # Find new surface temperature (LW is parametrized using cloud fraction) fun, surface_temperature, lw_radiation_in, lw_radiation_out, sensible_heat_flux, latent_heat_flux, \ ground_heat_flux, rain_heat_flux, rho, Lv, MOL, Cs_t, Cs_q, q0, q2 \ - = update_surface_temperature(GRID, dt, z, z0, T2[t], RH2[t], PRES[t], - sw_radiation_net, U2[t], RAIN, SLOPE, N=N[t], opt_dict=opt_dict) + = update_surface_temperature(GRID, dt, z, z0, T2[t], RH2[t], PRES[t], sw_radiation_net, \ + U2[t], RAIN, SLOPE, N=N[t], opt_dict=opt_dict) #-------------------------------------------- # Surface mass fluxes [m w.e.q.] #-------------------------------------------- if surface_temperature < zero_temperature: - sublimation = min(latent_heat_flux / (water_density * lat_heat_sublimation), 0) * dt - deposition = max(latent_heat_flux / (water_density * lat_heat_sublimation), 0) * dt - evaporation = 0 - condensation = 0 + sublimation = min(latent_heat_flux / (water_density * lat_heat_sublimation), 0.) * dt + deposition = max(latent_heat_flux / (water_density * lat_heat_sublimation), 0.) * dt + evaporation = 0. + condensation = 0. else: - sublimation = 0 - deposition = 0 - evaporation = min(latent_heat_flux / (water_density * lat_heat_vaporize), 0) * dt - condensation = max(latent_heat_flux / (water_density * lat_heat_vaporize), 0) * dt + sublimation = 0. + deposition = 0. + evaporation = min(latent_heat_flux / (water_density * lat_heat_vaporize), 0.) * dt + condensation = max(latent_heat_flux / (water_density * lat_heat_vaporize), 0.) * dt #-------------------------------------------- # Melt process - mass changes of snowpack (melting, sublimation, deposition, evaporation, condensation) @@ -336,7 +389,7 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat #-------------------------------------------- # Calculate new density to densification #-------------------------------------------- - densification(GRID, SLOPE, dt, opt_dict) + densification(GRID, SLOPE, dt) #-------------------------------------------- # Calculate mass balance @@ -356,46 +409,47 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat # Cumulative mass balance for stake evaluation MB_cum = MB_cum + mass_balance - + # Store cumulative MB in pandas frame for validation if stake_names: - if (DATA.isel(time=t).time.values in stake_data.index): + if DATA.isel(time=t).time.values in stake_data.index: _df['mb'].loc[DATA.isel(time=t).time.values] = MB_cum - _df['snowheight'].loc[DATA.isel(time=t).time.values] = GRID.get_total_snowheight() - - # Save results - _RAIN[t] = RAIN - _SNOWFALL[t] = SNOWFALL * (density_fresh_snow/water_density) - _LWin[t] = lw_radiation_in - _LWout[t] = lw_radiation_out - _H[t] = sensible_heat_flux - _LE[t] = latent_heat_flux - _B[t] = ground_heat_flux - _QRR[t] = rain_heat_flux - _MB[t] = mass_balance - _surfMB[t] = surface_mass_balance - _Q[t] = Q - _SNOWHEIGHT[t] = GRID.get_total_snowheight() - _TOTALHEIGHT[t] = GRID.get_total_height() - _TS[t] = surface_temperature - _ALBEDO[t] = alpha - _NLAYERS[t] = GRID.get_number_layers() - _ME[t] = melt_energy - _intMB[t] = internal_mass_balance - _EVAPORATION[t] = evaporation - _SUBLIMATION[t] = sublimation - _CONDENSATION[t] = condensation - _DEPOSITION[t] = deposition - _REFREEZE[t] = water_refreezed - _subM[t] = subsurface_melt - _Z0[t] = z0 - _surfM[t] = melt - _MOL[t] = MOL - - if Config.full_field: - if GRID.get_number_layers()>max_layers: - print('Maximum number of layers reached') - else: + _df['snowheight'].loc[DATA.isel(time=t).time.values] = GRID.get_total_snowheight() + + # Save results -- standalone cosipy case + if not WRF_X_CSPY: + _RAIN[t] = RAIN + _SNOWFALL[t] = SNOWFALL * (density_fresh_snow/water_density) + _LWin[t] = lw_radiation_in + _LWout[t] = lw_radiation_out + _H[t] = sensible_heat_flux + _LE[t] = latent_heat_flux + _B[t] = ground_heat_flux + _QRR[t] = rain_heat_flux + _MB[t] = mass_balance + _surfMB[t] = surface_mass_balance + _Q[t] = Q + _SNOWHEIGHT[t] = GRID.get_total_snowheight() + _TOTALHEIGHT[t] = GRID.get_total_height() + _TS[t] = surface_temperature + _ALBEDO[t] = alpha + _NLAYERS[t] = GRID.get_number_layers() + _ME[t] = melt_energy + _intMB[t] = internal_mass_balance + _EVAPORATION[t] = evaporation + _SUBLIMATION[t] = sublimation + _CONDENSATION[t] = condensation + _DEPOSITION[t] = deposition + _REFREEZE[t] = water_refreezed + _subM[t] = subsurface_melt + _Z0[t] = z0 + _surfM[t] = melt + _MOL[t] = MOL + _new_snow_height[t], _new_snow_timestamp[t], _old_snow_timestamp[t] = GRID.get_fresh_snow_props() + + if Config.full_field: + if GRID.get_number_layers()>max_layers: + raise ValueError('Maximum number of layers reached') _LAYER_HEIGHT[t, 0:GRID.get_number_layers()] = GRID.get_height() _LAYER_RHO[t, 0:GRID.get_number_layers()] = GRID.get_density() _LAYER_T[t, 0:GRID.get_number_layers()] = GRID.get_temperature() @@ -405,40 +459,68 @@ def cosipy_core(DATA, indY, indX, GRID_RESTART=None, stake_names=None, stake_dat _LAYER_ICE_FRACTION[t, 0:GRID.get_number_layers()] = GRID.get_ice_fraction() _LAYER_IRREDUCIBLE_WATER[t, 0:GRID.get_number_layers()] = GRID.get_irreducible_water_content() _LAYER_REFREEZE[t, 0:GRID.get_number_layers()] = GRID.get_refreeze() + else: + _LAYER_HEIGHT = None + _LAYER_RHO = None + _LAYER_T = None + _LAYER_LWC = None + _LAYER_CC = None + _LAYER_POROSITY = None + _LAYER_ICE_FRACTION = None + _LAYER_IRREDUCIBLE_WATER = None + _LAYER_REFREEZE = None + + # Save results -- WRF_X_CSPY case else: - _LAYER_HEIGHT = None - _LAYER_RHO = None - _LAYER_T = None - _LAYER_LWC = None - _LAYER_CC = None - _LAYER_POROSITY = None - _LAYER_ICE_FRACTION = None - _LAYER_IRREDUCIBLE_WATER = None - _LAYER_REFREEZE = None + _SNOWHEIGHT = GRID.get_total_snowheight() + _TOTALHEIGHT = GRID.get_total_height() + _NLAYERS = GRID.get_number_layers() + _new_snow_height, _new_snow_timestamp, _old_snow_timestamp = GRID.get_fresh_snow_props() + + _LAYER_HEIGHT = np.array(max_layers * [np.nan]) + _LAYER_RHO = np.array(max_layers * [np.nan]) + _LAYER_T = np.array(max_layers * [np.nan]) + _LAYER_LWC = np.array(max_layers * [np.nan]) + _LAYER_ICE_FRACTION = np.array(max_layers * [np.nan]) + if GRID.get_number_layers()>max_layers: + raise ValueError('Maximum number of layers reached') + _LAYER_HEIGHT[0:GRID.get_number_layers()] = GRID.get_height() + _LAYER_RHO[0:GRID.get_number_layers()] = GRID.get_density() + _LAYER_T[0:GRID.get_number_layers()] = GRID.get_temperature() + _LAYER_LWC[0:GRID.get_number_layers()] = GRID.get_liquid_water_content() + _LAYER_ICE_FRACTION[0:GRID.get_number_layers()] = GRID.get_ice_fraction() + + return (None,None,None,None,None,None,lw_radiation_out,sensible_heat_flux,latent_heat_flux, \ + ground_heat_flux,rain_heat_flux,mass_balance,surface_mass_balance,Q,_SNOWHEIGHT, \ + _TOTALHEIGHT,surface_temperature,alpha,_NLAYERS,melt_energy,internal_mass_balance, \ + evaporation,sublimation,condensation,deposition,water_refreezed,subsurface_melt,z0, \ + melt,_new_snow_height,_new_snow_timestamp,_old_snow_timestamp,MOL,_LAYER_HEIGHT, \ + _LAYER_RHO,_LAYER_T,_LAYER_LWC,None,None,_LAYER_ICE_FRACTION,None,None,None,None,None) + + if not WRF_X_CSPY: + if Config.stake_evaluation: + # Evaluate stakes + _stat = evaluate(stake_names, stake_data, _df) - if Config.stake_evaluation: - # Evaluate stakes - _stat = evaluate(stake_names, stake_data, _df) - else: - _stat = None - _df = None - # Restart - if not Config.WRF_X_CSPY: - new_snow_height, new_snow_timestamp, old_snow_timestamp = GRID.get_fresh_snow_props() + else: + _stat = None + _df = None + + # Restart RESTART.NLAYERS.values[:] = GRID.get_number_layers() - RESTART.NEWSNOWHEIGHT.values[:] = new_snow_height - RESTART.NEWSNOWTIMESTAMP.values[:] = new_snow_timestamp - RESTART.OLDSNOWTIMESTAMP.values[:] = old_snow_timestamp + RESTART.NEWSNOWHEIGHT.values[:] = _new_snow_height[t] + RESTART.NEWSNOWTIMESTAMP.values[:] = _new_snow_timestamp[t] + RESTART.OLDSNOWTIMESTAMP.values[:] = _old_snow_timestamp[t] RESTART.LAYER_HEIGHT[0:GRID.get_number_layers()] = GRID.get_height() RESTART.LAYER_RHO[0:GRID.get_number_layers()] = GRID.get_density() RESTART.LAYER_T[0:GRID.get_number_layers()] = GRID.get_temperature() RESTART.LAYER_LWC[0:GRID.get_number_layers()] = GRID.get_liquid_water_content() RESTART.LAYER_IF[0:GRID.get_number_layers()] = GRID.get_ice_fraction() - return (indY,indX,RESTART,_RAIN,_SNOWFALL,_LWin,_LWout,_H,_LE,_B, _QRR, \ + return (indY,indX,RESTART,_RAIN,_SNOWFALL,_LWin,_LWout,_H,_LE,_B,_QRR, \ _MB,_surfMB,_Q,_SNOWHEIGHT,_TOTALHEIGHT,_TS,_ALBEDO,_NLAYERS, \ _ME,_intMB,_EVAPORATION,_SUBLIMATION,_CONDENSATION,_DEPOSITION,_REFREEZE, \ - _subM,_Z0,_surfM, _MOL, \ + _subM,_Z0,_surfM,_new_snow_height,_new_snow_timestamp,_old_snow_timestamp,_MOL, \ _LAYER_HEIGHT,_LAYER_RHO,_LAYER_T,_LAYER_LWC,_LAYER_CC,_LAYER_POROSITY,_LAYER_ICE_FRACTION, \ _LAYER_IRREDUCIBLE_WATER,_LAYER_REFREEZE,stake_names,_stat,_df) diff --git a/cosipy/cpkernel/grid.py b/cosipy/cpkernel/grid.py index 247f8ba..94072a1 100644 --- a/cosipy/cpkernel/grid.py +++ b/cosipy/cpkernel/grid.py @@ -29,6 +29,7 @@ temperature_threshold_merging = Constants.temperature_threshold_merging density_threshold_merging = Constants.density_threshold_merging merge_max = Constants.merge_max +minimum_snowfall = Constants.minimum_snowfall minimum_snow_layer_height = Constants.minimum_snow_layer_height remesh_method = Constants.remesh_method ice_density = Constants.ice_density @@ -128,7 +129,7 @@ def init_grid(self): ) def add_fresh_snow( - self, height, density, temperature, liquid_water_content + self, height, density, temperature, liquid_water_content, dt ): """Add a fresh snow layer (node). @@ -141,6 +142,7 @@ def add_fresh_snow( temperature (float): Layer temperature [K]. liquid_water_content (float): Liquid water content of the layer [|m w.e.|]. + dt (int): Integration time [s]. """ # Add new node @@ -151,9 +153,12 @@ def add_fresh_snow( # Increase node counter self.number_nodes += 1 - # Set fresh snow properties for albedo calculation (height and - # timestamp) - self.set_fresh_snow_props(height) + if height < minimum_snowfall: + # Ignore impact of small snowfall on fresh snow layer properties + self.set_fresh_snow_props_update_time(dt) + else: + # Set the fresh snow properties for albedo calculation (height and timestamp) + self.set_fresh_snow_props(height) def remove_node(self, idx: list = None): """Remove a layer (node) from the grid (node list). @@ -427,11 +432,11 @@ def adaptive_profile(self): merged if: (1) the density difference between one layer and the next is - smaller than the user defined threshold. + smaller than the user defined threshold. (2) the temperature difference is smaller than the user defined - threshold. + threshold. (3) the number of merges per time step does not exceed the user - defined threshold. + defined threshold. The thresholds are defined by ``temperature_threshold_merging``, ``density_threshold_merging``, and ``merge_max`` in @@ -463,7 +468,21 @@ def adaptive_profile(self): else: idx += 1 - # Correct first layer + # Remesh ice + # remeshing layer 0 done by correct_layer above + min_ice_idx = max(1, self.get_number_snow_layers()) + # Ensure top ice layer has first_layer_height when thin snow layers will be removed in update_grid + if (min_ice_idx == 1) & (self.get_node_height(0) < first_layer_height): + self.correct_layer(min_ice_idx, first_layer_height) + min_ice_idx += 1 + + idx = min_ice_idx + while idx < self.get_number_layers() - 1: + # Correct first layer + if self.get_node_height(idx) < minimum_snow_layer_height: + self.merge_nodes(idx) + else: + idx += 1 self.correct_layer(0, first_layer_height) def split_node(self, pos: int): @@ -584,12 +603,12 @@ def merge_snow_with_glacier(self, idx: int): self.get_node_density(idx + 1) >= snow_ice_threshold ): # Update node properties - first_layer_height = self.get_node_height(idx) * ( + surface_layer_height = self.get_node_height(idx) * ( self.get_node_density(idx) / ice_density ) self.update_node( idx + 1, - self.get_node_height(idx + 1) + first_layer_height, + self.get_node_height(idx + 1) + surface_layer_height, self.get_node_temperature(idx + 1), self.get_node_ice_fraction(idx + 1), 0.0, @@ -599,7 +618,7 @@ def merge_snow_with_glacier(self, idx: int): # self.check('Merge snow with glacier function') - def remove_melt_weq(self, melt: float, idx: int=0): + def remove_melt_weq(self, melt: float, idx: int = 0) -> float: """Remove mass from a layer. Reduces the mass/height of layer ``idx`` by the available melt @@ -609,13 +628,13 @@ def remove_melt_weq(self, melt: float, idx: int=0): melt: Snow water equivalent of melt [|m w.e.|]. idx: Index of the layer. If no value is given, the function acts on the first layer. - + Returns: - float: Liquid water content from removed layers. + Liquid water content from removed layers. """ - lwc_from_layers = 0 + lwc_from_layers = 0.0 - while melt > 0: + while melt > 0.0 and idx < self.number_nodes: # Get SWE of layer SWE = self.get_node_height(idx) * ( self.get_node_density(idx) / water_density @@ -673,7 +692,7 @@ def set_fresh_snow_props_update_time(self, seconds: float): """Update the timestamp of the snow properties. Args: - seconds : seconds without snowfall [s]. + seconds: seconds without snowfall [s]. """ self.old_snow_timestamp = self.old_snow_timestamp + seconds # Set the timestamp to zero @@ -687,11 +706,16 @@ def set_fresh_snow_props_height(self, height: float): """ self.new_snow_height = height - def get_fresh_snow_props(self): + def get_fresh_snow_props(self) -> tuple: """Get the first snow layer's properties. This is used internally to track the albedo properties of the first snow layer. + + Returns: + First snow layer's updated height, time elapsed since the + last snowfall, and the time elapsed between the last and + penultimate snowfall. """ return ( self.new_snow_height, @@ -757,7 +781,7 @@ def set_refreeze(self, refreeze: np.ndarray): for idx in range(self.number_nodes): self.grid[idx].set_layer_refreeze(refreeze[idx]) - def get_temperature(self): + def get_temperature(self) -> list: """Get the temperature profile.""" return [ self.grid[idx].get_layer_temperature() @@ -768,7 +792,7 @@ def get_node_temperature(self, idx: int): """Get a node's temperature.""" return self.grid[idx].get_layer_temperature() - def get_specific_heat(self): + def get_specific_heat(self) -> list: """Get the specific heat capacity profile (air+water+ice).""" return [ self.grid[idx].get_layer_specific_heat() @@ -779,21 +803,21 @@ def get_node_specific_heat(self, idx: int): """Get a node's specific heat capacity (air+water+ice).""" return self.grid[idx].get_layer_specific_heat() - def get_height(self): + def get_height(self) -> list: """Get the heights of all the layers.""" return [ self.grid[idx].get_layer_height() for idx in range(self.number_nodes) ] - def get_snow_heights(self): + def get_snow_heights(self) -> list: """Get the heights of the snow layers.""" return [ self.grid[idx].get_layer_height() for idx in range(self.get_number_snow_layers()) ] - def get_ice_heights(self): + def get_ice_heights(self) -> list: """Get the heights of the ice layers.""" return [ self.grid[idx].get_layer_height() @@ -809,7 +833,7 @@ def get_node_density(self, idx: int): """Get a node's density.""" return self.grid[idx].get_layer_density() - def get_density(self): + def get_density(self) -> list: """Get the density profile.""" return [ self.grid[idx].get_layer_density() @@ -820,7 +844,7 @@ def get_node_liquid_water_content(self, idx: int): """Get a node's liquid water content.""" return self.grid[idx].get_layer_liquid_water_content() - def get_liquid_water_content(self): + def get_liquid_water_content(self) -> list: """Get a profile of the liquid water content.""" return [ self.grid[idx].get_layer_liquid_water_content() @@ -831,7 +855,7 @@ def get_node_ice_fraction(self, idx: int): """Get a node's ice fraction.""" return self.grid[idx].get_layer_ice_fraction() - def get_ice_fraction(self): + def get_ice_fraction(self) -> list: """Get a profile of the ice fraction.""" return [ self.grid[idx].get_layer_ice_fraction() @@ -842,7 +866,7 @@ def get_node_irreducible_water_content(self, idx: int): """Get a node's irreducible water content.""" return self.grid[idx].get_layer_irreducible_water_content() - def get_irreducible_water_content(self): + def get_irreducible_water_content(self) -> list: """Get a profile of the irreducible water content.""" return [ self.grid[idx].get_layer_irreducible_water_content() @@ -853,7 +877,7 @@ def get_node_cold_content(self, idx: int): """Get a node's cold content.""" return self.grid[idx].get_layer_cold_content() - def get_cold_content(self): + def get_cold_content(self) -> list: """Get the cold content profile.""" return [ self.grid[idx].get_layer_cold_content() @@ -864,7 +888,7 @@ def get_node_porosity(self, idx: int): """Get a node's porosity.""" return self.grid[idx].get_layer_porosity() - def get_porosity(self): + def get_porosity(self) -> list: """Get the porosity profile.""" return [ self.grid[idx].get_layer_porosity() @@ -875,7 +899,7 @@ def get_node_thermal_conductivity(self, idx: int): """Get a node's thermal conductivity.""" return self.grid[idx].get_layer_thermal_conductivity() - def get_thermal_conductivity(self): + def get_thermal_conductivity(self) -> list: """Get the thermal conductivity profile.""" return [ self.grid[idx].get_layer_thermal_conductivity() @@ -886,7 +910,7 @@ def get_node_thermal_diffusivity(self, idx: int): """Get a node's thermal diffusivity.""" return self.grid[idx].get_layer_thermal_diffusivity() - def get_thermal_diffusivity(self): + def get_thermal_diffusivity(self) -> list: """Get the thermal diffusivity profile.""" return [ self.grid[idx].get_layer_thermal_diffusivity() @@ -897,7 +921,7 @@ def get_node_refreeze(self, idx: int): """Get the amount of refrozen water in a node.""" return self.grid[idx].get_layer_refreeze() - def get_refreeze(self): + def get_refreeze(self) -> list: """Get the profile of refrozen water.""" return [ self.grid[idx].get_layer_refreeze() @@ -906,21 +930,20 @@ def get_refreeze(self): def get_node_depth(self, idx: int): """Get a node's depth relative to the surface.""" - d = 0 - for i in range(idx + 1): - if i == 0: - d = d + self.get_node_height(i) / 2.0 - else: - d = ( - d - + self.get_node_height(i - 1) / 2.0 - + self.get_node_height(i) / 2.0 - ) + d = self.get_node_height(idx) / 2.0 + if idx != 0: + for i in range(idx): + d = d + self.get_node_height(i) return d - def get_depth(self): + def get_depth(self) -> list: """Get a depth profile.""" - return [self.get_node_depth(idx) for idx in range(self.number_nodes)] + h = np.array(self.get_height()) + z = np.empty_like(h) # faster than copy + z[0] = 0.5 * h[0] + z[1:] = np.cumsum(h[:-1]) + (0.5 * h[1:]) + + return [z[idx] for idx in range(self.number_nodes)] def get_total_snowheight(self, verbose=False): """Get the total snowheight (density maximum or np.nanmin(property) < minimum: + if ( + np.nanmax(layer_property) > maximum + or np.nanmin(layer_property) < minimum + ): print( str.capitalize(name), "max", - np.nanmax(property), + np.nanmax(layer_property), "min", - np.nanmin(property), + np.nanmin(layer_property), ) os._exit() diff --git a/cosipy/cpkernel/init.py b/cosipy/cpkernel/init.py index cd83728..264b5ec 100644 --- a/cosipy/cpkernel/init.py +++ b/cosipy/cpkernel/init.py @@ -1,14 +1,16 @@ -import sys - import numpy as np +from numba import njit from cosipy.constants import Constants from cosipy.cpkernel.grid import Grid -#from cosipy.utils.options import read_opt -def init_snowpack(DATA, opt_dict=None): + +def init_snowpack(DATA): """Initialise the snowpack. + Args: + DATA (xarray.Dataset): Glacier data for a single grid point. + Returns: Grid: Initialised glacier data structure with snowpack. """ @@ -21,81 +23,96 @@ def init_snowpack(DATA, opt_dict=None): initial_top_density_snowpack = Constants.initial_top_density_snowpack initial_bottom_density_snowpack = Constants.initial_bottom_density_snowpack ice_density = Constants.ice_density - # Read and set options - - #if opt_dict is not None: - # mult_factor_RRR = opt_dict[0] - # albedo_ice = opt_dict[1] - # albedo_fresh_snow = opt_dict[2] - # albedo_firn = opt_dict[3] - # albedo_mod_snow_aging = opt_dict[4] - # albedo_mod_snow_depth = opt_dict[5] - # center_snow_transfer_function = opt_dict[6] - # spread_snow_transfer_function = opt_dict[7] - # roughness_fresh_snow = opt_dict[8] - # roughness_ice = opt_dict[9] - # roughness_firn = opt_dict[10] - #read_opt(opt_dict, globals()) - - ##-------------------------------------- - ## Check for WRF data - ##-------------------------------------- - if ('SNOWHEIGHT' in DATA): + + # Check for WRF data + if "SNOWHEIGHT" in DATA: initial_snowheight = DATA.SNOWHEIGHT.values if np.isnan(initial_snowheight): initial_snowheight = 0.0 - else: + else: initial_snowheight = initial_snowheight_constant - temperature_top = np.minimum(DATA.T2.values[0], 273.16) + temperature_top = np.minimum(DATA.T2.values[0], Constants.zero_temperature) - #-------------------------------------- # Do the vertical interpolation - #-------------------------------------- layer_heights = [] layer_densities = [] layer_T = [] layer_liquid_water = [] - - if (initial_snowheight > 0.0): - optimal_height = 0.1 # 10 cm + + if initial_snowheight > 0.0: + optimal_height = 0.1 # 10 cm nlayers = int(min(initial_snowheight / optimal_height, 5)) - dT = (temperature_top-temperature_bottom)/(initial_snowheight+initial_glacier_height) + dT = (temperature_top - temperature_bottom) / ( + initial_snowheight + initial_glacier_height + ) if nlayers == 0: layer_heights = np.array([initial_snowheight]) layer_densities = np.array([initial_top_density_snowpack]) layer_T = np.array([temperature_top]) layer_liquid_water = np.array([0.0]) elif nlayers > 0: - drho = (initial_top_density_snowpack-initial_bottom_density_snowpack)/initial_snowheight - layer_heights = np.ones(nlayers) * (initial_snowheight/nlayers) - layer_densities = np.array([initial_top_density_snowpack-drho*(initial_snowheight/nlayers)*i for i in range(1,nlayers+1)]) - layer_T = np.array([temperature_top-dT*(initial_snowheight/nlayers)*i for i in range(1,nlayers+1)]) + drho = ( + initial_top_density_snowpack - initial_bottom_density_snowpack + ) / initial_snowheight + layer_heights = np.ones(nlayers) * (initial_snowheight / nlayers) + layer_densities = np.array( + [ + initial_top_density_snowpack + - drho * (initial_snowheight / nlayers) * i + for i in range(1, nlayers + 1) + ] + ) + layer_T = np.array( + [ + temperature_top - dT * (initial_snowheight / nlayers) * i + for i in range(1, nlayers + 1) + ] + ) layer_liquid_water = np.zeros(nlayers) # add glacier - nlayers = int(initial_glacier_height/initial_glacier_layer_heights) - layer_heights = np.append(layer_heights, np.ones(nlayers)*initial_glacier_layer_heights) - layer_densities = np.append(layer_densities, np.ones(nlayers)*ice_density) - layer_T = np.append(layer_T, [layer_T[-1]-dT*initial_glacier_layer_heights*i for i in range(1,nlayers+1)]) + nlayers = int(initial_glacier_height / initial_glacier_layer_heights) + layer_heights = np.append( + layer_heights, np.ones(nlayers) * initial_glacier_layer_heights + ) + layer_densities = np.append( + layer_densities, np.ones(nlayers) * ice_density + ) + layer_T = np.append( + layer_T, + [ + layer_T[-1] - dT * initial_glacier_layer_heights * i + for i in range(1, nlayers + 1) + ], + ) layer_liquid_water = np.append(layer_liquid_water, np.zeros(nlayers)) - else: - - # only glacier - nlayers = int(initial_glacier_height/initial_glacier_layer_heights) - dT = (temperature_top-temperature_bottom)/initial_glacier_height - layer_heights = np.ones(nlayers)*initial_glacier_layer_heights - layer_densities = np.ones(nlayers)*ice_density - layer_T = np.array([temperature_top-dT*initial_glacier_layer_heights*i for i in range(1,nlayers+1)]) + else: # only glacier + nlayers = int(initial_glacier_height / initial_glacier_layer_heights) + dT = (temperature_top - temperature_bottom) / initial_glacier_height + layer_heights = np.ones(nlayers) * initial_glacier_layer_heights + layer_densities = np.ones(nlayers) * ice_density + layer_T = np.array( + [ + temperature_top - dT * initial_glacier_layer_heights * i + for i in range(1, nlayers + 1) + ] + ) layer_liquid_water = np.zeros(nlayers) # Initialize grid, the grid class contains all relevant grid information - GRID = Grid(np.array(layer_heights, dtype=np.float64), np.array(layer_densities, dtype=np.float64), - np.array(layer_T, dtype=np.float64), np.array(layer_liquid_water, dtype=np.float64), - None, None, None, None) - - return GRID + GRID = create_grid_jitted( + np.array(layer_heights, dtype=np.float64), + np.array(layer_densities, dtype=np.float64), + np.array(layer_T, dtype=np.float64), + np.array(layer_liquid_water, dtype=np.float64), + None, + None, + None, + None, + ) + return GRID def load_snowpack(GRID_RESTART): @@ -103,7 +120,7 @@ def load_snowpack(GRID_RESTART): # Number of layers num_layers = int(GRID_RESTART.NLAYERS.values) - + # Init layer height # Weird slicing position to accommodate NestedNamespace in WRF_X_CSPY layer_heights = GRID_RESTART.LAYER_HEIGHT.values[0:num_layers] @@ -115,12 +132,51 @@ def load_snowpack(GRID_RESTART): new_snow_height = np.float64(GRID_RESTART.new_snow_height.values) new_snow_timestamp = np.float64(GRID_RESTART.new_snow_timestamp.values) old_snow_timestamp = np.float64(GRID_RESTART.old_snow_timestamp.values) - - GRID = Grid(layer_heights, layer_density, layer_T, layer_LWC, layer_IF, new_snow_height, - new_snow_timestamp, old_snow_timestamp) + + GRID = create_grid_jitted( + layer_heights, + layer_density, + layer_T, + layer_LWC, + layer_IF, + new_snow_height, + new_snow_timestamp, + old_snow_timestamp, + ) + + return GRID + + +@njit +def create_grid_jitted( + layer_heights: np.ndarray, + layer_density: np.ndarray, + layer_T: np.ndarray, + layer_LWC: np.ndarray, + layer_IF: np.ndarray, + new_snow_height: float, + new_snow_timestamp: float, + old_snow_timestamp: float, +): + """Create Grid with JIT. + + Returns: + Grid: Glacier data structure. + """ + + GRID = Grid( + layer_heights, + layer_density, + layer_T, + layer_LWC, + layer_IF, + new_snow_height, + new_snow_timestamp, + old_snow_timestamp, + ) if np.isnan(layer_T).any(): GRID.grid_info_screen() - sys.exit(1) - + raise ValueError("NaNs encountered in GRID creation.") + return GRID diff --git a/cosipy/cpkernel/io.py b/cosipy/cpkernel/io.py index 22c0840..8c0c56b 100644 --- a/cosipy/cpkernel/io.py +++ b/cosipy/cpkernel/io.py @@ -2,134 +2,190 @@ Read the input data (model forcing) and write the output to netCDF file. """ -import configparser -import inspect import os -import sys +import warnings +from datetime import datetime import numpy as np -import pandas as pd import xarray as xr from cosipy.config import Config from cosipy.constants import Constants + class IOClass: def __init__(self, DATA=None, opt_dict=None): - """Initialise the IO Class""" + """Initialise the IO Class. + Attributes: + DATA (xarray.Dataset or None): Meteorological data. + RESTART (xarray.Dataset or None): Restart file data. + RESULT (xarray.Dataset or None): Model result data. + opt_dict (tuple or None): Model parameter tuples. """ - # Read and set options - if opt_dict is not None: - mult_factor_RRR = opt_dict[0] - albedo_ice = opt_dict[1] - albedo_fresh_snow = opt_dict[2] - albedo_firn = opt_dict[3] - albedo_mod_snow_aging = opt_dict[4] - albedo_mod_snow_depth = opt_dict[5] - center_snow_transfer_function = opt_dict[6] - spread_snow_transfer_function = opt_dict[7] - roughness_fresh_snow = opt_dict[8] - roughness_ice = opt_dict[9] - roughness_firn = opt_dict[10] - else: - mult_factor_RRR = Constants.mult_factor_RRR - albedo_ice = Constants.albedo_ice - albedo_fresh_snow = Constants.albedo_fresh_snow - albedo_firn = Constants.albedo_firn - albedo_mod_snow_aging = Constants.albedo_mod_snow_aging - albedo_mod_snow_depth = Constants.albedo_mod_snow_depth - center_snow_transfer_function = Constants.center_snow_transfer_function - spread_snow_transfer_function = Constants.spread_snow_transfer_function - roughness_fresh_snow = Constants.roughness_fresh_snow - roughness_ice = Constants.roughness_ice - roughness_firn = Constants.roughness_firn - #the above needs to be put into attrs - #read_opt(opt_dict, globals()) - """ - output_vars = self.get_output_structure() - self.atm = output_vars['vars']['atm'] - self.internal = output_vars['vars']['internal'] - self.full = output_vars['vars']['full'] + + self.atm = self.get_output_variables(Config.output_atm) + self.internal = self.get_output_variables(Config.output_internal) + self.full = self.get_output_variables(Config.output_full) # Initialize data self.DATA = DATA self.RESTART = None self.RESULT = None - + # If local IO class is initialized we need to get the dimensions of the dataset if DATA is not None: - self.time = self.DATA.sizes['time'] + self.time = self.DATA.sizes["time"] - def get_output_structure(self): - """Get the model output variables. - - Returns: - Output variables for internal and full-field simulations. + def get_output_variables(self, variables: str) -> list: + """Get output variable names from configuration string.""" + # Sets are unordered -> same config, different checksums + if not variables: + return [] + else: + return [item for item in variables.split(",")] + + def init_atm_attribute(self, name: str): + """Initialise empty atm attribute.""" + if name in self.atm: + setattr(self, name, self.create_nan_array()) + + def init_internal_attribute(self, name: str): + """Initialise empty internal attribute.""" + if name in self.internal: + setattr(self, name, self.create_nan_array()) + + def init_full_field_attribute(self, name: str, max_layers: int): + """Initialise empty layer attribute.""" + if name in self.full: + setattr(self, f"LAYER_{name}", self.create_3d_nan_array(max_layers)) + + def set_atm_attribute(self, name: str, value: np.ndarray, x: int, y: int): + """Set atm attribute if it is a desired output variable. - Raises: - FileNotFoundError: If the "output" file is not found. + Args: + name: Output variable name. + value: Output variable data. + x: Slice x-coordinate. + y: Slice y-coordinate. + """ + if name in self.atm: + getattr(self, name)[:, y, x] = value + + def set_internal_attribute(self, name: str, value: np.ndarray, x: int, y: int): + """Set internal attribute if it is a desired output variable. + + Args: + name: Output variable name. + value: Output variable data. + x: Slice x-coordinate. + y: Slice y-coordinate. """ - # Package is not installed in working directory - filename = inspect.getfile(inspect.currentframe()) - filename = f"{filename[:-14]}output" - if not os.path.isfile(filename): - raise FileNotFoundError(f"{filename} not found.") + if name in self.internal: + getattr(self, name)[:, y, x] = value + + def set_full_field_attribute(self, name: str, value: np.ndarray, x: int, y: int): + """Set layer attribute if it is a desired output variable. - # Read variable list from file - output_structure = configparser.ConfigParser() - output_structure.read(filename) + Args: + name: Output variable name. + value: Output variable data. + x: Slice x-coordinate. + y: Slice y-coordinate. + """ + if name in self.full: + getattr(self, f"LAYER_{name}")[:, y, x, :] = value + + def get_datetime( + self, timestamp: str, use_np: bool = True, fmt: str = "%Y-%m-%dT%H:%M" + ): + """Get datetime object from a string. + + Args: + timestamp: Timestamp. + use_np: Convert to numpy datetime64. Default True. + fmt: Timestamp format. - return output_structure + Returns: + datetime|np.datetime64: Timestamp converted to datetime or + np.datetime64. + """ + if isinstance(timestamp, str): + if use_np: + return np.datetime64(timestamp) + else: + return datetime.strptime(timestamp, fmt) + else: + return timestamp - def create_data_file(self): + def create_data_file(self) -> xr.Dataset: """Create the input data and read the restart file if necessary. Returns: - xarray.Dataset: Model input data. + Model input data. """ - + if Config.restart: - print('--------------------------------------------------------------') - print('\t RESTART FROM PREVIOUS STATE') - print('-------------------------------------------------------------- \n') - + print(f"{'-'*62}\n\tRESTART FROM PREVIOUS STATE\n{'-'*62}\n") + # Load the restart file - timestamp = pd.to_datetime(Config.time_start).strftime('%Y-%m-%dT%H-%M') - if (os.path.isfile(os.path.join(Config.data_path, 'restart', 'restart_'+timestamp+'.nc')) & (Config.time_start != Config.time_end)): - self.GRID_RESTART = xr.open_dataset(os.path.join(Config.data_path, 'restart', 'restart_'+timestamp+'.nc')) - self.restart_date = self.GRID_RESTART.time+np.timedelta64(Constants.dt,'s') # Get time of the last calculation and add one time step - self.init_data_dataset() # Read data from the last date to the end of the data file - else: - print('No restart file available for the given date %s' % (timestamp)) # if there is a problem kill the program - print('OR start date %s equals end date %s \n' % (Config.time_start, Config.time_end)) - sys.exit(1) + time_start = Config.time_start + time_end = Config.time_end + start_timestamp = self.get_datetime(time_start) + end_timestamp = self.get_datetime(time_end) + timestamp = start_timestamp.strftime("%Y-%m-%dT%H-%M") + restart_path = os.path.join( + Config.data_path, "restart", f"restart_{timestamp}.nc" + ) + try: + if not os.path.isfile(restart_path): + raise FileNotFoundError + elif start_timestamp == end_timestamp: + raise IndexError + else: + self.GRID_RESTART = xr.open_dataset(restart_path) + """Get time of the last calculation and add one time + step. GRID_RESTART.time is an array of np.datetime64 + objects. + """ + self.restart_date = self.GRID_RESTART.time.values + np.timedelta64( + Constants.dt, "s" + ) + # Read data from the last date to the end of the data file + self.init_data_dataset() + except FileNotFoundError: + raise SystemExit(f"No restart file available for the given date: {timestamp}") + except IndexError: + raise SystemExit(f"Start date {time_start} equals end date {time_end}\n") else: + # If no restart, read data according to the dates defined in config file self.restart_date = None - self.init_data_dataset() # If no restart, read data according to the dates defined in the config.py + self.init_data_dataset() - #---------------------------------------------- - # Tile the data is desired - #---------------------------------------------- - if Config.tile: + if Config.tile: # Tile the data if desired if Config.WRF: - self.DATA = self.DATA.isel(south_north=slice(Config.ystart,Config.yend), west_east=slice(Config.xstart,Config.xend)) + self.DATA = self.DATA.isel( + south_north=slice(Config.ystart, Config.yend), + west_east=slice(Config.xstart, Config.xend), + ) else: - self.DATA = self.DATA.isel(lat=slice(Config.ystart,Config.yend), lon=slice(Config.xstart,Config.xend)) + self.DATA = self.DATA.isel( + lat=slice(Config.ystart, Config.yend), + lon=slice(Config.xstart, Config.xend), + ) self.ny = self.DATA.sizes[Config.northing] self.nx = self.DATA.sizes[Config.easting] - self.time = self.DATA.sizes['time'] + self.time = self.DATA.sizes["time"] return self.DATA - def create_result_file(self, opt_dict=None) -> xr.Dataset: """Create and initialise the RESULT dataset.""" self.init_result_dataset(opt_dict=opt_dict) return self.RESULT - + def create_restart_file(self) -> xr.Dataset: """Create and initialise the RESTART dataset.""" self.init_restart_dataset() @@ -143,10 +199,66 @@ def create_grid_restart(self): """ return self.GRID_RESTART + def create_nan_array(self) -> np.ndarray: + """Create and fill a NaN array with time, (x,y) dimensions. + + Returns: + Filled array with time and 2D spatial coordinates. + """ + + return np.full((self.time, self.ny, self.nx), np.nan) + + def create_3d_nan_array(self, max_layers: int) -> np.ndarray: + """Create and fill a NaN array with time, (x,y,z) dimensions. + + Args: + The maximum number of model layers. + + Returns: + Filled array with time and 3D spatial coordinates. + """ + + return np.full((self.time, self.ny, self.nx, max_layers), np.nan) + + def check_field(self, field, _max, _min) -> bool: + """Check the validity of the input data.""" + if np.nanmax(field) > _max or np.nanmin(field) < _min: + print( + f"Please check the input data, it seems they are out of range:" + f"\n{field.name} ... MAX: {np.nanmax(field):.2f} " + f"MIN: {np.nanmin(field):.2f}\n" + ) + return False + else: + return True + + def check_input_data(self) -> bool: + """Check the input data is within valid bounds.""" + print(f"{'-'*62}\nChecking input data ....\n") + + data_bounds = { + "T2": (313.16, 243.16), + "RH2": (100.0, 0.0), + "G": (1600.0, 0.0), + "U2": (50.0, 0.0), + "RRR": (20.0, 0.0), + "N": (1.0, 0.0), + "PRES": (1080.0, 400.0), + "LWin": (400.0, 200.0), + "SNOWFALL": (0.1, 0.0), + "SLOPE": (0.0, 90.0), + } + + for key, bounds in data_bounds.items(): + if key in self.DATA: + if self.check_field(self.DATA[key], bounds[0], bounds[1]): + print(f"{key} ... ok") + + return True def init_data_dataset(self): """Read and store the input netCDF data. - + The input data should contain the following variables: :PRES: Air pressure [hPa]. :N: Cloud cover fraction [-]. @@ -158,81 +270,68 @@ def init_data_dataset(self): :U2: Wind speed (magnitude) [|m s^-1|]. :HGT: Elevation [m]. """ - - # Open input dataset - self.DATA = xr.open_dataset(os.path.join(Config.data_path,'input',Config.input_netcdf)) - self.DATA['time'] = np.sort(self.DATA['time'].values) - start_interval=str(self.DATA.time.values[0])[0:16] - end_interval = str(self.DATA.time.values[-1])[0:16] - time_steps = str(self.DATA.sizes['time']) - print('\n Maximum available time interval from %s until %s. Time steps: %s \n\n' % (start_interval, end_interval, time_steps)) - - # Check if restart option is set - if self.restart_date is None: - print('--------------------------------------------------------------') - print('\t Integration from %s to %s' % (Config.time_start, Config.time_end)) - print('--------------------------------------------------------------\n') - self.DATA = self.DATA.sel(time=slice(Config.time_start, Config.time_end)) # Select dates from config.py + + try: + input_path = os.path.join(Config.data_path, "input", Config.input_netcdf) + self.DATA = xr.open_dataset(input_path) + except FileNotFoundError: + raise SystemExit(f"Input file not found at: {input_path}") + + + self.DATA["time"] = np.sort(self.DATA["time"].values) + minimum_time = str(self.DATA.time.values[0])[0:16] + maximum_time = str(self.DATA.time.values[-1])[0:16] + start_interval = self.get_datetime(minimum_time) + end_interval = self.get_datetime(maximum_time) + time_steps = str(self.DATA.sizes["time"]) + print( + f"\nMaximum available time interval from {minimum_time} " + f"until {maximum_time}. Time steps: {time_steps}\n\n" + ) + + time_start = Config.time_start # avoid repeat calls + time_end = Config.time_end + start_time = self.get_datetime(time_start) + end_time = self.get_datetime(time_end) + + if (start_time > end_interval) or (end_time < start_interval): + raise IndexError("Selected period not available in input data.\n") + if start_time < start_interval: + warnings.warn( + "\nWARNING! Selected startpoint before first timestep of input data\n", + ) + if end_time > end_interval: + warnings.warn( + "\nWARNING! Selected endpoint after last timestep of input data\n", + ) + + if self.restart_date is None: # Check if restart option is set + print( + f"{'-'*62}\n\tIntegration from {time_start} to {time_end}\n{'-'*62}\n" + ) + self.DATA = self.DATA.sel(time=slice(time_start, time_end)) else: # There is nothing to do if the dates are equal - if (self.restart_date==Config.time_end): - print('Start date equals end date ... no new data ... EXIT') - sys.exit(1) + if self.restart_date == end_time: + raise SystemExit("Start date equals end date ... no new data ... EXIT") else: # otherwise, run the model from the restart date to the defined end date - print('Starting from %s (from restart file) to %s (from config.py) \n' % (self.restart_date.values, Config.time_end)) - self.DATA = self.DATA.sel(time=slice(self.restart_date, Config.time_end)) - - if Config.time_start < start_interval: - print('\n WARNING! Selected startpoint before first timestep of input data\n') - if Config.time_end > end_interval: - print('\n WARNING! Selected endpoint after last timestep of input data\n') - if Config.time_start > end_interval or Config.time_end < start_interval: - print('\n ERROR! Selected period not available in input data\n') + print( + f"Starting from {self.restart_date} (from restart file) " + f"to {time_end} (from config file)\n" + ) + self.DATA = self.DATA.sel(time=slice(self.restart_date, time_end)) + self.check_input_data() + print(f"\nGlacier gridpoints: {np.nansum(self.DATA.MASK >= 1)} \n\n") - print('--------------------------------------------------------------') - print('Checking input data .... \n') - - # Define an auxiliary function to check the validity of the data - def check(field, max, min): - """Check the validity of the input data.""" - if np.nanmax(field) > max or np.nanmin(field) < min: - print('Please check the input data, its seems they are out of range %s MAX: %.2f MIN: %.2f \n' % (str.capitalize(field.name), np.nanmax(field), np.nanmin(field))) - # Check if data is within valid bounds - if ('T2' in self.DATA): - print('Temperature data (T2) ... ok ') - check(self.DATA.T2, 313.16, 243.16) - if ('RH2' in self.DATA): - print('Relative humidity data (RH2) ... ok ') - check(self.DATA.RH2, 100.0, 0.0) - if ('G' in self.DATA): - print('Shortwave data (G) ... ok ') - check(self.DATA.G, 1600.0, 0.0) - if ('U2' in self.DATA): - print('Wind velocity data (U2) ... ok ') - check(self.DATA.U2, 50.0, 0.0) - if ('RRR' in self.DATA): - print('Precipitation data (RRR) ... ok ') - check(self.DATA.RRR, 20.0, 0.0) - if ('N' in self.DATA): - print('Cloud cover data (N) ... ok ') - check(self.DATA.N, 1.0, 0.0) - if ('PRES' in self.DATA): - print('Pressure data (PRES) ... ok ') - check(self.DATA.PRES, 1080.0, 400.0) - if ('LWin' in self.DATA): - print('Incoming longwave data (LWin) ... ok ') - check(self.DATA.LWin, 400.0, 200.0) - if ('SNOWFALL' in self.DATA): - print('Snowfall data (SNOWFALL) ... ok ') - check(self.DATA.SNOWFALL, 0.1, 0.0) - - print('\n Glacier gridpoints: %s \n\n' %(np.nansum(self.DATA.MASK>=1))) + def get_input_metadata(self) -> tuple: + """Get input variable names and units. - - def get_result_metadata(self) -> tuple: - """Get variable names and units.""" + Returns: + tuple[dict, dict]: Metadata for spatial and spatiotemporal + variables, including netCDF keyname, unit, and long name. + """ metadata_spatial = { "HGT": ("m", "Elevation"), "MASK": ("boolean", "Glacier mask"), @@ -246,7 +345,7 @@ def get_result_metadata(self) -> tuple: "U2": ("m s\u207b\xb9", "Wind velocity at 2 m"), "PRES": ("hPa", "Atmospheric pressure"), "G": ("W m\u207b\xb2", "Incoming shortwave radiation"), - "RRR": ("mm", "Total precipiation"), + "RRR": ("mm", "Total precipitation"), "SNOWFALL": ("m", "Snowfall"), "N": ("-", "Cloud fraction"), "LWin": ("W m\u207b\xb2", "Incoming longwave radiation"), @@ -254,7 +353,90 @@ def get_result_metadata(self) -> tuple: return metadata_spatial, metadata_spatiotemporal - def init_result_dataset(self, opt_dict=None) -> xr.Dataset: + def get_full_field_metadata(self) -> dict: + """Get layer variable names and units. + + Returns: + Metadata for full field layer variables, including netCDF + keyname, unit, and long name. + """ + metadata = { + "LAYER_HEIGHT": ("m", "Layer height"), + "LAYER_RHO": ("kg m^-3", "Layer density"), + "LAYER_T": ("K", "Layer temperature"), + "LAYER_LWC": ("kg m^-2", "layer liquid water content"), + "LAYER_CC": ("J m^-2", "Cold content"), + "LAYER_POROSITY": ("-", "Porosity"), + "LAYER_ICE_FRACTION": ("-", "Layer ice fraction"), + "LAYER_IF": ("-", "Layer ice fraction"), # RESTART compatibility + "LAYER_IRREDUCIBLE_WATER": ("-", "Irreducible water"), + "LAYER_REFREEZE": ("m w.e.", "Refreezing"), + } + + return metadata + + def get_restart_metadata(self) -> dict: + + field_metadata = self.get_full_field_metadata() + restart_metadata = { + "new_snow_height": ("m .w.e", "New snow height"), + "new_snow_timestamp": ("s", "New snow timestamp"), + "old_snow_timestamp": ("s", "Old snow timestamp"), + "NLAYERS": ("-", "Number of layers"), + "NEWSNOWHEIGHT": ("m .w.e", "New snow height"), + "NEWSNOWTIMESTAMP": ("s", "New snow timestamp"), + "OLDSNOWTIMESTAMP": ("s", "Old snow timestamp"), + } + metadata = restart_metadata | field_metadata + + return metadata + + def get_result_metadata(self) -> dict: + """Get all variable names and units. + + Returns: + Metadata for all input and output variables, including + netCDF keyname, unit, and long name. + """ + metadata_spatial, metadata_spatiotemporal = self.get_input_metadata() + metadata_full = self.get_full_field_metadata() + metadata_result = { + "RAIN": ("mm", "Liquid precipitation"), + "SNOWFALL": ("m w.e.", "Snowfall"), + "LWin": ("W m\u207b\xb2", "Incoming longwave radiation"), + "LWout": ("W m\u207b\xb2", "Outgoing longwave radiation"), + "H": ("W m\u207b\xb2", "Sensible heat flux"), + "LE": ("W m\u207b\xb2", "Latent heat flux"), + "B": ("W m\u207b\xb2", "Ground heat flux"), + "QRR": ("W m\u207b\xb2", "Rain heat flux"), + "surfMB": ("m w.e.", "Surface mass balance"), + "MB": ("m w.e.", "Mass balance"), + "Q": ("m w.e.", "Runoff"), + "SNOWHEIGHT": ("m", "Snowheight"), + "TOTALHEIGHT": ("m", "Total domain height"), + "TS": ("K", "Surface temperature"), + "ALBEDO": ("-", "Albedo"), + "LAYERS": ("-", "Number of layers"), + "ME": ("W m\u207b\xb2", "Available melt energy"), + "intMB": ("m w.e.", "Internal mass balance"), + "EVAPORATION": ("m w.e.", "Evaporation"), + "SUBLIMATION": ("m w.e.", "Sublimation"), + "CONDENSATION": ("m w.e.", "Condensation"), + "DEPOSITION": ("m w.e.", "Deposition"), + "REFREEZE": ("m w.e.", "Refreezing"), + "subM": ("m w.e.", "Subsurface melt"), + "Z0": ("m", "Roughness length"), + "surfM": ("m w.e.", "Surface melt"), + "MOL": ("m", "Monin Obukhov length"), + } + # metadata_result overwrites items in previous union! + metadata = ( + metadata_spatial | metadata_spatiotemporal | metadata_full | metadata_result + ) + + return metadata + + def init_result_dataset(self, opt_dict=opt_dict) -> xr.Dataset: """Create the final dataset to aggregate and store the results. Aggregates results from individual COSIPY runs. After the @@ -296,86 +478,86 @@ def init_result_dataset(self, opt_dict=None) -> xr.Dataset: # Coordinates self.RESULT = xr.Dataset() - self.RESULT.coords['time'] = self.DATA.coords['time'] - self.RESULT.coords['lat'] = self.DATA.coords['lat'] - self.RESULT.coords['lon'] = self.DATA.coords['lon'] + self.RESULT.coords["time"] = self.DATA.coords["time"] + self.RESULT.coords["lat"] = self.DATA.coords["lat"] + self.RESULT.coords["lon"] = self.DATA.coords["lon"] # Global attributes from config.py - self.RESULT.attrs['Start_from_restart_file'] = str(Config.restart) - self.RESULT.attrs['Stake_evaluation'] = str(Config.stake_evaluation) - self.RESULT.attrs['WRF_simulation'] = str(Config.WRF) - self.RESULT.attrs['Compression_level'] = Config.compression_level - self.RESULT.attrs['Slurm_use'] = str(Config.slurm_use) - self.RESULT.attrs['Full_fiels'] = str(Config.full_field) - self.RESULT.attrs['Force_use_TP'] = str(Config.force_use_TP) - self.RESULT.attrs['Force_use_N'] = str(Config.force_use_N) - self.RESULT.attrs['Tile_of_glacier_of_interest'] = str(Config.tile) + self.RESULT.attrs["Start_from_restart_file"] = str(Config.restart) + self.RESULT.attrs["Stake_evaluation"] = str(Config.stake_evaluation) + self.RESULT.attrs["WRF_simulation"] = str(Config.WRF) + self.RESULT.attrs["Compression_level"] = Config.compression_level + self.RESULT.attrs["Slurm_use"] = str(Config.slurm_use) + self.RESULT.attrs["Full_field"] = str(Config.full_field) + self.RESULT.attrs["Force_use_TP"] = str(Config.force_use_TP) + self.RESULT.attrs["Force_use_N"] = str(Config.force_use_N) + self.RESULT.attrs["Tile_of_glacier_of_interest"] = str(Config.tile) # Global attributes from constants.py - self.RESULT.attrs['Time_step_input_file_seconds'] = Constants.dt - self.RESULT.attrs['Max_layers'] = Constants.max_layers - self.RESULT.attrs['Z_measurment_height'] = Constants.z - self.RESULT.attrs['Stability_correction'] = Constants.stability_correction - self.RESULT.attrs['Albedo_method'] = Constants.albedo_method - self.RESULT.attrs['Densification_method'] = Constants.densification_method - self.RESULT.attrs['Penetrating_method'] = Constants.penetrating_method - self.RESULT.attrs['Roughness_method'] = Constants.roughness_method - self.RESULT.attrs['Saturation_water_vapour_method'] = Constants.saturation_water_vapour_method - - self.RESULT.attrs['Initial_snowheight'] = Constants.initial_snowheight_constant - self.RESULT.attrs['Initial_snow_layer_heights'] = Constants.initial_snow_layer_heights - self.RESULT.attrs['Initial_glacier_height'] = Constants.initial_glacier_height - self.RESULT.attrs['Initial_glacier_layer_heights'] = Constants.initial_glacier_layer_heights - self.RESULT.attrs['Initial_top_density_snowpack'] = Constants.initial_top_density_snowpack - self.RESULT.attrs['Initial_bottom_density_snowpack'] = Constants.initial_bottom_density_snowpack - self.RESULT.attrs['Temperature_bottom'] = Constants.temperature_bottom - self.RESULT.attrs['Const_init_temp'] = Constants.const_init_temp - - self.RESULT.attrs['Center_snow_transfer_function'] = center_snow_transfer_function - self.RESULT.attrs['Spread_snow_transfer_function'] = spread_snow_transfer_function - self.RESULT.attrs['Multiplication_factor_for_RRR_or_SNOWFALL'] = mult_factor_RRR - self.RESULT.attrs['Minimum_snow_layer_height'] = Constants.minimum_snow_layer_height - self.RESULT.attrs['Minimum_snowfall'] = Constants.minimum_snowfall - - self.RESULT.attrs['Remesh_method'] = Constants.remesh_method - self.RESULT.attrs['First_layer_height_log_profile'] = Constants.first_layer_height - self.RESULT.attrs['Layer_stretching_log_profile'] = Constants.layer_stretching - - self.RESULT.attrs['Merge_max'] = Constants.merge_max - self.RESULT.attrs['Layer_stretching_log_profile'] = Constants.layer_stretching - self.RESULT.attrs['Density_threshold_merging'] = Constants.density_threshold_merging - self.RESULT.attrs['Temperature_threshold_merging'] = Constants.temperature_threshold_merging - - self.RESULT.attrs['Density_fresh_snow'] = Constants.constant_density - self.RESULT.attrs['Albedo_fresh_snow'] = albedo_fresh_snow - self.RESULT.attrs['Albedo_firn'] = albedo_firn - self.RESULT.attrs['Albedo_ice'] = albedo_ice - self.RESULT.attrs['Albedo_mod_snow_aging'] = albedo_mod_snow_aging - self.RESULT.attrs['Albedo_mod_snow_depth'] = albedo_mod_snow_depth - self.RESULT.attrs['Roughness_fresh_snow'] = roughness_fresh_snow - self.RESULT.attrs['Roughness_ice'] = roughness_ice - self.RESULT.attrs['Roughness_firn'] = roughness_firn - self.RESULT.attrs['Aging_factor_roughness'] = aging_factor_roughness - self.RESULT.attrs['Snow_ice_threshold'] = Constants.snow_ice_threshold - - self.RESULT.attrs['lat_heat_melting'] = Constants.lat_heat_melting - self.RESULT.attrs['lat_heat_vaporize'] = Constants.lat_heat_vaporize - self.RESULT.attrs['lat_heat_sublimation'] = Constants.lat_heat_sublimation - self.RESULT.attrs['spec_heat_air'] = Constants.spec_heat_air - self.RESULT.attrs['spec_heat_ice'] = Constants.spec_heat_ice - self.RESULT.attrs['spec_heat_water'] = Constants.spec_heat_water - self.RESULT.attrs['k_i'] = Constants.k_i - self.RESULT.attrs['k_w'] = Constants.k_w - self.RESULT.attrs['k_a'] = Constants.k_a - self.RESULT.attrs['water_density'] = Constants.water_density - self.RESULT.attrs['ice_density'] = Constants.ice_density - self.RESULT.attrs['air_density'] = Constants.air_density - self.RESULT.attrs['sigma'] = Constants.sigma - self.RESULT.attrs['zero_temperature'] = Constants.zero_temperature - self.RESULT.attrs['Surface_emission_coeff'] = Constants.surface_emission_coeff + self.RESULT.attrs["Time_step_input_file_seconds"] = Constants.dt + self.RESULT.attrs["Max_layers"] = Constants.max_layers + self.RESULT.attrs["Z_measurement_height"] = Constants.z + self.RESULT.attrs["Stability_correction"] = Constants.stability_correction + self.RESULT.attrs["Albedo_method"] = Constants.albedo_method + self.RESULT.attrs["Densification_method"] = Constants.densification_method + self.RESULT.attrs["Penetrating_method"] = Constants.penetrating_method + self.RESULT.attrs["Roughness_method"] = Constants.roughness_method + self.RESULT.attrs["Saturation_water_vapour_method"] = Constants.saturation_water_vapour_method + + self.RESULT.attrs["Initial_snowheight"] = Constants.initial_snowheight_constant + self.RESULT.attrs["Initial_snow_layer_heights"] = Constants.initial_snow_layer_heights + self.RESULT.attrs["Initial_glacier_height"] = Constants.initial_glacier_height + self.RESULT.attrs["Initial_glacier_layer_heights"] = Constants.initial_glacier_layer_heights + self.RESULT.attrs["Initial_top_density_snowpack"] = Constants.initial_top_density_snowpack + self.RESULT.attrs["Initial_bottom_density_snowpack"] = Constants.initial_bottom_density_snowpack + self.RESULT.attrs["Temperature_bottom"] = Constants.temperature_bottom + self.RESULT.attrs["Const_init_temp"] = Constants.const_init_temp + + self.RESULT.attrs["Center_snow_transfer_function"] = center_snow_transfer_function + self.RESULT.attrs["Spread_snow_transfer_function"] = spread_snow_transfer_function + self.RESULT.attrs["Multiplication_factor_for_RRR_or_SNOWFALL"] = mult_factor_RRR + self.RESULT.attrs["Minimum_snow_layer_height"] = Constants.minimum_snow_layer_height + self.RESULT.attrs["Minimum_snowfall"] = Constants.minimum_snowfall + + self.RESULT.attrs["Remesh_method"] = Constants.remesh_method + self.RESULT.attrs["First_layer_height_log_profile"] = Constants.first_layer_height + self.RESULT.attrs["Layer_stretching_log_profile"] = Constants.layer_stretching + + self.RESULT.attrs["Merge_max"] = Constants.merge_max + self.RESULT.attrs["Layer_stretching_log_profile"] = Constants.layer_stretching + self.RESULT.attrs["Density_threshold_merging"] = Constants.density_threshold_merging + self.RESULT.attrs["Temperature_threshold_merging"] = Constants.temperature_threshold_merging + + self.RESULT.attrs["Density_fresh_snow"] = Constants.constant_density + self.RESULT.attrs["Albedo_fresh_snow"] = albedo_fresh_snow + self.RESULT.attrs["Albedo_firn"] = albedo_firn + self.RESULT.attrs["Albedo_ice"] = albedo_ice + self.RESULT.attrs["Albedo_mod_snow_aging"] = albedo_mod_snow_aging + self.RESULT.attrs["Albedo_mod_snow_depth"] = albedo_mod_snow_depth + self.RESULT.attrs["Roughness_fresh_snow"] = roughness_fresh_snow + self.RESULT.attrs["Roughness_ice"] = roughness_ice + self.RESULT.attrs["Roughness_firn"] = roughness_firn + self.RESULT.attrs["Aging_factor_roughness"] = aging_factor_roughness + self.RESULT.attrs["Snow_ice_threshold"] = Constants.snow_ice_threshold + + self.RESULT.attrs["lat_heat_melting"] = Constants.lat_heat_melting + self.RESULT.attrs["lat_heat_vaporize"] = Constants.lat_heat_vaporize + self.RESULT.attrs["lat_heat_sublimation"] = Constants.lat_heat_sublimation + self.RESULT.attrs["spec_heat_air"] = Constants.spec_heat_air + self.RESULT.attrs["spec_heat_ice"] = Constants.spec_heat_ice + self.RESULT.attrs["spec_heat_water"] = Constants.spec_heat_water + self.RESULT.attrs["k_i"] = Constants.k_i + self.RESULT.attrs["k_w"] = Constants.k_w + self.RESULT.attrs["k_a"] = Constants.k_a + self.RESULT.attrs["water_density"] = Constants.water_density + self.RESULT.attrs["ice_density"] = Constants.ice_density + self.RESULT.attrs["air_density"] = Constants.air_density + self.RESULT.attrs["sigma"] = Constants.sigma + self.RESULT.attrs["zero_temperature"] = Constants.zero_temperature + self.RESULT.attrs["Surface_emission_coeff"] = Constants.surface_emission_coeff # Variables given by the input dataset - spatial, spatiotemporal = self.get_result_metadata() + spatial, spatiotemporal = self.get_input_metadata() for name, metadata in spatial.items(): if name in self.DATA: @@ -393,271 +575,164 @@ def init_result_dataset(self, opt_dict=None) -> xr.Dataset: self.RESULT, np.full_like(self.DATA.T2, np.nan), "RRR", - "mm", - "Total precipiation", + spatiotemporal["RRR"][1], + spatiotemporal["RRR"][0], ) if "N" not in self.DATA: self.add_variable_along_latlontime( self.RESULT, np.full_like(self.DATA.T2, np.nan), "N", - "-", - "Cloud fraction", + spatiotemporal["N"][1], + spatiotemporal["N"][0], ) - print("\n") - print("Output dataset ... ok") + print(f"\nOutput dataset ... ok") + return self.RESULT - def create_global_result_arrays(self): """Create the global numpy arrays to store each output variable. - Each global array will be filled with local results from the - workers. The arrays will then be assigned to the RESULT dataset - and stored to disk (see COSIPY.py). + Each global array is filled with local results from the workers. + The arrays are then assigned to the RESULT dataset and stored to + disk (see COSIPY.py). """ - if ('RAIN' in self.atm): - self.RAIN = np.full((self.time,self.ny,self.nx), np.nan) - if ('SNOWFALL' in self.atm): - self.SNOWFALL = np.full((self.time,self.ny,self.nx), np.nan) - if ('LWin' in self.atm): - self.LWin = np.full((self.time,self.ny,self.nx), np.nan) - if ('LWout' in self.atm): - self.LWout = np.full((self.time,self.ny,self.nx), np.nan) - if ('H' in self.atm): - self.H = np.full((self.time,self.ny,self.nx), np.nan) - if ('LE' in self.atm): - self.LE = np.full((self.time,self.ny,self.nx), np.nan) - if ('B' in self.atm): - self.B = np.full((self.time,self.ny,self.nx), np.nan) - if ('QRR' in self.atm): - self.QRR = np.full((self.time,self.ny,self.nx), np.nan) - if ('MB' in self.internal): - self.MB = np.full((self.time,self.ny,self.nx), np.nan) - if ('surfMB' in self.internal): - self.surfMB = np.full((self.time,self.ny,self.nx), np.nan) - if ('Q' in self.internal): - self.Q = np.full((self.time,self.ny,self.nx), np.nan) - if ('SNOWHEIGHT' in self.internal): - self.SNOWHEIGHT = np.full((self.time,self.ny,self.nx), np.nan) - if ('TOTALHEIGHT' in self.internal): - self.TOTALHEIGHT = np.full((self.time,self.ny,self.nx), np.nan) - if ('TS' in self.atm): - self.TS = np.full((self.time,self.ny,self.nx), np.nan) - if ('ALBEDO' in self.atm): - self.ALBEDO = np.full((self.time,self.ny,self.nx), np.nan) - if ('LAYERS' in self.internal): - self.LAYERS = np.full((self.time,self.ny,self.nx), np.nan) - if ('ME' in self.internal): - self.ME = np.full((self.time,self.ny,self.nx), np.nan) - if ('intMB' in self.internal): - self.intMB = np.full((self.time,self.ny,self.nx), np.nan) - if ('EVAPORATION' in self.internal): - self.EVAPORATION = np.full((self.time,self.ny,self.nx), np.nan) - if ('SUBLIMATION' in self.internal): - self.SUBLIMATION = np.full((self.time,self.ny,self.nx), np.nan) - if ('CONDENSATION' in self.internal): - self.CONDENSATION = np.full((self.time,self.ny,self.nx), np.nan) - if ('DEPOSITION' in self.internal): - self.DEPOSITION = np.full((self.time,self.ny,self.nx), np.nan) - if ('REFREEZE' in self.internal): - self.REFREEZE = np.full((self.time,self.ny,self.nx), np.nan) - if ('subM' in self.internal): - self.subM = np.full((self.time,self.ny,self.nx), np.nan) - if ('Z0' in self.atm): - self.Z0 = np.full((self.time,self.ny,self.nx), np.nan) - if ('surfM' in self.internal): - self.surfM = np.full((self.time,self.ny,self.nx), np.nan) - if ('MOL' in self.internal): - self.MOL = np.full((self.time,self.ny,self.nx), np.nan) - if Config.full_field: - max_layers = Constants.max_layers # faster lookup - if ('HEIGHT' in self.full): - self.LAYER_HEIGHT = np.full((self.time,self.ny,self.nx,max_layers), np.nan) - if ('RHO' in self.full): - self.LAYER_RHO = np.full((self.time,self.ny,self.nx,max_layers), np.nan) - if ('T' in self.full): - self.LAYER_T = np.full((self.time,self.ny,self.nx,max_layers), np.nan) - if ('LWC' in self.full): - self.LAYER_LWC = np.full((self.time,self.ny,self.nx,max_layers), np.nan) - if ('CC' in self.full): - self.LAYER_CC = np.full((self.time,self.ny,self.nx,max_layers), np.nan) - if ('POROSITY' in self.full): - self.LAYER_POROSITY = np.full((self.time,self.ny,self.nx,max_layers), np.nan) - if ('ICE_FRACTION' in self.full): - self.LAYER_ICE_FRACTION = np.full((self.time,self.ny,self.nx,max_layers), np.nan) - if ('IRREDUCIBLE_WATER' in self.full): - self.LAYER_IRREDUCIBLE_WATER = np.full((self.time,self.ny,self.nx,max_layers), np.nan) - if ('REFREEZE' in self.full): - self.LAYER_REFREEZE = np.full((self.time,self.ny,self.nx,max_layers), np.nan) - - - def copy_local_to_global(self,y,x,local_RAIN,local_SNOWFALL,local_LWin,local_LWout,local_H,local_LE,local_B,local_QRR, - local_MB, local_surfMB,local_Q,local_SNOWHEIGHT,local_TOTALHEIGHT,local_TS,local_ALBEDO, \ - local_LAYERS,local_ME,local_intMB,local_EVAPORATION,local_SUBLIMATION,local_CONDENSATION, \ - local_DEPOSITION,local_REFREEZE,local_subM,local_Z0,local_surfM,local_MOL,local_LAYER_HEIGHT, \ - local_LAYER_RHO,local_LAYER_T,local_LAYER_LWC,local_LAYER_CC,local_LAYER_POROSITY, \ - local_LAYER_ICE_FRACTION,local_LAYER_IRREDUCIBLE_WATER,local_LAYER_REFREEZE): - """Copy the local results from workers to global numpy arrays. + if self.atm: + for atm_var in self.atm: + self.init_atm_attribute(atm_var) - Args: - y: Latitude index. - x: Longitude index. - """ - if ('RAIN' in self.atm): - self.RAIN[:,y,x] = local_RAIN - if ('SNOWFALL' in self.atm): - self.SNOWFALL[:,y,x] = local_SNOWFALL - if ('LWin' in self.atm): - self.LWin[:,y,x] = local_LWin - if ('LWout' in self.atm): - self.LWout[:,y,x] = local_LWout - if ('H' in self.atm): - self.H[:,y,x] = local_H - if ('LE' in self.atm): - self.LE[:,y,x] = local_LE - if ('B' in self.atm): - self.B[:,y,x] = local_B - if ('QRR' in self.atm): - self.QRR[:,y,x] = local_QRR - if ('surfMB' in self.internal): - self.surfMB[:,y,x] = local_surfMB - if ('MB' in self.internal): - self.MB[:,y,x] = local_MB - if ('Q' in self.internal): - self.Q[:,y,x] = local_Q - if ('SNOWHEIGHT' in self.internal): - self.SNOWHEIGHT[:,y,x] = local_SNOWHEIGHT - if ('TOTALHEIGHT' in self.internal): - self.TOTALHEIGHT[:,y,x] = local_TOTALHEIGHT - if ('TS' in self.atm): - self.TS[:,y,x] = local_TS - if ('ALBEDO' in self.atm): - self.ALBEDO[:,y,x] = local_ALBEDO - if ('LAYERS' in self.internal): - self.LAYERS[:,y,x] = local_LAYERS - if ('ME' in self.internal): - self.ME[:,y,x] = local_ME - if ('intMB' in self.internal): - self.intMB[:,y,x] = local_intMB - if ('EVAPORATION' in self.internal): - self.EVAPORATION[:,y,x] = local_EVAPORATION - if ('SUBLIMATION' in self.internal): - self.SUBLIMATION[:,y,x] = local_SUBLIMATION - if ('CONDENSATION' in self.internal): - self.CONDENSATION[:,y,x] = local_CONDENSATION - if ('DEPOSITION' in self.internal): - self.DEPOSITION[:,y,x] = local_DEPOSITION - if ('REFREEZE' in self.internal): - self.REFREEZE[:,y,x] = local_REFREEZE - if ('subM' in self.internal): - self.subM[:,y,x] = local_subM - if ('Z0' in self.atm): - self.Z0[:,y,x] = local_Z0 - if ('surfM' in self.internal): - self.surfM[:,y,x] = local_surfM - if ('MOL' in self.internal): - self.MOL[:,y,x] = local_MOL + if self.internal: + for internal_var in self.internal: + self.init_internal_attribute(internal_var) - if Config.full_field: - if ('HEIGHT' in self.full): - self.LAYER_HEIGHT[:,y,x,:] = local_LAYER_HEIGHT - if ('RHO' in self.full): - self.LAYER_RHO[:,y,x,:] = local_LAYER_RHO - if ('T' in self.full): - self.LAYER_T[:,y,x,:] = local_LAYER_T - if ('LWC' in self.full): - self.LAYER_LWC[:,y,x,:] = local_LAYER_LWC - if ('CC' in self.full): - self.LAYER_CC[:,y,x,:] = local_LAYER_CC - if ('POROSITY' in self.full): - self.LAYER_POROSITY[:,y,x,:] = local_LAYER_POROSITY - if ('ICE_FRACTION' in self.full): - self.LAYER_ICE_FRACTION[:,y,x,:] = local_LAYER_ICE_FRACTION - if ('IRREDUCIBLE_WATER' in self.full): - self.LAYER_IRREDUCIBLE_WATER[:,y,x,:] = local_LAYER_IRREDUCIBLE_WATER - if ('REFREEZE' in self.full): - self.LAYER_REFREEZE[:,y,x,:] = local_LAYER_REFREEZE + if Config.full_field and self.full: + max_layers = Constants.max_layers # faster lookup + for full_field_var in self.full: + self.init_full_field_attribute(full_field_var, max_layers) + + def copy_local_to_global( + self, + y: int, + x: int, + local_RAIN: np.ndarray, + local_SNOWFALL: np.ndarray, + local_LWin: np.ndarray, + local_LWout: np.ndarray, + local_H: np.ndarray, + local_LE: np.ndarray, + local_B: np.ndarray, + local_QRR: np.ndarray, + local_MB: np.ndarray, + local_surfMB: np.ndarray, + local_Q: np.ndarray, + local_SNOWHEIGHT: np.ndarray, + local_TOTALHEIGHT: np.ndarray, + local_TS: np.ndarray, + local_ALBEDO: np.ndarray, + local_LAYERS: np.ndarray, + local_ME: np.ndarray, + local_intMB: np.ndarray, + local_EVAPORATION: np.ndarray, + local_SUBLIMATION: np.ndarray, + local_CONDENSATION: np.ndarray, + local_DEPOSITION: np.ndarray, + local_REFREEZE: np.ndarray, + local_subM: np.ndarray, + local_Z0: np.ndarray, + local_surfM: np.ndarray, + local_MOL: np.ndarray, + local_LAYER_HEIGHT: np.ndarray, + local_LAYER_RHO: np.ndarray, + local_LAYER_T: np.ndarray, + local_LAYER_LWC: np.ndarray, + local_LAYER_CC: np.ndarray, + local_LAYER_POROSITY: np.ndarray, + local_LAYER_ICE_FRACTION: np.ndarray, + local_LAYER_IRREDUCIBLE_WATER: np.ndarray, + local_LAYER_REFREEZE: np.ndarray, + ): + """Copy the local results from workers to global numpy arrays.""" + + self.set_atm_attribute("RAIN", local_RAIN, x, y) + self.set_atm_attribute("SNOWFALL", local_SNOWFALL, x, y) + self.set_atm_attribute("LWin", local_LWin, x, y) + self.set_atm_attribute("LWout", local_LWout, x, y) + self.set_atm_attribute("H", local_H, x, y) + self.set_atm_attribute("LE", local_LE, x, y) + self.set_atm_attribute("B", local_B, x, y) + self.set_atm_attribute("QRR", local_QRR, x, y) + self.set_atm_attribute("TS", local_TS, x, y) + self.set_atm_attribute("ALBEDO", local_ALBEDO, x, y) + self.set_atm_attribute("Z0", local_Z0, x, y) + + self.set_internal_attribute("surfMB", local_surfMB, x, y) + self.set_internal_attribute("MB", local_MB, x, y) + self.set_internal_attribute("Q", local_Q, x, y) + self.set_internal_attribute("SNOWHEIGHT", local_SNOWHEIGHT, x, y) + self.set_internal_attribute("TOTALHEIGHT", local_TOTALHEIGHT, x, y) + self.set_internal_attribute("LAYERS", local_LAYERS, x, y) + self.set_internal_attribute("ME", local_ME, x, y) + self.set_internal_attribute("intMB", local_intMB, x, y) + self.set_internal_attribute("EVAPORATION", local_EVAPORATION, x, y) + self.set_internal_attribute("SUBLIMATION", local_SUBLIMATION, x, y) + self.set_internal_attribute("CONDENSATION", local_CONDENSATION, x, y) + self.set_internal_attribute("DEPOSITION", local_DEPOSITION, x, y) + self.set_internal_attribute("REFREEZE", local_REFREEZE, x, y) + self.set_internal_attribute("subM", local_subM, x, y) + self.set_internal_attribute("surfM", local_surfM, x, y) + self.set_internal_attribute("MOL", local_MOL, x, y) + if Config.full_field: + self.set_full_field_attribute("HEIGHT", local_LAYER_HEIGHT, x, y) + self.set_full_field_attribute("RHO", local_LAYER_RHO, x, y) + self.set_full_field_attribute("T", local_LAYER_T, x, y) + self.set_full_field_attribute("LWC", local_LAYER_LWC, x, y) + self.set_full_field_attribute("CC", local_LAYER_CC, x, y) + self.set_full_field_attribute("POROSITY", local_LAYER_POROSITY, x, y) + self.set_full_field_attribute( + "ICE_FRACTION", local_LAYER_ICE_FRACTION, x, y + ) + self.set_full_field_attribute( + "IRREDUCIBLE_WATER", local_LAYER_IRREDUCIBLE_WATER, x, y + ) + self.set_full_field_attribute("REFREEZE", local_LAYER_REFREEZE, x, y) def write_results_to_file(self): """Add the global numpy arrays to the RESULT dataset.""" - if ('RAIN' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.RAIN, 'RAIN', 'mm', 'Liquid precipitation') - if ('SNOWFALL' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.SNOWFALL, 'SNOWFALL', 'm w.e.', 'Snowfall') - if ('LWin' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.LWin, 'LWin', 'W m\u207b\xb2', 'Incoming longwave radiation') - if ('LWout' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.LWout, 'LWout', 'W m\u207b\xb2', 'Outgoing longwave radiation') - if ('H' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.H, 'H', 'W m\u207b\xb2', 'Sensible heat flux') - if ('LE' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.LE, 'LE', 'W m\u207b\xb2', 'Latent heat flux') - if ('B' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.B, 'B', 'W m\u207b\xb2', 'Ground heat flux') - if ('QRR' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.QRR, 'QRR', 'W m\u207b\xb2', 'Rain heat flux') - if ('surfMB' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.surfMB, 'surfMB', 'm w.e.', 'Surface mass balance') - if ('MB' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.MB, 'MB', 'm w.e.', 'Mass balance') - if ('Q' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.Q, 'Q', 'm w.e.', 'Runoff') - if ('SNOWHEIGHT' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.SNOWHEIGHT, 'SNOWHEIGHT', 'm', 'Snowheight') - if ('TOTALHEIGHT' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.TOTALHEIGHT, 'TOTALHEIGHT', 'm', 'Total domain height') - if ('TS' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.TS, 'TS', 'K', 'Surface temperature') - if ('ALBEDO' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.ALBEDO, 'ALBEDO', '-', 'Albedo') - if ('LAYERS' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.LAYERS, 'LAYERS', '-', 'Number of layers') - if ('ME' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.ME, 'ME', 'W m\u207b\xb2', 'Available melt energy') - if ('intMB' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.intMB, 'intMB', 'm w.e.', 'Internal mass balance') - if ('EVAPORATION' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.EVAPORATION, 'EVAPORATION', 'm w.e.', 'Evaporation') - if ('SUBLIMATION' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.SUBLIMATION, 'SUBLIMATION', 'm w.e.', 'Sublimation') - if ('CONDENSATION' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.CONDENSATION, 'CONDENSATION', 'm w.e.', 'Condensation') - if ('DEPOSITION' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.DEPOSITION, 'DEPOSITION', 'm w.e.', 'Deposition') - if ('REFREEZE' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.REFREEZE, 'REFREEZE', 'm w.e.', 'Refreezing') - if ('subM' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.subM, 'subM', 'm w.e.', 'Subsurface melt') - if ('Z0' in self.atm): - self.add_variable_along_latlontime(self.RESULT, self.Z0, 'Z0', 'm', 'Roughness length') - if ('surfM' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.surfM, 'surfM', 'm w.e.', 'Surface melt') - if ('MOL' in self.internal): - self.add_variable_along_latlontime(self.RESULT, self.MOL, 'MOL', '', 'Monin Obukhov length') - if Config.full_field: - if ('HEIGHT' in self.full): - self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_HEIGHT, 'LAYER_HEIGHT', 'm', 'Layer height') - if ('RHO' in self.full): - self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_RHO, 'LAYER_RHO', 'kg m^-3', 'Layer density') - if ('T' in self.full): - self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_T, 'LAYER_T', 'K', 'Layer temperature') - if ('LWC' in self.full): - self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_LWC, 'LAYER_LWC', 'kg m^-2', 'Liquid water content') - if ('CC' in self.full): - self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_CC, 'LAYER_CC', 'J m^-2', 'Cold content') - if ('POROSITY' in self.full): - self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_POROSITY, 'LAYER_POROSITY', '-', 'Porosity') - if ('ICE_FRACTION' in self.full): - self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_ICE_FRACTION, 'LAYER_ICE_FRACTION', '-', 'Ice fraction') - if ('IRREDUCIBLE_WATER' in self.full): - self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_IRREDUCIBLE_WATER, 'LAYER_IRREDUCIBLE_WATER', '-', 'Irreducible water') - if ('REFREEZE' in self.full): - self.add_variable_along_latlonlayertime(self.RESULT, self.LAYER_REFREEZE, 'LAYER_REFREEZE', 'm w.e.', 'Refreezing') + metadata = self.get_result_metadata() + if self.atm: + for atm_var in self.atm: + self.add_variable_along_latlontime( + self.RESULT, + getattr(self, atm_var), + atm_var, + metadata[atm_var][0], + metadata[atm_var][1], + ) + + if self.internal: + for internal_var in self.internal: + self.add_variable_along_latlontime( + self.RESULT, + getattr(self, internal_var), + internal_var, + metadata[internal_var][0], + metadata[internal_var][1], + ) + + if Config.full_field and self.full: + for full_field_var in self.full: + layer_name = f"LAYER_{full_field_var}" + self.add_variable_along_latlonlayertime( + self.RESULT, + getattr(self, layer_name), + layer_name, + metadata[layer_name][0], + metadata[layer_name][1], + ) def create_empty_restart(self) -> xr.Dataset: """Create an empty dataset for the RESTART attribute. @@ -666,26 +741,23 @@ def create_empty_restart(self) -> xr.Dataset: Empty xarray dataset with coordinates from self.DATA. """ dataset = xr.Dataset() - dataset.coords['time'] = self.DATA.coords['time'][-1] - dataset.coords['lat'] = self.DATA.coords['lat'] - dataset.coords['lon'] = self.DATA.coords['lon'] - dataset.coords['layer'] = np.arange(Constants.max_layers) + dataset.coords["time"] = self.DATA.coords["time"][-1] + dataset.coords["lat"] = self.DATA.coords["lat"] + dataset.coords["lon"] = self.DATA.coords["lon"] + dataset.coords["layer"] = np.arange(Constants.max_layers) return dataset def init_restart_dataset(self) -> xr.Dataset: """Initialise the restart dataset. - + Returns: The empty restart dataset. """ self.RESTART = self.create_empty_restart() - - print('Restart dataset ... ok \n') - print('--------------------------------------------------------------\n') - + print(f"Restart dataset ... ok\n{'-'*62}\n") + return self.RESTART - def create_global_restart_arrays(self): """Initialise the global numpy arrays to store layer profiles. @@ -697,42 +769,58 @@ def create_global_restart_arrays(self): max_layers = Constants.max_layers # faster lookup - self.RES_NLAYERS = np.full((self.ny,self.nx), np.nan) - self.RES_NEWSNOWHEIGHT = np.full((self.ny, self.nx), np.nan) - self.RES_NEWSNOWTIMESTAMP = np.full((self.ny, self.nx), np.nan) - self.RES_OLDSNOWTIMESTAMP = np.full((self.ny, self.nx), np.nan) - self.RES_LAYER_HEIGHT = np.full((self.ny,self.nx,max_layers), np.nan) - self.RES_LAYER_RHO = np.full((self.ny,self.nx,max_layers), np.nan) - self.RES_LAYER_T = np.full((self.ny,self.nx,max_layers), np.nan) - self.RES_LAYER_LWC = np.full((self.ny,self.nx,max_layers), np.nan) - self.RES_LAYER_IF = np.full((self.ny,self.nx,max_layers), np.nan) - + for name in [ + "NLAYERS", + "NEWSNOWHEIGHT", + "NEWSNOWTIMESTAMP", + "OLDSNOWTIMESTAMP", + ]: + setattr(self, f"RES_{name}", np.full((self.ny, self.nx), np.nan)) + + for name in ["HEIGHT", "RHO", "T", "LWC", "IF"]: + setattr( + self, + f"RES_LAYER_{name}", + np.full((self.ny, self.nx, max_layers), np.nan), + ) def create_local_restart_dataset(self) -> xr.Dataset: """Create the result dataset for a single grid point. - + Returns: RESTART dataset initialised with layer profiles. """ - + self.RESTART = self.create_empty_restart() - - self.add_variable_along_scalar(self.RESTART, np.full((1), np.nan), 'NLAYERS', '-', 'Number of layers') - self.add_variable_along_scalar(self.RESTART, np.full((1), np.nan), 'NEWSNOWHEIGHT', 'm .w.e', 'New snow height') - self.add_variable_along_scalar(self.RESTART, np.full((1), np.nan), 'NEWSNOWTIMESTAMP', 's', 'New snow timestamp') - self.add_variable_along_scalar(self.RESTART, np.full((1), np.nan), 'OLDSNOWTIMESTAMP', 's', 'Old snow timestamp') - self.add_variable_along_layer(self.RESTART, np.full((self.RESTART.coords['layer'].shape[0]), np.nan), 'LAYER_HEIGHT', 'm', 'Layer height') - self.add_variable_along_layer(self.RESTART, np.full((self.RESTART.coords['layer'].shape[0]), np.nan), 'LAYER_RHO', 'kg m^-3', 'Density of layer') - self.add_variable_along_layer(self.RESTART, np.full((self.RESTART.coords['layer'].shape[0]), np.nan), 'LAYER_T', 'K', 'Layer temperature') - self.add_variable_along_layer(self.RESTART, np.full((self.RESTART.coords['layer'].shape[0]), np.nan), 'LAYER_LWC', '-', 'Layer liquid water content') - self.add_variable_along_layer(self.RESTART, np.full((self.RESTART.coords['layer'].shape[0]), np.nan), 'LAYER_IF', '-', 'Layer ice fraction') + metadata = self.get_restart_metadata() + for name in [ + "NLAYERS", + "NEWSNOWHEIGHT", + "NEWSNOWTIMESTAMP", + "OLDSNOWTIMESTAMP", + ]: + self.add_variable_along_scalar( + self.RESTART, + np.full((1), np.nan), + name, + metadata[name][0], + metadata[name][1], + ) + for layer_name in ["HEIGHT", "RHO", "T", "LWC", "IF"]: + keyname = f"LAYER_{layer_name}" + self.add_variable_along_layer( + self.RESTART, + np.full((self.RESTART.coords["layer"].shape[0]), np.nan), + keyname, + metadata[keyname][0], + metadata[keyname][1], + ) return self.RESTART - - def copy_local_restart_to_global(self,y,x,local_restart): + def copy_local_restart_to_global(self, y: int, x: int, local_restart: xr.Dataset): """Copy local restart data from workers to global numpy arrays. Args: @@ -740,29 +828,47 @@ def copy_local_restart_to_global(self,y,x,local_restart): x: Longitude index. local_restart: Local RESTART dataset. """ - self.RES_NLAYERS[y,x] = local_restart.NLAYERS - self.RES_NEWSNOWHEIGHT[y,x] = local_restart.NEWSNOWHEIGHT - self.RES_NEWSNOWTIMESTAMP[y,x] = local_restart.NEWSNOWTIMESTAMP - self.RES_OLDSNOWTIMESTAMP[y,x] = local_restart.OLDSNOWTIMESTAMP - self.RES_LAYER_HEIGHT[y,x,:] = local_restart.LAYER_HEIGHT - self.RES_LAYER_RHO[y,x,:] = local_restart.LAYER_RHO - self.RES_LAYER_T[y,x,:] = local_restart.LAYER_T - self.RES_LAYER_LWC[y,x,:] = local_restart.LAYER_LWC - self.RES_LAYER_IF[y,x,:] = local_restart.LAYER_IF + for name in [ + "NLAYERS", + "NEWSNOWHEIGHT", + "NEWSNOWTIMESTAMP", + "OLDSNOWTIMESTAMP", + ]: + getattr(self, f"RES_{name}")[y, x] = getattr(local_restart, name) + + for name in ["HEIGHT", "RHO", "T", "LWC", "IF"]: + getattr(self, f"RES_LAYER_{name}")[y, x, :] = getattr( + local_restart, f"LAYER_{name}" + ) def write_restart_to_file(self): """Add global numpy arrays to the RESTART dataset.""" - self.add_variable_along_latlon(self.RESTART, self.RES_NLAYERS, 'NLAYERS', '-', 'Number of layers') - self.add_variable_along_latlon(self.RESTART, self.RES_NEWSNOWHEIGHT, 'new_snow_height', 'm .w.e', 'New snow height') - self.add_variable_along_latlon(self.RESTART, self.RES_NEWSNOWTIMESTAMP, 'new_snow_timestamp', 's', 'New snow timestamp') - self.add_variable_along_latlon(self.RESTART, self.RES_OLDSNOWTIMESTAMP, 'old_snow_timestamp', 's', 'Old snow timestamp') - self.add_variable_along_latlonlayer(self.RESTART, self.RES_LAYER_HEIGHT, 'LAYER_HEIGHT', 'm', 'Height of each layer') - self.add_variable_along_latlonlayer(self.RESTART, self.RES_LAYER_RHO, 'LAYER_RHO', 'kg m^-3', 'Layer density') - self.add_variable_along_latlonlayer(self.RESTART, self.RES_LAYER_T, 'LAYER_T', 'K', 'Layer temperature') - self.add_variable_along_latlonlayer(self.RESTART, self.RES_LAYER_LWC, 'LAYER_LWC', '-', 'Layer liquid water content') - self.add_variable_along_latlonlayer(self.RESTART, self.RES_LAYER_IF, 'LAYER_IF', '-', 'Layer ice fraction') + metadata = self.get_restart_metadata() + for name in [ + "NLAYERS", + "NEWSNOWHEIGHT", + "NEWSNOWTIMESTAMP", + "OLDSNOWTIMESTAMP", + ]: + keyname = f"RES_{name}" + self.add_variable_along_latlon( + self.RESTART, + getattr(self, keyname), + name, + metadata[name][0], + metadata[name][1], + ) + for name in ["HEIGHT", "RHO", "T", "LWC", "IF"]: + keyname = f"LAYER_{name}" + self.add_variable_along_latlonlayer( + self.RESTART, + getattr(self, f"RES_{keyname}"), + keyname, + metadata[keyname][0], + metadata[keyname][1], + ) # TODO: Make it Pythonian - Finish the getter/setter functions @property @@ -792,8 +898,7 @@ def QRR(self): @property def MB(self): return self.__MB - - + @RAIN.setter def RAIN(self, x): self.__RAIN = x @@ -822,7 +927,6 @@ def QRR(self, x): def MB(self, x): self.__MB = x - def get_result(self) -> xr.Dataset: """Get the RESULT data structure.""" return self.RESULT @@ -838,26 +942,48 @@ def get_grid_restart(self) -> xr.Dataset: # ================================================================== # Auxiliary functions for writing variables to NetCDF files # ================================================================== - def add_variable_along_scalar(self, ds, var, name, units, long_name): + def set_variable_metadata( + self, data: xr.Variable, units: str, long_name: str, fill_value: int = -9999 + ) -> xr.Dataset: + """Add long name, units, and default fill value to variable. + + Args: + data: netCDF data structure. + units: Variable units. + long_name: Full name of variable. + fill_value: NaN fill value. Default "-9999". + + Returns: + Dataset with updated metadata for a specific variable. + """ + data.attrs["units"] = units + data.attrs["long_name"] = long_name + data.encoding["_FillValue"] = fill_value + + return data + + def add_variable_along_scalar( + self, ds: xr.Dataset, var: np.ndarray, name: str, units: str, long_name: str + ) -> xr.Dataset: """Add scalar data to a dataset. Args: - ds (xr.Dataset): Target data structure. - var (np.ndarray): New data. - name (str): The new variable's abbreviated name. - units (str): New variable units. - long_name (str): The new variable's full name. + ds: Target data structure. + var: New data. + name: The new variable's abbreviated name. + units: The new variable's units. + long_name: The new variable's full name. Returns: - xr.Dataset: Target dataset with the new scalar variable. + Target dataset with the new scalar variable. """ ds[name] = var.data - ds[name].attrs['units'] = units - ds[name].attrs['long_name'] = long_name - ds[name].encoding['_FillValue'] = -9999 + self.set_variable_metadata(ds[name], units, long_name) return ds - def add_variable_along_latlon(self, ds, var, name, units, long_name): + def add_variable_along_latlon( + self, ds: xr.Dataset, var: np.ndarray, name: str, units: str, long_name: str + ) -> xr.Dataset: """Add spatial data to a dataset. Args: @@ -868,125 +994,122 @@ def add_variable_along_latlon(self, ds, var, name, units, long_name): long_name (str): The new variable's full name. Returns: - xr.Dataset: Target dataset with the new spatial variable. + Target dataset with the new spatial variable. """ - ds[name] = ((Config.northing,Config.easting), var.data) - ds[name].attrs['units'] = units - ds[name].attrs['long_name'] = long_name - ds[name].encoding['_FillValue'] = -9999 + ds[name] = ((Config.northing, Config.easting), var.data) + self.set_variable_metadata(ds[name], units, long_name) return ds - - def add_variable_along_time(self, ds, var, name, units, long_name): + + def add_variable_along_time( + self, ds: xr.Dataset, var: np.ndarray, name: str, units: str, long_name: str + ) -> xr.Dataset: """Add temporal data to a dataset. Args: - ds (xr.Dataset): Target data structure. - var (np.ndarray): New temporal data. - name (str): The new variable's abbreviated name. - units (str): New variable units. - long_name (str): The new variable's full name. + ds: Target data structure. + var: New temporal data. + name: The new variable's abbreviated name. + units: The new variable's units. + long_name: The new variable's full name. Returns: - xr.Dataset: Target dataset with the new temporal variable. + Target dataset with the new temporal variable. """ - ds[name] = xr.DataArray(var.data, coords=[('time', ds.time)]) - ds[name].attrs['units'] = units - ds[name].attrs['long_name'] = long_name - ds[name].encoding['_FillValue'] = -9999 + ds[name] = xr.DataArray(var.data, coords=[("time", ds.time)]) + self.set_variable_metadata(ds[name], units, long_name) return ds - - def add_variable_along_latlontime(self, ds, var, name, units, long_name): + + def add_variable_along_latlontime( + self, ds: xr.Dataset, var: np.ndarray, name: str, units: str, long_name: str + ) -> xr.Dataset: """Add spatiotemporal data to a dataset. Args: - ds (xr.Dataset): Target data structure. - var (np.ndarray): New spatiotemporal data. - name (str): The new variable's abbreviated name. - units (str): New variable units. - long_name (str): The new variable's full name. + ds: Target data structure. + var: New spatiotemporal data. + name: The new variable's abbreviated name. + units: The new variable's units. + long_name: The new variable's full name. Returns: - xr.Dataset: Target dataset with the new spatiotemporal - variable. + Target dataset with the new spatiotemporal variable. """ - ds[name] = (('time',Config.northing,Config.easting), var.data) - ds[name].attrs['units'] = units - ds[name].attrs['long_name'] = long_name - ds[name].encoding['_FillValue'] = -9999 + ds[name] = (("time", Config.northing, Config.easting), var.data) + self.set_variable_metadata(ds[name], units, long_name) return ds - - def add_variable_along_latlonlayertime(self, ds, var, name, units, long_name): + + def add_variable_along_latlonlayertime( + self, ds: xr.Dataset, var: np.ndarray, name: str, units: str, long_name: str + ) -> xr.Dataset: """Add a spatiotemporal mesh to a dataset. Args: - ds (xr.Dataset): Target data structure. - var (np.ndarray): New spatiotemporal mesh data. - name (str): The new variable's abbreviated name. - units (str): New variable units. - long_name (str): The new variable's full name. + ds: Target data structure. + var: New spatiotemporal mesh data. + name: The new variable's abbreviated name. + units: The new variable's units. + long_name: The new variable's full name. Returns: - xr.Dataset: Target dataset with the new spatiotemporal mesh. + Target dataset with the new spatiotemporal mesh. """ - ds[name] = (('time',Config.northing,Config.easting,'layer'), var.data) - ds[name].attrs['units'] = units - ds[name].attrs['long_name'] = long_name - ds[name].encoding['_FillValue'] = -9999 + ds[name] = (("time", Config.northing, Config.easting, "layer"), var.data) + self.set_variable_metadata(ds[name], units, long_name) return ds - - def add_variable_along_latlonlayer(self, ds, var, name, units, long_name): + + def add_variable_along_latlonlayer( + self, ds: xr.Dataset, var: np.ndarray, name: str, units: str, long_name: str + ) -> xr.Dataset: """Add a spatial mesh to a dataset. Args: - ds (xr.Dataset): Target data structure. - var (np.ndarray): New spatial mesh. - name (str): The new variable's abbreviated name. - units (str): New variable units. - long_name (str): The new variable's full name. + ds: Target data structure. + var: New spatial mesh. + name: The new variable's abbreviated name. + units: The new variable's units. + long_name: The new variable's full name. Returns: - xr.Dataset: Target dataset with the new spatial mesh. + Target dataset with the new spatial mesh. """ - ds[name] = ((Config.northing,Config.easting,'layer'), var.data) - ds[name].attrs['units'] = units - ds[name].attrs['long_name'] = long_name - ds[name].encoding['_FillValue'] = -9999 + ds[name] = ((Config.northing, Config.easting, "layer"), var.data) + self.set_variable_metadata(ds[name], units, long_name) return ds - - def add_variable_along_layertime(self, ds, var, name, units, long_name): + + def add_variable_along_layertime( + self, ds: xr.Dataset, var: np.ndarray, name: str, units: str, long_name: str + ) -> xr.Dataset: """Add temporal layer data to a dataset. Args: - ds (xr.Dataset): Target data structure. - var (np.ndarray): New layer data with a time coordinate. - name (str): The new variable's abbreviated name. - units (str): New variable units. - long_name (str): The new variable's full name. + ds: Target data structure. + var: New layer data with a time coordinate. + name: The new variable's abbreviated name. + units: The new variable's units. + long_name: The new variable's full name. Returns: - xr.Dataset: Target dataset with the new layer data. + Target dataset with the new layer data. """ - ds[name] = (('time','layer'), var.data) - ds[name].attrs['units'] = units - ds[name].attrs['long_name'] = long_name - ds[name].encoding['_FillValue'] = -9999 + ds[name] = (("time", "layer"), var.data) + self.set_variable_metadata(ds[name], units, long_name) return ds - - def add_variable_along_layer(self, ds, var, name, units, long_name): + + def add_variable_along_layer( + self, ds: xr.Dataset, var: np.ndarray, name: str, units: str, long_name: str + ) -> xr.Dataset: """Add layer data to a dataset. Args: - ds (xr.Dataset): Target data structure. - var (np.ndarray): New layer data. - name (str): The new variable's abbreviated name. - units (str): New variable units. - long_name (str): The new variable's full name. + ds: Target data structure. + var: New layer data. + name: The new variable's abbreviated name. + units: The new variable's units. + long_name: The new variable's full name. Returns: - xr.Dataset: Target dataset with the new layer data. + Target dataset with the new layer data. """ - ds[name] = (('layer'), var.data) - ds[name].attrs['units'] = units - ds[name].attrs['long_name'] = long_name - ds[name].encoding['_FillValue'] = -9999 + ds[name] = ("layer", var.data) + self.set_variable_metadata(ds[name], units, long_name) return ds diff --git a/cosipy/cpkernel/node.py b/cosipy/cpkernel/node.py index a8fc08e..5527d17 100644 --- a/cosipy/cpkernel/node.py +++ b/cosipy/cpkernel/node.py @@ -128,7 +128,7 @@ def get_layer_air_porosity(self) -> float: return max(0.0, 1 - self.get_layer_liquid_water_content() - self.get_layer_ice_fraction()) def get_layer_specific_heat(self) -> float: - """Get the node's volumetric averaged specific heat capacity. + """Get the node's volume-weighted specific heat capacity. Returns: Specific heat capacity [|J kg^-1 K^-1|]. @@ -176,7 +176,7 @@ def get_layer_porosity(self) -> float: return 1-self.get_layer_ice_fraction()-self.get_layer_liquid_water_content() def get_layer_thermal_conductivity(self) -> float: - """Get the node's volumetric weighted thermal conductivity. + """Get the node's volume-weighted thermal conductivity. Returns: Thermal conductivity, |kappa| [|W m^-1 K^-1|]. @@ -186,12 +186,6 @@ def get_layer_thermal_conductivity(self) -> float: kappa = self.get_layer_ice_fraction()*k_i + self.get_layer_air_porosity()*k_a + self.get_layer_liquid_water_content()*k_w elif thermal_conductivity_method == 'empirical': kappa = 0.021 + 2.5 * np.power((self.get_layer_density()/1000),2) - elif thermal_conductivity_method == 'Sturm97': - kappa = 0.138 - 1.01e-3 * self.get_layer_density() + 3.23e-6 * np.power((self.get_layer_density()),2) - elif thermal_conductivity_method == 'Calonne19': - theta = 1 / (1 + np.exp(-2 * 0.02 * (self.get_layer_density() - 450))) - kappa = ((k_i * k_a / 2.107 *0.024) * ((1 - theta) * ((0.024 - (1.23e-4 * self.get_layer_density()) + (2.5e-6 * np.power(self.get_layer_density(),2)))))) + \ - ((k_i / 2.107) * (theta * (2.107 + 0.003618 * (self.get_layer_density() - ice_density)))) else: message = ("Thermal conductivity method =", f"{thermal_conductivity_method}",