From 2066f14bfd64b76b68d39e3d4b448eae05b886b5 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 18 Nov 2024 11:28:55 +0100 Subject: [PATCH 01/33] Refactored all names in the steps blending code from old to new --- pysteps/blending/steps.py | 847 ++++++++++++++++++++++++-------------- pysteps/nowcasts/steps.py | 1 - 2 files changed, 534 insertions(+), 314 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index a6ecc09f..6069cfca 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -616,28 +616,34 @@ def forecast( # Create an empty np array with shape [n_ens_members, rows, cols] # and fill it with the minimum value from precip (corresponding to # zero precipitation) - R_f_ = np.full( + precip_forecast = np.full( (n_ens_members, precip_shape[0], precip_shape[1]), np.nanmin(precip) ) if subtimestep_idx: if callback is not None: - if R_f_.shape[1] > 0: - callback(R_f_.squeeze()) + if precip_forecast.shape[1] > 0: + callback(precip_forecast.squeeze()) if return_output: for j in range(n_ens_members): - R_f[j].append(R_f_[j]) + R_f[j].append(precip_forecast[j]) - R_f_ = None + precip_forecast = None if measure_time: zero_precip_time = time.time() - starttime_init if return_output: - outarr = np.stack([np.stack(R_f[j]) for j in range(n_ens_members)]) + precip_forecast_all_members_all_times = np.stack( + [np.stack(R_f[j]) for j in range(n_ens_members)] + ) if measure_time: - return outarr, zero_precip_time, zero_precip_time + return ( + precip_forecast_all_members_all_times, + zero_precip_time, + zero_precip_time, + ) else: - return outarr + return precip_forecast_all_members_all_times else: return None @@ -699,7 +705,7 @@ def forecast( else: precip_noise_input = precip.copy() - pp, generate_noise, noise_std_coeffs = _init_noise( + generate_perturb, generate_noise, noise_std_coeffs = _init_noise( precip_noise_input, precip_thr, n_cascade_levels, @@ -734,18 +740,29 @@ def forecast( ) # 6. Initialize all the random generators and prepare for the forecast loop - randgen_prec, vps, generate_vel_noise = _init_random_generators( - velocity, - noise_method, - vel_pert_method, - vp_par, - vp_perp, - seed, - n_ens_members, - kmperpixel, - timestep, + randgen_precip, velocity_perturbations, generate_vel_noise = ( + _init_random_generators( + velocity, + noise_method, + vel_pert_method, + vp_par, + vp_perp, + seed, + n_ens_members, + kmperpixel, + timestep, + ) ) - D, D_Yn, D_pb, R_f, R_m, mask_rim, struct, fft_objs = _prepare_forecast_loop( + ( + previous_displacement, + previous_displacement_noise_cascade, + previous_displacement_prob_matching, + R_f, + precip_forecast_non_perturbed, + mask_rim, + struct, + fft_objs, + ) = _prepare_forecast_loop( precip_cascade, noise_method, fft_method, @@ -765,8 +782,8 @@ def forecast( n_cascade_levels=n_cascade_levels, generate_noise=generate_noise, decompositor=decompositor, - pp=pp, - randgen_prec=randgen_prec, + generate_perturb=generate_perturb, + randgen_precip=randgen_precip, fft_objs=fft_objs, bp_filter=bp_filter, domain=domain, @@ -779,8 +796,10 @@ def forecast( # 7. initizalize the current and previous extrapolation forecast scale # for the nowcasting component - rho_extr_prev = np.repeat(1.0, PHI.shape[0]) - rho_extr = PHI[:, 0] / (1.0 - PHI[:, 1]) # phi1 / (1 - phi2), see BPS2004 + rho_extrap_cascade_prev = np.repeat(1.0, PHI.shape[0]) + rho_extrap_cascade = PHI[:, 0] / ( + 1.0 - PHI[:, 1] + ) # phi1 / (1 - phi2), see BPS2004 if measure_time: init_time = time.time() - starttime_init @@ -795,11 +814,11 @@ def forecast( extrap_kwargs["return_displacement"] = True - forecast_prev = deepcopy(precip_cascade) - noise_prev = deepcopy(noise_cascade) + precip_forc_prev_subtimestep = deepcopy(precip_cascade) + noise_prev_subtimestep = deepcopy(noise_cascade) - t_prev = [0.0 for j in range(n_ens_members)] - t_total = [0.0 for j in range(n_ens_members)] + t_prev_timestep = [0.0 for j in range(n_ens_members)] + t_leadtime_since_start_forecast = [0.0 for j in range(n_ens_members)] # iterate each time step for t, subtimestep_idx in enumerate(timesteps): @@ -927,14 +946,16 @@ def forecast( # 8.1.3 Determine the skill of the components for lead time (t0 + t) # First for the extrapolation component. Only calculate it when t > 0. ( - rho_extr, - rho_extr_prev, + rho_extrap_cascade, + rho_extrap_cascade_prev, ) = blending.skill_scores.lt_dependent_cor_extrapolation( - PHI=PHI, correlations=rho_extr, correlations_prev=rho_extr_prev + PHI=PHI, + correlations=rho_extrap_cascade, + correlations_prev=rho_extrap_cascade_prev, ) # the nowcast iteration for each ensemble member - R_f_ = [None for _ in range(n_ens_members)] + precip_forecast = [None for _ in range(n_ens_members)] def worker(j): # 8.1.2 Determine the skill of the nwp components for lead time (t0 + t) @@ -951,8 +972,10 @@ def worker(j): for n_model in range(rho_nwp_models.shape[0]) ] rho_nwp_fc = np.stack(rho_nwp_fc) - # Concatenate rho_extr and rho_nwp - rho_fc = np.concatenate((rho_extr[None, :], rho_nwp_fc), axis=0) + # Concatenate rho_extrap_cascade and rho_nwp + rho_fc = np.concatenate( + (rho_extrap_cascade[None, :], rho_nwp_fc), axis=0 + ) else: rho_nwp_fc = blending.skill_scores.lt_dependent_cor_nwp( lt=(t * int(timestep)), @@ -961,9 +984,9 @@ def worker(j): n_model=n_model_indices[j], skill_kwargs=clim_kwargs, ) - # Concatenate rho_extr and rho_nwp + # Concatenate rho_extrap_cascade and rho_nwp rho_fc = np.concatenate( - (rho_extr[None, :], rho_nwp_fc[None, :]), axis=0 + (rho_extrap_cascade[None, :], rho_nwp_fc[None, :]), axis=0 ) # 8.2 Determine the weights per component @@ -991,7 +1014,7 @@ def worker(j): for i in range(n_cascade_levels): # Determine the normalized covariance matrix (containing) # the cross-correlations between the models - cov = np.corrcoef( + covariance_nwp_models = np.corrcoef( np.stack( [ precip_models_cascade_temp[ @@ -1005,7 +1028,8 @@ def worker(j): ) # Determine the weights for this cascade level weights_model_only[:, i] = calculate_weights_spn( - correlations=rho_fc[1:, i], cov=cov + correlations=rho_fc[1:, i], + covariance=covariance_nwp_models, ) else: # Same as correlation and noise is 1 - correlation @@ -1024,16 +1048,16 @@ def worker(j): # but spatially correlated noise if noise_method is not None: # generate noise field - EPS = generate_noise( - pp, - randstate=randgen_prec[j], + epsilon = generate_noise( + generate_perturb, + randstate=randgen_precip[j], fft_method=fft_objs[j], domain=domain, ) # decompose the noise field into a cascade - EPS = decompositor( - EPS, + epsilon_decomposed = decompositor( + epsilon, bp_filter, fft_method=fft_objs[j], input_domain=domain, @@ -1043,14 +1067,14 @@ def worker(j): compact_output=True, ) else: - EPS = None + epsilon_decomposed = None # 8.3.2 regress the extrapolation component to the subsequent time # step # iterate the AR(p) model for each cascade level for i in range(n_cascade_levels): # apply AR(p) process to extrapolation cascade level - if EPS is not None or vel_pert_method is not None: + if epsilon_decomposed is not None or vel_pert_method is not None: precip_cascade[j][i] = autoregression.iterate_ar_model( precip_cascade[j][i], PHI[i, :] ) @@ -1059,25 +1083,25 @@ def worker(j): else: # use the deterministic AR(p) model computed above if # perturbations are disabled - precip_cascade[j][i] = R_m[i] + precip_cascade[j][i] = precip_forecast_non_perturbed[i] # 8.3.3 regress the noise component to the subsequent time step # iterate the AR(p) model for each cascade level for i in range(n_cascade_levels): # normalize the noise cascade - if EPS is not None: - EPS_ = EPS["cascade_levels"][i] - EPS_ *= noise_std_coeffs[i] + if epsilon_decomposed is not None: + epsilon_temp = epsilon_decomposed["cascade_levels"][i] + epsilon_temp *= noise_std_coeffs[i] else: - EPS_ = None + epsilon_temp = None # apply AR(p) process to noise cascade level - # (Returns zero noise if EPS is None) + # (Returns zero noise if epsilon_decomposed is None) noise_cascade[j][i] = autoregression.iterate_ar_model( - noise_cascade[j][i], PHI[i, :], eps=EPS_ + noise_cascade[j][i], PHI[i, :], eps=epsilon_temp ) - EPS = None - EPS_ = None + epsilon_decomposed = None + epsilon_temp = None # 8.4 Perturb and blend the advection fields + advect the # extrapolation and noise cascade to the current time step @@ -1087,72 +1111,86 @@ def worker(j): extrap_kwargs_ = extrap_kwargs.copy() extrap_kwargs_noise = extrap_kwargs.copy() extrap_kwargs_pb = extrap_kwargs.copy() - velocity_pert = velocity - R_f_ep_out = [] - Yn_ep_out = [] - R_pm_ep = [] + velocity_perturbations_extrapolation = velocity + precip_forecast_extrapolated_decomp_done = [] + noise_extrapolated_decomp_done = [] + precip_forecast_extrapolated_probability_matching = [] # Extrapolate per sub time step for t_sub in subtimesteps: if t_sub > 0: - t_diff_prev_int = t_sub - int(t_sub) - if t_diff_prev_int > 0.0: - R_f_ip = [ - (1.0 - t_diff_prev_int) * forecast_prev[j][i][-1, :] - + t_diff_prev_int * precip_cascade[j][i][-1, :] + t_diff_prev_subtimestep_int = t_sub - int(t_sub) + if t_diff_prev_subtimestep_int > 0.0: + precip_forecast_cascade_subtimestep = [ + (1.0 - t_diff_prev_subtimestep_int) + * precip_forc_prev_subtimestep[j][i][-1, :] + + t_diff_prev_subtimestep_int + * precip_cascade[j][i][-1, :] for i in range(n_cascade_levels) ] - Yn_ip = [ - (1.0 - t_diff_prev_int) * noise_prev[j][i][-1, :] - + t_diff_prev_int * noise_cascade[j][i][-1, :] + noise_cascade_subtimestep = [ + (1.0 - t_diff_prev_subtimestep_int) + * noise_prev_subtimestep[j][i][-1, :] + + t_diff_prev_subtimestep_int + * noise_cascade[j][i][-1, :] for i in range(n_cascade_levels) ] else: - R_f_ip = [ - forecast_prev[j][i][-1, :] + precip_forecast_cascade_subtimestep = [ + precip_forc_prev_subtimestep[j][i][-1, :] for i in range(n_cascade_levels) ] - Yn_ip = [ - noise_prev[j][i][-1, :] for i in range(n_cascade_levels) + noise_cascade_subtimestep = [ + noise_prev_subtimestep[j][i][-1, :] + for i in range(n_cascade_levels) ] - R_f_ip = np.stack(R_f_ip) - Yn_ip = np.stack(Yn_ip) + precip_forecast_cascade_subtimestep = np.stack( + precip_forecast_cascade_subtimestep + ) + noise_cascade_subtimestep = np.stack(noise_cascade_subtimestep) - t_diff_prev = t_sub - t_prev[j] - t_total[j] += t_diff_prev + t_diff_prev_subtimestep = t_sub - t_prev_timestep[j] + t_leadtime_since_start_forecast[j] += t_diff_prev_subtimestep # compute the perturbed motion field - include the NWP # velocities and the weights. Note that we only perturb # the extrapolation velocity field, as the NWP velocity # field is present per time step if vel_pert_method is not None: - velocity_pert = velocity + generate_vel_noise( - vps[j], t_total[j] * timestep + velocity_perturbations_extrapolation = ( + velocity + + generate_vel_noise( + velocity_perturbations[j], + t_leadtime_since_start_forecast[j] * timestep, + ) ) # Stack the perturbed extrapolation and the NWP velocities if blend_nwp_members: - V_stack = np.concatenate( + velocity_stack_all = np.concatenate( ( - velocity_pert[None, :, :, :], + velocity_perturbations_extrapolation[None, :, :, :], velocity_models_temp, ), axis=0, ) else: - V_model_ = velocity_models_temp[j] - V_stack = np.concatenate( - (velocity_pert[None, :, :, :], V_model_[None, :, :, :]), + velocity_models = velocity_models_temp[j] + velocity_stack_all = np.concatenate( + ( + velocity_perturbations_extrapolation[None, :, :, :], + velocity_models[None, :, :, :], + ), axis=0, ) - V_model_ = None + velocity_models = None # Obtain a blended optical flow, using the weights of the # second cascade following eq. 24 in BPS2006 velocity_blended = blending.utils.blend_optical_flows( - flows=V_stack, + flows=velocity_stack_all, weights=weights[ :-1, 1 ], # [(extr_field, n_model_fields), cascade_level=2] @@ -1163,36 +1201,51 @@ def worker(j): # This is needed to remove the interpolation artifacts. # In addition, the number of extrapolations is greatly reduced # A. Radar Rain - R_f_ip_recomp = blending.utils.recompose_cascade( - combined_cascade=R_f_ip, - combined_mean=mu_extrapolation, - combined_sigma=sigma_extrapolation, + precip_forecast_recomp_subtimestep = ( + blending.utils.recompose_cascade( + combined_cascade=precip_forecast_cascade_subtimestep, + combined_mean=mu_extrapolation, + combined_sigma=sigma_extrapolation, + ) ) # Make sure we have values outside the mask if zero_precip_radar: - R_f_ip_recomp = np.nan_to_num( - R_f_ip_recomp, + precip_forecast_recomp_subtimestep = np.nan_to_num( + precip_forecast_recomp_subtimestep, copy=True, nan=zerovalue, posinf=zerovalue, neginf=zerovalue, ) # Put back the mask - R_f_ip_recomp[domain_mask] = np.nan - extrap_kwargs["displacement_prev"] = D[j] - R_f_ep_recomp_, D[j] = extrapolator( - R_f_ip_recomp, + precip_forecast_recomp_subtimestep[domain_mask] = np.nan + extrap_kwargs["displacement_prev"] = previous_displacement[j] + ( + precip_forecast_extrapolated_recomp_subtimestep_temp, + previous_displacement[j], + ) = extrapolator( + precip_forecast_recomp_subtimestep, velocity_blended, - [t_diff_prev], + [t_diff_prev_subtimestep], allow_nonfinite_values=True, **extrap_kwargs, ) - R_f_ep_recomp = R_f_ep_recomp_[0].copy() - temp_mask = ~np.isfinite(R_f_ep_recomp) + precip_forecast_extrapolated_recomp_subtimestep = ( + precip_forecast_extrapolated_recomp_subtimestep_temp[ + 0 + ].copy() + ) + temp_mask = ~np.isfinite( + precip_forecast_extrapolated_recomp_subtimestep + ) # TODO WHERE DO CAN I FIND THIS -15.0 - R_f_ep_recomp[~np.isfinite(R_f_ep_recomp)] = zerovalue - R_f_ep = decompositor( - R_f_ep_recomp, + precip_forecast_extrapolated_recomp_subtimestep[ + ~np.isfinite( + precip_forecast_extrapolated_recomp_subtimestep + ) + ] = zerovalue + precip_forecast_extrapolated_decomp = decompositor( + precip_forecast_extrapolated_recomp_subtimestep, bp_filter, mask=MASK_thr, fft_method=fft, @@ -1203,33 +1256,42 @@ def worker(j): )["cascade_levels"] # Make sure we have values outside the mask if zero_precip_radar: - R_f_ep = np.nan_to_num( - R_f_ep, + precip_forecast_extrapolated_decomp = np.nan_to_num( + precip_forecast_extrapolated_decomp, copy=True, - nan=np.nanmin(R_f_ip), - posinf=np.nanmin(R_f_ip), - neginf=np.nanmin(R_f_ip), + nan=np.nanmin(precip_forecast_cascade_subtimestep), + posinf=np.nanmin(precip_forecast_cascade_subtimestep), + neginf=np.nanmin(precip_forecast_cascade_subtimestep), ) for i in range(n_cascade_levels): - R_f_ep[i][temp_mask] = np.nan + precip_forecast_extrapolated_decomp[i][temp_mask] = np.nan # B. Noise - Yn_ip_recomp = blending.utils.recompose_cascade( - combined_cascade=Yn_ip, - combined_mean=mu_noise[j], - combined_sigma=sigma_noise[j], + noise_cascade_subtimestep_recomp = ( + blending.utils.recompose_cascade( + combined_cascade=noise_cascade_subtimestep, + combined_mean=mu_noise[j], + combined_sigma=sigma_noise[j], + ) + ) + extrap_kwargs_noise["displacement_prev"] = ( + previous_displacement_noise_cascade[j] ) - extrap_kwargs_noise["displacement_prev"] = D_Yn[j] extrap_kwargs_noise["map_coordinates_mode"] = "wrap" - Yn_ep_recomp_, D_Yn[j] = extrapolator( - Yn_ip_recomp, + ( + noise_extrapolated_recomp_temp, + previous_displacement_noise_cascade[j], + ) = extrapolator( + noise_cascade_subtimestep_recomp, velocity_blended, - [t_diff_prev], + [t_diff_prev_subtimestep], allow_nonfinite_values=True, **extrap_kwargs_noise, ) - Yn_ep_recomp = Yn_ep_recomp_[0].copy() - Yn_ep = decompositor( - Yn_ep_recomp, + noise_extrapolated_recomp = noise_extrapolated_recomp_temp[ + 0 + ].copy() + noise_extrapolated_decomp = decompositor( + noise_extrapolated_recomp, bp_filter, mask=MASK_thr, fft_method=fft, @@ -1239,77 +1301,104 @@ def worker(j): compact_output=True, )["cascade_levels"] for i in range(n_cascade_levels): - Yn_ep[i] *= noise_std_coeffs[i] + noise_extrapolated_decomp[i] *= noise_std_coeffs[i] # Append the results to the output lists - R_f_ep_out.append(R_f_ep.copy()) - Yn_ep_out.append(Yn_ep.copy()) - R_f_ip = None - R_f_ip_recomp = None - R_f_ep_recomp_ = None - R_f_ep_recomp = None - R_f_ep = None - Yn_ip = None - Yn_ip_recomp = None - Yn_ep_recomp_ = None - Yn_ep_recomp = None - Yn_ep = None + precip_forecast_extrapolated_decomp_done.append( + precip_forecast_extrapolated_decomp.copy() + ) + noise_extrapolated_decomp_done.append( + noise_extrapolated_decomp.copy() + ) + precip_forecast_cascade_subtimestep = None + precip_forecast_recomp_subtimestep = None + precip_forecast_extrapolated_recomp_subtimestep_temp = None + precip_forecast_extrapolated_recomp_subtimestep = None + precip_forecast_extrapolated_decomp = None + noise_cascade_subtimestep = None + noise_cascade_subtimestep_recomp = None + noise_extrapolated_recomp_temp = None + noise_extrapolated_recomp = None + noise_extrapolated_decomp = None # Finally, also extrapolate the initial radar rainfall # field. This will be blended with the rainfall field(s) # of the (NWP) model(s) for Lagrangian blended prob. matching # min_R = np.min(precip) - extrap_kwargs_pb["displacement_prev"] = D_pb[j] + extrap_kwargs_pb["displacement_prev"] = ( + previous_displacement_prob_matching[j] + ) # Apply the domain mask to the extrapolation component - R_ = precip.copy() - R_[domain_mask] = np.nan - R_pm_ep_, D_pb[j] = extrapolator( - R_, + precip_forecast = precip.copy() + precip_forecast[domain_mask] = np.nan + ( + precip_forecast_extrapolated_probability_matching_temp, + previous_displacement_prob_matching[j], + ) = extrapolator( + precip_forecast, velocity_blended, - [t_diff_prev], + [t_diff_prev_subtimestep], allow_nonfinite_values=True, **extrap_kwargs_pb, ) - R_pm_ep.append(R_pm_ep_[0]) + precip_forecast_extrapolated_probability_matching.append( + precip_forecast_extrapolated_probability_matching_temp[0] + ) - t_prev[j] = t_sub + t_prev_timestep[j] = t_sub - if len(R_f_ep_out) > 0: - R_f_ep_out = np.stack(R_f_ep_out) - Yn_ep_out = np.stack(Yn_ep_out) - R_pm_ep = np.stack(R_pm_ep) + if len(precip_forecast_extrapolated_decomp_done) > 0: + precip_forecast_extrapolated_decomp_done = np.stack( + precip_forecast_extrapolated_decomp_done + ) + noise_extrapolated_decomp_done = np.stack( + noise_extrapolated_decomp_done + ) + precip_forecast_extrapolated_probability_matching = np.stack( + precip_forecast_extrapolated_probability_matching + ) # advect the forecast field by one time step if no subtimesteps in the # current interval were found if not subtimesteps: - t_diff_prev = t + 1 - t_prev[j] - t_total[j] += t_diff_prev + t_diff_prev_subtimestep = t + 1 - t_prev_timestep[j] + t_leadtime_since_start_forecast[j] += t_diff_prev_subtimestep # compute the perturbed motion field - include the NWP # velocities and the weights if vel_pert_method is not None: - velocity_pert = velocity + generate_vel_noise( - vps[j], t_total[j] * timestep + velocity_perturbations_extrapolation = ( + velocity + + generate_vel_noise( + velocity_perturbations[j], + t_leadtime_since_start_forecast[j] * timestep, + ) ) # Stack the perturbed extrapolation and the NWP velocities if blend_nwp_members: - V_stack = np.concatenate( - (velocity_pert[None, :, :, :], velocity_models_temp), + velocity_stack_all = np.concatenate( + ( + velocity_perturbations_extrapolation[None, :, :, :], + velocity_models_temp, + ), axis=0, ) else: - V_model_ = velocity_models_temp[j] - V_stack = np.concatenate( - (velocity_pert[None, :, :, :], V_model_[None, :, :, :]), + velocity_models = velocity_models_temp[j] + velocity_stack_all = np.concatenate( + ( + velocity_perturbations_extrapolation[None, :, :, :], + velocity_models[None, :, :, :], + ), axis=0, ) - V_model_ = None + velocity_models = None # Obtain a blended optical flow, using the weights of the # second cascade following eq. 24 in BPS2006 velocity_blended = blending.utils.blend_optical_flows( - flows=V_stack, + flows=velocity_stack_all, weights=weights[ :-1, 1 ], # [(extr_field, n_model_fields), cascade_level=2] @@ -1317,44 +1406,48 @@ def worker(j): # Extrapolate the extrapolation and noise cascade - extrap_kwargs_["displacement_prev"] = D[j] - extrap_kwargs_noise["displacement_prev"] = D_Yn[j] + extrap_kwargs_["displacement_prev"] = previous_displacement[j] + extrap_kwargs_noise["displacement_prev"] = ( + previous_displacement_noise_cascade[j] + ) extrap_kwargs_noise["map_coordinates_mode"] = "wrap" - _, D[j] = extrapolator( + _, previous_displacement[j] = extrapolator( None, velocity_blended, - [t_diff_prev], + [t_diff_prev_subtimestep], allow_nonfinite_values=True, **extrap_kwargs_, ) - _, D_Yn[j] = extrapolator( + _, previous_displacement_noise_cascade[j] = extrapolator( None, velocity_blended, - [t_diff_prev], + [t_diff_prev_subtimestep], allow_nonfinite_values=True, **extrap_kwargs_noise, ) # Also extrapolate the radar observation, used for the probability # matching and post-processing steps - extrap_kwargs_pb["displacement_prev"] = D_pb[j] - _, D_pb[j] = extrapolator( + extrap_kwargs_pb["displacement_prev"] = ( + previous_displacement_prob_matching[j] + ) + _, previous_displacement_prob_matching[j] = extrapolator( None, velocity_blended, - [t_diff_prev], + [t_diff_prev_subtimestep], allow_nonfinite_values=True, **extrap_kwargs_pb, ) - t_prev[j] = t + 1 + t_prev_timestep[j] = t + 1 - forecast_prev[j] = precip_cascade[j] - noise_prev[j] = noise_cascade[j] + precip_forc_prev_subtimestep[j] = precip_cascade[j] + noise_prev_subtimestep[j] = noise_cascade[j] # 8.5 Blend the cascades - R_f_out = [] + final_blended_forecast = [] for t_sub in subtimesteps: # TODO: does it make sense to use sub time steps - check if it works? @@ -1363,11 +1456,13 @@ def worker(j): # First concatenate the cascades and the means and sigmas # precip_models = [n_models,timesteps,n_cascade_levels,m,n] if blend_nwp_members: - cascades_stacked = np.concatenate( + cascade_stack_all_components = np.concatenate( ( - R_f_ep_out[None, t_index], + precip_forecast_extrapolated_decomp_done[ + None, t_index + ], precip_models_cascade_temp, - Yn_ep_out[None, t_index], + noise_extrapolated_decomp_done[None, t_index], ), axis=0, ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] @@ -1379,11 +1474,13 @@ def worker(j): axis=0, ) else: - cascades_stacked = np.concatenate( + cascade_stack_all_components = np.concatenate( ( - R_f_ep_out[None, t_index], + precip_forecast_extrapolated_decomp_done[ + None, t_index + ], precip_models_cascade_temp[None, j], - Yn_ep_out[None, t_index], + noise_extrapolated_decomp_done[None, t_index], ), axis=0, ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] @@ -1403,36 +1500,46 @@ def worker(j): # weights for method bps have already been determined. if weights_method == "spn": weights = np.zeros( - (cascades_stacked.shape[0], n_cascade_levels) + ( + cascade_stack_all_components.shape[0], + n_cascade_levels, + ) ) for i in range(n_cascade_levels): # Determine the normalized covariance matrix (containing) # the cross-correlations between the models - cascades_stacked_ = np.stack( + cascade_stack_all_components_temp = np.stack( [ - cascades_stacked[n_model, i, :, :].flatten() + cascade_stack_all_components[ + n_model, i, :, : + ].flatten() for n_model in range( - cascades_stacked.shape[0] - 1 + cascade_stack_all_components.shape[0] - 1 ) ] ) # -1 to exclude the noise component - cov = np.ma.corrcoef( - np.ma.masked_invalid(cascades_stacked_) + covariance_nwp_models = np.ma.corrcoef( + np.ma.masked_invalid( + cascade_stack_all_components_temp + ) ) # Determine the weights for this cascade level weights[:, i] = calculate_weights_spn( - correlations=rho_fc[:, i], cov=cov + correlations=rho_fc[:, i], + covariance=covariance_nwp_models, ) # Blend the extrapolation, (NWP) model(s) and noise cascades - R_f_blended = blending.utils.blend_cascades( - cascades_norm=cascades_stacked, weights=weights + precip_forecast_blended = blending.utils.blend_cascades( + cascades_norm=cascade_stack_all_components, weights=weights ) # Also blend the cascade without the extrapolation component - R_f_blended_mod_only = blending.utils.blend_cascades( - cascades_norm=cascades_stacked[1:, :], - weights=weights_model_only, + precip_forecast_blended_mod_only = ( + blending.utils.blend_cascades( + cascades_norm=cascade_stack_all_components[1:, :], + weights=weights_model_only, + ) ) # Blend the means and standard deviations @@ -1451,24 +1558,30 @@ def worker(j): ) # 8.6 Recompose the cascade to a precipitation field - # (The function first normalizes the blended cascade, R_f_blended + # (The function first normalizes the blended cascade, precip_forecast_blended # again) - R_f_new = blending.utils.recompose_cascade( - combined_cascade=R_f_blended, + precip_forecast_recomposed = blending.utils.recompose_cascade( + combined_cascade=precip_forecast_blended, combined_mean=means_blended, combined_sigma=sigmas_blended, ) # The recomposed cascade without the extrapolation (for NaN filling # outside the radar domain) - R_f_new_mod_only = blending.utils.recompose_cascade( - combined_cascade=R_f_blended_mod_only, - combined_mean=means_blended_mod_only, - combined_sigma=sigmas_blended_mod_only, + precip_forecast_recomposed_mod_only = ( + blending.utils.recompose_cascade( + combined_cascade=precip_forecast_blended_mod_only, + combined_mean=means_blended_mod_only, + combined_sigma=sigmas_blended_mod_only, + ) ) if domain == "spectral": # TODO: Check this! (Only tested with domain == 'spatial') - R_f_new = fft_objs[j].irfft2(R_f_new) - R_f_new_mod_only = fft_objs[j].irfft2(R_f_new_mod_only) + precip_forecast_recomposed = fft_objs[j].irfft2( + precip_forecast_recomposed + ) + precip_forecast_recomposed_mod_only = fft_objs[j].irfft2( + precip_forecast_recomposed_mod_only + ) # 8.7 Post-processing steps - use the mask and fill no data with # the blended NWP forecast. Probability matching following @@ -1480,20 +1593,28 @@ def worker(j): # that is only used for post-processing steps) with the NWP # rainfall forecast for this time step using the weights # at scale level 2. - weights_pm = weights[:-1, 1] # Weights without noise, level 2 - weights_pm_normalized = weights_pm / np.sum(weights_pm) + weights_probability_matching = weights[ + :-1, 1 + ] # Weights without noise, level 2 + weights_probability_matching_normalized = ( + weights_probability_matching + / np.sum(weights_probability_matching) + ) # And the weights for outside the radar domain - weights_pm_mod_only = weights_model_only[ + weights_probability_matching_mod_only = weights_model_only[ :-1, 1 ] # Weights without noise, level 2 - weights_pm_normalized_mod_only = weights_pm_mod_only / np.sum( - weights_pm_mod_only + weights_probability_matching_normalized_mod_only = ( + weights_probability_matching_mod_only + / np.sum(weights_probability_matching_mod_only) ) # Stack the fields if blend_nwp_members: R_pm_stacked = np.concatenate( ( - R_pm_ep[None, t_index], + precip_forecast_extrapolated_probability_matching[ + None, t_index + ], precip_models_temp, ), axis=0, @@ -1501,29 +1622,37 @@ def worker(j): else: R_pm_stacked = np.concatenate( ( - R_pm_ep[None, t_index], + precip_forecast_extrapolated_probability_matching[ + None, t_index + ], precip_models_temp[None, j], ), axis=0, ) # Blend it - R_pm_blended = np.sum( - weights_pm_normalized.reshape( - weights_pm_normalized.shape[0], 1, 1 + precip_forecast_probability_matching_blended = np.sum( + weights_probability_matching_normalized.reshape( + weights_probability_matching_normalized.shape[0], 1, 1 ) * R_pm_stacked, axis=0, ) if blend_nwp_members: - R_pm_blended_mod_only = np.sum( - weights_pm_normalized_mod_only.reshape( - weights_pm_normalized_mod_only.shape[0], 1, 1 + precip_forecast_probability_matching_blended_mod_only = np.sum( + weights_probability_matching_normalized_mod_only.reshape( + weights_probability_matching_normalized_mod_only.shape[ + 0 + ], + 1, + 1, ) * precip_models_temp, axis=0, ) else: - R_pm_blended_mod_only = precip_models_temp[j] + precip_forecast_probability_matching_blended_mod_only = ( + precip_models_temp[j] + ) # The extrapolation components are NaN outside the advected # radar domain. This results in NaN values in the blended @@ -1531,7 +1660,7 @@ def worker(j): # areas with the "..._mod_only" blended forecasts, consisting # of the NWP and noise components. - nan_indices = np.isnan(R_f_new) + nan_indices = np.isnan(precip_forecast_recomposed) if smooth_radar_mask_range != 0: # Compute the smooth dilated mask new_mask = blending.utils.compute_smooth_dilated_mask( @@ -1543,49 +1672,68 @@ def worker(j): mask_model = np.clip(new_mask, 0, 1) mask_radar = np.clip(1 - new_mask, 0, 1) - # Handle NaNs in R_f_new and R_f_new_mod_only by setting NaNs to 0 in the blending step - R_f_new_mod_only_no_nan = np.nan_to_num( - R_f_new_mod_only, nan=0 + # Handle NaNs in precip_forecast_new and precip_forecast_new_mod_only by setting NaNs to 0 in the blending step + precip_forecast_recomposed_mod_only_no_nan = np.nan_to_num( + precip_forecast_recomposed_mod_only, nan=0 + ) + precip_forecast_recomposed_no_nan = np.nan_to_num( + precip_forecast_recomposed, nan=0 ) - R_f_new_no_nan = np.nan_to_num(R_f_new, nan=0) # Perform the blending of radar and model inside the radar domain using a weighted combination - R_f_new = np.nansum( + precip_forecast_recomposed = np.nansum( [ - mask_model * R_f_new_mod_only_no_nan, - mask_radar * R_f_new_no_nan, + mask_model + * precip_forecast_recomposed_mod_only_no_nan, + mask_radar * precip_forecast_recomposed_no_nan, ], axis=0, ) - nan_indices = np.isnan(R_pm_blended) - R_pm_blended = np.nansum( + nan_indices = np.isnan( + precip_forecast_probability_matching_blended + ) + precip_forecast_probability_matching_blended = np.nansum( [ - R_pm_blended * mask_radar, - R_pm_blended_mod_only * mask_model, + precip_forecast_probability_matching_blended + * mask_radar, + precip_forecast_probability_matching_blended_mod_only + * mask_model, ], axis=0, ) else: - R_f_new[nan_indices] = R_f_new_mod_only[nan_indices] - nan_indices = np.isnan(R_pm_blended) - R_pm_blended[nan_indices] = R_pm_blended_mod_only[ + precip_forecast_recomposed[nan_indices] = ( + precip_forecast_recomposed_mod_only[nan_indices] + ) + nan_indices = np.isnan( + precip_forecast_probability_matching_blended + ) + precip_forecast_probability_matching_blended[ + nan_indices + ] = precip_forecast_probability_matching_blended_mod_only[ nan_indices ] # Finally, fill the remaining nan values, if present, with # the minimum value in the forecast - nan_indices = np.isnan(R_f_new) - R_f_new[nan_indices] = np.nanmin(R_f_new) - nan_indices = np.isnan(R_pm_blended) - R_pm_blended[nan_indices] = np.nanmin(R_pm_blended) + nan_indices = np.isnan(precip_forecast_recomposed) + precip_forecast_recomposed[nan_indices] = np.nanmin( + precip_forecast_recomposed + ) + nan_indices = np.isnan( + precip_forecast_probability_matching_blended + ) + precip_forecast_probability_matching_blended[nan_indices] = ( + np.nanmin(precip_forecast_probability_matching_blended) + ) # 8.7.2. Apply the masking and prob. matching if mask_method is not None: # apply the precipitation mask to prevent generation of new # precipitation into areas where it was not originally # observed - R_cmin = R_f_new.min() + precip_forecast_min_value = precip_forecast_recomposed.min() if mask_method == "incremental": # The incremental mask is slightly different from # the implementation in the non-blended steps.py, as @@ -1593,60 +1741,100 @@ def worker(j): # on R_pm_blended. Therefore, the buffer does not # increase over time. # Get the mask for this forecast - MASK_prec = R_pm_blended >= precip_thr + precip_field_mask = ( + precip_forecast_probability_matching_blended + >= precip_thr + ) # Buffer the mask - MASK_prec = _compute_incremental_mask( - MASK_prec, struct, mask_rim + precip_field_mask = _compute_incremental_mask( + precip_field_mask, struct, mask_rim ) # Get the final mask - R_f_new = R_cmin + (R_f_new - R_cmin) * MASK_prec - MASK_prec_ = R_f_new > R_cmin + precip_forecast_recomposed = ( + precip_forecast_min_value + + ( + precip_forecast_recomposed + - precip_forecast_min_value + ) + * precip_field_mask + ) + precip_field_mask_temp = ( + precip_forecast_recomposed + > precip_forecast_min_value + ) elif mask_method == "obs": # The mask equals the most recent benchmark # rainfall field - MASK_prec_ = R_pm_blended >= precip_thr + precip_field_mask_temp = ( + precip_forecast_probability_matching_blended + >= precip_thr + ) # Set to min value outside of mask - R_f_new[~MASK_prec_] = R_cmin + precip_forecast_recomposed[~precip_field_mask_temp] = ( + precip_forecast_min_value + ) # If probmatching_method is not None, resample the distribution from # both the extrapolation cascade and the model (NWP) cascade and use # that for the probability matching. if probmatching_method is not None and resample_distribution: - arr1 = R_pm_ep[t_index] + arr1 = precip_forecast_extrapolated_probability_matching[ + t_index + ] arr2 = precip_models_temp[j] # resample weights based on cascade level 2. # Areas where one of the fields is nan are not included. R_pm_resampled = probmatching.resample_distributions( first_array=arr1, second_array=arr2, - probability_first_array=weights_pm_normalized[0], + probability_first_array=weights_probability_matching_normalized[ + 0 + ], ) else: - R_pm_resampled = R_pm_blended.copy() + R_pm_resampled = ( + precip_forecast_probability_matching_blended.copy() + ) if probmatching_method == "cdf": # nan indices in the extrapolation nowcast - nan_indices = np.isnan(R_pm_ep[t_index]) + nan_indices = np.isnan( + precip_forecast_extrapolated_probability_matching[ + t_index + ] + ) # Adjust the CDF of the forecast to match the resampled distribution combined from # extrapolation and model fields. # Rainfall outside the pure extrapolation domain is not taken into account. - if np.any(np.isfinite(R_f_new)): - R_f_new = probmatching.nonparam_match_empirical_cdf( - R_f_new, R_pm_resampled, nan_indices + if np.any(np.isfinite(precip_forecast_recomposed)): + precip_forecast_recomposed = ( + probmatching.nonparam_match_empirical_cdf( + precip_forecast_recomposed, + R_pm_resampled, + nan_indices, + ) ) R_pm_resampled = None elif probmatching_method == "mean": # Use R_pm_blended as benchmark field and - mu_0 = np.mean(R_pm_resampled[R_pm_resampled >= precip_thr]) - MASK = R_f_new >= precip_thr - mu_fct = np.mean(R_f_new[MASK]) - R_f_new[MASK] = R_f_new[MASK] - mu_fct + mu_0 + mean_probabiltity_matching_forecast = np.mean( + R_pm_resampled[R_pm_resampled >= precip_thr] + ) + no_rain_mask = precip_forecast_recomposed >= precip_thr + mean_precip_forecast = np.mean( + precip_forecast_recomposed[no_rain_mask] + ) + precip_forecast_recomposed[no_rain_mask] = ( + precip_forecast_recomposed[no_rain_mask] + - mean_precip_forecast + + mean_probabiltity_matching_forecast + ) R_pm_resampled = None - R_f_out.append(R_f_new) + final_blended_forecast.append(precip_forecast_recomposed) - R_f_[j] = R_f_out + precip_forecast[j] = final_blended_forecast res = [] @@ -1667,25 +1855,27 @@ def worker(j): print("done.") if callback is not None: - R_f_stacked = np.stack(R_f_) - if R_f_stacked.shape[1] > 0: - callback(R_f_stacked.squeeze()) + precip_forecast_final = np.stack(precip_forecast) + if precip_forecast_final.shape[1] > 0: + callback(precip_forecast_final.squeeze()) if return_output: for j in range(n_ens_members): - R_f[j].extend(R_f_[j]) + R_f[j].extend(precip_forecast[j]) - R_f_ = None + precip_forecast = None if measure_time: mainloop_time = time.time() - starttime_mainloop if return_output: - outarr = np.stack([np.stack(R_f[j]) for j in range(n_ens_members)]) + precip_forecast_all_members_all_times = np.stack( + [np.stack(R_f[j]) for j in range(n_ens_members)] + ) if measure_time: - return outarr, init_time, mainloop_time + return precip_forecast_all_members_all_times, init_time, mainloop_time else: - return outarr + return precip_forecast_all_members_all_times else: return None @@ -1773,7 +1963,7 @@ def calculate_weights_bps(correlations): return weights -def calculate_weights_spn(correlations, cov): +def calculate_weights_spn(correlations, covariance): """Calculate SPN blending weights for STEPS blending from correlation. Parameters @@ -1781,7 +1971,7 @@ def calculate_weights_spn(correlations, cov): correlations : array-like Array of shape [n_components] containing correlation (skills) for each component (NWP models and nowcast). - cov : array-like + covariance : array-like Array of shape [n_components, n_components] containing the covariance matrix of the models that will be blended. If cov is set to None and correlations only contains one model, the weight equals the correlation @@ -1801,21 +1991,21 @@ def calculate_weights_spn(correlations, cov): # Check if the correlations are positive, otherwise rho = 10e-5 correlations = np.where(correlations < 10e-5, 10e-5, correlations) - if correlations.shape[0] > 1 and len(cov) > 1: - if isinstance(cov, type(None)): + if correlations.shape[0] > 1 and len(covariance) > 1: + if isinstance(covariance, type(None)): raise ValueError("cov must contain a covariance matrix") else: # Make a numpy array out of cov and get the inverse - cov = np.where(cov == 0.0, 10e-5, cov) + covariance = np.where(covariance == 0.0, 10e-5, covariance) # Make sure the determinant of the matrix is not zero, otherwise # subtract 10e-5 from the cross-correlations between the models - if np.linalg.det(cov) == 0.0: - cov = cov - 10e-5 + if np.linalg.det(covariance) == 0.0: + covariance = covariance - 10e-5 # Ensure the correlation of the model with itself is always 1.0 - for i, _ in enumerate(cov): - cov[i][i] = 1.0 + for i, _ in enumerate(covariance): + covariance[i][i] = 1.0 # Use a numpy array instead of a matrix - cov_matrix = np.array(cov) + cov_matrix = np.array(covariance) # Get the inverse of the matrix using scipy's inv function cov_matrix_inv = inv(cov_matrix) # The component weights are the dot product between cov_matrix_inv and cor_vec @@ -2041,21 +2231,21 @@ def _init_noise( init_noise, generate_noise = noise.get_method(noise_method) # initialize the perturbation generator for the precipitation field - pp = init_noise(precip, fft_method=fft, **noise_kwargs) + generate_perturb = init_noise(precip, fft_method=fft, **noise_kwargs) if noise_stddev_adj == "auto": print("Computing noise adjustment coefficients... ", end="", flush=True) if measure_time: starttime = time.time() - R_min = np.min(precip) + precip_forecast_min = np.min(precip) noise_std_coeffs = noise.utils.compute_noise_stddev_adjs( precip[-1, :, :], precip_thr, - R_min, + precip_forecast_min, bp_filter, decompositor, - pp, + generate_perturb, generate_noise, 20, conditional=True, @@ -2076,7 +2266,7 @@ def _init_noise( if noise_stddev_adj is not None: print(f"noise std. dev. coeffs: {noise_std_coeffs}") - return pp, generate_noise, noise_std_coeffs + return generate_perturb, generate_noise, noise_std_coeffs def _compute_cascade_decomposition_radar( @@ -2091,9 +2281,9 @@ def _compute_cascade_decomposition_radar( fft, ): """Compute the cascade decompositions of the input precipitation fields.""" - R_d = [] + precip_forecast_decomp = [] for i in range(ar_order + 1): - R_ = decompositor( + precip_forecast = decompositor( precip[i, :, :], bp_filter, mask=MASK_thr, @@ -2103,17 +2293,21 @@ def _compute_cascade_decomposition_radar( compute_stats=True, compact_output=True, ) - R_d.append(R_) + precip_forecast_decomp.append(precip_forecast) # Rearrange the cascaded into a four-dimensional array of shape # (n_cascade_levels,ar_order+1,m,n) for the autoregressive model - R_c = nowcast_utils.stack_cascades(R_d, n_cascade_levels) + precip_forecast_cascades = nowcast_utils.stack_cascades( + precip_forecast_decomp, n_cascade_levels + ) - R_d = R_d[-1] - mu_extrapolation = np.array(R_d["means"]) - sigma_extrapolation = np.array(R_d["stds"]) - R_d = [R_d.copy() for j in range(n_ens_members)] - return R_c, mu_extrapolation, sigma_extrapolation + precip_forecast_decomp = precip_forecast_decomp[-1] + mu_extrapolation = np.array(precip_forecast_decomp["means"]) + sigma_extrapolation = np.array(precip_forecast_decomp["stds"]) + precip_forecast_decomp = [ + precip_forecast_decomp.copy() for j in range(n_ens_members) + ] + return precip_forecast_cascades, mu_extrapolation, sigma_extrapolation def _compute_cascade_recomposition_nwp(precip_models_cascade, recompositor): @@ -2137,7 +2331,7 @@ def _compute_cascade_recomposition_nwp(precip_models_cascade, recompositor): def _estimate_ar_parameters_radar( - R_c, ar_order, n_cascade_levels, MASK_thr, zero_precip_radar + precip_forecast_cascades, ar_order, n_cascade_levels, MASK_thr, zero_precip_radar ): """Estimate AR parameters for the radar rainfall field.""" # If there are values in the radar fields, compute the autocorrelations @@ -2145,7 +2339,9 @@ def _estimate_ar_parameters_radar( if not zero_precip_radar: # compute lag-l temporal autocorrelation coefficients for each cascade level for i in range(n_cascade_levels): - GAMMA[i, :] = correlation.temporal_autocorrelation(R_c[i], mask=MASK_thr) + GAMMA[i, :] = correlation.temporal_autocorrelation( + precip_forecast_cascades[i], mask=MASK_thr + ) # Else, use standard values for the autocorrelations else: @@ -2204,7 +2400,7 @@ def _estimate_ar_parameters_radar( def _find_nwp_combination( precip_models, - R_models_pm, + precip_forecast_probability_matching, velocity_models, mu_models, sigma_models, @@ -2258,7 +2454,9 @@ def _find_nwp_combination( sigma_models = np.repeat(sigma_models, n_ens_members_max, axis=0) velocity_models = np.repeat(velocity_models, n_ens_members_max, axis=0) # For the prob. matching - R_models_pm = np.repeat(R_models_pm, n_ens_members_max, axis=0) + precip_forecast_probability_matching = np.repeat( + precip_forecast_probability_matching, n_ens_members_max, axis=0 + ) # Finally, for the model indices n_model_indices = np.repeat(n_model_indices, n_ens_members_max, axis=0) @@ -2273,13 +2471,15 @@ def _find_nwp_combination( sigma_models = np.repeat(sigma_models, repeats, axis=0) velocity_models = np.repeat(velocity_models, repeats, axis=0) # For the prob. matching - R_models_pm = np.repeat(R_models_pm, repeats, axis=0) + precip_forecast_probability_matching = np.repeat( + precip_forecast_probability_matching, repeats, axis=0 + ) # Finally, for the model indices n_model_indices = np.repeat(n_model_indices, repeats, axis=0) return ( precip_models, - R_models_pm, + precip_forecast_probability_matching, velocity_models, mu_models, sigma_models, @@ -2300,11 +2500,11 @@ def _init_random_generators( ): """Initialize all the random generators.""" if noise_method is not None: - randgen_prec = [] + randgen_precip = [] randgen_motion = [] for j in range(n_ens_members): rs = np.random.RandomState(seed) - randgen_prec.append(rs) + randgen_precip.append(rs) seed = rs.randint(0, high=1e9) rs = np.random.RandomState(seed) randgen_motion.append(rs) @@ -2314,7 +2514,7 @@ def _init_random_generators( init_vel_noise, generate_vel_noise = noise.get_method(vel_pert_method) # initialize the perturbation generators for the motion field - vps = [] + velocity_perturbations = [] for j in range(n_ens_members): kwargs = { "randstate": randgen_motion[j], @@ -2322,15 +2522,15 @@ def _init_random_generators( "p_perp": vp_perp, } vp_ = init_vel_noise(velocity, 1.0 / kmperpixel, timestep, **kwargs) - vps.append(vp_) + velocity_perturbations.append(vp_) else: - vps, generate_vel_noise = None, None + velocity_perturbations, generate_vel_noise = None, None - return randgen_prec, vps, generate_vel_noise + return randgen_precip, velocity_perturbations, generate_vel_noise def _prepare_forecast_loop( - R_c, + precip_forecast_cascades, noise_method, fft_method, n_cascade_levels, @@ -2342,9 +2542,9 @@ def _prepare_forecast_loop( ): """Prepare for the forecast loop.""" # Empty arrays for the previous displacements and the forecast cascade - D = np.stack([None for j in range(n_ens_members)]) - D_Yn = np.stack([None for j in range(n_ens_members)]) - D_pb = np.stack([None for j in range(n_ens_members)]) + previous_displacement = np.stack([None for j in range(n_ens_members)]) + previous_displacement_noise_cascade = np.stack([None for j in range(n_ens_members)]) + previous_displacement_prob_matching = np.stack([None for j in range(n_ens_members)]) R_f = [[] for j in range(n_ens_members)] if mask_method == "incremental": @@ -2360,24 +2560,42 @@ def _prepare_forecast_loop( mask_rim, struct = None, None if noise_method is None: - R_m = [R_c[0][i].copy() for i in range(n_cascade_levels)] + precip_forecast_non_perturbed = [ + precip_forecast_cascades[0][i].copy() for i in range(n_cascade_levels) + ] else: - R_m = None + precip_forecast_non_perturbed = None fft_objs = [] for i in range(n_ens_members): - fft_objs.append(utils.get_method(fft_method, shape=R_c.shape[-2:])) + fft_objs.append( + utils.get_method(fft_method, shape=precip_forecast_cascades.shape[-2:]) + ) - return D, D_Yn, D_pb, R_f, R_m, mask_rim, struct, fft_objs + return ( + previous_displacement, + previous_displacement_noise_cascade, + previous_displacement_prob_matching, + R_f, + precip_forecast_non_perturbed, + mask_rim, + struct, + fft_objs, + ) def _compute_initial_nwp_skill( - R_c, precip_models, domain_mask, issuetime, outdir_path_skill, clim_kwargs + precip_forecast_cascades, + precip_models, + domain_mask, + issuetime, + outdir_path_skill, + clim_kwargs, ): """Calculate the initial skill of the (NWP) model forecasts at t=0.""" rho_nwp_models = [ blending.skill_scores.spatial_correlation( - obs=R_c[0, :, -1, :, :].copy(), + obs=precip_forecast_cascades[0, :, -1, :, :].copy(), mod=precip_models[n_model, :, :, :].copy(), domain_mask=domain_mask, ) @@ -2408,8 +2626,8 @@ def _init_noise_cascade( n_cascade_levels, generate_noise, decompositor, - pp, - randgen_prec, + generate_perturb, + randgen_precip, fft_objs, bp_filter, domain, @@ -2426,11 +2644,14 @@ def _init_noise_cascade( sigma_noise = np.zeros((n_ens_members, n_cascade_levels)) if noise_method: for j in range(n_ens_members): - EPS = generate_noise( - pp, randstate=randgen_prec[j], fft_method=fft_objs[j], domain=domain + epsilon = generate_noise( + generate_perturb, + randstate=randgen_precip[j], + fft_method=fft_objs[j], + domain=domain, ) - EPS = decompositor( - EPS, + epsilon_decomposed = decompositor( + epsilon, bp_filter, fft_method=fft_objs[j], input_domain=domain, @@ -2439,15 +2660,15 @@ def _init_noise_cascade( normalize=True, compact_output=True, ) - mu_noise[j] = EPS["means"] - sigma_noise[j] = EPS["stds"] + mu_noise[j] = epsilon_decomposed["means"] + sigma_noise[j] = epsilon_decomposed["stds"] for i in range(n_cascade_levels): - EPS_ = EPS["cascade_levels"][i] - EPS_ *= noise_std_coeffs[i] + epsilon_temp = epsilon_decomposed["cascade_levels"][i] + epsilon_temp *= noise_std_coeffs[i] for n in range(ar_order): - noise_cascade[j][i][n] = EPS_ - EPS = None - EPS_ = None + noise_cascade[j][i][n] = epsilon_temp + epsilon_decomposed = None + epsilon_temp = None return noise_cascade, mu_noise, sigma_noise diff --git a/pysteps/nowcasts/steps.py b/pysteps/nowcasts/steps.py index 806c4082..0f0b27b2 100644 --- a/pysteps/nowcasts/steps.py +++ b/pysteps/nowcasts/steps.py @@ -1214,7 +1214,6 @@ def reset_states_and_params(self): # Wrapper function to preserve backward compatibility -@deprecate_args({"R": "precip", "V": "velocity", "R_thr": "precip_thr"}, "1.8.0") def forecast( precip, velocity, From 72d0fbc102a1492f09e0be71f1b697961235dc6d Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 18 Nov 2024 14:32:43 +0100 Subject: [PATCH 02/33] Made some name changes but test still do not pass --- pysteps/blending/steps.py | 59 ++++++++++++++++++++------------------- 1 file changed, 31 insertions(+), 28 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 6069cfca..f873d463 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -606,7 +606,7 @@ def forecast( ) print("The resulting forecast will contain only zeros") # Create the output list - R_f = [[] for j in range(n_ens_members)] + precip_forecast = [[] for j in range(n_ens_members)] # Save per time step to ensure the array does not become too large if # no return_output is requested and callback is not None. @@ -616,25 +616,25 @@ def forecast( # Create an empty np array with shape [n_ens_members, rows, cols] # and fill it with the minimum value from precip (corresponding to # zero precipitation) - precip_forecast = np.full( + precip_forecast_temp = np.full( (n_ens_members, precip_shape[0], precip_shape[1]), np.nanmin(precip) ) if subtimestep_idx: if callback is not None: - if precip_forecast.shape[1] > 0: - callback(precip_forecast.squeeze()) + if precip_forecast_temp.shape[1] > 0: + callback(precip_forecast_temp.squeeze()) if return_output: for j in range(n_ens_members): - R_f[j].append(precip_forecast[j]) + precip_forecast[j].append(precip_forecast_temp[j]) - precip_forecast = None + precip_forecast_temp = None if measure_time: zero_precip_time = time.time() - starttime_init if return_output: precip_forecast_all_members_all_times = np.stack( - [np.stack(R_f[j]) for j in range(n_ens_members)] + [np.stack(precip_forecast[j]) for j in range(n_ens_members)] ) if measure_time: return ( @@ -757,7 +757,7 @@ def forecast( previous_displacement, previous_displacement_noise_cascade, previous_displacement_prob_matching, - R_f, + precip_forecast, precip_forecast_non_perturbed, mask_rim, struct, @@ -955,7 +955,7 @@ def forecast( ) # the nowcast iteration for each ensemble member - precip_forecast = [None for _ in range(n_ens_members)] + precip_forecast_temp = [None for _ in range(n_ens_members)] def worker(j): # 8.1.2 Determine the skill of the nwp components for lead time (t0 + t) @@ -1329,13 +1329,13 @@ def worker(j): previous_displacement_prob_matching[j] ) # Apply the domain mask to the extrapolation component - precip_forecast = precip.copy() - precip_forecast[domain_mask] = np.nan + precip_forecast_temp = precip.copy() + precip_forecast_temp[domain_mask] = np.nan ( precip_forecast_extrapolated_probability_matching_temp, previous_displacement_prob_matching[j], ) = extrapolator( - precip_forecast, + precip_forecast_temp, velocity_blended, [t_diff_prev_subtimestep], allow_nonfinite_values=True, @@ -1610,7 +1610,7 @@ def worker(j): ) # Stack the fields if blend_nwp_members: - R_pm_stacked = np.concatenate( + precip_forecast_probability_matching_final = np.concatenate( ( precip_forecast_extrapolated_probability_matching[ None, t_index @@ -1620,7 +1620,7 @@ def worker(j): axis=0, ) else: - R_pm_stacked = np.concatenate( + precip_forecast_probability_matching_final = np.concatenate( ( precip_forecast_extrapolated_probability_matching[ None, t_index @@ -1634,7 +1634,7 @@ def worker(j): weights_probability_matching_normalized.reshape( weights_probability_matching_normalized.shape[0], 1, 1 ) - * R_pm_stacked, + * precip_forecast_probability_matching_final, axis=0, ) if blend_nwp_members: @@ -1785,7 +1785,7 @@ def worker(j): arr2 = precip_models_temp[j] # resample weights based on cascade level 2. # Areas where one of the fields is nan are not included. - R_pm_resampled = probmatching.resample_distributions( + precip_forecast_probability_matching_resampled = probmatching.resample_distributions( first_array=arr1, second_array=arr2, probability_first_array=weights_probability_matching_normalized[ @@ -1793,7 +1793,7 @@ def worker(j): ], ) else: - R_pm_resampled = ( + precip_forecast_probability_matching_resampled = ( precip_forecast_probability_matching_blended.copy() ) @@ -1811,15 +1811,18 @@ def worker(j): precip_forecast_recomposed = ( probmatching.nonparam_match_empirical_cdf( precip_forecast_recomposed, - R_pm_resampled, + precip_forecast_probability_matching_resampled, nan_indices, ) ) - R_pm_resampled = None + precip_forecast_probability_matching_resampled = None elif probmatching_method == "mean": # Use R_pm_blended as benchmark field and mean_probabiltity_matching_forecast = np.mean( - R_pm_resampled[R_pm_resampled >= precip_thr] + precip_forecast_probability_matching_resampled[ + precip_forecast_probability_matching_resampled + >= precip_thr + ] ) no_rain_mask = precip_forecast_recomposed >= precip_thr mean_precip_forecast = np.mean( @@ -1830,11 +1833,11 @@ def worker(j): - mean_precip_forecast + mean_probabiltity_matching_forecast ) - R_pm_resampled = None + precip_forecast_probability_matching_resampled = None final_blended_forecast.append(precip_forecast_recomposed) - precip_forecast[j] = final_blended_forecast + precip_forecast_temp[j] = final_blended_forecast res = [] @@ -1855,22 +1858,22 @@ def worker(j): print("done.") if callback is not None: - precip_forecast_final = np.stack(precip_forecast) + precip_forecast_final = np.stack(precip_forecast_temp) if precip_forecast_final.shape[1] > 0: callback(precip_forecast_final.squeeze()) if return_output: for j in range(n_ens_members): - R_f[j].extend(precip_forecast[j]) + precip_forecast[j].extend(precip_forecast_temp[j]) - precip_forecast = None + precip_forecast_temp = None if measure_time: mainloop_time = time.time() - starttime_mainloop if return_output: precip_forecast_all_members_all_times = np.stack( - [np.stack(R_f[j]) for j in range(n_ens_members)] + [np.stack(precip_forecast[j]) for j in range(n_ens_members)] ) if measure_time: return precip_forecast_all_members_all_times, init_time, mainloop_time @@ -2545,7 +2548,7 @@ def _prepare_forecast_loop( previous_displacement = np.stack([None for j in range(n_ens_members)]) previous_displacement_noise_cascade = np.stack([None for j in range(n_ens_members)]) previous_displacement_prob_matching = np.stack([None for j in range(n_ens_members)]) - R_f = [[] for j in range(n_ens_members)] + precip_forecast = [[] for j in range(n_ens_members)] if mask_method == "incremental": # get mask parameters @@ -2576,7 +2579,7 @@ def _prepare_forecast_loop( previous_displacement, previous_displacement_noise_cascade, previous_displacement_prob_matching, - R_f, + precip_forecast, precip_forecast_non_perturbed, mask_rim, struct, From 1ce563e4275ef687bcb3da9e90f5db7224ba6185 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 18 Nov 2024 15:30:36 +0100 Subject: [PATCH 03/33] Fixed naming changes, now the tests pass --- .gitignore | 3 +++ pysteps/blending/steps.py | 28 +++++++++++++++------------- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/.gitignore b/.gitignore index e12dc8bf..4588187d 100644 --- a/.gitignore +++ b/.gitignore @@ -91,3 +91,6 @@ venv.bak/ # Mac OS Stuff .DS_Store + +# Running lcoal tests +/tmp diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index f873d463..fc4ffa34 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -616,18 +616,18 @@ def forecast( # Create an empty np array with shape [n_ens_members, rows, cols] # and fill it with the minimum value from precip (corresponding to # zero precipitation) - precip_forecast_temp = np.full( + precip_forecast_workers = np.full( (n_ens_members, precip_shape[0], precip_shape[1]), np.nanmin(precip) ) if subtimestep_idx: if callback is not None: - if precip_forecast_temp.shape[1] > 0: - callback(precip_forecast_temp.squeeze()) + if precip_forecast_workers.shape[1] > 0: + callback(precip_forecast_workers.squeeze()) if return_output: for j in range(n_ens_members): - precip_forecast[j].append(precip_forecast_temp[j]) + precip_forecast[j].append(precip_forecast_workers[j]) - precip_forecast_temp = None + precip_forecast_workers = None if measure_time: zero_precip_time = time.time() - starttime_init @@ -955,7 +955,7 @@ def forecast( ) # the nowcast iteration for each ensemble member - precip_forecast_temp = [None for _ in range(n_ens_members)] + precip_forecast_workers = [None for _ in range(n_ens_members)] def worker(j): # 8.1.2 Determine the skill of the nwp components for lead time (t0 + t) @@ -1329,13 +1329,15 @@ def worker(j): previous_displacement_prob_matching[j] ) # Apply the domain mask to the extrapolation component - precip_forecast_temp = precip.copy() - precip_forecast_temp[domain_mask] = np.nan + precip_forecast_temp_for_probability_matching = precip.copy() + precip_forecast_temp_for_probability_matching[domain_mask] = ( + np.nan + ) ( precip_forecast_extrapolated_probability_matching_temp, previous_displacement_prob_matching[j], ) = extrapolator( - precip_forecast_temp, + precip_forecast_temp_for_probability_matching, velocity_blended, [t_diff_prev_subtimestep], allow_nonfinite_values=True, @@ -1837,7 +1839,7 @@ def worker(j): final_blended_forecast.append(precip_forecast_recomposed) - precip_forecast_temp[j] = final_blended_forecast + precip_forecast_workers[j] = final_blended_forecast res = [] @@ -1858,15 +1860,15 @@ def worker(j): print("done.") if callback is not None: - precip_forecast_final = np.stack(precip_forecast_temp) + precip_forecast_final = np.stack(precip_forecast_workers) if precip_forecast_final.shape[1] > 0: callback(precip_forecast_final.squeeze()) if return_output: for j in range(n_ens_members): - precip_forecast[j].extend(precip_forecast_temp[j]) + precip_forecast[j].extend(precip_forecast_workers[j]) - precip_forecast_temp = None + precip_forecast_workers = None if measure_time: mainloop_time = time.time() - starttime_mainloop From fbe551b5f511900e64ca775e81337fccbb3de0bc Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 18 Nov 2024 16:33:01 +0100 Subject: [PATCH 04/33] Built the rough scaffolding for the blending class --- pysteps/blending/steps.py | 218 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 218 insertions(+) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index fc4ffa34..2fd9e3b3 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -66,6 +66,224 @@ DASK_IMPORTED = False +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + + +@dataclass +class StepsBlendingConfig: + # Configuration parameters + n_ens_members: int + n_cascade_levels: int + ar_order: int + timestep: float + blend_nwp_members: bool + precip_thr: float + norain_thr: float + kmperpixel: float + seed: Optional[int] + num_workers: int + measure_time: bool + domain: str = "spatial" + fft_method: str = "numpy" + extrap_method: str = "semilagrangian" + extrap_kwargs: Dict[str, Any] = field(default_factory=dict) + decomp_method: str = "fft" + bandpass_filter_method: str = "gaussian" + filter_kwargs: Dict[str, Any] = field(default_factory=dict) + noise_method: Optional[str] = "nonparametric" + noise_stddev_adj: Optional[str] = "auto" + noise_kwargs: Dict[str, Any] = field(default_factory=dict) + vel_pert_method: Optional[str] = "bps" + vel_pert_kwargs: Dict[str, Any] = field(default_factory=dict) + weights_method: str = "bps" + mask_method: Optional[str] = "incremental" + mask_kwargs: Dict[str, Any] = field(default_factory=dict) + probmatching_method: Optional[str] = "cdf" + resample_distribution: bool = True + smooth_radar_mask_range: int = 0 + outdir_path_skill: str = "./tmp/" + clim_kwargs: Dict[str, Any] = field(default_factory=dict) + callback: Optional[Any] = None + return_output: bool = True + # Additional configuration parameters as needed + + +@dataclass +class StepsBlendingState: + issuetime: Any # Replace with appropriate type, e.g., datetime.datetime + # Precomputed or intermediate data + domain_mask: Optional[np.ndarray] = None + MASK_thr: Optional[np.ndarray] = None + forecast_output: Optional[np.ndarray] = None + # Additional state variables as needed + + +@dataclass +class StepsBlendingParams: + # Parameters and variables calculated during initialization or processing + fft_method: Any = None + bandpass_filter: Any = None + decomposer: Any = None + recomposer: Any = None + generate_perturb: Any = None + generate_noise: Any = None + noise_std_coeffs: Optional[np.ndarray] = None + PHI: Optional[np.ndarray] = None + randgen_precip: Optional[List[np.random.RandomState]] = None + velocity_perturbations: Optional[List[Any]] = None + generate_vel_noise: Any = None + previous_displacement: Optional[np.ndarray] = None + previous_displacement_noise_cascade: Optional[np.ndarray] = None + previous_displacement_prob_matching: Optional[np.ndarray] = None + precip_cascade: Optional[np.ndarray] = None + mu: Optional[np.ndarray] = None + sigma: Optional[np.ndarray] = None + noise_cascade: Optional[np.ndarray] = None + mask_rim: Optional[int] = None + struct: Any = None + fft_objs: Optional[List[Any]] = None + # Additional parameters and variables as needed + + +class BlendingEngine: + def __init__( + self, + precip, + precip_models, + velocity, + velocity_models, + time_steps, + steps_blending_config: StepsBlendingConfig, + ): + # Store inputs and optional parameters + self.__precip = precip + self.__precip_models = precip_models + self.__velocity = velocity + self.__velocity_models = velocity_models + self.__time_steps = time_steps + + # Store the config data: + self.__config = steps_blending_config + + # Store the state and params data: + self.__state = StepsBlendingState() + self.__params = StepsBlendingParams() + + def forecast(self): + """Main method to perform the forecast.""" + self._check_inputs() + self._initialize() + self._prepare_data() + self._initialize_noise() + self._estimate_ar_parameters() + self._init_random_generators() + self._prepare_forecast_loop() + self._compute_forecast() + return self.state.forecast_output + + # Private methods for internal processing + def _initialize_methods(self): + """Set up methods for extrapolation, decomposition, etc.""" + pass + + def _initialize_bandpass_filter(self): + """Initialize the bandpass filter.""" + pass + + def _prepare_data(self): + """Transform data into Lagrangian coordinates and perform initial decomposition.""" + pass + + def _initialize_noise(self): + """Set up the noise generation mechanism.""" + pass + + def _estimate_ar_parameters(self): + """Estimate autoregressive model parameters.""" + self._estimate_ar_parameters_radar() + + def _init_random_generators(self): + """Initialize random number generators.""" + self._init_random_generators_for_noise() + if self.config.vel_pert_method: + self._init_velocity_perturbations() + + def _prepare_forecast_loop(self): + """Set up variables and structures for the forecasting loop.""" + self._initialize_forecast_variables() + + def _compute_forecast(self): + """Main loop to compute the forecast over the specified time steps.""" + self._run_forecast_loop() + + # Methods for specific functionalities + def _check_inputs(self): + """Validate input data and configurations.""" + # Implement input checks as needed + pass + + def _initialize(self): + """Perform any additional initialization steps.""" + # Initialize variables in self.params as needed + pass + + def _transform_to_lagrangian(self): + """Transform precipitation data to Lagrangian coordinates.""" + # Use self.state and self.params as needed + pass + + def _compute_cascade_decomposition_radar(self): + """Compute the cascade decomposition for radar data.""" + # Update self.params with computed cascades + pass + + def _estimate_ar_parameters_radar(self): + """Estimate AR parameters for radar data.""" + # Update self.params.PHI + pass + + def _init_random_generators_for_noise(self): + """Initialize random generators for noise.""" + # Update self.params.randgen_precip + pass + + def _init_velocity_perturbations(self): + """Initialize velocity perturbations if required.""" + # Update self.params.velocity_perturbations, generate_vel_noise + pass + + def _initialize_forecast_variables(self): + """Set up variables needed for the forecast loop.""" + # Initialize variables in self.params + pass + + def _run_forecast_loop(self): + """Run the main forecast loop over time steps.""" + # Use self.state and self.params as needed + pass + + # Additional helper methods as needed + + +class SkillScoreManager: + def __init__(self, config: StepsBlendingConfig): + self.outdir_path_skill = config.outdir_path_skill + self.clim_kwargs = config.clim_kwargs + + def compute_initial_skill(self, observed_cascades, model_cascades, domain_mask): + """Calculate the initial skill of NWP models at t=0.""" + pass # Implement as needed + + def update_skill(self, lead_time, correlations, model_indices): + """Update the skill scores based on lead time.""" + pass # Implement as needed + + def save_skill(self, current_skill, validtime): + """Save the skill scores to disk.""" + pass # Implement as needed + + def forecast( precip, precip_models, From 46a93e53084fd4506540a0b14823e072166ef151 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Wed, 27 Nov 2024 16:31:02 +0100 Subject: [PATCH 05/33] Refactored untill no rain case --- pysteps/blending/steps.py | 861 +++++++++++++++++++++++++++----------- pysteps/nowcasts/steps.py | 24 +- 2 files changed, 640 insertions(+), 245 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 2fd9e3b3..be3eee06 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -65,88 +65,111 @@ except ImportError: DASK_IMPORTED = False - from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional +from typing import Optional, List, Dict, Any, Callable @dataclass class StepsBlendingConfig: - # Configuration parameters + precip_threshold: Optional[float] + norain_threshold: float + kmperpixel: float + timestep: float n_ens_members: int n_cascade_levels: int - ar_order: int - timestep: float blend_nwp_members: bool - precip_thr: float - norain_thr: float - kmperpixel: float + extrapolation_method: str + decomposition_method: str + bandpass_filter_method: str + noise_method: Optional[str] + noise_stddev_adj: Optional[str] + ar_order: int + vel_pert_method: Optional[str] + weights_method: str + conditional: bool + probmatching_method: Optional[str] + mask_method: Optional[str] + resample_distribution: bool + smooth_radar_mask_range: int seed: Optional[int] num_workers: int - measure_time: bool - domain: str = "spatial" - fft_method: str = "numpy" - extrap_method: str = "semilagrangian" + fft_method: str + domain: str + outdir_path_skill: str extrap_kwargs: Dict[str, Any] = field(default_factory=dict) - decomp_method: str = "fft" - bandpass_filter_method: str = "gaussian" filter_kwargs: Dict[str, Any] = field(default_factory=dict) - noise_method: Optional[str] = "nonparametric" - noise_stddev_adj: Optional[str] = "auto" noise_kwargs: Dict[str, Any] = field(default_factory=dict) - vel_pert_method: Optional[str] = "bps" vel_pert_kwargs: Dict[str, Any] = field(default_factory=dict) - weights_method: str = "bps" - mask_method: Optional[str] = "incremental" - mask_kwargs: Dict[str, Any] = field(default_factory=dict) - probmatching_method: Optional[str] = "cdf" - resample_distribution: bool = True - smooth_radar_mask_range: int = 0 - outdir_path_skill: str = "./tmp/" clim_kwargs: Dict[str, Any] = field(default_factory=dict) + mask_kwargs: Dict[str, Any] = field(default_factory=dict) + measure_time: bool = False callback: Optional[Any] = None return_output: bool = True - # Additional configuration parameters as needed @dataclass -class StepsBlendingState: - issuetime: Any # Replace with appropriate type, e.g., datetime.datetime - # Precomputed or intermediate data - domain_mask: Optional[np.ndarray] = None - MASK_thr: Optional[np.ndarray] = None - forecast_output: Optional[np.ndarray] = None - # Additional state variables as needed +class StepsBlendingParams: + PHI: np.ndarray # AR(p) model parameters + noise_std_coeffs: np.ndarray # Noise standard deviation coefficients + mu_extrapolation: np.ndarray # Means of extrapolated cascades + sigma_extrapolation: np.ndarray # Std devs of extrapolated cascades + bandpass_filter: Any # Band-pass filter object + fft: Any # FFT method object + generate_perturb: Callable # Perturbation generator + generate_noise: Callable # Noise generator + generate_vel_noise: Optional[Callable] # Velocity noise generator + extrapolation_method: Any = None + decomposition_method: Any = None + recomposition_method: Any = None + vp_par: Optional[np.ndarray] = None # Velocity perturbation parameters (parallel) + vp_perp: Optional[np.ndarray] = ( + None # Velocity perturbation parameters (perpendicular) + ) + fft_objs: List[Any] = field( + default_factory=list + ) # FFT objects for ensemble members + mask_rim: Optional[int] = None # Rim size for masking + struct: Optional[np.ndarray] = None # Structuring element for mask + n_model_indices: Optional[np.ndarray] = None # NWP model indices + noise_method: Optional[str] = None # Noise method used + ar_order: int = 2 # Order of the AR model + seed: Optional[int] = None # Random seed for reproducibility + time_steps_is_list: bool = False # Time steps is a list + precip_models_provided_is_cascade: bool = False # Precip models are decomposed + xy_coordinates: np.ndarray | None = None + precip_zerovalue: Any = None + mask_threshold: Any = None + zero_precip_radar: bool = False + zero_precip_model_fields: bool = False @dataclass -class StepsBlendingParams: - # Parameters and variables calculated during initialization or processing - fft_method: Any = None - bandpass_filter: Any = None - decomposer: Any = None - recomposer: Any = None +class StepsBlendingState: + precip_cascade: Any = None + mu_extrapolation: Any = None + sigma_extrapolation: Any = None + precip_models_cascade: Any = None generate_perturb: Any = None generate_noise: Any = None - noise_std_coeffs: Optional[np.ndarray] = None - PHI: Optional[np.ndarray] = None - randgen_precip: Optional[List[np.random.RandomState]] = None - velocity_perturbations: Optional[List[Any]] = None + noise_std_coeffs: Any = None + PHI: Any = None + randgen_precip: Any = None + velocity_perturbations: Any = None generate_vel_noise: Any = None - previous_displacement: Optional[np.ndarray] = None - previous_displacement_noise_cascade: Optional[np.ndarray] = None - previous_displacement_prob_matching: Optional[np.ndarray] = None - precip_cascade: Optional[np.ndarray] = None - mu: Optional[np.ndarray] = None - sigma: Optional[np.ndarray] = None - noise_cascade: Optional[np.ndarray] = None - mask_rim: Optional[int] = None + previous_displacement: Any = None + previous_displacement_noise_cascade: Any = None + previous_displacement_prob_matching: Any = None + precip_forecast: Any = None + precip_forecast_non_perturbed: Any = None + mask_rim: Any = None struct: Any = None - fft_objs: Optional[List[Any]] = None - # Additional parameters and variables as needed + fft_objs: Any = None + t_prev_timestep: Any = None + t_leadtime_since_start_forecast: Any = None + # Add more state variables as needed -class BlendingEngine: +class StepsBlendingNowcaster: def __init__( self, precip, @@ -154,134 +177,586 @@ def __init__( velocity, velocity_models, time_steps, + issue_time, steps_blending_config: StepsBlendingConfig, ): - # Store inputs and optional parameters + """Initializes the StepsBlendingNowcaster with inputs and configurations.""" + # Store inputs self.__precip = precip self.__precip_models = precip_models self.__velocity = velocity self.__velocity_models = velocity_models - self.__time_steps = time_steps + self.__timesteps = time_steps + self.__issuetime = issue_time - # Store the config data: self.__config = steps_blending_config - # Store the state and params data: - self.__state = StepsBlendingState() + # Initialize Params and State self.__params = StepsBlendingParams() + self.__state = StepsBlendingState() - def forecast(self): - """Main method to perform the forecast.""" - self._check_inputs() - self._initialize() - self._prepare_data() - self._initialize_noise() - self._estimate_ar_parameters() - self._init_random_generators() - self._prepare_forecast_loop() - self._compute_forecast() - return self.state.forecast_output - - # Private methods for internal processing - def _initialize_methods(self): - """Set up methods for extrapolation, decomposition, etc.""" - pass + # Perform input validation + self.__check_inputs() - def _initialize_bandpass_filter(self): - """Initialize the bandpass filter.""" - pass + # Initialize nowcast components and parameters + self.__initialize_nowcast_components() - def _prepare_data(self): - """Transform data into Lagrangian coordinates and perform initial decomposition.""" - pass + # Additional variables for time measurement + self.__start_time_init = None + self.__init_time = None + self.__mainloop_time = None - def _initialize_noise(self): - """Set up the noise generation mechanism.""" + def compute_forecast(self): pass - def _estimate_ar_parameters(self): - """Estimate autoregressive model parameters.""" - self._estimate_ar_parameters_radar() + def __nowcast_main(self): + self.__check_inputs() + self.__print_forecast_info() + # Measure time for initialization + if self.__config.measure_time: + self.__start_time_init = time.time() + + # Slice the precipitation field to only use the last ar_order + 1 fields + self.__precip = self.__precip[-(self.__config.ar_order + 1) :, :, :].copy() + self.__initialize_nowcast_components() + self.__prepare_radar_and_NWP_fields() + if self.__params.zero_precip_radar and self.__params.zero_precip_model_fields: + self.__zero_precipitation_forecast() + else: + pass + + def __check_inputs(self): + """Validates the inputs and determines if the user provided raw forecasts or decomposed forecasts.""" + # Check dimensions of precip + if self.__precip.ndim != 3: + raise ValueError( + "precip must be a three-dimensional array of shape (ar_order + 1, m, n)" + ) + if self.__precip.shape[0] < self.__config.ar_order + 1: + raise ValueError( + f"precip must have at least {self.__config.ar_order + 1} time steps in the first dimension " + f"to match the autoregressive order (ar_order={self.__config.ar_order})" + ) + + # Check dimensions of velocity + if self.__velocity.ndim != 3: + raise ValueError( + "velocity must be a three-dimensional array of shape (2, m, n)" + ) + if self.__velocity_models.ndim != 5: + raise ValueError( + "velocity_models must be a five-dimensional array of shape (n_models, timestep, 2, m, n)" + ) + if self.__velocity.shape[0] != 2 or self.__velocity_models.shape[2] != 2: + raise ValueError( + "velocity and velocity_models must have an x- and y-component, check the shape" + ) + + # Check that spatial dimensions match between precip and velocity + if self.__precip.shape[1:3] != self.__velocity.shape[1:3]: + raise ValueError( + f"Spatial dimensions of precip and velocity do not match: " + f"{self.__precip.shape[1:3]} vs {self.__velocity.shape[1:3]}" + ) + # Check if the number of members in the precipitation models and velocity models match + if self.__precip_models.shape[0] != self.__velocity_models.shape[0]: + raise ValueError( + "The number of members in the precipitation models and velocity models must match" + ) + + if isinstance(self.__timesteps, list): + self.__params.time_steps_is_list = True + original_timesteps = [0] + list(self.__timesteps) + self.__timesteps = nowcast_utils.binned_timesteps(original_timesteps) + if not sorted(self.__timesteps) == self.__timesteps: + raise ValueError("timesteps is not in ascending order") + if self.__precip_models.shape[1] != math.ceil(self.__timesteps[-1]) + 1: + raise ValueError( + "precip_models does not contain sufficient lead times for this forecast" + ) + else: + self.__params.time_steps_is_list = False + self.__timesteps = list(range(self.__timesteps + 1)) + if self.__precip_models.shape[1] != self.__timesteps + 1: + raise ValueError( + "precip_models does not contain sufficient lead times for this forecast" + ) + + precip_nwp_dim = self.__precip_models.ndim + if precip_nwp_dim == 2: + if isinstance(self.__precip_models[0], dict): + # It's a 2D array of dictionaries with decomposed cascades + self.__params.precip_models_provided_is_cascade = True + else: + raise ValueError( + "When precip_models has ndim == 2, it must contain dictionaries with decomposed cascades." + ) + elif precip_nwp_dim == 4: + self.__params.precip_models_provided_is_cascade = False + else: + raise ValueError( + "precip_models must be either a two-dimensional array containing dictionaries with decomposed model fields" + "or a four-dimensional array containing the original (NWP) model forecasts" + ) + + if self.__config.extrap_kwargs is None: + self.__config.extrap_kwargs = dict() + + if self.__config.filter_kwargs is None: + self.__config.filter_kwargs = dict() + + if self.__config.noise_kwargs is None: + self.__config.noise_kwargs = dict() + + if self.__config.vel_pert_kwargs is None: + self.__config.vel_pert_kwargs = dict() + + if not self.__params.precip_models_provided_is_cascade: + if self.__config.clim_kwargs is None: + # Make sure clim_kwargs at least contains the number of models + self.__config.clim_kwargs = dict( + {"n_models": self.__precip_models.shape[0]} + ) + + if self.__config.mask_kwargs is None: + mask_kwargs = dict() + + if np.any(~np.isfinite(self.__velocity)): + raise ValueError("velocity contains non-finite values") + + if self.__config.mask_method not in ["obs", "incremental", None]: + raise ValueError( + "unknown mask method %s: must be 'obs', 'incremental' or None" + % self.__config.mask_method + ) + + if self.__config.conditional and self.__config.precip_threshold is None: + raise ValueError("conditional=True but precip_thr is not set") + + if ( + self.__config.mask_method is not None + and self.__config.precip_threshold is None + ): + raise ValueError("mask_method!=None but precip_thr=None") + + if self.__config.noise_stddev_adj not in ["auto", "fixed", None]: + raise ValueError( + "unknown noise_std_dev_adj method %s: must be 'auto', 'fixed', or None" + % self.__config.noise_stddev_adj + ) + + if self.__config.kmperpixel is None: + if self.__config.vel_pert_method is not None: + raise ValueError("vel_pert_method is set but kmperpixel=None") + if self.__config.mask_method == "incremental": + raise ValueError("mask_method='incremental' but kmperpixel=None") + + if self.__config.timestep is None: + if self.__config.vel_pert_method is not None: + raise ValueError("vel_pert_method is set but timestep=None") + if self.__config.mask_method == "incremental": + raise ValueError("mask_method='incremental' but timestep=None") + + def __print_forecast_info(self): + print("STEPS blending") + print("==============") + print("") + + print("Inputs") + print("------") + print(f"forecast issue time: {self.__issuetime.isoformat()}") + print( + f"input dimensions: {self.__precip.shape[1]}x{self.__precip.shape[2]}" + ) + if self.__config.kmperpixel is not None: + print(f"km/pixel: {self.__config.kmperpixel}") + if self.__config.timestep is not None: + print(f"time step: {self.__config.timestep} minutes") + print("") + + print("NWP and blending inputs") + print("-----------------------") + print(f"number of (NWP) models: {self.__precip_models.shape[0]}") + print(f"blend (NWP) model members: {self.__config.blend_nwp_members}") + print( + f"decompose (NWP) models: {'yes' if self.__precip_models.ndim == 4 else 'no'}" + ) + print("") + + print("Methods") + print("-------") + print(f"extrapolation: {self.__config.extrapolation_method}") + print(f"bandpass filter: {self.__config.bandpass_filter_method}") + print(f"decomposition: {self.__config.decomposition_method}") + print(f"noise generator: {self.__config.noise_method}") + print( + f"noise adjustment: {'yes' if self.__config.noise_stddev_adj else 'no'}" + ) + print(f"velocity perturbator: {self.__config.vel_pert_method}") + print(f"blending weights method: {self.__config.weights_method}") + print( + f"conditional statistics: {'yes' if self.__config.conditional else 'no'}" + ) + print(f"precip. mask method: {self.__config.mask_method}") + print(f"probability matching: {self.__config.probmatching_method}") + print(f"FFT method: {self.__config.fft_method}") + print(f"domain: {self.__config.domain}") + print("") + + print("Parameters") + print("----------") + if isinstance(self.__timesteps, int): + print(f"number of time steps: {self.__timesteps}") + else: + print(f"time steps: {self.__timesteps}") + print(f"ensemble size: {self.__config.n_ens_members}") + print(f"parallel threads: {self.__config.num_workers}") + print(f"number of cascade levels: {self.__config.n_cascade_levels}") + print(f"order of the AR(p) model: {self.__config.ar_order}") + if self.__config.vel_pert_method == "bps": + vp_par = self.__config.vel_pert_kwargs.get( + "p_par", noise.motion.get_default_params_bps_par() + ) + vp_perp = self.__config.vel_pert_kwargs.get( + "p_perp", noise.motion.get_default_params_bps_perp() + ) + print(f"vel. pert., parallel: {vp_par[0]},{vp_par[1]},{vp_par[2]}") + print( + f"vel. pert., perpendicular: {vp_perp[0]},{vp_perp[1]},{vp_perp[2]}" + ) + else: + vp_par, vp_perp = None, None + + if self.__config.conditional or self.__config.mask_method is not None: + print(f"precip. intensity threshold: {self.__config.precip_threshold}") + print(f"no-rain fraction threshold for radar: {self.__config.norain_threshold}") + print("") + + def __initialize_nowcast_components(self): + """ + Initialize the FFT, bandpass filters, decomposition methods, and extrapolation method. + """ + # Initialize number of ensemble workers + self.__params.num_ensemble_workers = min( + self.__config.n_ens_members, self.__config.num_workers + ) + + M, N = self.__precip.shape[1:] # Extract the spatial dimensions (height, width) + + # Initialize FFT method + self.__params.fft = utils.get_method( + self.__config.fft_method, shape=(M, N), n_threads=self.__config.num_workers + ) + + # Initialize the band-pass filter for the cascade decomposition + filter_method = cascade.get_method(self.__config.bandpass_filter_method) + self.__params.bandpass_filter = filter_method( + (M, N), + self.__config.n_cascade_levels, + **(self.__config.filter_kwargs or {}), + ) - def _init_random_generators(self): - """Initialize random number generators.""" - self._init_random_generators_for_noise() - if self.config.vel_pert_method: - self._init_velocity_perturbations() + # Get the decomposition method (e.g., FFT) + ( + self.__params.decomposition_method, + self.__params.recomposition_method, + ) = cascade.get_method(self.__config.decomposition_method) + + # Get the extrapolation method (e.g., semilagrangian) + self.__params.extrapolation_method = extrapolation.get_method( + self.__config.extrapolation_method + ) - def _prepare_forecast_loop(self): - """Set up variables and structures for the forecasting loop.""" - self._initialize_forecast_variables() + # Generate the mesh grid for spatial coordinates + x_values, y_values = np.meshgrid(np.arange(N), np.arange(M)) + self.__params.xy_coordinates = np.stack([x_values, y_values]) + + precip_copy = self.__precip[-(self.__config.ar_order + 1) :, :, :].copy() + # Determine the domain mask from non-finite values in the precipitation data + self.__params.domain_mask = np.logical_or.reduce( + [~np.isfinite(precip_copy[i, :]) for i in range(precip_copy.shape[0])] + ) + + print("Blended nowcast components initialized successfully.") + + def __prepare_radar_and_NWP_fields(self): + # determine the precipitation threshold mask + if self.__config.conditional: + self.__params.mask_threshold = np.logical_and.reduce( + [ + self.__precip[i, :, :] >= self.__config.precip_threshold + for i in range(self.__precip.shape[0]) + ] + ) + else: + self.__params.mask_threshold = None + + # we need to know the zerovalue of precip to replace the mask when decomposing after + # extrapolation + self.__params.precip_zerovalue = np.nanmin(self.__precip) + + # 1. Start with the radar rainfall fields. We want the fields in a + # Lagrangian space + self.__precip = _transform_to_lagrangian( + self.__precip, + self.__velocity, + self.__config.ar_order, + self.__params.xy_coordinates, + self.__params.extrapolation_method, + self.__config.extrap_kwargs, + self.__config.num_workers, + ) + + # 2. Perform the cascade decomposition for the input precip fields and + # and, if necessary, for the (NWP) model fields + # 2.1 Compute the cascade decompositions of the input precipitation fields + """Compute the cascade decompositions of the input precipitation fields.""" + precip_forecast_decomp = [] + for i in range(self.__config.ar_order + 1): + precip_forecast = self.__params.extrapolation_method( + self.__precip[i, :, :], + self.__params.bandpass_filter, + mask=self.__params.mask_threshold, + fft_method=self.__params.fft, + output_domain=self.__config.domain, + normalize=True, + compute_stats=True, + compact_output=True, + ) + precip_forecast_decomp.append(precip_forecast) + + # Rearrange the cascaded into a four-dimensional array of shape + # (n_cascade_levels,ar_order+1,m,n) for the autoregressive model + self.__state.precip_cascade = nowcast_utils.stack_cascades( + precip_forecast_decomp, self.__config.n_cascade_levels + ) - def _compute_forecast(self): - """Main loop to compute the forecast over the specified time steps.""" - self._run_forecast_loop() + precip_forecast_decomp = precip_forecast_decomp[-1] + self.__state.mu_extrapolation = np.array(precip_forecast_decomp["means"]) + self.__state.sigma_extrapolation = np.array(precip_forecast_decomp["stds"]) - # Methods for specific functionalities - def _check_inputs(self): - """Validate input data and configurations.""" - # Implement input checks as needed + # 2.2 If necessary, recompose (NWP) model forecasts + self.__state.precip_models_cascade = None + + if self.__precip_models.ndim != 4: + self.__state.precip_models_cascade = self.__precip_models + self.__precip_models = _compute_cascade_recomposition_nwp( + self.__precip_models, self.__params.recomposition_method + ) + + # 2.3 Check for zero input fields in the radar and NWP data. + self.__params.zero_precip_radar = blending.utils.check_norain( + self.__precip, + self.__config.precip_threshold, + self.__config.norain_threshold, + ) + # The norain fraction threshold used for nwp is the default value of 0.0, + # since nwp does not suffer from clutter. + self.__params.zero_precip_model_fields = blending.utils.check_norain( + self.__precip_models, + self.__config.precip_threshold, + self.__config.norain_threshold, + ) + + def __zero_precipitation_forecast(self): + print( + "No precipitation above the threshold found in both the radar and NWP fields" + ) + print("The resulting forecast will contain only zeros") + # Create the output list + precip_forecast = [[] for j in range(self.__config.n_ens_members)] + + # Save per time step to ensure the array does not become too large if + # no return_output is requested and callback is not None. + for t, subtimestep_idx in enumerate(self.__timesteps): + # If the timestep is not the first one, we need to provide the zero forecast + if t > 0: + # Create an empty np array with shape [n_ens_members, rows, cols] + # and fill it with the minimum value from precip (corresponding to + # zero precipitation) + N, M = self.__precip.shape + precip_forecast_workers = np.full( + (self.__config.n_ens_members, N, M), self.__params.precip_zerovalue + ) + if subtimestep_idx: + if self.__config.callback is not None: + if precip_forecast_workers.shape[1] > 0: + self.__config.callback(precip_forecast_workers.squeeze()) + if self.__config.return_output: + for j in range(self.__config.n_ens_members): + precip_forecast[j].append(precip_forecast_workers[j]) + + precip_forecast_workers = None + + if self.__config.measure_time: + zero_precip_time = time.time() - self.__start_time_init + + if self.__config.return_output: + precip_forecast_all_members_all_times = np.stack( + [ + np.stack(precip_forecast[j]) + for j in range(self.__config.n_ens_members) + ] + ) + if self.__config.measure_time: + return ( + precip_forecast_all_members_all_times, + zero_precip_time, + zero_precip_time, + ) + else: + return precip_forecast_all_members_all_times + else: + return None + + def __perform_extrapolation(self): + pass + + def __apply_noise_and_ar_model(self): + pass + + def __initialize_velocity_perturbations(self): pass - def _initialize(self): - """Perform any additional initialization steps.""" - # Initialize variables in self.params as needed + def __initialize_precipitation_mask(self): pass - def _transform_to_lagrangian(self): - """Transform precipitation data to Lagrangian coordinates.""" - # Use self.state and self.params as needed + def __initialize_fft_objects(self): pass - def _compute_cascade_decomposition_radar(self): - """Compute the cascade decomposition for radar data.""" - # Update self.params with computed cascades + def __return_state_dict(self): pass - def _estimate_ar_parameters_radar(self): - """Estimate AR parameters for radar data.""" - # Update self.params.PHI + def __return_params_dict(self): pass - def _init_random_generators_for_noise(self): - """Initialize random generators for noise.""" - # Update self.params.randgen_precip + def __update_state(self, state, params): pass - def _init_velocity_perturbations(self): - """Initialize velocity perturbations if required.""" - # Update self.params.velocity_perturbations, generate_vel_noise + def __update_deterministic_ar_model(self, state, params): pass - def _initialize_forecast_variables(self): - """Set up variables needed for the forecast loop.""" - # Initialize variables in self.params + def __apply_ar_model_to_cascades(self, j, state, params): pass - def _run_forecast_loop(self): - """Run the main forecast loop over time steps.""" - # Use self.state and self.params as needed + def __generate_and_decompose_noise(self, j, state, params): pass - # Additional helper methods as needed + def __recompose_and_apply_mask(self, j, state, params): + pass + + def __apply_precipitation_mask(self, precip_forecast, j, state, params): + pass + + def __measure_time(self, label, start_time): + """ + Measure and print the time taken for a specific part of the process. + + Parameters: + - label: A description of the part of the process being measured. + - start_time: The timestamp when the process started (from time.time()). + """ + if self.__config.measure_time: + elapsed_time = time.time() - start_time + print(f"{label} took {elapsed_time:.2f} seconds.") + + def reset_states_and_params(self): + """ + Reset the internal state and parameters of the nowcaster to allow multiple forecasts. + This method resets the state and params to their initial conditions without reinitializing + the inputs like precip, velocity, time_steps, or config. + """ + # Re-initialize the state and parameters + self.__state = StepsBlendingState() + self.__params = StepsBlendingParams() + + # Reset time measurement variables + self.__start_time_init = None + self.__init_time = None + self.__mainloop_time = None + + +def calculate_ratios(correlations): + """Calculate explained variance ratios from correlation. + + Parameters + ---------- + Array of shape [component, scale_level, ...] + containing correlation (skills) for each component (NWP and nowcast), + scale level, and optionally along [y, x] dimensions. + + Returns + ------- + out : numpy array + An array containing the ratios of explain variance for each + component, scale level, ... + """ + # correlations: [component, scale, ...] + square_corrs = np.square(correlations) + # Calculate the ratio of the explained variance to the unexplained + # variance of the nowcast and NWP model components + out = square_corrs / (1 - square_corrs) + # out: [component, scale, ...] + return out -class SkillScoreManager: - def __init__(self, config: StepsBlendingConfig): - self.outdir_path_skill = config.outdir_path_skill - self.clim_kwargs = config.clim_kwargs +def calculate_weights_bps(correlations): + """Calculate BPS blending weights for STEPS blending from correlation. - def compute_initial_skill(self, observed_cascades, model_cascades, domain_mask): - """Calculate the initial skill of NWP models at t=0.""" - pass # Implement as needed + Parameters + ---------- + correlations : array-like + Array of shape [component, scale_level, ...] + containing correlation (skills) for each component (NWP and nowcast), + scale level, and optionally along [y, x] dimensions. - def update_skill(self, lead_time, correlations, model_indices): - """Update the skill scores based on lead time.""" - pass # Implement as needed + Returns + ------- + weights : array-like + Array of shape [component+1, scale_level, ...] + containing the weights to be used in STEPS blending for + each original component plus an addtional noise component, scale level, + and optionally along [y, x] dimensions. - def save_skill(self, current_skill, validtime): - """Save the skill scores to disk.""" - pass # Implement as needed + References + ---------- + :cite:`BPS2006` + + Notes + ----- + The weights in the BPS method can sum op to more than 1.0. + """ + # correlations: [component, scale, ...] + # Check if the correlations are positive, otherwise rho = 10e-5 + correlations = np.where(correlations < 10e-5, 10e-5, correlations) + + # If we merge more than one component with the noise cascade, we follow + # the weights impolementation in either :cite:`BPS2006` or :cite:`SPN2013`. + if correlations.shape[0] > 1: + # Calculate weights for each source + ratios = calculate_ratios(correlations) + # ratios: [component, scale, ...] + total_ratios = np.sum(ratios, axis=0) + # total_ratios: [scale, ...] - the denominator of eq. 11 & 12 in BPS2006 + weights = correlations * np.sqrt(ratios / total_ratios) + # weights: [component, scale, ...] + # Calculate the weight of the noise component. + # Original BPS2006 method in the following two lines (eq. 13) + total_square_weights = np.sum(np.square(weights), axis=0) + noise_weight = np.sqrt(1.0 - total_square_weights) + # Finally, add the noise_weights to the weights variable. + weights = np.concatenate((weights, noise_weight[None, ...]), axis=0) + + # Otherwise, the weight equals the correlation on that scale level and + # the noise component weight equals 1 - this weight. This only occurs for + # the weights calculation outside the radar domain where in the case of 1 + # NWP model or ensemble member, no blending of multiple models has to take + # place + else: + noise_weight = 1.0 - correlations + weights = np.concatenate((correlations, noise_weight), axis=0) + + return weights def forecast( @@ -380,7 +855,6 @@ def forecast( Time step of the motion vectors (minutes). Required if vel_pert_method is not None or mask_method is 'incremental'. issuetime: datetime - Datetime object containing the date and time for which the forecast is issued. n_ens_members: int The number of ensemble members to generate. This number should always be @@ -958,18 +1432,20 @@ def forecast( ) # 6. Initialize all the random generators and prepare for the forecast loop - randgen_precip, velocity_perturbations, generate_vel_noise = ( - _init_random_generators( - velocity, - noise_method, - vel_pert_method, - vp_par, - vp_perp, - seed, - n_ens_members, - kmperpixel, - timestep, - ) + ( + randgen_precip, + velocity_perturbations, + generate_vel_noise, + ) = _init_random_generators( + velocity, + noise_method, + vel_pert_method, + vp_par, + vp_perp, + seed, + n_ens_members, + kmperpixel, + timestep, ) ( previous_displacement, @@ -2103,89 +2579,6 @@ def worker(j): return None -def calculate_ratios(correlations): - """Calculate explained variance ratios from correlation. - - Parameters - ---------- - Array of shape [component, scale_level, ...] - containing correlation (skills) for each component (NWP and nowcast), - scale level, and optionally along [y, x] dimensions. - - Returns - ------- - out : numpy array - An array containing the ratios of explain variance for each - component, scale level, ... - """ - # correlations: [component, scale, ...] - square_corrs = np.square(correlations) - # Calculate the ratio of the explained variance to the unexplained - # variance of the nowcast and NWP model components - out = square_corrs / (1 - square_corrs) - # out: [component, scale, ...] - return out - - -def calculate_weights_bps(correlations): - """Calculate BPS blending weights for STEPS blending from correlation. - - Parameters - ---------- - correlations : array-like - Array of shape [component, scale_level, ...] - containing correlation (skills) for each component (NWP and nowcast), - scale level, and optionally along [y, x] dimensions. - - Returns - ------- - weights : array-like - Array of shape [component+1, scale_level, ...] - containing the weights to be used in STEPS blending for - each original component plus an addtional noise component, scale level, - and optionally along [y, x] dimensions. - - References - ---------- - :cite:`BPS2006` - - Notes - ----- - The weights in the BPS method can sum op to more than 1.0. - """ - # correlations: [component, scale, ...] - # Check if the correlations are positive, otherwise rho = 10e-5 - correlations = np.where(correlations < 10e-5, 10e-5, correlations) - - # If we merge more than one component with the noise cascade, we follow - # the weights impolementation in either :cite:`BPS2006` or :cite:`SPN2013`. - if correlations.shape[0] > 1: - # Calculate weights for each source - ratios = calculate_ratios(correlations) - # ratios: [component, scale, ...] - total_ratios = np.sum(ratios, axis=0) - # total_ratios: [scale, ...] - the denominator of eq. 11 & 12 in BPS2006 - weights = correlations * np.sqrt(ratios / total_ratios) - # weights: [component, scale, ...] - # Calculate the weight of the noise component. - # Original BPS2006 method in the following two lines (eq. 13) - total_square_weights = np.sum(np.square(weights), axis=0) - noise_weight = np.sqrt(1.0 - total_square_weights) - # Finally, add the noise_weights to the weights variable. - weights = np.concatenate((weights, noise_weight[None, ...]), axis=0) - - # Otherwise, the weight equals the correlation on that scale level and - # the noise component weight equals 1 - this weight. This only occurs for - # the weights calculation outside the radar domain where in the case of 1 - # NWP model or ensemble member, no blending of multiple models has to take - # place - else: - noise_weight = 1.0 - correlations - weights = np.concatenate((correlations, noise_weight), axis=0) - - return weights - - def calculate_weights_spn(correlations, covariance): """Calculate SPN blending weights for STEPS blending from correlation. diff --git a/pysteps/nowcasts/steps.py b/pysteps/nowcasts/steps.py index 0f0b27b2..b04e8275 100644 --- a/pysteps/nowcasts/steps.py +++ b/pysteps/nowcasts/steps.py @@ -341,9 +341,9 @@ def compute_forecast(self): if self.__config.measure_time: self.__start_time_init = time.time() - self.__initialize_nowcast_components() # Slice the precipitation field to only use the last ar_order + 1 fields self.__precip = self.__precip[-(self.__config.ar_order + 1) :, :, :].copy() + self.__initialize_nowcast_components() self.__perform_extrapolation() self.__apply_noise_and_ar_model() @@ -358,9 +358,10 @@ def compute_forecast(self): self.__nowcast_main() if self.__config.measure_time: - self.__state.precip_forecast, self.__mainloop_time = ( - self.__state.precip_forecast - ) + ( + self.__state.precip_forecast, + self.__mainloop_time, + ) = self.__state.precip_forecast # Stack and return the forecast output if self.__config.return_output: @@ -392,8 +393,8 @@ def __nowcast_main(self): ] # Extract the last available precipitation field # Prepare state and params dictionaries, these need to be formatted a specific way for the nowcast_main_loop - state = self.__initialize_state() - params = self.__initialize_params(precip) + state = self.__return_state_dict() + params = self.__return_params_dict(precip) print("Starting nowcast computation.") @@ -589,9 +590,10 @@ def __initialize_nowcast_components(self): ) # Get the decomposition method (e.g., FFT) - self.__params.decomposition_method, self.__params.recomposition_method = ( - cascade.get_method(self.__config.decomposition_method) - ) + ( + self.__params.decomposition_method, + self.__params.recomposition_method, + ) = cascade.get_method(self.__config.decomposition_method) # Get the extrapolation method (e.g., semilagrangian) self.__params.extrapolation_method = extrapolation.get_method( @@ -957,7 +959,7 @@ def __initialize_fft_objects(self): self.__state.fft_objs.append(fft_obj) print("FFT objects initialized successfully.") - def __initialize_state(self): + def __return_state_dict(self): """ Initialize the state dictionary used during the nowcast iteration. """ @@ -971,7 +973,7 @@ def __initialize_state(self): "randgen_prec": self.__state.random_generator_precip, } - def __initialize_params(self, precip): + def __return_params_dict(self, precip): """ Initialize the params dictionary used during the nowcast iteration. """ From 1eede39494d9d8ca6f265bc3c564ae3123cadbd2 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 28 Nov 2024 17:58:44 +0100 Subject: [PATCH 06/33] Added code to estimation of ar parameters of radar --- pysteps/blending/steps.py | 168 +++++++++++++++++++++++++++++++++++++- 1 file changed, 165 insertions(+), 3 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index be3eee06..302a2ea9 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -141,6 +141,7 @@ class StepsBlendingParams: mask_threshold: Any = None zero_precip_radar: bool = False zero_precip_model_fields: bool = False + PHI: Any = None @dataclass @@ -166,6 +167,7 @@ class StepsBlendingState: fft_objs: Any = None t_prev_timestep: Any = None t_leadtime_since_start_forecast: Any = None + precip_noise_input: Any = None # Add more state variables as needed @@ -220,10 +222,17 @@ def __nowcast_main(self): self.__precip = self.__precip[-(self.__config.ar_order + 1) :, :, :].copy() self.__initialize_nowcast_components() self.__prepare_radar_and_NWP_fields() + + # Determine if rain is present in both radar and NWP fields if self.__params.zero_precip_radar and self.__params.zero_precip_model_fields: self.__zero_precipitation_forecast() else: - pass + # Prepare the data for the zero precipitation radar case and initialize the noise correctly + if self.__params.zero_precip_radar: + self.__prepare_nowcast_for_zero_radar() + else: + self.__state.precip_noise_input = self.__precip.copy() + self.__estimate_ar_parameters_radar() def __check_inputs(self): """Validates the inputs and determines if the user provided raw forecasts or decomposed forecasts.""" @@ -482,6 +491,7 @@ def __initialize_nowcast_components(self): def __prepare_radar_and_NWP_fields(self): # determine the precipitation threshold mask if self.__config.conditional: + # TODO: is this logical_and correct here? Now only those places where precip is in all images is saved? self.__params.mask_threshold = np.logical_and.reduce( [ self.__precip[i, :, :] >= self.__config.precip_threshold @@ -507,8 +517,8 @@ def __prepare_radar_and_NWP_fields(self): self.__config.num_workers, ) - # 2. Perform the cascade decomposition for the input precip fields and - # and, if necessary, for the (NWP) model fields + # 2. Perform the cascade decomposition for the input precip fields and, + # if necessary, for the (NWP) model fields # 2.1 Compute the cascade decompositions of the input precipitation fields """Compute the cascade decompositions of the input precipitation fields.""" precip_forecast_decomp = [] @@ -609,6 +619,158 @@ def __zero_precipitation_forecast(self): else: return None + def __prepare_nowcast_for_zero_radar(self): + # 2.3.3 If zero_precip_radar, make sure that precip_cascade does not contain + # only nans or infs. If so, fill it with the zero value. + + # Look for a timestep and member with rain so that we have a sensible decomposition + done = False + for t in self.__timesteps: + if done: + break + for j in range(self.__precip_models.shape[0]): + if not blending.utils.check_norain( + self.__precip_models[j, t], + self.__config.precip_threshold, + self.__config.norain_threshold, + ): + if self.__state.precip_models_cascade is not None: + self.__state.precip_cascade[ + ~np.isfinite(self.__state.precip_cascade) + ] = np.nanmin( + self.__state.precip_models_cascade[j, t]["cascade_levels"] + ) + continue + precip_models_cascade_temp = self.__params.decomposition_method( + self.__precip_models[j, t, :, :], + bp_filter=self.__params.bandpass_filter, + fft_method=self.__params.fft, + output_domain=self.__config.domain, + normalize=True, + compute_stats=True, + compact_output=True, + )["cascade_levels"] + self.__state.precip_cascade[ + ~np.isfinite(self.__state.precip_cascade) + ] = np.nanmin(precip_models_cascade_temp) + done = True + break + + # 2.3.5 If zero_precip_radar is True, only use the velocity field of the NWP + # forecast. I.e., velocity (radar) equals velocity_model at the first time + # step. + # Use the velocity from velocity_models at time step 0 + self.__velocity = self.__velocity_models[:, 0, :, :, :].astype( + np.float64, copy=False + ) + # Take the average over the first axis, which corresponds to n_models + # (hence, the model average) + self.__velocity = np.mean(self.__velocity, axis=0) + + # 3. Initialize the noise method. + # If zero_precip_radar is True, initialize noise based on the NWP field time + # step where the fraction of rainy cells is highest (because other lead times + # might be zero as well). Else, initialize the noise with the radar + # rainfall data + """Initialize noise based on the NWP field time step where the fraction of rainy cells is highest""" + if self.__config.precip_threshold is None: + self.__config.precip_threshold = np.nanmin(self.__precip_models) + + max_rain_pixels = -1 + max_rain_pixels_j = -1 + max_rain_pixels_t = -1 + for j in range(self.__precip_models.shape[0]): + for t in self.__timesteps: + rain_pixels = self.__precip_models[j][t][ + self.__precip_models[j][t] > self.__config.precip_threshold + ].size + if rain_pixels > max_rain_pixels: + max_rain_pixels = rain_pixels + max_rain_pixels_j = j + max_rain_pixels_t = t + self.__state.precip_noise_input = self.__precip_models[max_rain_pixels_j][ + max_rain_pixels_t + ] + + # Make sure precip_noise_input is three-dimensional + if len(self.__state.precip_noise_input.shape) != 3: + self.__state.precip_noise_input = self.__state.precip_noise_input[ + np.newaxis, :, : + ] + + def __estimate_ar_parameters_radar(self): + # 4. Estimate AR parameters for the radar rainfall field + """Estimate AR parameters for the radar rainfall field.""" + # If there are values in the radar fields, compute the auto-correlations + GAMMA = np.empty((self.__config.n_cascade_levels, self.__config.ar_order)) + if not self.__params.zero_precip_radar: + # compute lag-l temporal auto-correlation coefficients for each cascade level + for i in range(self.__config.n_cascade_levels): + GAMMA[i, :] = correlation.temporal_autocorrelation( + self.__state.precip_cascade[i], mask=self.__params.mask_threshold + ) + + # Else, use standard values for the auto-correlations + else: + # Get the climatological lag-1 and lag-2 auto-correlation values from Table 2 + # in `BPS2004`. + # Hard coded, change to own (climatological) values when present. + # TODO: add user warning here so users can be aware of this without reading the code? + GAMMA = np.array( + [ + [0.99805, 0.9925, 0.9776, 0.9297, 0.796, 0.482, 0.079, 0.0006], + [0.9933, 0.9752, 0.923, 0.750, 0.367, 0.069, 0.0018, 0.0014], + ] + ) + + # Check whether the number of cascade_levels is correct + if GAMMA.shape[1] > self.__config.n_cascade_levels: + GAMMA = GAMMA[:, 0 : self.__config.n_cascade_levels] + elif GAMMA.shape[1] < self.__config.n_cascade_levels: + # Get the number of cascade levels that is missing + n_extra_lev = self.__config.n_cascade_levels - GAMMA.shape[1] + # Append the array with correlation values of 10e-4 + GAMMA = np.append( + GAMMA, + [np.repeat(0.0006, n_extra_lev), np.repeat(0.0014, n_extra_lev)], + axis=1, + ) + + # Finally base GAMMA.shape[0] on the AR-level + if self.__config.ar_order == 1: + GAMMA = GAMMA[0, :] + if self.__config.ar_order > 2: + for repeat_index in range(self.__config.ar_order - 2): + GAMMA = np.vstack((GAMMA, GAMMA[1, :])) + + # Finally, transpose GAMMA to ensure that the shape is the same as np.empty((n_cascade_levels, ar_order)) + GAMMA = GAMMA.transpose() + assert GAMMA.shape == ( + self.__config.n_cascade_levels, + self.__config.ar_order, + ) + + # Print the GAMMA value + nowcast_utils.print_corrcoefs(GAMMA) + + if self.__config.ar_order == 2: + # adjust the lag-2 correlation coefficient to ensure that the AR(p) + # process is stationary + for i in range(self.__config.n_cascade_levels): + GAMMA[i, 1] = autoregression.adjust_lag2_corrcoef2( + GAMMA[i, 0], GAMMA[i, 1] + ) + + # estimate the parameters of the AR(p) model from the auto-correlation + # coefficients + self.__params.PHI = np.empty( + (self.__config.n_cascade_levels, self.__config.ar_order + 1) + ) + for i in range(self.__config.n_cascade_levels): + self.__params.PHI[i, :] = autoregression.estimate_ar_params_yw(GAMMA[i, :]) + + nowcast_utils.print_ar_params(self.__params.PHI) + def __perform_extrapolation(self): pass From a18f1f62d374b65838efd30268107fcb4c4a58d7 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Fri, 29 Nov 2024 16:23:20 +0100 Subject: [PATCH 07/33] Next go, start with forecast loop #7 --- pysteps/blending/steps.py | 286 +++++++++++++++++++++++++++++++++----- 1 file changed, 252 insertions(+), 34 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 302a2ea9..6209bb6b 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -84,7 +84,7 @@ class StepsBlendingConfig: noise_method: Optional[str] noise_stddev_adj: Optional[str] ar_order: int - vel_pert_method: Optional[str] + velocity_perturbation_method: Optional[str] weights_method: str conditional: bool probmatching_method: Optional[str] @@ -115,14 +115,16 @@ class StepsBlendingParams: sigma_extrapolation: np.ndarray # Std devs of extrapolated cascades bandpass_filter: Any # Band-pass filter object fft: Any # FFT method object - generate_perturb: Callable # Perturbation generator - generate_noise: Callable # Noise generator + perturbation_generator: Callable # Perturbation generator + noise_generator: Callable # Noise generator generate_vel_noise: Optional[Callable] # Velocity noise generator extrapolation_method: Any = None decomposition_method: Any = None recomposition_method: Any = None - vp_par: Optional[np.ndarray] = None # Velocity perturbation parameters (parallel) - vp_perp: Optional[np.ndarray] = ( + velocity_perturbations_parallel: Optional[np.ndarray] = ( + None # Velocity perturbation parameters (parallel) + ) + velocity_perturbations_perpendicular: Optional[np.ndarray] = ( None # Velocity perturbation parameters (perpendicular) ) fft_objs: List[Any] = field( @@ -146,17 +148,14 @@ class StepsBlendingParams: @dataclass class StepsBlendingState: - precip_cascade: Any = None + precip_cascades: Any = None mu_extrapolation: Any = None sigma_extrapolation: Any = None - precip_models_cascade: Any = None - generate_perturb: Any = None - generate_noise: Any = None - noise_std_coeffs: Any = None + precip_models_cascades: Any = None PHI: Any = None randgen_precip: Any = None velocity_perturbations: Any = None - generate_vel_noise: Any = None + generate_velocity_noise: Any = None previous_displacement: Any = None previous_displacement_noise_cascade: Any = None previous_displacement_prob_matching: Any = None @@ -168,6 +167,9 @@ class StepsBlendingState: t_prev_timestep: Any = None t_leadtime_since_start_forecast: Any = None precip_noise_input: Any = None + precip_noise_cascade: Any = None + precip_mean_noise: Any = None + precip_std_noise: Any = None # Add more state variables as needed @@ -232,7 +234,12 @@ def __nowcast_main(self): self.__prepare_nowcast_for_zero_radar() else: self.__state.precip_noise_input = self.__precip.copy() + self.__initialize_noise() self.__estimate_ar_parameters_radar() + self.__multiply_precip_cascade_to_match_ensemble_members() + self.__initialize_random_generators() + self.__prepare_forecast_loop() + self.__initialize_noise_cascade() def __check_inputs(self): """Validates the inputs and determines if the user provided raw forecasts or decomposed forecasts.""" @@ -355,14 +362,18 @@ def __check_inputs(self): ) if self.__config.kmperpixel is None: - if self.__config.vel_pert_method is not None: - raise ValueError("vel_pert_method is set but kmperpixel=None") + if self.__config.velocity_perturbation_method is not None: + raise ValueError( + "velocity_perturbation_method is set but kmperpixel=None" + ) if self.__config.mask_method == "incremental": raise ValueError("mask_method='incremental' but kmperpixel=None") if self.__config.timestep is None: - if self.__config.vel_pert_method is not None: - raise ValueError("vel_pert_method is set but timestep=None") + if self.__config.velocity_perturbation_method is not None: + raise ValueError( + "velocity_perturbation_method is set but timestep=None" + ) if self.__config.mask_method == "incremental": raise ValueError("mask_method='incremental' but timestep=None") @@ -401,7 +412,9 @@ def __print_forecast_info(self): print( f"noise adjustment: {'yes' if self.__config.noise_stddev_adj else 'no'}" ) - print(f"velocity perturbator: {self.__config.vel_pert_method}") + print( + f"velocity perturbator: {self.__config.velocity_perturbation_method}" + ) print(f"blending weights method: {self.__config.weights_method}") print( f"conditional statistics: {'yes' if self.__config.conditional else 'no'}" @@ -422,19 +435,28 @@ def __print_forecast_info(self): print(f"parallel threads: {self.__config.num_workers}") print(f"number of cascade levels: {self.__config.n_cascade_levels}") print(f"order of the AR(p) model: {self.__config.ar_order}") - if self.__config.vel_pert_method == "bps": - vp_par = self.__config.vel_pert_kwargs.get( - "p_par", noise.motion.get_default_params_bps_par() + if self.__config.velocity_perturbation_method == "bps": + self.__params.velocity_perturbations_parallel = ( + self.__config.vel_pert_kwargs.get( + "p_par", noise.motion.get_default_params_bps_par() + ) ) - vp_perp = self.__config.vel_pert_kwargs.get( - "p_perp", noise.motion.get_default_params_bps_perp() + self.__params.velocity_perturbations_perpendicular = ( + self.__config.vel_pert_kwargs.get( + "p_perp", noise.motion.get_default_params_bps_perp() + ) + ) + print( + f"vel. pert., parallel: {self.__params.velocity_perturbations_parallel[0]},{self.__params.velocity_perturbations_parallel[1]},{self.__params.velocity_perturbations_parallel[2]}" ) - print(f"vel. pert., parallel: {vp_par[0]},{vp_par[1]},{vp_par[2]}") print( - f"vel. pert., perpendicular: {vp_perp[0]},{vp_perp[1]},{vp_perp[2]}" + f"vel. pert., perpendicular: {self.__params.velocity_perturbations_perpendicular[0]},{self.__params.velocity_perturbations_perpendicular[1]},{self.__params.velocity_perturbations_perpendicular[2]}" ) else: - vp_par, vp_perp = None, None + ( + self.__params.velocity_perturbations_parallel, + self.__params.velocity_perturbations_perpendicular, + ) = (None, None) if self.__config.conditional or self.__config.mask_method is not None: print(f"precip. intensity threshold: {self.__config.precip_threshold}") @@ -537,7 +559,7 @@ def __prepare_radar_and_NWP_fields(self): # Rearrange the cascaded into a four-dimensional array of shape # (n_cascade_levels,ar_order+1,m,n) for the autoregressive model - self.__state.precip_cascade = nowcast_utils.stack_cascades( + self.__state.precip_cascades = nowcast_utils.stack_cascades( precip_forecast_decomp, self.__config.n_cascade_levels ) @@ -546,10 +568,10 @@ def __prepare_radar_and_NWP_fields(self): self.__state.sigma_extrapolation = np.array(precip_forecast_decomp["stds"]) # 2.2 If necessary, recompose (NWP) model forecasts - self.__state.precip_models_cascade = None + self.__state.precip_models_cascades = None if self.__precip_models.ndim != 4: - self.__state.precip_models_cascade = self.__precip_models + self.__state.precip_models_cascades = self.__precip_models self.__precip_models = _compute_cascade_recomposition_nwp( self.__precip_models, self.__params.recomposition_method ) @@ -634,11 +656,11 @@ def __prepare_nowcast_for_zero_radar(self): self.__config.precip_threshold, self.__config.norain_threshold, ): - if self.__state.precip_models_cascade is not None: - self.__state.precip_cascade[ - ~np.isfinite(self.__state.precip_cascade) + if self.__state.precip_models_cascades is not None: + self.__state.precip_cascades[ + ~np.isfinite(self.__state.precip_cascades) ] = np.nanmin( - self.__state.precip_models_cascade[j, t]["cascade_levels"] + self.__state.precip_models_cascades[j, t]["cascade_levels"] ) continue precip_models_cascade_temp = self.__params.decomposition_method( @@ -650,8 +672,8 @@ def __prepare_nowcast_for_zero_radar(self): compute_stats=True, compact_output=True, )["cascade_levels"] - self.__state.precip_cascade[ - ~np.isfinite(self.__state.precip_cascade) + self.__state.precip_cascades[ + ~np.isfinite(self.__state.precip_cascades) ] = np.nanmin(precip_models_cascade_temp) done = True break @@ -698,6 +720,61 @@ def __prepare_nowcast_for_zero_radar(self): np.newaxis, :, : ] + def __initialize_noise(self): + """Initialize the noise method.""" + if self.__config.noise_method is not None: + # get methods for perturbations + init_noise, self.__params.noise_generator = noise.get_method( + self.__config.noise_method + ) + + # initialize the perturbation generator for the precipitation field + self.__params.perturbation_generator = init_noise( + self.__precip, + fft_method=self.__params.fft, + **self.__config.noise_kwargs, + ) + + if self.__config.noise_stddev_adj == "auto": + print("Computing noise adjustment coefficients... ", end="", flush=True) + if self.__config.measure_time: + starttime = time.time() + + precip_forecast_min = np.min(self.__precip) + self.__params.noise_std_coeffs = noise.utils.compute_noise_stddev_adjs( + self.__precip[-1, :, :], + self.__config.precip_threshold, + precip_forecast_min, + self.__params.bandpass_filter, + self.__params.decomposition_method, + self.__params.perturbation_generator, + self.__params.noise_generator, + 20, + conditional=True, + num_workers=self.__config.num_workers, + seed=self.__config.seed, + ) + + if self.__config.measure_time: + print(f"{time.time() - starttime:.2f} seconds.") + else: + print("done.") + elif self.__config.noise_stddev_adj == "fixed": + f = lambda k: 1.0 / (0.75 + 0.09 * k) + self.__params.noise_std_coeffs = [ + f(k) for k in range(1, self.__config.n_cascade_levels + 1) + ] + else: + self.__params.noise_std_coeffs = np.ones(self.__config.n_cascade_levels) + + if self.__params.noise_stddev_adj is not None: + print(f"noise std. dev. coeffs: {self.__params.noise_std_coeffs}") + + else: + self.__params.perturbation_generator = None + self.__params.noise_generator = None + self.__params.noise_std_coeffs = None + def __estimate_ar_parameters_radar(self): # 4. Estimate AR parameters for the radar rainfall field """Estimate AR parameters for the radar rainfall field.""" @@ -707,7 +784,7 @@ def __estimate_ar_parameters_radar(self): # compute lag-l temporal auto-correlation coefficients for each cascade level for i in range(self.__config.n_cascade_levels): GAMMA[i, :] = correlation.temporal_autocorrelation( - self.__state.precip_cascade[i], mask=self.__params.mask_threshold + self.__state.precip_cascades[i], mask=self.__params.mask_threshold ) # Else, use standard values for the auto-correlations @@ -771,6 +848,147 @@ def __estimate_ar_parameters_radar(self): nowcast_utils.print_ar_params(self.__params.PHI) + def __multiply_precip_cascade_to_match_ensemble_members(self): + # 5. Repeat precip_cascade for n ensemble members + # First, discard all except the p-1 last cascades because they are not needed + # for the AR(p) model + + self.__state.precip_cascades = np.stack( + [ + [ + self.__state.precip_cascades[i][-self.__config.ar_order :].copy() + for i in range(self.__config.n_cascade_levels) + ] + ] + * self.__config.n_ens_members + ) + + def __initialize_random_generators(self): + # 6. Initialize all the random generators and prepare for the forecast loop + """Initialize all the random generators.""" + # TODO: randgen_motion and randgen_precip are not defined if no noise method is given? Should we end the program in that case? + if self.__config.noise_method is not None: + self.__state.randgen_precip = [] + randgen_motion = [] + for j in range(self.__config.n_ens_members): + rs = np.random.RandomState(self.__config.seed) + self.__state.randgen_precip.append(rs) + seed = rs.randint(0, high=1e9) + rs = np.random.RandomState(seed) + randgen_motion.append(rs) + seed = rs.randint(0, high=1e9) + + if self.__config.velocity_perturbation_method is not None: + ( + init_velocity_noise, + self.__state.generate_velocity_noise, + ) = noise.get_method(self.__config.velocity_perturbation_method) + + # initialize the perturbation generators for the motion field + self.__state.velocity_perturbations = [] + for j in range(self.__config.n_ens_members): + kwargs = { + "randstate": randgen_motion[j], + "p_par": self.__params.velocity_perturbations_parallel, + "p_perp": self.__params.velocity_perturbations_perpendicular, + } + vp_ = init_velocity_noise( + self.__velocity, + 1.0 / self.__config.kmperpixel, + self.__config.timestep, + **kwargs, + ) + self.__state.velocity_perturbations.append(vp_) + else: + ( + self.__state.velocity_perturbations, + self.__state.generate_velocity_noise, + ) = (None, None) + + def __prepare_forecast_loop(self): + """Prepare for the forecast loop.""" + # Empty arrays for the previous displacements and the forecast cascade + self.__state.previous_displacement = np.stack( + [None for j in range(self.__config.n_ens_members)] + ) + self.__state.previous_displacement_noise_cascade = np.stack( + [None for j in range(self.__config.n_ens_members)] + ) + self.__state.previous_displacement_prob_matching = np.stack( + [None for j in range(self.__config.n_ens_members)] + ) + self.__state.precip_forecast = [[] for j in range(self.__config.n_ens_members)] + + if self.__config.mask_method == "incremental": + # get mask parameters + self.__state.mask_rim = self.__config.mask_kwargs.get("mask_rim", 10) + mask_f = self.__config.mask_kwargs.get("mask_f", 1.0) + # initialize the structuring element + struct = generate_binary_structure(2, 1) + # iterate it to expand it nxn + n = mask_f * self.__config.timestep / self.__config.kmperpixel + self.__state.struct = iterate_structure(struct, int((n - 1) / 2.0)) + else: + self.__state.mask_rim, self.__state.struct = None, None + + if self.__config.noise_method is None: + self.__state.precip_forecast_non_perturbed = [ + self.__state.precip_cascades[0][i].copy() + for i in range(self.__config.n_cascade_levels) + ] + else: + self.__state.precip_forecast_non_perturbed = None + + self.__state.fft_objs = [] + for i in range(self.__config.n_ens_members): + self.__state.fft_objs.append( + utils.get_method( + self.__config.fft_method, + shape=self.__state.precip_cascades.shape[-2:], + ) + ) + + def __initialize_noise_cascade(self): + """Initialize the noise cascade with identical noise for all AR(n) steps + We also need to return the mean and standard deviations of the noise + for the recombination of the noise before advecting it. + """ + self.__state.precip_noise_cascade = np.zeros(self.__state.precip_cascades.shape) + self.__state.precip_mean_noise = np.zeros( + (self.__config.n_ens_members, self.__config.n_cascade_levels) + ) + self.__state.precip_std_noise = np.zeros( + (self.__config.n_ens_members, self.__config.n_cascade_levels) + ) + if self.__config.noise_method: + for j in range(self.__config.n_ens_members): + # TODO: check rest later, starts at #3 so should look above what these terms match to + epsilon = self.__params.noise_generator( + self.__params.perturbation_generator, + randstate=self.__state.randgen_precip[j], + fft_method=self.__state.fft_objs[j], + domain=self.__config.domain, + ) + epsilon_decomposed = self.__params.decomposition_method( + epsilon, + self.__params.bandpass_filter, + fft_method=self.__state.fft_objs[j], + input_domain=self.__config.domain, + output_domain=self.__config.domain, + compute_stats=True, + normalize=True, + compact_output=True, + ) + self.__state.precip_mean_noise[j] = epsilon_decomposed["means"] + self.__state.precip_std_noise[j] = epsilon_decomposed["stds"] + for i in range(self.__config.n_cascade_levels): + epsilon_temp = epsilon_decomposed["cascade_levels"][i] + epsilon_temp *= self.__params.noise_std_coeffs[i] + for n in range(self.__config.ar_order): + self.__state.precip_noise_cascade[j][i][n] = epsilon_temp + epsilon_decomposed = None + epsilon_temp = None + def __perform_extrapolation(self): pass From 8d16c11043e224c8e22126f5d0c45d32c676dc17 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 2 Dec 2024 13:47:17 +0100 Subject: [PATCH 08/33] Added some uniformity between nowcast and blending steps. Now at # 8.4 for the refactoring --- pysteps/blending/steps.py | 552 +++++++++++++++++++++++++++++++++++--- pysteps/nowcasts/steps.py | 11 +- 2 files changed, 518 insertions(+), 45 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 6209bb6b..1f314a75 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -107,6 +107,7 @@ class StepsBlendingConfig: return_output: bool = True +# TODO: typing could be improved here @dataclass class StepsBlendingParams: PHI: np.ndarray # AR(p) model parameters @@ -132,7 +133,6 @@ class StepsBlendingParams: ) # FFT objects for ensemble members mask_rim: Optional[int] = None # Rim size for masking struct: Optional[np.ndarray] = None # Structuring element for mask - n_model_indices: Optional[np.ndarray] = None # NWP model indices noise_method: Optional[str] = None # Noise method used ar_order: int = 2 # Order of the AR model seed: Optional[int] = None # Random seed for reproducibility @@ -144,13 +144,18 @@ class StepsBlendingParams: zero_precip_radar: bool = False zero_precip_model_fields: bool = False PHI: Any = None + original_timesteps: Any = None + num_ensemble_workers: int = None + rho_nwp_models: Any = None + domain_mask: Any = None +# TODO: typing could be improved here @dataclass class StepsBlendingState: precip_cascades: Any = None - mu_extrapolation: Any = None - sigma_extrapolation: Any = None + mean_extrapolation: Any = None + std_extrapolation: Any = None precip_models_cascades: Any = None PHI: Any = None randgen_precip: Any = None @@ -167,10 +172,23 @@ class StepsBlendingState: t_prev_timestep: Any = None t_leadtime_since_start_forecast: Any = None precip_noise_input: Any = None - precip_noise_cascade: Any = None + precip_noise_cascades: Any = None precip_mean_noise: Any = None precip_std_noise: Any = None - # Add more state variables as needed + rho_extrap_cascade_prev: Any = None + rho_extrap_cascade: Any = None + subtimesteps: Any = None + is_nowcast_time_step: bool = None + # Variables to save data over (sub)time steps + precip_models_cascades_temp: Any = None + precip_models_temp: Any = None + mean_models_temp: Any = None + std_models_temp: Any = None + velocity_models_temp: Any = None + n_model_indices: Optional[np.ndarray] = None # NWP model indices + rho_forecast: Any = None + weights: Any = None + weights_model_only: Any = None class StepsBlendingNowcaster: @@ -207,13 +225,11 @@ def __init__( # Additional variables for time measurement self.__start_time_init = None + self.__zero_precip_time = None self.__init_time = None self.__mainloop_time = None def compute_forecast(self): - pass - - def __nowcast_main(self): self.__check_inputs() self.__print_forecast_info() # Measure time for initialization @@ -240,6 +256,76 @@ def __nowcast_main(self): self.__initialize_random_generators() self.__prepare_forecast_loop() self.__initialize_noise_cascade() + if self.__config.measure_time: + self.__init_time = self.__measure_time( + "initialization", self.__start_time_init + ) + + self.__blended_nowcast_main() + # Stack and return the forecast output + if self.__config.return_output: + self.__state.precip_forecast = np.stack( + [ + np.stack(self.__state.precip_forecast[j]) + for j in range(self.__config.n_ens_members) + ] + ) + if self.__config.measure_time: + return ( + self.__state.precip_forecast, + self.__init_time, + self.__mainloop_time, + ) + else: + return self.__state.precip_forecast + else: + return None + + def __blended_nowcast_main(self): + """ + Main nowcast loop that iterates through the ensemble members and time steps + to generate forecasts. + """ + ### + # 8. Start the forecasting loop + ### + # Isolate the last time slice of observed precipitation + precip = self.__precip[-1, :, :] + print("Starting blended nowcast computation.") + + if self.__config.measure_time: + starttime_mainloop = time.time() + + self.__config.extrap_kwargs["return_displacement"] = True + + precip_forc_prev_subtimestep = deepcopy(self.__state.precip_cascades) + noise_prev_subtimestep = deepcopy(self.__state.precip_noise_cascades) + + t_prev_timestep = [0.0 for j in range(self.__config.n_ens_members)] + t_leadtime_since_start_forecast = [ + 0.0 for j in range(self.__config.n_ens_members) + ] + + # iterate each time step + for t, subtimestep_idx in enumerate(self.__timesteps): + self.__determine_subtimesteps_and_nowcast_time_step(t, subtimestep_idx) + if self.__config.measure_time: + starttime = time.time() + self.__decompose_nwp_if_needed_and_fill_nans_in_nwp(t) + self.__find_nowcast_NWP_combination(t) + self.__determine_skill_for_current_timestep(t) + # the nowcast iteration for each ensemble member + precip_forecast_workers = [None for _ in range(self.__config.n_ens_members)] + + def worker(j): + self.__determine_skill_for_next_timestep(t, j) + self.__determine_weights_per_component() + self.__regress_extrapolation_and_noise_cascades(j) + + # Perturb and blend the advection fields + advect the extrapolation and noise cascade to the current time step + # Blend the cascades + + pass def __check_inputs(self): """Validates the inputs and determines if the user provided raw forecasts or decomposed forecasts.""" @@ -282,8 +368,10 @@ def __check_inputs(self): if isinstance(self.__timesteps, list): self.__params.time_steps_is_list = True - original_timesteps = [0] + list(self.__timesteps) - self.__timesteps = nowcast_utils.binned_timesteps(original_timesteps) + self.__params.original_timesteps = [0] + list(self.__timesteps) + self.__timesteps = nowcast_utils.binned_timesteps( + self.__params.original_timesteps + ) if not sorted(self.__timesteps) == self.__timesteps: raise ValueError("timesteps is not in ascending order") if self.__precip_models.shape[1] != math.ceil(self.__timesteps[-1]) + 1: @@ -564,8 +652,8 @@ def __prepare_radar_and_NWP_fields(self): ) precip_forecast_decomp = precip_forecast_decomp[-1] - self.__state.mu_extrapolation = np.array(precip_forecast_decomp["means"]) - self.__state.sigma_extrapolation = np.array(precip_forecast_decomp["stds"]) + self.__state.mean_extrapolation = np.array(precip_forecast_decomp["means"]) + self.__state.std_extrapolation = np.array(precip_forecast_decomp["stds"]) # 2.2 If necessary, recompose (NWP) model forecasts self.__state.precip_models_cascades = None @@ -756,7 +844,7 @@ def __initialize_noise(self): ) if self.__config.measure_time: - print(f"{time.time() - starttime:.2f} seconds.") + __ = self.__measure_time("Initialize noise", starttime) else: print("done.") elif self.__config.noise_stddev_adj == "fixed": @@ -947,13 +1035,24 @@ def __prepare_forecast_loop(self): shape=self.__state.precip_cascades.shape[-2:], ) ) + # TODO: moved this from # 7 to here as it seems to fit better here. The only parameter used and needed is PHI, this is its last use untill # 7 + # initizalize the current and previous extrapolation forecast scale for the nowcasting component + # phi1 / (1 - phi2), see BPS2004 + self.__state.rho_extrap_cascade_prev = np.repeat( + 1.0, self.__params.PHI.shape[0] + ) + self.__state.rho_extrap_cascade = self.__params.PHI[:, 0] / ( + 1.0 - self.__params.PHI[:, 1] + ) def __initialize_noise_cascade(self): """Initialize the noise cascade with identical noise for all AR(n) steps We also need to return the mean and standard deviations of the noise for the recombination of the noise before advecting it. """ - self.__state.precip_noise_cascade = np.zeros(self.__state.precip_cascades.shape) + self.__state.precip_noise_cascades = np.zeros( + self.__state.precip_cascades.shape + ) self.__state.precip_mean_noise = np.zeros( (self.__config.n_ens_members, self.__config.n_cascade_levels) ) @@ -962,7 +1061,6 @@ def __initialize_noise_cascade(self): ) if self.__config.noise_method: for j in range(self.__config.n_ens_members): - # TODO: check rest later, starts at #3 so should look above what these terms match to epsilon = self.__params.noise_generator( self.__params.perturbation_generator, randstate=self.__state.randgen_precip[j], @@ -985,48 +1083,416 @@ def __initialize_noise_cascade(self): epsilon_temp = epsilon_decomposed["cascade_levels"][i] epsilon_temp *= self.__params.noise_std_coeffs[i] for n in range(self.__config.ar_order): - self.__state.precip_noise_cascade[j][i][n] = epsilon_temp + self.__state.precip_noise_cascades[j][i][n] = epsilon_temp epsilon_decomposed = None epsilon_temp = None - def __perform_extrapolation(self): - pass + def __determine_subtimesteps_and_nowcast_time_step(self, t, subtimestep_idx): + if self.__params.time_steps_is_list: + self.__state.subtimesteps = [ + self.__params.original_timesteps[t_] for t_ in subtimestep_idx + ] + else: + self.__state.subtimesteps = [t] - def __apply_noise_and_ar_model(self): - pass + if (self.__params.time_steps_is_list and self.__state.subtimesteps) or ( + not self.__params.time_steps_is_list and t > 0 + ): + self.__state.is_nowcast_time_step = True + else: + self.__state.is_nowcast_time_step = False - def __initialize_velocity_perturbations(self): - pass + if self.__state.is_nowcast_time_step: + print( + f"Computing nowcast for time step {t}... ", + end="", + flush=True, + ) - def __initialize_precipitation_mask(self): - pass + def __decompose_nwp_if_needed_and_fill_nans_in_nwp(self, t): + if self.__state.precip_models_cascades is not None: + decomp_precip_models = list(self.__state.precip_models_cascades[:, t]) - def __initialize_fft_objects(self): - pass + else: + if self.__precip_models.shape[0] == 1: + decomp_precip_models = [ + self.__params.decomposition_method( + self.__precip_models[0, t, :, :], + bp_filter=self.__params.bandpass_filter, + fft_method=self.__params.fft, + output_domain=self.__config.domain, + normalize=True, + compute_stats=True, + compact_output=True, + ) + ] + else: + with ThreadPool(self.__config.num_workers) as pool: + decomp_precip_models = pool.map( + partial( + self.__params.decomposition_method, + bp_filter=self.__params.bandpass_filter, + fft_method=self.__params.fft, + output_domain=self.__config.domain, + normalize=True, + compute_stats=True, + compact_output=True, + ), + list(self.__precip_models[:, t, :, :]), + ) - def __return_state_dict(self): - pass + self.__state.precip_models_cascades_temp = np.array( + [decomp["cascade_levels"] for decomp in decomp_precip_models] + ) + self.__state.mean_models_temp = np.array( + [decomp["means"] for decomp in decomp_precip_models] + ) + self.__state.std_models_temp = np.array( + [decomp["stds"] for decomp in decomp_precip_models] + ) - def __return_params_dict(self): - pass + # 2.3.4 Check if the NWP fields contain nans or infinite numbers. If so, + # fill these with the minimum value present in precip (corresponding to + # zero rainfall in the radar observations) - def __update_state(self, state, params): - pass + """Ensure that the NWP cascade and fields do no contain any nans or infinite number""" + # Fill nans and infinite numbers with the minimum value present in precip + self.__state.precip_models_temp = self.__precip_models[:, t, :, :].astype( + np.float64, copy=False + ) # (corresponding to zero rainfall in the radar observations) + min_cascade = np.nanmin(self.__state.precip_cascades) + min_precip = np.nanmin(self.__precip) + self.__state.precip_models_cascades_temp[ + ~np.isfinite(self.__state.precip_models_cascades_temp) + ] = min_cascade + self.__state.precip_models_temp[ + ~np.isfinite(self.__state.precip_models_temp) + ] = min_precip + # Also set any nans or infs in the mean and sigma of the cascade to + # respectively 0.0 and 1.0 + self.__state.mean_models_temp[~np.isfinite(self.__state.mean_models_temp)] = 0.0 + self.__state.std_models_temp[~np.isfinite(self.__state.std_models_temp)] = 0.0 + + def __find_nowcast_NWP_combination(self, t): + # 8.1.1 Before calling the worker for the forecast loop, determine which (NWP) + # models will be combined with which nowcast ensemble members. With the + # way it is implemented at this moment: n_ens_members of the output equals + # the maximum number of (ensemble) members in the input (either the nowcasts or NWP). + + """Determine which (NWP) models will be combined with which nowcast ensemble members. + With the way it is implemented at this moment: n_ens_members of the output equals + the maximum number of (ensemble) members in the input (either the nowcasts or NWP). + """ + self.__state.velocity_models_temp = self.__velocity_models[ + :, t, :, :, : + ].astype(np.float64, copy=False) + # Make sure the number of model members is not larger than or equal to n_ens_members + n_model_members = self.__state.precip_models_cascades_temp.shape[0] + if n_model_members > self.__config.n_ens_members: + raise ValueError( + "The number of NWP model members is larger than the given number of ensemble members. n_model_members <= n_ens_members." + ) - def __update_deterministic_ar_model(self, state, params): - pass + # Check if NWP models/members should be used individually, or if all of + # them are blended together per nowcast ensemble member. + if self.__config.blend_nwp_members: + self.__state.n_model_indices = None - def __apply_ar_model_to_cascades(self, j, state, params): - pass + else: + # Start with determining the maximum and mimimum number of members/models + # in both input products + n_ens_members_max = max(self.__config.n_ens_members, n_model_members) + n_ens_members_min = min(self.__config.n_ens_members, n_model_members) + # Also make a list of the model index numbers. These indices are needed + # for indexing the right climatological skill file when pysteps calculates + # the blended forecast in parallel. + if n_model_members > 1: + self.__state.n_model_indices = np.arange(n_model_members) + else: + self.__state.n_model_indices = [0] + + # Now, repeat the nowcast ensemble members or the nwp models/members until + # it has the same amount of members as n_ens_members_max. For instance, if + # you have 10 ensemble nowcasts members and 3 NWP members, the output will + # be an ensemble of 10 members. Hence, the three NWP members are blended + # with the first three members of the nowcast (member one with member one, + # two with two, etc.), subsequently, the same NWP members are blended with + # the next three members (NWP member one with member 4, NWP member 2 with + # member 5, etc.), until 10 is reached. + if n_ens_members_min != n_ens_members_max: + if n_model_members == 1: + self.__state.precip_models_cascades_temp = np.repeat( + self.__state.precip_models_cascades_temp, + n_ens_members_max, + axis=0, + ) + self.__state.mean_models_temp = np.repeat( + self.__state.mean_models_temp, n_ens_members_max, axis=0 + ) + self.__state.std_models_temp = np.repeat( + self.__state.std_models_temp, n_ens_members_max, axis=0 + ) + self.__state.velocity_models_temp = np.repeat( + self.__state.velocity_models_temp, n_ens_members_max, axis=0 + ) + # For the prob. matching + self.__state.precip_models_temp = np.repeat( + self.__state.precip_models_temp, n_ens_members_max, axis=0 + ) + # Finally, for the model indices + self.__state.n_model_indices = np.repeat( + self.__state.n_model_indices, n_ens_members_max, axis=0 + ) + + elif n_model_members == n_ens_members_min: + repeats = [ + (n_ens_members_max + i) // n_ens_members_min + for i in range(n_ens_members_min) + ] + if n_model_members == n_ens_members_min: + self.__state.precip_models_cascades_temp = np.repeat( + self.__state.precip_models_cascades_temp, repeats, axis=0 + ) + self.__state.mean_models_temp = np.repeat( + self.__state.mean_models_temp, repeats, axis=0 + ) + self.__state.std_models_temp = np.repeat( + self.__state.std_models_temp, repeats, axis=0 + ) + self.__state.velocity_models_temp = np.repeat( + self.__state.velocity_models_temp, repeats, axis=0 + ) + # For the prob. matching + self.__state.precip_models_temp = np.repeat( + self.__state.precip_models_temp, repeats, axis=0 + ) + # Finally, for the model indices + self.__state.n_model_indices = np.repeat( + self.__state.n_model_indices, repeats, axis=0 + ) + + # TODO: is this not duplicate from part 2.3.5? + # If zero_precip_radar is True, set the velocity field equal to the NWP + # velocity field for the current time step (velocity_models_temp). + if self.__params.zero_precip_radar: + # Use the velocity from velocity_models and take the average over + # n_models (axis=0) + self.__velocity = np.mean(self.__state.velocity_models_temp, axis=0) + + def __determine_skill_for_current_timestep(self, t): + if t == 0: + """Calculate the initial skill of the (NWP) model forecasts at t=0.""" + # TODO: n_model is not defined here, how does this work? + self.__params.rho_nwp_models = [ + blending.skill_scores.spatial_correlation( + obs=self.__state.precip_cascades[0, :, -1, :, :].copy(), + mod=self.__state.precip_models_cascades_temp[ + n_model, :, :, : + ].copy(), + domain_mask=self.__params.domain_mask, + ) + for n_model in range(self.__state.precip_models_cascades_temp.shape[0]) + ] + self.__params.rho_nwp_models = np.stack(self.__params.rho_nwp_models) + + # Ensure that the model skill decreases with increasing scale level. + for n_model in range(self.__state.precip_models_cascades_temp.shape[0]): + for i in range(1, self.__state.precip_models_cascades_temp.shape[1]): + if ( + self.__params.rho_nwp_models[n_model, i] + > self.__params.rho_nwp_models[n_model, i - 1] + ): + # Set it equal to the previous scale level + self.__params.rho_nwp_models[n_model, i] = ( + self.__params.rho_nwp_models[n_model, i - 1] + ) + + # Save this in the climatological skill file + blending.clim.save_skill( + current_skill=self.__params.rho_nwp_models, + validtime=self.__issuetime, + outdir_path=self.__config.outdir_path_skill, + **self.__config.clim_kwargs, + ) + if t > 0: + # 8.1.3 Determine the skill of the components for lead time (t0 + t) + # First for the extrapolation component. Only calculate it when t > 0. + ( + self.__state.rho_extrap_cascade, + self.__state.rho_extrap_cascade_prev, + ) = blending.skill_scores.lt_dependent_cor_extrapolation( + PHI=self.__params.PHI, + correlations=self.__state.rho_extrap_cascade, + correlations_prev=self.__state.rho_extrap_cascade_prev, + ) - def __generate_and_decompose_noise(self, j, state, params): - pass + def __determine_skill_for_next_timestep(self, t, j): + # 8.1.2 Determine the skill of the nwp components for lead time (t0 + t) + # Then for the model components + if self.__config.blend_nwp_members: + rho_nwp_forecast = [ + blending.skill_scores.lt_dependent_cor_nwp( + lt=(t * int(self.__config.timestep)), + correlations=self.__params.rho_nwp_models[n_model], + outdir_path=self.__config.outdir_path_skill, + n_model=n_model, + skill_kwargs=self.__config.clim_kwargs, + ) + for n_model in range(self.__params.rho_nwp_models.shape[0]) + ] + rho_nwp_forecast = np.stack(rho_nwp_forecast) + # Concatenate rho_extrap_cascade and rho_nwp + self.__state.rho_forecast = np.concatenate( + (self.__state.rho_extrap_cascade[None, :], rho_nwp_forecast), axis=0 + ) + else: + rho_nwp_forecast = blending.skill_scores.lt_dependent_cor_nwp( + lt=(t * int(self.__config.timestep)), + correlations=self.__params.rho_nwp_models[j], + outdir_path=self.__config.outdir_path_skill, + n_model=self.__params.n_model_indices[j], + skill_kwargs=self.__config.clim_kwargs, + ) + # Concatenate rho_extrap_cascade and rho_nwp + self.__state.rho_forecast = np.concatenate( + (self.__state.rho_extrap_cascade[None, :], rho_nwp_forecast[None, :]), + axis=0, + ) - def __recompose_and_apply_mask(self, j, state, params): - pass + def __determine_weights_per_component(self): + # 8.2 Determine the weights per component + + # Weights following the bps method. These are needed for the velocity + # weights prior to the advection step. If weights method spn is + # selected, weights will be overwritten with those weights prior to + # blending step. + # weight = [(extr_field, n_model_fields, noise), n_cascade_levels, ...] + self.__state.weights = calculate_weights_bps(self.__state.rho_forecast) + + # The model only weights + if self.__config.weights_method == "bps": + # Determine the weights of the components without the extrapolation + # cascade, in case this is no data or outside the mask. + self.__state.weights_model_only = calculate_weights_bps( + self.__state.rho_forecast[1:, :] + ) + elif self.__config.weights_method == "spn": + # Only the weights of the components without the extrapolation + # cascade will be determined here. The full set of weights are + # determined after the extrapolation step in this method. + if ( + self.__config.blend_nwp_members + and self.__state.precip_models_cascades_temp.shape[0] > 1 + ): + self.__state.weights_model_only = np.zeros( + ( + self.__state.precip_models_cascades_temp.shape[0] + 1, + self.__config.n_cascade_levels, + ) + ) + for i in range(self.__config.n_cascade_levels): + # Determine the normalized covariance matrix (containing) + # the cross-correlations between the models + covariance_nwp_models = np.corrcoef( + np.stack( + [ + self.__state.precip_models_cascades_temp[ + n_model, i, :, : + ].flatten() + for n_model in range( + self.__state.precip_models_cascades_temp.shape[0] + ) + ] + ) + ) + # Determine the weights for this cascade level + self.__state.weights_model_only[:, i] = calculate_weights_spn( + correlations=self.__state.rho_forecast[1:, i], + covariance=covariance_nwp_models, + ) + else: + # Same as correlation and noise is 1 - correlation + self.__state.weights_model_only = calculate_weights_bps( + self.__state.rho_forecast[1:, :] + ) + else: + raise ValueError( + "Unknown weights method %s: must be 'bps' or 'spn'" + % self.__config.weights_method + ) + + def __regress_extrapolation_and_noise_cascades(self, j): + # 8.3 Determine the noise cascade and regress this to the subsequent + # time step + regress the extrapolation component to the subsequent + # time step + + # 8.3.1 Determine the epsilon, a cascade of temporally independent + # but spatially correlated noise + if self.__config.noise_method is not None: + # generate noise field + epsilon = self.__params.noise_generator( + self.__params.perturbation_generator, + randstate=self.__state.randgen_precip[j], + fft_method=self.__state.fft_objs[j], + domain=self.__config.domain, + ) + + # decompose the noise field into a cascade + epsilon_decomposed = self.__params.decomposition_method( + epsilon, + self.__params.bandpass_filter, + fft_method=self.__state.fft_objs[j], + input_domain=self.__config.domain, + output_domain=self.__config.domain, + compute_stats=True, + normalize=True, + compact_output=True, + ) + else: + epsilon_decomposed = None + + # 8.3.2 regress the extrapolation component to the subsequent time + # step + # iterate the AR(p) model for each cascade level + for i in range(self.__config.n_cascade_levels): + # apply AR(p) process to extrapolation cascade level + if ( + epsilon_decomposed is not None + or self.__config.velocity_perturbation_method is not None + ): + self.__state.precip_cascade[j][i] = autoregression.iterate_ar_model( + self.__state.precip_cascade[j][i], self.__params.PHI[i, :] + ) + # Renormalize the cascade + self.__state.precip_cascade[j][i][1] /= np.std( + self.__state.precip_cascade[j][i][1] + ) + else: + # use the deterministic AR(p) model computed above if + # perturbations are disabled + self.__state.precip_cascade[j][i] = ( + self.__state.precip_forecast_non_perturbed[i] + ) + + # 8.3.3 regress the noise component to the subsequent time step + # iterate the AR(p) model for each cascade level + for i in range(self.__config.n_cascade_levels): + # normalize the noise cascade + if epsilon_decomposed is not None: + epsilon_temp = epsilon_decomposed["cascade_levels"][i] + epsilon_temp *= self.__paramsnoise_std_coeffs[i] + else: + epsilon_temp = None + # apply AR(p) process to noise cascade level + # (Returns zero noise if epsilon_decomposed is None) + self.__state.precip_noise_cascades[j][i] = autoregression.iterate_ar_model( + self.__state.precip_noise_cascades[j][i], + self.__params.PHI[i, :], + eps=epsilon_temp, + ) - def __apply_precipitation_mask(self, precip_forecast, j, state, params): - pass + epsilon_decomposed = None + epsilon_temp = None def __measure_time(self, label, start_time): """ @@ -1039,6 +1505,8 @@ def __measure_time(self, label, start_time): if self.__config.measure_time: elapsed_time = time.time() - start_time print(f"{label} took {elapsed_time:.2f} seconds.") + return elapsed_time + return None def reset_states_and_params(self): """ diff --git a/pysteps/nowcasts/steps.py b/pysteps/nowcasts/steps.py index b04e8275..6efd9758 100644 --- a/pysteps/nowcasts/steps.py +++ b/pysteps/nowcasts/steps.py @@ -352,11 +352,14 @@ def compute_forecast(self): self.__initialize_fft_objects() # Measure and print initialization time if self.__config.measure_time: - self.__measure_time("Initialization", self.__start_time_init) + self.__init_time = self.__measure_time( + "Initialization", self.__start_time_init + ) # Run the main nowcast loop self.__nowcast_main() + # Unstack nowcast output if return_output is True if self.__config.measure_time: ( self.__state.precip_forecast, @@ -387,7 +390,7 @@ def __nowcast_main(self): Main nowcast loop that iterates through the ensemble members and time steps to generate forecasts. """ - # Isolate the last time slice of precipitation + # Isolate the last time slice of observed precipitation precip = self.__precip[ -1, :, : ] # Extract the last available precipitation field @@ -717,7 +720,7 @@ def __apply_noise_and_ar_model(self): # Measure and print time taken if self.__config.measure_time: - self.__measure_time( + __ = self.__measure_time( "Noise adjustment coefficient computation", starttime ) else: @@ -1198,6 +1201,8 @@ def __measure_time(self, label, start_time): if self.__config.measure_time: elapsed_time = time.time() - start_time print(f"{label} took {elapsed_time:.2f} seconds.") + return elapsed_time + return None def reset_states_and_params(self): """ From 88df97dbfd9a8bdfbe43626e01b1939dc94929f1 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 2 Dec 2024 14:03:16 +0100 Subject: [PATCH 09/33] Small changes since prev commit --- pysteps/blending/steps.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 1f314a75..5e3cb43e 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -255,7 +255,7 @@ def compute_forecast(self): self.__multiply_precip_cascade_to_match_ensemble_members() self.__initialize_random_generators() self.__prepare_forecast_loop() - self.__initialize_noise_cascade() + self.__initialize_noise_cascades() if self.__config.measure_time: self.__init_time = self.__measure_time( "initialization", self.__start_time_init @@ -1045,7 +1045,7 @@ def __prepare_forecast_loop(self): 1.0 - self.__params.PHI[:, 1] ) - def __initialize_noise_cascade(self): + def __initialize_noise_cascades(self): """Initialize the noise cascade with identical noise for all AR(n) steps We also need to return the mean and standard deviations of the noise for the recombination of the noise before advecting it. From 7ee0020d5e5863e04557d7d4b57f1578501f9551 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 2 Dec 2024 16:57:33 +0100 Subject: [PATCH 10/33] All code is tranfered. Last part of the main loop needs to be refactored --- pysteps/blending/steps.py | 778 +++++++++++++++++++++++++++++++++++++- 1 file changed, 759 insertions(+), 19 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 5e3cb43e..15ac25b9 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -148,6 +148,8 @@ class StepsBlendingParams: num_ensemble_workers: int = None rho_nwp_models: Any = None domain_mask: Any = None + velocity_perturbations: Any = None + generate_velocity_noise: Any = None # TODO: typing could be improved here @@ -159,8 +161,6 @@ class StepsBlendingState: precip_models_cascades: Any = None PHI: Any = None randgen_precip: Any = None - velocity_perturbations: Any = None - generate_velocity_noise: Any = None previous_displacement: Any = None previous_displacement_noise_cascade: Any = None previous_displacement_prob_matching: Any = None @@ -189,6 +189,11 @@ class StepsBlendingState: rho_forecast: Any = None weights: Any = None weights_model_only: Any = None + precip_forecast_extrapolated_decomp_done: Any = None + noise_extrapolated_decomp_done: Any = None + precip_forecast_extrapolated_probability_matching: Any = None + precip_forecast_prev_subtimestep: Any = None + noise_prev_subtimestep: Any = None class StepsBlendingNowcaster: @@ -295,14 +300,19 @@ def __blended_nowcast_main(self): if self.__config.measure_time: starttime_mainloop = time.time() - + # TODO: problem with the config here! This variable changes over time... + # extrap_kwargs is in config but by adding info to it, the next run of a blended forecast will have issues! self.__config.extrap_kwargs["return_displacement"] = True - precip_forc_prev_subtimestep = deepcopy(self.__state.precip_cascades) - noise_prev_subtimestep = deepcopy(self.__state.precip_noise_cascades) + self.__state.precip_forecast_prev_subtimestep = deepcopy( + self.__state.precip_cascades + ) + self.__state.noise_prev_subtimestep = deepcopy( + self.__state.precip_noise_cascades + ) - t_prev_timestep = [0.0 for j in range(self.__config.n_ens_members)] - t_leadtime_since_start_forecast = [ + self.__state.t_prev_timestep = [0.0 for j in range(self.__config.n_ens_members)] + self.__state.t_leadtime_since_start_forecast = [ 0.0 for j in range(self.__config.n_ens_members) ] @@ -321,11 +331,49 @@ def worker(j): self.__determine_skill_for_next_timestep(t, j) self.__determine_weights_per_component() self.__regress_extrapolation_and_noise_cascades(j) + self.__perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( + t, j + ) + # 8.5 Blend the cascades + final_blended_forecast_single_member = [] + for t_sub in self.__state.subtimesteps: + # TODO: does it make sense to use sub time steps - check if it works? + if t_sub > 0: + self.__blend_cascades() + self.__recompose_cascade_to_rainfall_field() + self.__post_process_output(final_blended_forecast_single_member) + precip_forecast_workers[j] = final_blended_forecast_single_member + + result = [] + + if DASK_IMPORTED and self.__config.n_ens_members > 1: + for j in range(self.__config.n_ens_members): + result.append(dask.delayed(worker)(j)) + dask.compute(*result, num_workers=self.__params.num_ensemble_workers) + else: + for j in range(self.__config.n_ens_members): + worker(j) - # Perturb and blend the advection fields + advect the extrapolation and noise cascade to the current time step - # Blend the cascades + result = None - pass + if self.__state.is_nowcast_time_step: + if self.__config.measure_time: + __ = self.__measure_time("subtimestep", starttime) + else: + print("done.") + + if self.__config.callback is not None: + precip_forecast_final = np.stack(precip_forecast_workers) + if precip_forecast_final.shape[1] > 0: + self.__config.callback(precip_forecast_final.squeeze()) + + if self.__config.return_output: + for j in range(self.__config.n_ens_members): + self.__state.precip_forecast[j].extend(precip_forecast_workers[j]) + + precip_forecast_workers = None + if self.__config.measure_time: + self.__mainloop_time = time.time() - starttime_mainloop def __check_inputs(self): """Validates the inputs and determines if the user provided raw forecasts or decomposed forecasts.""" @@ -969,7 +1017,7 @@ def __initialize_random_generators(self): if self.__config.velocity_perturbation_method is not None: ( init_velocity_noise, - self.__state.generate_velocity_noise, + self.__params.generate_velocity_noise, ) = noise.get_method(self.__config.velocity_perturbation_method) # initialize the perturbation generators for the motion field @@ -989,8 +1037,8 @@ def __initialize_random_generators(self): self.__state.velocity_perturbations.append(vp_) else: ( - self.__state.velocity_perturbations, - self.__state.generate_velocity_noise, + self.__params.velocity_perturbations, + self.__params.generate_velocity_noise, ) = (None, None) def __prepare_forecast_loop(self): @@ -1460,17 +1508,17 @@ def __regress_extrapolation_and_noise_cascades(self, j): epsilon_decomposed is not None or self.__config.velocity_perturbation_method is not None ): - self.__state.precip_cascade[j][i] = autoregression.iterate_ar_model( - self.__state.precip_cascade[j][i], self.__params.PHI[i, :] + self.__state.precip_cascades[j][i] = autoregression.iterate_ar_model( + self.__state.precip_cascades[j][i], self.__params.PHI[i, :] ) # Renormalize the cascade - self.__state.precip_cascade[j][i][1] /= np.std( - self.__state.precip_cascade[j][i][1] + self.__state.precip_cascades[j][i][1] /= np.std( + self.__state.precip_cascades[j][i][1] ) else: # use the deterministic AR(p) model computed above if # perturbations are disabled - self.__state.precip_cascade[j][i] = ( + self.__state.precip_cascades[j][i] = ( self.__state.precip_forecast_non_perturbed[i] ) @@ -1480,7 +1528,7 @@ def __regress_extrapolation_and_noise_cascades(self, j): # normalize the noise cascade if epsilon_decomposed is not None: epsilon_temp = epsilon_decomposed["cascade_levels"][i] - epsilon_temp *= self.__paramsnoise_std_coeffs[i] + epsilon_temp *= self.__params.noise_std_coeffs[i] else: epsilon_temp = None # apply AR(p) process to noise cascade level @@ -1494,6 +1542,698 @@ def __regress_extrapolation_and_noise_cascades(self, j): epsilon_decomposed = None epsilon_temp = None + def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( + self, t, j + ): + # 8.4 Perturb and blend the advection fields + advect the + # extrapolation and noise cascade to the current time step + # (or subtimesteps if non-integer time steps are given) + + # Settings and initialize the output + extrap_kwargs_ = self.__config.extrap_kwargs.copy() + extrap_kwargs_noise = self.__config.extrap_kwargs.copy() + extrap_kwargs_pb = self.__config.extrap_kwargs.copy() + velocity_perturbations_extrapolation = self.__velocity + # The following should be accesseble after this function + self.__state.precip_forecast_extrapolated_decomp_done = [] + self.__state.noise_extrapolated_decomp_done = [] + self.__state.precip_forecast_extrapolated_probability_matching = [] + + # Extrapolate per sub time step + for t_sub in self.__state.subtimesteps: + if t_sub > 0: + t_diff_prev_subtimestep_int = t_sub - int(t_sub) + if t_diff_prev_subtimestep_int > 0.0: + precip_forecast_cascade_subtimestep = [ + (1.0 - t_diff_prev_subtimestep_int) + * self.__state.precip_forecast_prev_subtimestep[j][i][-1, :] + + t_diff_prev_subtimestep_int + * self.__state.precip_cascades[j][i][-1, :] + for i in range(self.__config.n_cascade_levels) + ] + noise_cascade_subtimestep = [ + (1.0 - t_diff_prev_subtimestep_int) + * self.__state.noise_prev_subtimestep[j][i][-1, :] + + t_diff_prev_subtimestep_int + * self.__state.precip_noise_cascades[j][i][-1, :] + for i in range(self.__config.n_cascade_levels) + ] + + else: + precip_forecast_cascade_subtimestep = [ + self.__state.precip_forecast_prev_subtimestep[j][i][-1, :] + for i in range(self.__config.n_cascade_levels) + ] + noise_cascade_subtimestep = [ + self.__state.noise_prev_subtimestep[j][i][-1, :] + for i in range(self.__config.n_cascade_levels) + ] + + precip_forecast_cascade_subtimestep = np.stack( + precip_forecast_cascade_subtimestep + ) + noise_cascade_subtimestep = np.stack(noise_cascade_subtimestep) + + t_diff_prev_subtimestep = t_sub - self.__state.t_prev_timestep[j] + self.__state.t_leadtime_since_start_forecast[ + j + ] += t_diff_prev_subtimestep + + # compute the perturbed motion field - include the NWP + # velocities and the weights. Note that we only perturb + # the extrapolation velocity field, as the NWP velocity + # field is present per time step + if self.__config.velocity_perturbation_method is not None: + velocity_perturbations_extrapolation = ( + self.__velocity + + self.__params.generate_velocity_noise( + self.__params.velocity_perturbations[j], + self.__state.t_leadtime_since_start_forecast[j] + * self.__config.timestep, + ) + ) + + # Stack the perturbed extrapolation and the NWP velocities + if self.__config.blend_nwp_members: + velocity_stack_all = np.concatenate( + ( + velocity_perturbations_extrapolation[None, :, :, :], + self.__state.velocity_models_temp, + ), + axis=0, + ) + else: + velocity_models = self.__state.velocity_models_temp[j] + velocity_stack_all = np.concatenate( + ( + velocity_perturbations_extrapolation[None, :, :, :], + velocity_models[None, :, :, :], + ), + axis=0, + ) + velocity_models = None + + # Obtain a blended optical flow, using the weights of the + # second cascade following eq. 24 in BPS2006 + velocity_blended = blending.utils.blend_optical_flows( + flows=velocity_stack_all, + weights=self.__state.weights[ + :-1, 1 + ], # [(extr_field, n_model_fields), cascade_level=2] + ) + + # Extrapolate both cascades to the next time step + # First recompose the cascade, advect it and decompose it again + # This is needed to remove the interpolation artifacts. + # In addition, the number of extrapolations is greatly reduced + # A. Radar Rain + precip_forecast_recomp_subtimestep = blending.utils.recompose_cascade( + combined_cascade=precip_forecast_cascade_subtimestep, + combined_mean=self.__state.mean_extrapolation, + combined_sigma=self.__state.std_extrapolation, + ) + # Make sure we have values outside the mask + if self.__params.zero_precip_radar: + precip_forecast_recomp_subtimestep = np.nan_to_num( + precip_forecast_recomp_subtimestep, + copy=True, + nan=self.__params.precip_zerovalue, + posinf=self.__params.precip_zerovalue, + neginf=self.__params.precip_zerovalue, + ) + # Put back the mask + precip_forecast_recomp_subtimestep[self.__params.domain_mask] = np.nan + # TODO: problem with the config here! This variable changes over time... + # extrap_kwargs is in config but by adding info to it, the next run of a blended forecast will have issues! + self.__config.extrap_kwargs["displacement_prev"] = ( + self.__state.previous_displacement[j] + ) + ( + precip_forecast_extrapolated_recomp_subtimestep_temp, + self.__state.previous_displacement[j], + ) = self.__params.extrapolation_method( + precip_forecast_recomp_subtimestep, + velocity_blended, + [t_diff_prev_subtimestep], + allow_nonfinite_values=True, + **self.__config.extrap_kwargs, + ) + precip_forecast_extrapolated_recomp_subtimestep = ( + precip_forecast_extrapolated_recomp_subtimestep_temp[0].copy() + ) + temp_mask = ~np.isfinite( + precip_forecast_extrapolated_recomp_subtimestep + ) + # TODO: WHERE DO CAN I FIND THIS -15.0 + precip_forecast_extrapolated_recomp_subtimestep[ + ~np.isfinite(precip_forecast_extrapolated_recomp_subtimestep) + ] = self.__params.precip_zerovalue + precip_forecast_extrapolated_decomp = ( + self.__params.decomposition_method( + precip_forecast_extrapolated_recomp_subtimestep, + self.__params.bandpass_filter, + mask=self.__params.mask_threshold, + fft_method=self.__params.fft, + output_domain=self.__config.domain, + normalize=True, + compute_stats=True, + compact_output=True, + )["cascade_levels"] + ) + # Make sure we have values outside the mask + if self.__params.zero_precip_radar: + precip_forecast_extrapolated_decomp = np.nan_to_num( + precip_forecast_extrapolated_decomp, + copy=True, + nan=np.nanmin(precip_forecast_cascade_subtimestep), + posinf=np.nanmin(precip_forecast_cascade_subtimestep), + neginf=np.nanmin(precip_forecast_cascade_subtimestep), + ) + for i in range(self.__config.n_cascade_levels): + precip_forecast_extrapolated_decomp[i][temp_mask] = np.nan + # B. Noise + noise_cascade_subtimestep_recomp = blending.utils.recompose_cascade( + combined_cascade=noise_cascade_subtimestep, + combined_mean=self.__state.precip_mean_noise[j], + combined_sigma=self.__state.precip_std_noise[j], + ) + extrap_kwargs_noise["displacement_prev"] = ( + self.__state.previous_displacement_noise_cascade[j] + ) + extrap_kwargs_noise["map_coordinates_mode"] = "wrap" + ( + noise_extrapolated_recomp_temp, + self.__state.previous_displacement_noise_cascade[j], + ) = self.__params.extrapolation_method( + noise_cascade_subtimestep_recomp, + velocity_blended, + [t_diff_prev_subtimestep], + allow_nonfinite_values=True, + **extrap_kwargs_noise, + ) + noise_extrapolated_recomp = noise_extrapolated_recomp_temp[0].copy() + noise_extrapolated_decomp = self.__params.decomposition_method( + noise_extrapolated_recomp, + self.__params.bandpass_filter, + mask=self.__params.mask_threshold, + fft_method=self.__params.fft, + output_domain=self.__config.domain, + normalize=True, + compute_stats=True, + compact_output=True, + )["cascade_levels"] + for i in range(self.__config.n_cascade_levels): + noise_extrapolated_decomp[i] *= self.__params.noise_std_coeffs[i] + + # Append the results to the output lists + self.__state.precip_forecast_extrapolated_decomp_done.append( + precip_forecast_extrapolated_decomp.copy() + ) + self.__state.noise_extrapolated_decomp_done.append( + noise_extrapolated_decomp.copy() + ) + precip_forecast_cascade_subtimestep = None + precip_forecast_recomp_subtimestep = None + precip_forecast_extrapolated_recomp_subtimestep_temp = None + precip_forecast_extrapolated_recomp_subtimestep = None + precip_forecast_extrapolated_decomp = None + noise_cascade_subtimestep = None + noise_cascade_subtimestep_recomp = None + noise_extrapolated_recomp_temp = None + noise_extrapolated_recomp = None + noise_extrapolated_decomp = None + + # Finally, also extrapolate the initial radar rainfall + # field. This will be blended with the rainfall field(s) + # of the (NWP) model(s) for Lagrangian blended prob. matching + # min_R = np.min(precip) + extrap_kwargs_pb["displacement_prev"] = ( + self.__state.previous_displacement_prob_matching[j] + ) + # Apply the domain mask to the extrapolation component + precip_forecast_temp_for_probability_matching = self.__precip.copy() + precip_forecast_temp_for_probability_matching[ + self.__params.domain_mask + ] = np.nan + ( + precip_forecast_extrapolated_probability_matching_temp, + self.__state.previous_displacement_prob_matching[j], + ) = self.__params.extrapolation_method( + precip_forecast_temp_for_probability_matching, + velocity_blended, + [t_diff_prev_subtimestep], + allow_nonfinite_values=True, + **extrap_kwargs_pb, + ) + self.__state.precip_forecast_extrapolated_probability_matching.append( + precip_forecast_extrapolated_probability_matching_temp[0] + ) + + self.__state.t_prev_timestep[j] = t_sub + + if len(self.__state.precip_forecast_extrapolated_decomp_done) > 0: + self.__state.precip_forecast_extrapolated_decomp_done = np.stack( + self.__state.precip_forecast_extrapolated_decomp_done + ) + self.__state.noise_extrapolated_decomp_done = np.stack( + self.__state.noise_extrapolated_decomp_done + ) + self.__state.precip_forecast_extrapolated_probability_matching = np.stack( + self.__state.precip_forecast_extrapolated_probability_matching + ) + + # advect the forecast field by one time step if no subtimesteps in the + # current interval were found + if not self.__state.subtimesteps: + t_diff_prev_subtimestep = t + 1 - self.__state.t_prev_timestep[j] + self.__state.t_leadtime_since_start_forecast[j] += t_diff_prev_subtimestep + + # compute the perturbed motion field - include the NWP + # velocities and the weights + if self.__config.velocity_perturbation_method is not None: + velocity_perturbations_extrapolation = ( + self.__velocity + + self.__params.generate_velocity_noise( + self.__params.velocity_perturbations[j], + self.__state.t_leadtime_since_start_forecast[j] + * self.__config.timestep, + ) + ) + + # Stack the perturbed extrapolation and the NWP velocities + if self.__config.blend_nwp_members: + velocity_stack_all = np.concatenate( + ( + velocity_perturbations_extrapolation[None, :, :, :], + self.__state.velocity_models_temp, + ), + axis=0, + ) + else: + velocity_models = self.__state.velocity_models_temp[j] + velocity_stack_all = np.concatenate( + ( + velocity_perturbations_extrapolation[None, :, :, :], + velocity_models[None, :, :, :], + ), + axis=0, + ) + velocity_models = None + + # Obtain a blended optical flow, using the weights of the + # second cascade following eq. 24 in BPS2006 + velocity_blended = blending.utils.blend_optical_flows( + flows=velocity_stack_all, + weights=self.__state.weights[ + :-1, 1 + ], # [(extr_field, n_model_fields), cascade_level=2] + ) + + # Extrapolate the extrapolation and noise cascade + + extrap_kwargs_["displacement_prev"] = self.__state.previous_displacement[j] + extrap_kwargs_noise["displacement_prev"] = ( + self.__state.previous_displacement_noise_cascade[j] + ) + extrap_kwargs_noise["map_coordinates_mode"] = "wrap" + + _, self.__state.previous_displacement[j] = ( + self.__params.extrapolation_method( + None, + velocity_blended, + [t_diff_prev_subtimestep], + allow_nonfinite_values=True, + **extrap_kwargs_, + ) + ) + + _, self.__state.previous_displacement_noise_cascade[j] = ( + self.__params.extrapolation_method( + None, + velocity_blended, + [t_diff_prev_subtimestep], + allow_nonfinite_values=True, + **extrap_kwargs_noise, + ) + ) + + # Also extrapolate the radar observation, used for the probability + # matching and post-processing steps + extrap_kwargs_pb["displacement_prev"] = ( + self.__state.previous_displacement_prob_matching[j] + ) + _, self.__state.previous_displacement_prob_matching[j] = ( + self.__params.extrapolation_method( + None, + velocity_blended, + [t_diff_prev_subtimestep], + allow_nonfinite_values=True, + **extrap_kwargs_pb, + ) + ) + + self.__state.t_prev_timestep[j] = t + 1 + + self.__state.precip_forecast_prev_subtimestep[j] = self.__state.precip_cascades[ + j + ] + self.__state.noise_prev_subtimestep[j] = self.__state.precip_noise_cascades[j] + + def __blend_cascades(self): + t_index = np.where(np.array(subtimesteps) == t_sub)[0][0] + # First concatenate the cascades and the means and sigmas + # precip_models = [n_models,timesteps,n_cascade_levels,m,n] + if blend_nwp_members: + cascade_stack_all_components = np.concatenate( + ( + precip_forecast_extrapolated_decomp_done[None, t_index], + precip_models_cascade_temp, + noise_extrapolated_decomp_done[None, t_index], + ), + axis=0, + ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] + means_stacked = np.concatenate( + (mu_extrapolation[None, :], mu_models_temp), axis=0 + ) + sigmas_stacked = np.concatenate( + (sigma_extrapolation[None, :], sigma_models_temp), + axis=0, + ) + else: + cascade_stack_all_components = np.concatenate( + ( + precip_forecast_extrapolated_decomp_done[None, t_index], + precip_models_cascade_temp[None, j], + noise_extrapolated_decomp_done[None, t_index], + ), + axis=0, + ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] + means_stacked = np.concatenate( + (mu_extrapolation[None, :], mu_models_temp[None, j]), + axis=0, + ) + sigmas_stacked = np.concatenate( + ( + sigma_extrapolation[None, :], + sigma_models_temp[None, j], + ), + axis=0, + ) + + # First determine the blending weights if method is spn. The + # weights for method bps have already been determined. + if weights_method == "spn": + weights = np.zeros( + ( + cascade_stack_all_components.shape[0], + n_cascade_levels, + ) + ) + for i in range(n_cascade_levels): + # Determine the normalized covariance matrix (containing) + # the cross-correlations between the models + cascade_stack_all_components_temp = np.stack( + [ + cascade_stack_all_components[n_model, i, :, :].flatten() + for n_model in range(cascade_stack_all_components.shape[0] - 1) + ] + ) # -1 to exclude the noise component + covariance_nwp_models = np.ma.corrcoef( + np.ma.masked_invalid(cascade_stack_all_components_temp) + ) + # Determine the weights for this cascade level + weights[:, i] = calculate_weights_spn( + correlations=rho_fc[:, i], + covariance=covariance_nwp_models, + ) + + # Blend the extrapolation, (NWP) model(s) and noise cascades + precip_forecast_blended = blending.utils.blend_cascades( + cascades_norm=cascade_stack_all_components, weights=weights + ) + + # Also blend the cascade without the extrapolation component + precip_forecast_blended_mod_only = blending.utils.blend_cascades( + cascades_norm=cascade_stack_all_components[1:, :], + weights=weights_model_only, + ) + + # Blend the means and standard deviations + # Input is array of shape [number_components, scale_level, ...] + means_blended, sigmas_blended = blend_means_sigmas( + means=means_stacked, sigmas=sigmas_stacked, weights=weights + ) + # Also blend the means and sigmas for the cascade without extrapolation + ( + means_blended_mod_only, + sigmas_blended_mod_only, + ) = blend_means_sigmas( + means=means_stacked[1:, :], + sigmas=sigmas_stacked[1:, :], + weights=weights_model_only, + ) + + def __recompose_cascade_to_rainfall_field(self): + # 8.6 Recompose the cascade to a precipitation field + # (The function first normalizes the blended cascade, precip_forecast_blended + # again) + precip_forecast_recomposed = blending.utils.recompose_cascade( + combined_cascade=precip_forecast_blended, + combined_mean=means_blended, + combined_sigma=sigmas_blended, + ) + # The recomposed cascade without the extrapolation (for NaN filling + # outside the radar domain) + precip_forecast_recomposed_mod_only = blending.utils.recompose_cascade( + combined_cascade=precip_forecast_blended_mod_only, + combined_mean=means_blended_mod_only, + combined_sigma=sigmas_blended_mod_only, + ) + if domain == "spectral": + # TODO: Check this! (Only tested with domain == 'spatial') + precip_forecast_recomposed = fft_objs[j].irfft2(precip_forecast_recomposed) + precip_forecast_recomposed_mod_only = fft_objs[j].irfft2( + precip_forecast_recomposed_mod_only + ) + + def __post_process_output(self, final_blended_forecast_single_member): + # 8.7 Post-processing steps - use the mask and fill no data with + # the blended NWP forecast. Probability matching following + # Lagrangian blended probability matching which uses the + # latest extrapolated radar rainfall field blended with the + # nwp model(s) rainfall forecast fields as 'benchmark'. + + # 8.7.1 first blend the extrapolated rainfall field (the field + # that is only used for post-processing steps) with the NWP + # rainfall forecast for this time step using the weights + # at scale level 2. + weights_probability_matching = weights[:-1, 1] # Weights without noise, level 2 + weights_probability_matching_normalized = weights_probability_matching / np.sum( + weights_probability_matching + ) + # And the weights for outside the radar domain + weights_probability_matching_mod_only = weights_model_only[ + :-1, 1 + ] # Weights without noise, level 2 + weights_probability_matching_normalized_mod_only = ( + weights_probability_matching_mod_only + / np.sum(weights_probability_matching_mod_only) + ) + # Stack the fields + if blend_nwp_members: + precip_forecast_probability_matching_final = np.concatenate( + ( + precip_forecast_extrapolated_probability_matching[None, t_index], + precip_models_temp, + ), + axis=0, + ) + else: + precip_forecast_probability_matching_final = np.concatenate( + ( + precip_forecast_extrapolated_probability_matching[None, t_index], + precip_models_temp[None, j], + ), + axis=0, + ) + # Blend it + precip_forecast_probability_matching_blended = np.sum( + weights_probability_matching_normalized.reshape( + weights_probability_matching_normalized.shape[0], 1, 1 + ) + * precip_forecast_probability_matching_final, + axis=0, + ) + if blend_nwp_members: + precip_forecast_probability_matching_blended_mod_only = np.sum( + weights_probability_matching_normalized_mod_only.reshape( + weights_probability_matching_normalized_mod_only.shape[0], + 1, + 1, + ) + * precip_models_temp, + axis=0, + ) + else: + precip_forecast_probability_matching_blended_mod_only = precip_models_temp[ + j + ] + + # The extrapolation components are NaN outside the advected + # radar domain. This results in NaN values in the blended + # forecast outside the radar domain. Therefore, fill these + # areas with the "..._mod_only" blended forecasts, consisting + # of the NWP and noise components. + + nan_indices = np.isnan(precip_forecast_recomposed) + if smooth_radar_mask_range != 0: + # Compute the smooth dilated mask + new_mask = blending.utils.compute_smooth_dilated_mask( + nan_indices, + max_padding_size_in_px=smooth_radar_mask_range, + ) + + # Ensure mask values are between 0 and 1 + mask_model = np.clip(new_mask, 0, 1) + mask_radar = np.clip(1 - new_mask, 0, 1) + + # Handle NaNs in precip_forecast_new and precip_forecast_new_mod_only by setting NaNs to 0 in the blending step + precip_forecast_recomposed_mod_only_no_nan = np.nan_to_num( + precip_forecast_recomposed_mod_only, nan=0 + ) + precip_forecast_recomposed_no_nan = np.nan_to_num( + precip_forecast_recomposed, nan=0 + ) + + # Perform the blending of radar and model inside the radar domain using a weighted combination + precip_forecast_recomposed = np.nansum( + [ + mask_model * precip_forecast_recomposed_mod_only_no_nan, + mask_radar * precip_forecast_recomposed_no_nan, + ], + axis=0, + ) + + nan_indices = np.isnan(precip_forecast_probability_matching_blended) + precip_forecast_probability_matching_blended = np.nansum( + [ + precip_forecast_probability_matching_blended * mask_radar, + precip_forecast_probability_matching_blended_mod_only * mask_model, + ], + axis=0, + ) + else: + precip_forecast_recomposed[nan_indices] = ( + precip_forecast_recomposed_mod_only[nan_indices] + ) + nan_indices = np.isnan(precip_forecast_probability_matching_blended) + precip_forecast_probability_matching_blended[nan_indices] = ( + precip_forecast_probability_matching_blended_mod_only[nan_indices] + ) + + # Finally, fill the remaining nan values, if present, with + # the minimum value in the forecast + nan_indices = np.isnan(precip_forecast_recomposed) + precip_forecast_recomposed[nan_indices] = np.nanmin(precip_forecast_recomposed) + nan_indices = np.isnan(precip_forecast_probability_matching_blended) + precip_forecast_probability_matching_blended[nan_indices] = np.nanmin( + precip_forecast_probability_matching_blended + ) + + # 8.7.2. Apply the masking and prob. matching + if mask_method is not None: + # apply the precipitation mask to prevent generation of new + # precipitation into areas where it was not originally + # observed + precip_forecast_min_value = precip_forecast_recomposed.min() + if mask_method == "incremental": + # The incremental mask is slightly different from + # the implementation in the non-blended steps.py, as + # it is not based on the last forecast, but instead + # on R_pm_blended. Therefore, the buffer does not + # increase over time. + # Get the mask for this forecast + precip_field_mask = ( + precip_forecast_probability_matching_blended >= precip_thr + ) + # Buffer the mask + precip_field_mask = _compute_incremental_mask( + precip_field_mask, struct, mask_rim + ) + # Get the final mask + precip_forecast_recomposed = ( + precip_forecast_min_value + + (precip_forecast_recomposed - precip_forecast_min_value) + * precip_field_mask + ) + precip_field_mask_temp = ( + precip_forecast_recomposed > precip_forecast_min_value + ) + elif mask_method == "obs": + # The mask equals the most recent benchmark + # rainfall field + precip_field_mask_temp = ( + precip_forecast_probability_matching_blended >= precip_thr + ) + + # Set to min value outside of mask + precip_forecast_recomposed[~precip_field_mask_temp] = ( + precip_forecast_min_value + ) + + # If probmatching_method is not None, resample the distribution from + # both the extrapolation cascade and the model (NWP) cascade and use + # that for the probability matching. + if probmatching_method is not None and resample_distribution: + arr1 = precip_forecast_extrapolated_probability_matching[t_index] + arr2 = precip_models_temp[j] + # resample weights based on cascade level 2. + # Areas where one of the fields is nan are not included. + precip_forecast_probability_matching_resampled = ( + probmatching.resample_distributions( + first_array=arr1, + second_array=arr2, + probability_first_array=weights_probability_matching_normalized[0], + ) + ) + else: + precip_forecast_probability_matching_resampled = ( + precip_forecast_probability_matching_blended.copy() + ) + + if probmatching_method == "cdf": + # nan indices in the extrapolation nowcast + nan_indices = np.isnan( + precip_forecast_extrapolated_probability_matching[t_index] + ) + # Adjust the CDF of the forecast to match the resampled distribution combined from + # extrapolation and model fields. + # Rainfall outside the pure extrapolation domain is not taken into account. + if np.any(np.isfinite(precip_forecast_recomposed)): + precip_forecast_recomposed = probmatching.nonparam_match_empirical_cdf( + precip_forecast_recomposed, + precip_forecast_probability_matching_resampled, + nan_indices, + ) + precip_forecast_probability_matching_resampled = None + elif probmatching_method == "mean": + # Use R_pm_blended as benchmark field and + mean_probabiltity_matching_forecast = np.mean( + precip_forecast_probability_matching_resampled[ + precip_forecast_probability_matching_resampled >= precip_thr + ] + ) + no_rain_mask = precip_forecast_recomposed >= precip_thr + mean_precip_forecast = np.mean(precip_forecast_recomposed[no_rain_mask]) + precip_forecast_recomposed[no_rain_mask] = ( + precip_forecast_recomposed[no_rain_mask] + - mean_precip_forecast + + mean_probabiltity_matching_forecast + ) + precip_forecast_probability_matching_resampled = None + + final_blended_forecast_single_member.append(precip_forecast_recomposed) + def __measure_time(self, label, start_time): """ Measure and print the time taken for a specific part of the process. From f3879812de7a79c7c93ecd52f12b5215af038d76 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 5 Dec 2024 10:37:16 +0100 Subject: [PATCH 11/33] Everything is refactored, no test ran as of yet --- pysteps/blending/steps.py | 395 +++++++++++++++++++++++--------------- pysteps/nowcasts/steps.py | 1 - 2 files changed, 242 insertions(+), 154 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 15ac25b9..e866d45c 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -180,11 +180,11 @@ class StepsBlendingState: subtimesteps: Any = None is_nowcast_time_step: bool = None # Variables to save data over (sub)time steps - precip_models_cascades_temp: Any = None - precip_models_temp: Any = None - mean_models_temp: Any = None - std_models_temp: Any = None - velocity_models_temp: Any = None + precip_models_cascades_timestep: Any = None + precip_models_timestep: Any = None + mean_models_timestep: Any = None + std_models_timestep: Any = None + velocity_models_timestep: Any = None n_model_indices: Optional[np.ndarray] = None # NWP model indices rho_forecast: Any = None weights: Any = None @@ -194,6 +194,16 @@ class StepsBlendingState: precip_forecast_extrapolated_probability_matching: Any = None precip_forecast_prev_subtimestep: Any = None noise_prev_subtimestep: Any = None + final_blended_forecast_single_member: Any = None + means_blended: Any = None + sigmas_blended: Any = None + means_blended_mod_only: Any = None + sigmas_blended_mod_only: Any = None + precip_forecast_blended: Any = None + precip_forecast_blended_mod_only: Any = None + precip_forecast_recomposed: Any = None + precip_forecast_recomposed_mod_only: Any = None + t_index: Any = None class StepsBlendingNowcaster: @@ -295,7 +305,8 @@ def __blended_nowcast_main(self): # 8. Start the forecasting loop ### # Isolate the last time slice of observed precipitation - precip = self.__precip[-1, :, :] + # TODO: This precip was "precip = self.__precip[-1, :, :]", changed to self.__precip = self.__precip[-1, :, :]. Might need to chage again and user local variable precip in all following functions + self.__precip = self.__precip[-1, :, :] print("Starting blended nowcast computation.") if self.__config.measure_time: @@ -335,14 +346,17 @@ def worker(j): t, j ) # 8.5 Blend the cascades - final_blended_forecast_single_member = [] + self.__state.final_blended_forecast_single_member = [] for t_sub in self.__state.subtimesteps: # TODO: does it make sense to use sub time steps - check if it works? if t_sub > 0: - self.__blend_cascades() - self.__recompose_cascade_to_rainfall_field() - self.__post_process_output(final_blended_forecast_single_member) - precip_forecast_workers[j] = final_blended_forecast_single_member + self.__blend_cascades(t_sub, j) + self.__recompose_cascade_to_rainfall_field(j) + # TODO: could be I need to return and ave final_blended_forecast_single_member + self.__post_process_output(j) + precip_forecast_workers[j] = ( + self.__state.final_blended_forecast_single_member + ) result = [] @@ -706,6 +720,7 @@ def __prepare_radar_and_NWP_fields(self): # 2.2 If necessary, recompose (NWP) model forecasts self.__state.precip_models_cascades = None + # TODO: This type of check needs to be changed when going to xarray if self.__precip_models.ndim != 4: self.__state.precip_models_cascades = self.__precip_models self.__precip_models = _compute_cascade_recomposition_nwp( @@ -799,7 +814,7 @@ def __prepare_nowcast_for_zero_radar(self): self.__state.precip_models_cascades[j, t]["cascade_levels"] ) continue - precip_models_cascade_temp = self.__params.decomposition_method( + precip_models_cascade_timestep = self.__params.decomposition_method( self.__precip_models[j, t, :, :], bp_filter=self.__params.bandpass_filter, fft_method=self.__params.fft, @@ -810,7 +825,7 @@ def __prepare_nowcast_for_zero_radar(self): )["cascade_levels"] self.__state.precip_cascades[ ~np.isfinite(self.__state.precip_cascades) - ] = np.nanmin(precip_models_cascade_temp) + ] = np.nanmin(precip_models_cascade_timestep) done = True break @@ -903,7 +918,7 @@ def __initialize_noise(self): else: self.__params.noise_std_coeffs = np.ones(self.__config.n_cascade_levels) - if self.__params.noise_stddev_adj is not None: + if self.__config.noise_stddev_adj is not None: print(f"noise std. dev. coeffs: {self.__params.noise_std_coeffs}") else: @@ -1189,13 +1204,13 @@ def __decompose_nwp_if_needed_and_fill_nans_in_nwp(self, t): list(self.__precip_models[:, t, :, :]), ) - self.__state.precip_models_cascades_temp = np.array( + self.__state.precip_models_cascades_timestep = np.array( [decomp["cascade_levels"] for decomp in decomp_precip_models] ) - self.__state.mean_models_temp = np.array( + self.__state.mean_models_timestep = np.array( [decomp["means"] for decomp in decomp_precip_models] ) - self.__state.std_models_temp = np.array( + self.__state.std_models_timestep = np.array( [decomp["stds"] for decomp in decomp_precip_models] ) @@ -1205,21 +1220,25 @@ def __decompose_nwp_if_needed_and_fill_nans_in_nwp(self, t): """Ensure that the NWP cascade and fields do no contain any nans or infinite number""" # Fill nans and infinite numbers with the minimum value present in precip - self.__state.precip_models_temp = self.__precip_models[:, t, :, :].astype( + self.__state.precip_models_timestep = self.__precip_models[:, t, :, :].astype( np.float64, copy=False ) # (corresponding to zero rainfall in the radar observations) min_cascade = np.nanmin(self.__state.precip_cascades) min_precip = np.nanmin(self.__precip) - self.__state.precip_models_cascades_temp[ - ~np.isfinite(self.__state.precip_models_cascades_temp) + self.__state.precip_models_cascades_timestep[ + ~np.isfinite(self.__state.precip_models_cascades_timestep) ] = min_cascade - self.__state.precip_models_temp[ - ~np.isfinite(self.__state.precip_models_temp) + self.__state.precip_models_timestep[ + ~np.isfinite(self.__state.precip_models_timestep) ] = min_precip # Also set any nans or infs in the mean and sigma of the cascade to # respectively 0.0 and 1.0 - self.__state.mean_models_temp[~np.isfinite(self.__state.mean_models_temp)] = 0.0 - self.__state.std_models_temp[~np.isfinite(self.__state.std_models_temp)] = 0.0 + self.__state.mean_models_timestep[ + ~np.isfinite(self.__state.mean_models_timestep) + ] = 0.0 + self.__state.std_models_timestep[ + ~np.isfinite(self.__state.std_models_timestep) + ] = 0.0 def __find_nowcast_NWP_combination(self, t): # 8.1.1 Before calling the worker for the forecast loop, determine which (NWP) @@ -1231,11 +1250,11 @@ def __find_nowcast_NWP_combination(self, t): With the way it is implemented at this moment: n_ens_members of the output equals the maximum number of (ensemble) members in the input (either the nowcasts or NWP). """ - self.__state.velocity_models_temp = self.__velocity_models[ + self.__state.velocity_models_timestep = self.__velocity_models[ :, t, :, :, : ].astype(np.float64, copy=False) # Make sure the number of model members is not larger than or equal to n_ens_members - n_model_members = self.__state.precip_models_cascades_temp.shape[0] + n_model_members = self.__state.precip_models_cascades_timestep.shape[0] if n_model_members > self.__config.n_ens_members: raise ValueError( "The number of NWP model members is larger than the given number of ensemble members. n_model_members <= n_ens_members." @@ -1269,23 +1288,23 @@ def __find_nowcast_NWP_combination(self, t): # member 5, etc.), until 10 is reached. if n_ens_members_min != n_ens_members_max: if n_model_members == 1: - self.__state.precip_models_cascades_temp = np.repeat( - self.__state.precip_models_cascades_temp, + self.__state.precip_models_cascades_timestep = np.repeat( + self.__state.precip_models_cascades_timestep, n_ens_members_max, axis=0, ) - self.__state.mean_models_temp = np.repeat( - self.__state.mean_models_temp, n_ens_members_max, axis=0 + self.__state.mean_models_timestep = np.repeat( + self.__state.mean_models_timestep, n_ens_members_max, axis=0 ) - self.__state.std_models_temp = np.repeat( - self.__state.std_models_temp, n_ens_members_max, axis=0 + self.__state.std_models_timestep = np.repeat( + self.__state.std_models_timestep, n_ens_members_max, axis=0 ) - self.__state.velocity_models_temp = np.repeat( - self.__state.velocity_models_temp, n_ens_members_max, axis=0 + self.__state.velocity_models_timestep = np.repeat( + self.__state.velocity_models_timestep, n_ens_members_max, axis=0 ) # For the prob. matching - self.__state.precip_models_temp = np.repeat( - self.__state.precip_models_temp, n_ens_members_max, axis=0 + self.__state.precip_models_timestep = np.repeat( + self.__state.precip_models_timestep, n_ens_members_max, axis=0 ) # Finally, for the model indices self.__state.n_model_indices = np.repeat( @@ -1298,21 +1317,23 @@ def __find_nowcast_NWP_combination(self, t): for i in range(n_ens_members_min) ] if n_model_members == n_ens_members_min: - self.__state.precip_models_cascades_temp = np.repeat( - self.__state.precip_models_cascades_temp, repeats, axis=0 + self.__state.precip_models_cascades_timestep = np.repeat( + self.__state.precip_models_cascades_timestep, + repeats, + axis=0, ) - self.__state.mean_models_temp = np.repeat( - self.__state.mean_models_temp, repeats, axis=0 + self.__state.mean_models_timestep = np.repeat( + self.__state.mean_models_timestep, repeats, axis=0 ) - self.__state.std_models_temp = np.repeat( - self.__state.std_models_temp, repeats, axis=0 + self.__state.std_models_timestep = np.repeat( + self.__state.std_models_timestep, repeats, axis=0 ) - self.__state.velocity_models_temp = np.repeat( - self.__state.velocity_models_temp, repeats, axis=0 + self.__state.velocity_models_timestep = np.repeat( + self.__state.velocity_models_timestep, repeats, axis=0 ) # For the prob. matching - self.__state.precip_models_temp = np.repeat( - self.__state.precip_models_temp, repeats, axis=0 + self.__state.precip_models_timestep = np.repeat( + self.__state.precip_models_timestep, repeats, axis=0 ) # Finally, for the model indices self.__state.n_model_indices = np.repeat( @@ -1325,7 +1346,7 @@ def __find_nowcast_NWP_combination(self, t): if self.__params.zero_precip_radar: # Use the velocity from velocity_models and take the average over # n_models (axis=0) - self.__velocity = np.mean(self.__state.velocity_models_temp, axis=0) + self.__velocity = np.mean(self.__state.velocity_models_timestep, axis=0) def __determine_skill_for_current_timestep(self, t): if t == 0: @@ -1334,18 +1355,22 @@ def __determine_skill_for_current_timestep(self, t): self.__params.rho_nwp_models = [ blending.skill_scores.spatial_correlation( obs=self.__state.precip_cascades[0, :, -1, :, :].copy(), - mod=self.__state.precip_models_cascades_temp[ + mod=self.__state.precip_models_cascades_timestep[ n_model, :, :, : ].copy(), domain_mask=self.__params.domain_mask, ) - for n_model in range(self.__state.precip_models_cascades_temp.shape[0]) + for n_model in range( + self.__state.precip_models_cascades_timestep.shape[0] + ) ] self.__params.rho_nwp_models = np.stack(self.__params.rho_nwp_models) # Ensure that the model skill decreases with increasing scale level. - for n_model in range(self.__state.precip_models_cascades_temp.shape[0]): - for i in range(1, self.__state.precip_models_cascades_temp.shape[1]): + for n_model in range(self.__state.precip_models_cascades_timestep.shape[0]): + for i in range( + 1, self.__state.precip_models_cascades_timestep.shape[1] + ): if ( self.__params.rho_nwp_models[n_model, i] > self.__params.rho_nwp_models[n_model, i - 1] @@ -1398,7 +1423,7 @@ def __determine_skill_for_next_timestep(self, t, j): lt=(t * int(self.__config.timestep)), correlations=self.__params.rho_nwp_models[j], outdir_path=self.__config.outdir_path_skill, - n_model=self.__params.n_model_indices[j], + n_model=self.__state.n_model_indices[j], skill_kwargs=self.__config.clim_kwargs, ) # Concatenate rho_extrap_cascade and rho_nwp @@ -1430,11 +1455,11 @@ def __determine_weights_per_component(self): # determined after the extrapolation step in this method. if ( self.__config.blend_nwp_members - and self.__state.precip_models_cascades_temp.shape[0] > 1 + and self.__state.precip_models_cascades_timestep.shape[0] > 1 ): self.__state.weights_model_only = np.zeros( ( - self.__state.precip_models_cascades_temp.shape[0] + 1, + self.__state.precip_models_cascades_timestep.shape[0] + 1, self.__config.n_cascade_levels, ) ) @@ -1444,11 +1469,13 @@ def __determine_weights_per_component(self): covariance_nwp_models = np.corrcoef( np.stack( [ - self.__state.precip_models_cascades_temp[ + self.__state.precip_models_cascades_timestep[ n_model, i, :, : ].flatten() for n_model in range( - self.__state.precip_models_cascades_temp.shape[0] + self.__state.precip_models_cascades_timestep.shape[ + 0 + ] ) ] ) @@ -1618,12 +1645,12 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( velocity_stack_all = np.concatenate( ( velocity_perturbations_extrapolation[None, :, :, :], - self.__state.velocity_models_temp, + self.__state.velocity_models_timestep, ), axis=0, ) else: - velocity_models = self.__state.velocity_models_temp[j] + velocity_models = self.__state.velocity_models_timestep[j] velocity_stack_all = np.concatenate( ( velocity_perturbations_extrapolation[None, :, :, :], @@ -1825,12 +1852,12 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( velocity_stack_all = np.concatenate( ( velocity_perturbations_extrapolation[None, :, :, :], - self.__state.velocity_models_temp, + self.__state.velocity_models_timestep, ), axis=0, ) else: - velocity_models = self.__state.velocity_models_temp[j] + velocity_models = self.__state.velocity_models_timestep[j] velocity_stack_all = np.concatenate( ( velocity_perturbations_extrapolation[None, :, :, :], @@ -1899,57 +1926,82 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( ] self.__state.noise_prev_subtimestep[j] = self.__state.precip_noise_cascades[j] - def __blend_cascades(self): - t_index = np.where(np.array(subtimesteps) == t_sub)[0][0] + def __blend_cascades(self, t_sub, j): + self.__state.t_index = np.where(np.array(self.__state.subtimesteps) == t_sub)[ + 0 + ][0] # First concatenate the cascades and the means and sigmas # precip_models = [n_models,timesteps,n_cascade_levels,m,n] - if blend_nwp_members: + if self.__config.blend_nwp_members: cascade_stack_all_components = np.concatenate( ( - precip_forecast_extrapolated_decomp_done[None, t_index], - precip_models_cascade_temp, - noise_extrapolated_decomp_done[None, t_index], + self.__state.precip_forecast_extrapolated_decomp_done[ + None, self.__state.t_index + ], + self.__state.precip_models_cascades_timestep, + self.__state.noise_extrapolated_decomp_done[ + None, self.__state.t_index + ], ), axis=0, ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] means_stacked = np.concatenate( - (mu_extrapolation[None, :], mu_models_temp), axis=0 + ( + self.__state.mean_extrapolation[None, :], + self.__state.mean_models_timestep, + ), + axis=0, ) sigmas_stacked = np.concatenate( - (sigma_extrapolation[None, :], sigma_models_temp), + ( + self.__state.std_extrapolation[None, :], + self.__state.std_models_timestep, + ), axis=0, ) else: cascade_stack_all_components = np.concatenate( ( - precip_forecast_extrapolated_decomp_done[None, t_index], - precip_models_cascade_temp[None, j], - noise_extrapolated_decomp_done[None, t_index], + self.__state.precip_forecast_extrapolated_decomp_done[ + None, self.__state.t_index + ], + self.__state.precip_models_cascades_timestep[None, j], + self.__state.noise_extrapolated_decomp_done[ + None, self.__state.t_index + ], ), axis=0, ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] means_stacked = np.concatenate( - (mu_extrapolation[None, :], mu_models_temp[None, j]), + ( + self.__state.mean_extrapolation[None, :], + self.__state.mean_models_timestep[None, j], + ), axis=0, ) sigmas_stacked = np.concatenate( ( - sigma_extrapolation[None, :], - sigma_models_temp[None, j], + self.__state.std_extrapolation[None, :], + self.__state.std_models_timestep[None, j], ), axis=0, ) # First determine the blending weights if method is spn. The # weights for method bps have already been determined. - if weights_method == "spn": - weights = np.zeros( + # TODO: no other weight method is possible, should we not al least give a user warning if a different weight + # method is given? Or does this mean that in all other circumstances the weights + # have been calculated in a different way? + + # TODO: changed weights to self.__state.weights + if self.__config.weights_method == "spn": + self.__state.weights = np.zeros( ( cascade_stack_all_components.shape[0], - n_cascade_levels, + self.__config.n_cascade_levels, ) ) - for i in range(n_cascade_levels): + for i in range(self.__config.n_cascade_levels): # Determine the normalized covariance matrix (containing) # the cross-correlations between the models cascade_stack_all_components_temp = np.stack( @@ -1962,61 +2014,68 @@ def __blend_cascades(self): np.ma.masked_invalid(cascade_stack_all_components_temp) ) # Determine the weights for this cascade level - weights[:, i] = calculate_weights_spn( - correlations=rho_fc[:, i], + self.__state.weights[:, i] = calculate_weights_spn( + correlations=self.__state.rho_forecast[:, i], covariance=covariance_nwp_models, ) # Blend the extrapolation, (NWP) model(s) and noise cascades - precip_forecast_blended = blending.utils.blend_cascades( - cascades_norm=cascade_stack_all_components, weights=weights + self.__state.precip_forecast_blended = blending.utils.blend_cascades( + cascades_norm=cascade_stack_all_components, weights=self.__state.weights ) # Also blend the cascade without the extrapolation component - precip_forecast_blended_mod_only = blending.utils.blend_cascades( + self.__state.precip_forecast_blended_mod_only = blending.utils.blend_cascades( cascades_norm=cascade_stack_all_components[1:, :], - weights=weights_model_only, + weights=self.__state.weights_model_only, ) # Blend the means and standard deviations # Input is array of shape [number_components, scale_level, ...] - means_blended, sigmas_blended = blend_means_sigmas( - means=means_stacked, sigmas=sigmas_stacked, weights=weights + self.__state.means_blended, self.__state.sigmas_blended = blend_means_sigmas( + means=means_stacked, sigmas=sigmas_stacked, weights=self.__state.weights ) # Also blend the means and sigmas for the cascade without extrapolation + ( - means_blended_mod_only, - sigmas_blended_mod_only, + self.__state.means_blended_mod_only, + self.__state.sigmas_blended_mod_only, ) = blend_means_sigmas( means=means_stacked[1:, :], sigmas=sigmas_stacked[1:, :], - weights=weights_model_only, + weights=self.__state.weights_model_only, ) - def __recompose_cascade_to_rainfall_field(self): + def __recompose_cascade_to_rainfall_field(self, j): # 8.6 Recompose the cascade to a precipitation field # (The function first normalizes the blended cascade, precip_forecast_blended # again) - precip_forecast_recomposed = blending.utils.recompose_cascade( - combined_cascade=precip_forecast_blended, - combined_mean=means_blended, - combined_sigma=sigmas_blended, + self.__state.precip_forecast_recomposed = blending.utils.recompose_cascade( + combined_cascade=self.__state.precip_forecast_blended, + combined_mean=self.__state.means_blended, + combined_sigma=self.__state.sigmas_blended, ) # The recomposed cascade without the extrapolation (for NaN filling # outside the radar domain) - precip_forecast_recomposed_mod_only = blending.utils.recompose_cascade( - combined_cascade=precip_forecast_blended_mod_only, - combined_mean=means_blended_mod_only, - combined_sigma=sigmas_blended_mod_only, + self.__state.precip_forecast_recomposed_mod_only = ( + blending.utils.recompose_cascade( + combined_cascade=self.__state.precip_forecast_blended_mod_only, + combined_mean=self.__state.means_blended_mod_only, + combined_sigma=self.__state.sigmas_blended_mod_only, + ) ) - if domain == "spectral": + if self.__config.domain == "spectral": # TODO: Check this! (Only tested with domain == 'spatial') - precip_forecast_recomposed = fft_objs[j].irfft2(precip_forecast_recomposed) - precip_forecast_recomposed_mod_only = fft_objs[j].irfft2( - precip_forecast_recomposed_mod_only + + # TODO: what needs to happen with above TODO? + self.__state.precip_forecast_recomposed = self.__state.fft_objs[j].irfft2( + self.__state.precip_forecast_recomposed ) + self.__state.precip_forecast_recomposed_mod_only = self.__state.fft_objs[ + j + ].irfft2(self.__state.precip_forecast_recomposed_mod_only) - def __post_process_output(self, final_blended_forecast_single_member): + def __post_process_output(self, j): # 8.7 Post-processing steps - use the mask and fill no data with # the blended NWP forecast. Probability matching following # Lagrangian blended probability matching which uses the @@ -2027,12 +2086,14 @@ def __post_process_output(self, final_blended_forecast_single_member): # that is only used for post-processing steps) with the NWP # rainfall forecast for this time step using the weights # at scale level 2. - weights_probability_matching = weights[:-1, 1] # Weights without noise, level 2 + weights_probability_matching = self.__state.weights[ + :-1, 1 + ] # Weights without noise, level 2 weights_probability_matching_normalized = weights_probability_matching / np.sum( weights_probability_matching ) # And the weights for outside the radar domain - weights_probability_matching_mod_only = weights_model_only[ + weights_probability_matching_mod_only = self.__state.weights_model_only[ :-1, 1 ] # Weights without noise, level 2 weights_probability_matching_normalized_mod_only = ( @@ -2040,19 +2101,23 @@ def __post_process_output(self, final_blended_forecast_single_member): / np.sum(weights_probability_matching_mod_only) ) # Stack the fields - if blend_nwp_members: + if self.__config.blend_nwp_members: precip_forecast_probability_matching_final = np.concatenate( ( - precip_forecast_extrapolated_probability_matching[None, t_index], - precip_models_temp, + self.__state.precip_forecast_extrapolated_probability_matching[ + None, self.__state.t_index + ], + self.__state.precip_models_timestep, ), axis=0, ) else: precip_forecast_probability_matching_final = np.concatenate( ( - precip_forecast_extrapolated_probability_matching[None, t_index], - precip_models_temp[None, j], + self.__state.precip_forecast_extrapolated_probability_matching[ + None, self.__state.t_index + ], + self.__state.precip_models_timestep[None, j], ), axis=0, ) @@ -2064,20 +2129,20 @@ def __post_process_output(self, final_blended_forecast_single_member): * precip_forecast_probability_matching_final, axis=0, ) - if blend_nwp_members: + if self.__config.blend_nwp_members: precip_forecast_probability_matching_blended_mod_only = np.sum( weights_probability_matching_normalized_mod_only.reshape( weights_probability_matching_normalized_mod_only.shape[0], 1, 1, ) - * precip_models_temp, + * self.__state.precip_models_timestep, axis=0, ) else: - precip_forecast_probability_matching_blended_mod_only = precip_models_temp[ - j - ] + precip_forecast_probability_matching_blended_mod_only = ( + self.__state.precip_models_timestep[j] + ) # The extrapolation components are NaN outside the advected # radar domain. This results in NaN values in the blended @@ -2085,12 +2150,12 @@ def __post_process_output(self, final_blended_forecast_single_member): # areas with the "..._mod_only" blended forecasts, consisting # of the NWP and noise components. - nan_indices = np.isnan(precip_forecast_recomposed) - if smooth_radar_mask_range != 0: + nan_indices = np.isnan(self.__state.precip_forecast_recomposed) + if self.__config.smooth_radar_mask_range != 0: # Compute the smooth dilated mask new_mask = blending.utils.compute_smooth_dilated_mask( nan_indices, - max_padding_size_in_px=smooth_radar_mask_range, + max_padding_size_in_px=self.__config.smooth_radar_mask_range, ) # Ensure mask values are between 0 and 1 @@ -2099,10 +2164,10 @@ def __post_process_output(self, final_blended_forecast_single_member): # Handle NaNs in precip_forecast_new and precip_forecast_new_mod_only by setting NaNs to 0 in the blending step precip_forecast_recomposed_mod_only_no_nan = np.nan_to_num( - precip_forecast_recomposed_mod_only, nan=0 + self.__state.precip_forecast_recomposed_mod_only, nan=0 ) precip_forecast_recomposed_no_nan = np.nan_to_num( - precip_forecast_recomposed, nan=0 + self.__state.precip_forecast_recomposed, nan=0 ) # Perform the blending of radar and model inside the radar domain using a weighted combination @@ -2114,7 +2179,6 @@ def __post_process_output(self, final_blended_forecast_single_member): axis=0, ) - nan_indices = np.isnan(precip_forecast_probability_matching_blended) precip_forecast_probability_matching_blended = np.nansum( [ precip_forecast_probability_matching_blended * mask_radar, @@ -2123,8 +2187,8 @@ def __post_process_output(self, final_blended_forecast_single_member): axis=0, ) else: - precip_forecast_recomposed[nan_indices] = ( - precip_forecast_recomposed_mod_only[nan_indices] + self.__state.precip_forecast_recomposed[nan_indices] = ( + self.__state.precip_forecast_recomposed_mod_only[nan_indices] ) nan_indices = np.isnan(precip_forecast_probability_matching_blended) precip_forecast_probability_matching_blended[nan_indices] = ( @@ -2133,20 +2197,23 @@ def __post_process_output(self, final_blended_forecast_single_member): # Finally, fill the remaining nan values, if present, with # the minimum value in the forecast - nan_indices = np.isnan(precip_forecast_recomposed) - precip_forecast_recomposed[nan_indices] = np.nanmin(precip_forecast_recomposed) + nan_indices = np.isnan(self.__state.precip_forecast_recomposed) + self.__state.precip_forecast_recomposed[nan_indices] = np.nanmin( + self.__state.precip_forecast_recomposed + ) nan_indices = np.isnan(precip_forecast_probability_matching_blended) precip_forecast_probability_matching_blended[nan_indices] = np.nanmin( precip_forecast_probability_matching_blended ) # 8.7.2. Apply the masking and prob. matching - if mask_method is not None: + precip_field_mask_temp = None + if self.__config.mask_method is not None: # apply the precipitation mask to prevent generation of new # precipitation into areas where it was not originally # observed - precip_forecast_min_value = precip_forecast_recomposed.min() - if mask_method == "incremental": + precip_forecast_min_value = self.__state.precip_forecast_recomposed.min() + if self.__config.mask_method == "incremental": # The incremental mask is slightly different from # the implementation in the non-blended steps.py, as # it is not based on the last forecast, but instead @@ -2154,39 +2221,49 @@ def __post_process_output(self, final_blended_forecast_single_member): # increase over time. # Get the mask for this forecast precip_field_mask = ( - precip_forecast_probability_matching_blended >= precip_thr + precip_forecast_probability_matching_blended + >= self.__config.precip_threshold ) # Buffer the mask precip_field_mask = _compute_incremental_mask( - precip_field_mask, struct, mask_rim + precip_field_mask, self.__params.struct, self.__params.mask_rim ) # Get the final mask - precip_forecast_recomposed = ( + self.__state.precip_forecast_recomposed = ( precip_forecast_min_value - + (precip_forecast_recomposed - precip_forecast_min_value) + + ( + self.__state.precip_forecast_recomposed + - precip_forecast_min_value + ) * precip_field_mask ) precip_field_mask_temp = ( - precip_forecast_recomposed > precip_forecast_min_value + self.__state.precip_forecast_recomposed > precip_forecast_min_value ) - elif mask_method == "obs": + elif self.__config.mask_method == "obs": # The mask equals the most recent benchmark # rainfall field precip_field_mask_temp = ( - precip_forecast_probability_matching_blended >= precip_thr + precip_forecast_probability_matching_blended + >= self.__config.precip_threshold ) # Set to min value outside of mask - precip_forecast_recomposed[~precip_field_mask_temp] = ( + self.__state.precip_forecast_recomposed[~precip_field_mask_temp] = ( precip_forecast_min_value ) # If probmatching_method is not None, resample the distribution from # both the extrapolation cascade and the model (NWP) cascade and use # that for the probability matching. - if probmatching_method is not None and resample_distribution: - arr1 = precip_forecast_extrapolated_probability_matching[t_index] - arr2 = precip_models_temp[j] + if ( + self.__config.probmatching_method is not None + and self.__config.resample_distribution + ): + arr1 = self.__state.precip_forecast_extrapolated_probability_matching[ + self.__state.t_index + ] + arr2 = self.__state.precip_models_timestep[j] # resample weights based on cascade level 2. # Areas where one of the fields is nan are not included. precip_forecast_probability_matching_resampled = ( @@ -2201,38 +2278,50 @@ def __post_process_output(self, final_blended_forecast_single_member): precip_forecast_probability_matching_blended.copy() ) - if probmatching_method == "cdf": + if self.__config.probmatching_method == "cdf": # nan indices in the extrapolation nowcast nan_indices = np.isnan( - precip_forecast_extrapolated_probability_matching[t_index] + self.__state.precip_forecast_extrapolated_probability_matching[ + self.__state.t_index + ] ) # Adjust the CDF of the forecast to match the resampled distribution combined from # extrapolation and model fields. # Rainfall outside the pure extrapolation domain is not taken into account. - if np.any(np.isfinite(precip_forecast_recomposed)): - precip_forecast_recomposed = probmatching.nonparam_match_empirical_cdf( - precip_forecast_recomposed, - precip_forecast_probability_matching_resampled, - nan_indices, + if np.any(np.isfinite(self.__state.precip_forecast_recomposed)): + self.__state.precip_forecast_recomposed = ( + probmatching.nonparam_match_empirical_cdf( + self.__state.precip_forecast_recomposed, + precip_forecast_probability_matching_resampled, + nan_indices, + ) ) precip_forecast_probability_matching_resampled = None - elif probmatching_method == "mean": + elif self.__config.probmatching_method == "mean": # Use R_pm_blended as benchmark field and mean_probabiltity_matching_forecast = np.mean( precip_forecast_probability_matching_resampled[ - precip_forecast_probability_matching_resampled >= precip_thr + precip_forecast_probability_matching_resampled + >= self.__config.precip_threshold ] ) - no_rain_mask = precip_forecast_recomposed >= precip_thr - mean_precip_forecast = np.mean(precip_forecast_recomposed[no_rain_mask]) - precip_forecast_recomposed[no_rain_mask] = ( - precip_forecast_recomposed[no_rain_mask] + no_rain_mask = ( + self.__state.precip_forecast_recomposed + >= self.__config.precip_threshold + ) + mean_precip_forecast = np.mean( + self.__state.precip_forecast_recomposed[no_rain_mask] + ) + self.__state.precip_forecast_recomposed[no_rain_mask] = ( + self.__state.precip_forecast_recomposed[no_rain_mask] - mean_precip_forecast + mean_probabiltity_matching_forecast ) precip_forecast_probability_matching_resampled = None - final_blended_forecast_single_member.append(precip_forecast_recomposed) + self.__state.final_blended_forecast_single_member.append( + self.__state.precip_forecast_recomposed + ) def __measure_time(self, label, start_time): """ diff --git a/pysteps/nowcasts/steps.py b/pysteps/nowcasts/steps.py index 6efd9758..b61ee8e7 100644 --- a/pysteps/nowcasts/steps.py +++ b/pysteps/nowcasts/steps.py @@ -19,7 +19,6 @@ from pysteps import extrapolation from pysteps import noise from pysteps import utils -from pysteps.decorators import deprecate_args from pysteps.nowcasts import utils as nowcast_utils from pysteps.postprocessing import probmatching from pysteps.timeseries import autoregression, correlation From 760c1850474c983b7d206585a479170461dc43a7 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 5 Dec 2024 10:46:27 +0100 Subject: [PATCH 12/33] Old forecast function is updated to fit newly refactored code --- pysteps/blending/steps.py | 1562 ++----------------------------------- 1 file changed, 55 insertions(+), 1507 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index e866d45c..0e229a43 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -96,10 +96,10 @@ class StepsBlendingConfig: fft_method: str domain: str outdir_path_skill: str - extrap_kwargs: Dict[str, Any] = field(default_factory=dict) + extrapolation_kwargs: Dict[str, Any] = field(default_factory=dict) filter_kwargs: Dict[str, Any] = field(default_factory=dict) noise_kwargs: Dict[str, Any] = field(default_factory=dict) - vel_pert_kwargs: Dict[str, Any] = field(default_factory=dict) + velocity_perturbation_kwargs: Dict[str, Any] = field(default_factory=dict) clim_kwargs: Dict[str, Any] = field(default_factory=dict) mask_kwargs: Dict[str, Any] = field(default_factory=dict) measure_time: bool = False @@ -313,7 +313,7 @@ def __blended_nowcast_main(self): starttime_mainloop = time.time() # TODO: problem with the config here! This variable changes over time... # extrap_kwargs is in config but by adding info to it, the next run of a blended forecast will have issues! - self.__config.extrap_kwargs["return_displacement"] = True + self.__config.extrapolation_kwargs["return_displacement"] = True self.__state.precip_forecast_prev_subtimestep = deepcopy( self.__state.precip_cascades @@ -465,8 +465,8 @@ def __check_inputs(self): "or a four-dimensional array containing the original (NWP) model forecasts" ) - if self.__config.extrap_kwargs is None: - self.__config.extrap_kwargs = dict() + if self.__config.extrapolation_kwargs is None: + self.__config.extrapolation_kwargs = dict() if self.__config.filter_kwargs is None: self.__config.filter_kwargs = dict() @@ -474,8 +474,8 @@ def __check_inputs(self): if self.__config.noise_kwargs is None: self.__config.noise_kwargs = dict() - if self.__config.vel_pert_kwargs is None: - self.__config.vel_pert_kwargs = dict() + if self.__config.velocity_perturbation_kwargs is None: + self.__config.velocity_perturbation_kwargs = dict() if not self.__params.precip_models_provided_is_cascade: if self.__config.clim_kwargs is None: @@ -587,12 +587,12 @@ def __print_forecast_info(self): print(f"order of the AR(p) model: {self.__config.ar_order}") if self.__config.velocity_perturbation_method == "bps": self.__params.velocity_perturbations_parallel = ( - self.__config.vel_pert_kwargs.get( + self.__config.velocity_perturbation_kwargs.get( "p_par", noise.motion.get_default_params_bps_par() ) ) self.__params.velocity_perturbations_perpendicular = ( - self.__config.vel_pert_kwargs.get( + self.__config.velocity_perturbation_kwargs.get( "p_perp", noise.motion.get_default_params_bps_perp() ) ) @@ -685,7 +685,7 @@ def __prepare_radar_and_NWP_fields(self): self.__config.ar_order, self.__params.xy_coordinates, self.__params.extrapolation_method, - self.__config.extrap_kwargs, + self.__config.extrapolation_kwargs, self.__config.num_workers, ) @@ -1577,9 +1577,9 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( # (or subtimesteps if non-integer time steps are given) # Settings and initialize the output - extrap_kwargs_ = self.__config.extrap_kwargs.copy() - extrap_kwargs_noise = self.__config.extrap_kwargs.copy() - extrap_kwargs_pb = self.__config.extrap_kwargs.copy() + extrap_kwargs_ = self.__config.extrapolation_kwargs.copy() + extrap_kwargs_noise = self.__config.extrapolation_kwargs.copy() + extrap_kwargs_pb = self.__config.extrapolation_kwargs.copy() velocity_perturbations_extrapolation = self.__velocity # The following should be accesseble after this function self.__state.precip_forecast_extrapolated_decomp_done = [] @@ -1692,7 +1692,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( precip_forecast_recomp_subtimestep[self.__params.domain_mask] = np.nan # TODO: problem with the config here! This variable changes over time... # extrap_kwargs is in config but by adding info to it, the next run of a blended forecast will have issues! - self.__config.extrap_kwargs["displacement_prev"] = ( + self.__config.extrapolation_kwargs["displacement_prev"] = ( self.__state.previous_displacement[j] ) ( @@ -1703,7 +1703,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( velocity_blended, [t_diff_prev_subtimestep], allow_nonfinite_values=True, - **self.__config.extrap_kwargs, + **self.__config.extrapolation_kwargs, ) precip_forecast_extrapolated_recomp_subtimestep = ( precip_forecast_extrapolated_recomp_subtimestep_temp[0].copy() @@ -2757,1503 +2757,51 @@ def forecast( turns out to be a warranted functionality. """ - # 0.1 Start with some checks - _check_inputs(precip, precip_models, velocity, velocity_models, timesteps, ar_order) - - if extrap_kwargs is None: - extrap_kwargs = dict() - - if filter_kwargs is None: - filter_kwargs = dict() - - if noise_kwargs is None: - noise_kwargs = dict() - - if vel_pert_kwargs is None: - vel_pert_kwargs = dict() - - if clim_kwargs is None: - # Make sure clim_kwargs at least contains the number of models - clim_kwargs = dict({"n_models": precip_models.shape[0]}) - - if mask_kwargs is None: - mask_kwargs = dict() - - if np.any(~np.isfinite(velocity)): - raise ValueError("velocity contains non-finite values") - - if mask_method not in ["obs", "incremental", None]: - raise ValueError( - "unknown mask method %s: must be 'obs', 'incremental' or None" % mask_method - ) - - if conditional and precip_thr is None: - raise ValueError("conditional=True but precip_thr is not set") - - if mask_method is not None and precip_thr is None: - raise ValueError("mask_method!=None but precip_thr=None") - - if noise_stddev_adj not in ["auto", "fixed", None]: - raise ValueError( - "unknown noise_std_dev_adj method %s: must be 'auto', 'fixed', or None" - % noise_stddev_adj - ) - - if kmperpixel is None: - if vel_pert_method is not None: - raise ValueError("vel_pert_method is set but kmperpixel=None") - if mask_method == "incremental": - raise ValueError("mask_method='incremental' but kmperpixel=None") - - if timestep is None: - if vel_pert_method is not None: - raise ValueError("vel_pert_method is set but timestep=None") - if mask_method == "incremental": - raise ValueError("mask_method='incremental' but timestep=None") - - # 0.2 Log some settings - print("STEPS blending") - print("==============") - print("") - - print("Inputs") - print("------") - print(f"forecast issue time: {issuetime.isoformat()}") - print(f"input dimensions: {precip.shape[1]}x{precip.shape[2]}") - if kmperpixel is not None: - print(f"km/pixel: {kmperpixel}") - if timestep is not None: - print(f"time step: {timestep} minutes") - print("") - - print("NWP and blending inputs") - print("-----------------------") - print(f"number of (NWP) models: {precip_models.shape[0]}") - print(f"blend (NWP) model members: {blend_nwp_members}") - print(f"decompose (NWP) models: {'yes' if precip_models.ndim == 4 else 'no'}") - print("") - - print("Methods") - print("-------") - print(f"extrapolation: {extrap_method}") - print(f"bandpass filter: {bandpass_filter_method}") - print(f"decomposition: {decomp_method}") - print(f"noise generator: {noise_method}") - print(f"noise adjustment: {'yes' if noise_stddev_adj else 'no'}") - print(f"velocity perturbator: {vel_pert_method}") - print(f"blending weights method: {weights_method}") - print(f"conditional statistics: {'yes' if conditional else 'no'}") - print(f"precip. mask method: {mask_method}") - print(f"probability matching: {probmatching_method}") - print(f"FFT method: {fft_method}") - print(f"domain: {domain}") - print("") - - print("Parameters") - print("----------") - if isinstance(timesteps, int): - print(f"number of time steps: {timesteps}") - else: - print(f"time steps: {timesteps}") - print(f"ensemble size: {n_ens_members}") - print(f"parallel threads: {num_workers}") - print(f"number of cascade levels: {n_cascade_levels}") - print(f"order of the AR(p) model: {ar_order}") - if vel_pert_method == "bps": - vp_par = vel_pert_kwargs.get("p_par", noise.motion.get_default_params_bps_par()) - vp_perp = vel_pert_kwargs.get( - "p_perp", noise.motion.get_default_params_bps_perp() - ) - print(f"vel. pert., parallel: {vp_par[0]},{vp_par[1]},{vp_par[2]}") - print(f"vel. pert., perpendicular: {vp_perp[0]},{vp_perp[1]},{vp_perp[2]}") - else: - vp_par, vp_perp = None, None - - if conditional or mask_method is not None: - print(f"precip. intensity threshold: {precip_thr}") - print(f"no-rain fraction threshold for radar: {norain_thr}") - print("") - - # 0.3 Get the methods that will be used - num_ensemble_workers = n_ens_members if num_workers > n_ens_members else num_workers - - if measure_time: - starttime_init = time.time() - - fft = utils.get_method(fft_method, shape=precip.shape[1:], n_threads=num_workers) - - precip_shape = precip.shape[1:] - - # initialize the band-pass filter - filter_method = cascade.get_method(bandpass_filter_method) - bp_filter = filter_method(precip_shape, n_cascade_levels, **filter_kwargs) - - decompositor, recompositor = cascade.get_method(decomp_method) - - extrapolator = extrapolation.get_method(extrap_method) - - x_values, y_values = np.meshgrid( - np.arange(precip.shape[2]), np.arange(precip.shape[1]) - ) - - xy_coords = np.stack([x_values, y_values]) - - precip = precip[-(ar_order + 1) :, :, :].copy() - - # determine the domain mask from non-finite values - domain_mask = np.logical_or.reduce( - [~np.isfinite(precip[i, :]) for i in range(precip.shape[0])] - ) - - # determine the precipitation threshold mask - if conditional: - MASK_thr = np.logical_and.reduce( - [precip[i, :, :] >= precip_thr for i in range(precip.shape[0])] - ) - else: - MASK_thr = None - - # we need to know the zerovalue of precip to replace the mask when decomposing after extrapolation - zerovalue = np.nanmin(precip) - - # 1. Start with the radar rainfall fields. We want the fields in a - # Lagrangian space - precip = _transform_to_lagrangian( - precip, velocity, ar_order, xy_coords, extrapolator, extrap_kwargs, num_workers + blending_config = StepsBlendingConfig( + n_ens_members=n_ens_members, + n_cascade_levels=n_cascade_levels, + precip_threshold=precip_thr, + kmperpixel=kmperpixel, + timestep=timestep, + extrapolation_method=extrap_method, + decomposition_method=decomp_method, + bandpass_filter_method=bandpass_filter_method, + noise_method=noise_method, + noise_stddev_adj=noise_stddev_adj, + ar_order=ar_order, + velocity_perturbation_method=vel_pert_method, + conditional=conditional, + probmatching_method=probmatching_method, + mask_method=mask_method, + seed=seed, + num_workers=num_workers, + fft_method=fft_method, + domain=domain, + extrapolation_kwargs=extrap_kwargs, + filter_kwargs=filter_kwargs, + noise_kwargs=noise_kwargs, + velocity_perturbation_kwargs=vel_pert_kwargs, + mask_kwargs=mask_kwargs, + measure_time=measure_time, + callback=callback, + return_output=return_output, ) - # 2. Perform the cascade decomposition for the input precip fields and - # and, if necessary, for the (NWP) model fields - # 2.1 Compute the cascade decompositions of the input precipitation fields - ( - precip_cascade, - mu_extrapolation, - sigma_extrapolation, - ) = _compute_cascade_decomposition_radar( + # Create an instance of the new class with all the provided arguments + blended_nowcaster = StepsBlendingNowcaster( precip, - ar_order, - n_cascade_levels, - n_ens_members, - MASK_thr, - domain, - bp_filter, - decompositor, - fft, - ) - - # 2.2 If necessary, recompose (NWP) model forecasts - precip_models_cascade = None - - if precip_models.ndim != 4: - precip_models_cascade = precip_models - precip_models = _compute_cascade_recomposition_nwp(precip_models, recompositor) - - # 2.3 Check for zero input fields in the radar and NWP data. - zero_precip_radar = blending.utils.check_norain(precip, precip_thr, norain_thr) - # The norain fraction threshold used for nwp is the default value of 0.0, - # since nwp does not suffer from clutter. - zero_model_fields = blending.utils.check_norain( - precip_models, precip_thr, norain_thr + precip_models, + velocity, + velocity_models, + timesteps, + issuetime, + blending_config, ) - if isinstance(timesteps, int): - timesteps = list(range(timesteps + 1)) - timestep_type = "int" - else: - original_timesteps = [0] + list(timesteps) - timesteps = nowcast_utils.binned_timesteps(original_timesteps) - timestep_type = "list" - - # 2.3.1 If precip is below the norain threshold and precip_models is zero, - # we consider it as no rain in the domain. - # The forecast will directly return an array filled with the minimum - # value present in precip (which equals zero rainfall in the used - # transformation) - if zero_precip_radar and zero_model_fields: - print( - "No precipitation above the threshold found in both the radar and NWP fields" - ) - print("The resulting forecast will contain only zeros") - # Create the output list - precip_forecast = [[] for j in range(n_ens_members)] - - # Save per time step to ensure the array does not become too large if - # no return_output is requested and callback is not None. - for t, subtimestep_idx in enumerate(timesteps): - # If the timestep is not the first one, we need to provide the zero forecast - if t > 0: - # Create an empty np array with shape [n_ens_members, rows, cols] - # and fill it with the minimum value from precip (corresponding to - # zero precipitation) - precip_forecast_workers = np.full( - (n_ens_members, precip_shape[0], precip_shape[1]), np.nanmin(precip) - ) - if subtimestep_idx: - if callback is not None: - if precip_forecast_workers.shape[1] > 0: - callback(precip_forecast_workers.squeeze()) - if return_output: - for j in range(n_ens_members): - precip_forecast[j].append(precip_forecast_workers[j]) - - precip_forecast_workers = None - - if measure_time: - zero_precip_time = time.time() - starttime_init - - if return_output: - precip_forecast_all_members_all_times = np.stack( - [np.stack(precip_forecast[j]) for j in range(n_ens_members)] - ) - if measure_time: - return ( - precip_forecast_all_members_all_times, - zero_precip_time, - zero_precip_time, - ) - else: - return precip_forecast_all_members_all_times - else: - return None - - else: - # 2.3.3 If zero_precip_radar, make sure that precip_cascade does not contain - # only nans or infs. If so, fill it with the zero value. - if zero_precip_radar: - # Look for a timestep and member with rain so that we have a sensible decomposition - done = False - for t in timesteps: - if done: - break - for j in range(precip_models.shape[0]): - if not blending.utils.check_norain( - precip_models[j, t], precip_thr, norain_thr - ): - if precip_models_cascade is not None: - precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin( - precip_models_cascade[j, t]["cascade_levels"] - ) - continue - precip_models_cascade_temp = decompositor( - precip_models[j, t, :, :], - bp_filter=bp_filter, - fft_method=fft, - output_domain=domain, - normalize=True, - compute_stats=True, - compact_output=True, - )["cascade_levels"] - precip_cascade[~np.isfinite(precip_cascade)] = np.nanmin( - precip_models_cascade_temp - ) - done = True - break - - # 2.3.5 If zero_precip_radar is True, only use the velocity field of the NWP - # forecast. I.e., velocity (radar) equals velocity_model at the first time - # step. - if zero_precip_radar: - # Use the velocity from velocity_models at time step 0 - velocity = velocity_models[:, 0, :, :, :].astype(np.float64, copy=False) - # Take the average over the first axis, which corresponds to n_models - # (hence, the model average) - velocity = np.mean(velocity, axis=0) - - # 3. Initialize the noise method. - # If zero_precip_radar is True, initialize noise based on the NWP field time - # step where the fraction of rainy cells is highest (because other lead times - # might be zero as well). Else, initialize the noise with the radar - # rainfall data - if zero_precip_radar: - precip_noise_input = _determine_max_nr_rainy_cells_nwp( - precip_models, precip_thr, precip_models.shape[0], timesteps - ) - # Make sure precip_noise_input is three dimensional - if len(precip_noise_input.shape) != 3: - precip_noise_input = precip_noise_input[np.newaxis, :, :] - else: - precip_noise_input = precip.copy() - - generate_perturb, generate_noise, noise_std_coeffs = _init_noise( - precip_noise_input, - precip_thr, - n_cascade_levels, - bp_filter, - decompositor, - fft, - noise_method, - noise_kwargs, - noise_stddev_adj, - measure_time, - num_workers, - seed, - ) - precip_noise_input = None - - # 4. Estimate AR parameters for the radar rainfall field - PHI = _estimate_ar_parameters_radar( - precip_cascade, - ar_order, - n_cascade_levels, - MASK_thr, - zero_precip_radar, - ) - - # 5. Repeat precip_cascade for n ensemble members - # First, discard all except the p-1 last cascades because they are not needed - # for the AR(p) model - - precip_cascade = np.stack( - [[precip_cascade[i][-ar_order:].copy() for i in range(n_cascade_levels)]] - * n_ens_members - ) - - # 6. Initialize all the random generators and prepare for the forecast loop - ( - randgen_precip, - velocity_perturbations, - generate_vel_noise, - ) = _init_random_generators( - velocity, - noise_method, - vel_pert_method, - vp_par, - vp_perp, - seed, - n_ens_members, - kmperpixel, - timestep, - ) - ( - previous_displacement, - previous_displacement_noise_cascade, - previous_displacement_prob_matching, - precip_forecast, - precip_forecast_non_perturbed, - mask_rim, - struct, - fft_objs, - ) = _prepare_forecast_loop( - precip_cascade, - noise_method, - fft_method, - n_cascade_levels, - n_ens_members, - mask_method, - mask_kwargs, - timestep, - kmperpixel, - ) - - # Also initialize the cascade of temporally correlated noise, which has the - # same shape as precip_cascade, but starts random noise. - noise_cascade, mu_noise, sigma_noise = _init_noise_cascade( - shape=precip_cascade.shape, - n_ens_members=n_ens_members, - n_cascade_levels=n_cascade_levels, - generate_noise=generate_noise, - decompositor=decompositor, - generate_perturb=generate_perturb, - randgen_precip=randgen_precip, - fft_objs=fft_objs, - bp_filter=bp_filter, - domain=domain, - noise_method=noise_method, - noise_std_coeffs=noise_std_coeffs, - ar_order=ar_order, - ) - - precip = precip[-1, :, :] - - # 7. initizalize the current and previous extrapolation forecast scale - # for the nowcasting component - rho_extrap_cascade_prev = np.repeat(1.0, PHI.shape[0]) - rho_extrap_cascade = PHI[:, 0] / ( - 1.0 - PHI[:, 1] - ) # phi1 / (1 - phi2), see BPS2004 - - if measure_time: - init_time = time.time() - starttime_init - - ### - # 8. Start the forecasting loop - ### - print("Starting blended nowcast computation.") - - if measure_time: - starttime_mainloop = time.time() - - extrap_kwargs["return_displacement"] = True - - precip_forc_prev_subtimestep = deepcopy(precip_cascade) - noise_prev_subtimestep = deepcopy(noise_cascade) - - t_prev_timestep = [0.0 for j in range(n_ens_members)] - t_leadtime_since_start_forecast = [0.0 for j in range(n_ens_members)] - - # iterate each time step - for t, subtimestep_idx in enumerate(timesteps): - if timestep_type == "list": - subtimesteps = [original_timesteps[t_] for t_ in subtimestep_idx] - else: - subtimesteps = [t] - - if (timestep_type == "list" and subtimesteps) or ( - timestep_type == "int" and t > 0 - ): - is_nowcast_time_step = True - else: - is_nowcast_time_step = False - - if is_nowcast_time_step: - print( - f"Computing nowcast for time step {t}... ", - end="", - flush=True, - ) - - if measure_time: - starttime = time.time() - - if precip_models_cascade is not None: - decomp_precip_models = list(precip_models_cascade[:, t]) - else: - if precip_models.shape[0] == 1: - decomp_precip_models = [ - decompositor( - precip_models[0, t, :, :], - bp_filter=bp_filter, - fft_method=fft, - output_domain=domain, - normalize=True, - compute_stats=True, - compact_output=True, - ) - ] - else: - with ThreadPool(num_workers) as pool: - decomp_precip_models = pool.map( - partial( - decompositor, - bp_filter=bp_filter, - fft_method=fft, - output_domain=domain, - normalize=True, - compute_stats=True, - compact_output=True, - ), - list(precip_models[:, t, :, :]), - ) - - precip_models_cascade_temp = np.array( - [decomp["cascade_levels"] for decomp in decomp_precip_models] - ) - mu_models_temp = np.array( - [decomp["means"] for decomp in decomp_precip_models] - ) - sigma_models_temp = np.array( - [decomp["stds"] for decomp in decomp_precip_models] - ) - - # 2.3.4 Check if the NWP fields contain nans or infinite numbers. If so, - # fill these with the minimum value present in precip (corresponding to - # zero rainfall in the radar observations) - ( - precip_models_cascade_temp, - precip_models_temp, - mu_models_temp, - sigma_models_temp, - ) = _fill_nans_infs_nwp_cascade( - precip_models_cascade_temp, - precip_models[:, t, :, :].astype(np.float64, copy=False), - precip_cascade, - precip, - mu_models_temp, - sigma_models_temp, - ) - - # 8.1.1 Before calling the worker for the forecast loop, determine which (NWP) - # models will be combined with which nowcast ensemble members. With the - # way it is implemented at this moment: n_ens_members of the output equals - # the maximum number of (ensemble) members in the input (either the nowcasts or NWP). - ( - precip_models_cascade_temp, - precip_models_temp, - velocity_models_temp, - mu_models_temp, - sigma_models_temp, - n_model_indices, - ) = _find_nwp_combination( - precip_models_cascade_temp, - precip_models_temp, - velocity_models[:, t, :, :, :].astype(np.float64, copy=False), - mu_models_temp, - sigma_models_temp, - n_ens_members, - ar_order, - n_cascade_levels, - blend_nwp_members, - ) - - # If zero_precip_radar is True, set the velocity field equal to the NWP - # velocity field for the current time step (velocity_models_temp). - if zero_precip_radar: - # Use the velocity from velocity_models and take the average over - # n_models (axis=0) - velocity = np.mean(velocity_models_temp, axis=0) - - if t == 0: - # 8.1.2 Calculate the initial skill of the (NWP) model forecasts at t=0 - rho_nwp_models = _compute_initial_nwp_skill( - precip_cascade, - precip_models_cascade_temp, - domain_mask, - issuetime, - outdir_path_skill, - clim_kwargs, - ) - - if t > 0: - # 8.1.3 Determine the skill of the components for lead time (t0 + t) - # First for the extrapolation component. Only calculate it when t > 0. - ( - rho_extrap_cascade, - rho_extrap_cascade_prev, - ) = blending.skill_scores.lt_dependent_cor_extrapolation( - PHI=PHI, - correlations=rho_extrap_cascade, - correlations_prev=rho_extrap_cascade_prev, - ) - - # the nowcast iteration for each ensemble member - precip_forecast_workers = [None for _ in range(n_ens_members)] - - def worker(j): - # 8.1.2 Determine the skill of the nwp components for lead time (t0 + t) - # Then for the model components - if blend_nwp_members: - rho_nwp_fc = [ - blending.skill_scores.lt_dependent_cor_nwp( - lt=(t * int(timestep)), - correlations=rho_nwp_models[n_model], - outdir_path=outdir_path_skill, - n_model=n_model, - skill_kwargs=clim_kwargs, - ) - for n_model in range(rho_nwp_models.shape[0]) - ] - rho_nwp_fc = np.stack(rho_nwp_fc) - # Concatenate rho_extrap_cascade and rho_nwp - rho_fc = np.concatenate( - (rho_extrap_cascade[None, :], rho_nwp_fc), axis=0 - ) - else: - rho_nwp_fc = blending.skill_scores.lt_dependent_cor_nwp( - lt=(t * int(timestep)), - correlations=rho_nwp_models[j], - outdir_path=outdir_path_skill, - n_model=n_model_indices[j], - skill_kwargs=clim_kwargs, - ) - # Concatenate rho_extrap_cascade and rho_nwp - rho_fc = np.concatenate( - (rho_extrap_cascade[None, :], rho_nwp_fc[None, :]), axis=0 - ) - - # 8.2 Determine the weights per component - - # Weights following the bps method. These are needed for the velocity - # weights prior to the advection step. If weights method spn is - # selected, weights will be overwritten with those weights prior to - # blending step. - # weight = [(extr_field, n_model_fields, noise), n_cascade_levels, ...] - weights = calculate_weights_bps(rho_fc) - - # The model only weights - if weights_method == "bps": - # Determine the weights of the components without the extrapolation - # cascade, in case this is no data or outside the mask. - weights_model_only = calculate_weights_bps(rho_fc[1:, :]) - elif weights_method == "spn": - # Only the weights of the components without the extrapolation - # cascade will be determined here. The full set of weights are - # determined after the extrapolation step in this method. - if blend_nwp_members and precip_models_cascade_temp.shape[0] > 1: - weights_model_only = np.zeros( - (precip_models_cascade_temp.shape[0] + 1, n_cascade_levels) - ) - for i in range(n_cascade_levels): - # Determine the normalized covariance matrix (containing) - # the cross-correlations between the models - covariance_nwp_models = np.corrcoef( - np.stack( - [ - precip_models_cascade_temp[ - n_model, i, :, : - ].flatten() - for n_model in range( - precip_models_cascade_temp.shape[0] - ) - ] - ) - ) - # Determine the weights for this cascade level - weights_model_only[:, i] = calculate_weights_spn( - correlations=rho_fc[1:, i], - covariance=covariance_nwp_models, - ) - else: - # Same as correlation and noise is 1 - correlation - weights_model_only = calculate_weights_bps(rho_fc[1:, :]) - else: - raise ValueError( - "Unknown weights method %s: must be 'bps' or 'spn'" - % weights_method - ) - - # 8.3 Determine the noise cascade and regress this to the subsequent - # time step + regress the extrapolation component to the subsequent - # time step - - # 8.3.1 Determine the epsilon, a cascade of temporally independent - # but spatially correlated noise - if noise_method is not None: - # generate noise field - epsilon = generate_noise( - generate_perturb, - randstate=randgen_precip[j], - fft_method=fft_objs[j], - domain=domain, - ) - - # decompose the noise field into a cascade - epsilon_decomposed = decompositor( - epsilon, - bp_filter, - fft_method=fft_objs[j], - input_domain=domain, - output_domain=domain, - compute_stats=True, - normalize=True, - compact_output=True, - ) - else: - epsilon_decomposed = None - - # 8.3.2 regress the extrapolation component to the subsequent time - # step - # iterate the AR(p) model for each cascade level - for i in range(n_cascade_levels): - # apply AR(p) process to extrapolation cascade level - if epsilon_decomposed is not None or vel_pert_method is not None: - precip_cascade[j][i] = autoregression.iterate_ar_model( - precip_cascade[j][i], PHI[i, :] - ) - # Renormalize the cascade - precip_cascade[j][i][1] /= np.std(precip_cascade[j][i][1]) - else: - # use the deterministic AR(p) model computed above if - # perturbations are disabled - precip_cascade[j][i] = precip_forecast_non_perturbed[i] - - # 8.3.3 regress the noise component to the subsequent time step - # iterate the AR(p) model for each cascade level - for i in range(n_cascade_levels): - # normalize the noise cascade - if epsilon_decomposed is not None: - epsilon_temp = epsilon_decomposed["cascade_levels"][i] - epsilon_temp *= noise_std_coeffs[i] - else: - epsilon_temp = None - # apply AR(p) process to noise cascade level - # (Returns zero noise if epsilon_decomposed is None) - noise_cascade[j][i] = autoregression.iterate_ar_model( - noise_cascade[j][i], PHI[i, :], eps=epsilon_temp - ) - - epsilon_decomposed = None - epsilon_temp = None - - # 8.4 Perturb and blend the advection fields + advect the - # extrapolation and noise cascade to the current time step - # (or subtimesteps if non-integer time steps are given) - - # Settings and initialize the output - extrap_kwargs_ = extrap_kwargs.copy() - extrap_kwargs_noise = extrap_kwargs.copy() - extrap_kwargs_pb = extrap_kwargs.copy() - velocity_perturbations_extrapolation = velocity - precip_forecast_extrapolated_decomp_done = [] - noise_extrapolated_decomp_done = [] - precip_forecast_extrapolated_probability_matching = [] - - # Extrapolate per sub time step - for t_sub in subtimesteps: - if t_sub > 0: - t_diff_prev_subtimestep_int = t_sub - int(t_sub) - if t_diff_prev_subtimestep_int > 0.0: - precip_forecast_cascade_subtimestep = [ - (1.0 - t_diff_prev_subtimestep_int) - * precip_forc_prev_subtimestep[j][i][-1, :] - + t_diff_prev_subtimestep_int - * precip_cascade[j][i][-1, :] - for i in range(n_cascade_levels) - ] - noise_cascade_subtimestep = [ - (1.0 - t_diff_prev_subtimestep_int) - * noise_prev_subtimestep[j][i][-1, :] - + t_diff_prev_subtimestep_int - * noise_cascade[j][i][-1, :] - for i in range(n_cascade_levels) - ] - - else: - precip_forecast_cascade_subtimestep = [ - precip_forc_prev_subtimestep[j][i][-1, :] - for i in range(n_cascade_levels) - ] - noise_cascade_subtimestep = [ - noise_prev_subtimestep[j][i][-1, :] - for i in range(n_cascade_levels) - ] - - precip_forecast_cascade_subtimestep = np.stack( - precip_forecast_cascade_subtimestep - ) - noise_cascade_subtimestep = np.stack(noise_cascade_subtimestep) - - t_diff_prev_subtimestep = t_sub - t_prev_timestep[j] - t_leadtime_since_start_forecast[j] += t_diff_prev_subtimestep - - # compute the perturbed motion field - include the NWP - # velocities and the weights. Note that we only perturb - # the extrapolation velocity field, as the NWP velocity - # field is present per time step - if vel_pert_method is not None: - velocity_perturbations_extrapolation = ( - velocity - + generate_vel_noise( - velocity_perturbations[j], - t_leadtime_since_start_forecast[j] * timestep, - ) - ) - - # Stack the perturbed extrapolation and the NWP velocities - if blend_nwp_members: - velocity_stack_all = np.concatenate( - ( - velocity_perturbations_extrapolation[None, :, :, :], - velocity_models_temp, - ), - axis=0, - ) - else: - velocity_models = velocity_models_temp[j] - velocity_stack_all = np.concatenate( - ( - velocity_perturbations_extrapolation[None, :, :, :], - velocity_models[None, :, :, :], - ), - axis=0, - ) - velocity_models = None - - # Obtain a blended optical flow, using the weights of the - # second cascade following eq. 24 in BPS2006 - velocity_blended = blending.utils.blend_optical_flows( - flows=velocity_stack_all, - weights=weights[ - :-1, 1 - ], # [(extr_field, n_model_fields), cascade_level=2] - ) - - # Extrapolate both cascades to the next time step - # First recompose the cascade, advect it and decompose it again - # This is needed to remove the interpolation artifacts. - # In addition, the number of extrapolations is greatly reduced - # A. Radar Rain - precip_forecast_recomp_subtimestep = ( - blending.utils.recompose_cascade( - combined_cascade=precip_forecast_cascade_subtimestep, - combined_mean=mu_extrapolation, - combined_sigma=sigma_extrapolation, - ) - ) - # Make sure we have values outside the mask - if zero_precip_radar: - precip_forecast_recomp_subtimestep = np.nan_to_num( - precip_forecast_recomp_subtimestep, - copy=True, - nan=zerovalue, - posinf=zerovalue, - neginf=zerovalue, - ) - # Put back the mask - precip_forecast_recomp_subtimestep[domain_mask] = np.nan - extrap_kwargs["displacement_prev"] = previous_displacement[j] - ( - precip_forecast_extrapolated_recomp_subtimestep_temp, - previous_displacement[j], - ) = extrapolator( - precip_forecast_recomp_subtimestep, - velocity_blended, - [t_diff_prev_subtimestep], - allow_nonfinite_values=True, - **extrap_kwargs, - ) - precip_forecast_extrapolated_recomp_subtimestep = ( - precip_forecast_extrapolated_recomp_subtimestep_temp[ - 0 - ].copy() - ) - temp_mask = ~np.isfinite( - precip_forecast_extrapolated_recomp_subtimestep - ) - # TODO WHERE DO CAN I FIND THIS -15.0 - precip_forecast_extrapolated_recomp_subtimestep[ - ~np.isfinite( - precip_forecast_extrapolated_recomp_subtimestep - ) - ] = zerovalue - precip_forecast_extrapolated_decomp = decompositor( - precip_forecast_extrapolated_recomp_subtimestep, - bp_filter, - mask=MASK_thr, - fft_method=fft, - output_domain=domain, - normalize=True, - compute_stats=True, - compact_output=True, - )["cascade_levels"] - # Make sure we have values outside the mask - if zero_precip_radar: - precip_forecast_extrapolated_decomp = np.nan_to_num( - precip_forecast_extrapolated_decomp, - copy=True, - nan=np.nanmin(precip_forecast_cascade_subtimestep), - posinf=np.nanmin(precip_forecast_cascade_subtimestep), - neginf=np.nanmin(precip_forecast_cascade_subtimestep), - ) - for i in range(n_cascade_levels): - precip_forecast_extrapolated_decomp[i][temp_mask] = np.nan - # B. Noise - noise_cascade_subtimestep_recomp = ( - blending.utils.recompose_cascade( - combined_cascade=noise_cascade_subtimestep, - combined_mean=mu_noise[j], - combined_sigma=sigma_noise[j], - ) - ) - extrap_kwargs_noise["displacement_prev"] = ( - previous_displacement_noise_cascade[j] - ) - extrap_kwargs_noise["map_coordinates_mode"] = "wrap" - ( - noise_extrapolated_recomp_temp, - previous_displacement_noise_cascade[j], - ) = extrapolator( - noise_cascade_subtimestep_recomp, - velocity_blended, - [t_diff_prev_subtimestep], - allow_nonfinite_values=True, - **extrap_kwargs_noise, - ) - noise_extrapolated_recomp = noise_extrapolated_recomp_temp[ - 0 - ].copy() - noise_extrapolated_decomp = decompositor( - noise_extrapolated_recomp, - bp_filter, - mask=MASK_thr, - fft_method=fft, - output_domain=domain, - normalize=True, - compute_stats=True, - compact_output=True, - )["cascade_levels"] - for i in range(n_cascade_levels): - noise_extrapolated_decomp[i] *= noise_std_coeffs[i] - - # Append the results to the output lists - precip_forecast_extrapolated_decomp_done.append( - precip_forecast_extrapolated_decomp.copy() - ) - noise_extrapolated_decomp_done.append( - noise_extrapolated_decomp.copy() - ) - precip_forecast_cascade_subtimestep = None - precip_forecast_recomp_subtimestep = None - precip_forecast_extrapolated_recomp_subtimestep_temp = None - precip_forecast_extrapolated_recomp_subtimestep = None - precip_forecast_extrapolated_decomp = None - noise_cascade_subtimestep = None - noise_cascade_subtimestep_recomp = None - noise_extrapolated_recomp_temp = None - noise_extrapolated_recomp = None - noise_extrapolated_decomp = None - - # Finally, also extrapolate the initial radar rainfall - # field. This will be blended with the rainfall field(s) - # of the (NWP) model(s) for Lagrangian blended prob. matching - # min_R = np.min(precip) - extrap_kwargs_pb["displacement_prev"] = ( - previous_displacement_prob_matching[j] - ) - # Apply the domain mask to the extrapolation component - precip_forecast_temp_for_probability_matching = precip.copy() - precip_forecast_temp_for_probability_matching[domain_mask] = ( - np.nan - ) - ( - precip_forecast_extrapolated_probability_matching_temp, - previous_displacement_prob_matching[j], - ) = extrapolator( - precip_forecast_temp_for_probability_matching, - velocity_blended, - [t_diff_prev_subtimestep], - allow_nonfinite_values=True, - **extrap_kwargs_pb, - ) - precip_forecast_extrapolated_probability_matching.append( - precip_forecast_extrapolated_probability_matching_temp[0] - ) - - t_prev_timestep[j] = t_sub - - if len(precip_forecast_extrapolated_decomp_done) > 0: - precip_forecast_extrapolated_decomp_done = np.stack( - precip_forecast_extrapolated_decomp_done - ) - noise_extrapolated_decomp_done = np.stack( - noise_extrapolated_decomp_done - ) - precip_forecast_extrapolated_probability_matching = np.stack( - precip_forecast_extrapolated_probability_matching - ) - - # advect the forecast field by one time step if no subtimesteps in the - # current interval were found - if not subtimesteps: - t_diff_prev_subtimestep = t + 1 - t_prev_timestep[j] - t_leadtime_since_start_forecast[j] += t_diff_prev_subtimestep - - # compute the perturbed motion field - include the NWP - # velocities and the weights - if vel_pert_method is not None: - velocity_perturbations_extrapolation = ( - velocity - + generate_vel_noise( - velocity_perturbations[j], - t_leadtime_since_start_forecast[j] * timestep, - ) - ) - - # Stack the perturbed extrapolation and the NWP velocities - if blend_nwp_members: - velocity_stack_all = np.concatenate( - ( - velocity_perturbations_extrapolation[None, :, :, :], - velocity_models_temp, - ), - axis=0, - ) - else: - velocity_models = velocity_models_temp[j] - velocity_stack_all = np.concatenate( - ( - velocity_perturbations_extrapolation[None, :, :, :], - velocity_models[None, :, :, :], - ), - axis=0, - ) - velocity_models = None - - # Obtain a blended optical flow, using the weights of the - # second cascade following eq. 24 in BPS2006 - velocity_blended = blending.utils.blend_optical_flows( - flows=velocity_stack_all, - weights=weights[ - :-1, 1 - ], # [(extr_field, n_model_fields), cascade_level=2] - ) - - # Extrapolate the extrapolation and noise cascade - - extrap_kwargs_["displacement_prev"] = previous_displacement[j] - extrap_kwargs_noise["displacement_prev"] = ( - previous_displacement_noise_cascade[j] - ) - extrap_kwargs_noise["map_coordinates_mode"] = "wrap" - - _, previous_displacement[j] = extrapolator( - None, - velocity_blended, - [t_diff_prev_subtimestep], - allow_nonfinite_values=True, - **extrap_kwargs_, - ) - - _, previous_displacement_noise_cascade[j] = extrapolator( - None, - velocity_blended, - [t_diff_prev_subtimestep], - allow_nonfinite_values=True, - **extrap_kwargs_noise, - ) - - # Also extrapolate the radar observation, used for the probability - # matching and post-processing steps - extrap_kwargs_pb["displacement_prev"] = ( - previous_displacement_prob_matching[j] - ) - _, previous_displacement_prob_matching[j] = extrapolator( - None, - velocity_blended, - [t_diff_prev_subtimestep], - allow_nonfinite_values=True, - **extrap_kwargs_pb, - ) - - t_prev_timestep[j] = t + 1 - - precip_forc_prev_subtimestep[j] = precip_cascade[j] - noise_prev_subtimestep[j] = noise_cascade[j] - - # 8.5 Blend the cascades - final_blended_forecast = [] - - for t_sub in subtimesteps: - # TODO: does it make sense to use sub time steps - check if it works? - if t_sub > 0: - t_index = np.where(np.array(subtimesteps) == t_sub)[0][0] - # First concatenate the cascades and the means and sigmas - # precip_models = [n_models,timesteps,n_cascade_levels,m,n] - if blend_nwp_members: - cascade_stack_all_components = np.concatenate( - ( - precip_forecast_extrapolated_decomp_done[ - None, t_index - ], - precip_models_cascade_temp, - noise_extrapolated_decomp_done[None, t_index], - ), - axis=0, - ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] - means_stacked = np.concatenate( - (mu_extrapolation[None, :], mu_models_temp), axis=0 - ) - sigmas_stacked = np.concatenate( - (sigma_extrapolation[None, :], sigma_models_temp), - axis=0, - ) - else: - cascade_stack_all_components = np.concatenate( - ( - precip_forecast_extrapolated_decomp_done[ - None, t_index - ], - precip_models_cascade_temp[None, j], - noise_extrapolated_decomp_done[None, t_index], - ), - axis=0, - ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] - means_stacked = np.concatenate( - (mu_extrapolation[None, :], mu_models_temp[None, j]), - axis=0, - ) - sigmas_stacked = np.concatenate( - ( - sigma_extrapolation[None, :], - sigma_models_temp[None, j], - ), - axis=0, - ) - - # First determine the blending weights if method is spn. The - # weights for method bps have already been determined. - if weights_method == "spn": - weights = np.zeros( - ( - cascade_stack_all_components.shape[0], - n_cascade_levels, - ) - ) - for i in range(n_cascade_levels): - # Determine the normalized covariance matrix (containing) - # the cross-correlations between the models - cascade_stack_all_components_temp = np.stack( - [ - cascade_stack_all_components[ - n_model, i, :, : - ].flatten() - for n_model in range( - cascade_stack_all_components.shape[0] - 1 - ) - ] - ) # -1 to exclude the noise component - covariance_nwp_models = np.ma.corrcoef( - np.ma.masked_invalid( - cascade_stack_all_components_temp - ) - ) - # Determine the weights for this cascade level - weights[:, i] = calculate_weights_spn( - correlations=rho_fc[:, i], - covariance=covariance_nwp_models, - ) - - # Blend the extrapolation, (NWP) model(s) and noise cascades - precip_forecast_blended = blending.utils.blend_cascades( - cascades_norm=cascade_stack_all_components, weights=weights - ) - - # Also blend the cascade without the extrapolation component - precip_forecast_blended_mod_only = ( - blending.utils.blend_cascades( - cascades_norm=cascade_stack_all_components[1:, :], - weights=weights_model_only, - ) - ) - - # Blend the means and standard deviations - # Input is array of shape [number_components, scale_level, ...] - means_blended, sigmas_blended = blend_means_sigmas( - means=means_stacked, sigmas=sigmas_stacked, weights=weights - ) - # Also blend the means and sigmas for the cascade without extrapolation - ( - means_blended_mod_only, - sigmas_blended_mod_only, - ) = blend_means_sigmas( - means=means_stacked[1:, :], - sigmas=sigmas_stacked[1:, :], - weights=weights_model_only, - ) - - # 8.6 Recompose the cascade to a precipitation field - # (The function first normalizes the blended cascade, precip_forecast_blended - # again) - precip_forecast_recomposed = blending.utils.recompose_cascade( - combined_cascade=precip_forecast_blended, - combined_mean=means_blended, - combined_sigma=sigmas_blended, - ) - # The recomposed cascade without the extrapolation (for NaN filling - # outside the radar domain) - precip_forecast_recomposed_mod_only = ( - blending.utils.recompose_cascade( - combined_cascade=precip_forecast_blended_mod_only, - combined_mean=means_blended_mod_only, - combined_sigma=sigmas_blended_mod_only, - ) - ) - if domain == "spectral": - # TODO: Check this! (Only tested with domain == 'spatial') - precip_forecast_recomposed = fft_objs[j].irfft2( - precip_forecast_recomposed - ) - precip_forecast_recomposed_mod_only = fft_objs[j].irfft2( - precip_forecast_recomposed_mod_only - ) - - # 8.7 Post-processing steps - use the mask and fill no data with - # the blended NWP forecast. Probability matching following - # Lagrangian blended probability matching which uses the - # latest extrapolated radar rainfall field blended with the - # nwp model(s) rainfall forecast fields as 'benchmark'. - - # 8.7.1 first blend the extrapolated rainfall field (the field - # that is only used for post-processing steps) with the NWP - # rainfall forecast for this time step using the weights - # at scale level 2. - weights_probability_matching = weights[ - :-1, 1 - ] # Weights without noise, level 2 - weights_probability_matching_normalized = ( - weights_probability_matching - / np.sum(weights_probability_matching) - ) - # And the weights for outside the radar domain - weights_probability_matching_mod_only = weights_model_only[ - :-1, 1 - ] # Weights without noise, level 2 - weights_probability_matching_normalized_mod_only = ( - weights_probability_matching_mod_only - / np.sum(weights_probability_matching_mod_only) - ) - # Stack the fields - if blend_nwp_members: - precip_forecast_probability_matching_final = np.concatenate( - ( - precip_forecast_extrapolated_probability_matching[ - None, t_index - ], - precip_models_temp, - ), - axis=0, - ) - else: - precip_forecast_probability_matching_final = np.concatenate( - ( - precip_forecast_extrapolated_probability_matching[ - None, t_index - ], - precip_models_temp[None, j], - ), - axis=0, - ) - # Blend it - precip_forecast_probability_matching_blended = np.sum( - weights_probability_matching_normalized.reshape( - weights_probability_matching_normalized.shape[0], 1, 1 - ) - * precip_forecast_probability_matching_final, - axis=0, - ) - if blend_nwp_members: - precip_forecast_probability_matching_blended_mod_only = np.sum( - weights_probability_matching_normalized_mod_only.reshape( - weights_probability_matching_normalized_mod_only.shape[ - 0 - ], - 1, - 1, - ) - * precip_models_temp, - axis=0, - ) - else: - precip_forecast_probability_matching_blended_mod_only = ( - precip_models_temp[j] - ) - - # The extrapolation components are NaN outside the advected - # radar domain. This results in NaN values in the blended - # forecast outside the radar domain. Therefore, fill these - # areas with the "..._mod_only" blended forecasts, consisting - # of the NWP and noise components. - - nan_indices = np.isnan(precip_forecast_recomposed) - if smooth_radar_mask_range != 0: - # Compute the smooth dilated mask - new_mask = blending.utils.compute_smooth_dilated_mask( - nan_indices, - max_padding_size_in_px=smooth_radar_mask_range, - ) - - # Ensure mask values are between 0 and 1 - mask_model = np.clip(new_mask, 0, 1) - mask_radar = np.clip(1 - new_mask, 0, 1) - - # Handle NaNs in precip_forecast_new and precip_forecast_new_mod_only by setting NaNs to 0 in the blending step - precip_forecast_recomposed_mod_only_no_nan = np.nan_to_num( - precip_forecast_recomposed_mod_only, nan=0 - ) - precip_forecast_recomposed_no_nan = np.nan_to_num( - precip_forecast_recomposed, nan=0 - ) - - # Perform the blending of radar and model inside the radar domain using a weighted combination - precip_forecast_recomposed = np.nansum( - [ - mask_model - * precip_forecast_recomposed_mod_only_no_nan, - mask_radar * precip_forecast_recomposed_no_nan, - ], - axis=0, - ) - - nan_indices = np.isnan( - precip_forecast_probability_matching_blended - ) - precip_forecast_probability_matching_blended = np.nansum( - [ - precip_forecast_probability_matching_blended - * mask_radar, - precip_forecast_probability_matching_blended_mod_only - * mask_model, - ], - axis=0, - ) - else: - precip_forecast_recomposed[nan_indices] = ( - precip_forecast_recomposed_mod_only[nan_indices] - ) - nan_indices = np.isnan( - precip_forecast_probability_matching_blended - ) - precip_forecast_probability_matching_blended[ - nan_indices - ] = precip_forecast_probability_matching_blended_mod_only[ - nan_indices - ] - - # Finally, fill the remaining nan values, if present, with - # the minimum value in the forecast - nan_indices = np.isnan(precip_forecast_recomposed) - precip_forecast_recomposed[nan_indices] = np.nanmin( - precip_forecast_recomposed - ) - nan_indices = np.isnan( - precip_forecast_probability_matching_blended - ) - precip_forecast_probability_matching_blended[nan_indices] = ( - np.nanmin(precip_forecast_probability_matching_blended) - ) - - # 8.7.2. Apply the masking and prob. matching - if mask_method is not None: - # apply the precipitation mask to prevent generation of new - # precipitation into areas where it was not originally - # observed - precip_forecast_min_value = precip_forecast_recomposed.min() - if mask_method == "incremental": - # The incremental mask is slightly different from - # the implementation in the non-blended steps.py, as - # it is not based on the last forecast, but instead - # on R_pm_blended. Therefore, the buffer does not - # increase over time. - # Get the mask for this forecast - precip_field_mask = ( - precip_forecast_probability_matching_blended - >= precip_thr - ) - # Buffer the mask - precip_field_mask = _compute_incremental_mask( - precip_field_mask, struct, mask_rim - ) - # Get the final mask - precip_forecast_recomposed = ( - precip_forecast_min_value - + ( - precip_forecast_recomposed - - precip_forecast_min_value - ) - * precip_field_mask - ) - precip_field_mask_temp = ( - precip_forecast_recomposed - > precip_forecast_min_value - ) - elif mask_method == "obs": - # The mask equals the most recent benchmark - # rainfall field - precip_field_mask_temp = ( - precip_forecast_probability_matching_blended - >= precip_thr - ) - - # Set to min value outside of mask - precip_forecast_recomposed[~precip_field_mask_temp] = ( - precip_forecast_min_value - ) - - # If probmatching_method is not None, resample the distribution from - # both the extrapolation cascade and the model (NWP) cascade and use - # that for the probability matching. - if probmatching_method is not None and resample_distribution: - arr1 = precip_forecast_extrapolated_probability_matching[ - t_index - ] - arr2 = precip_models_temp[j] - # resample weights based on cascade level 2. - # Areas where one of the fields is nan are not included. - precip_forecast_probability_matching_resampled = probmatching.resample_distributions( - first_array=arr1, - second_array=arr2, - probability_first_array=weights_probability_matching_normalized[ - 0 - ], - ) - else: - precip_forecast_probability_matching_resampled = ( - precip_forecast_probability_matching_blended.copy() - ) - - if probmatching_method == "cdf": - # nan indices in the extrapolation nowcast - nan_indices = np.isnan( - precip_forecast_extrapolated_probability_matching[ - t_index - ] - ) - # Adjust the CDF of the forecast to match the resampled distribution combined from - # extrapolation and model fields. - # Rainfall outside the pure extrapolation domain is not taken into account. - if np.any(np.isfinite(precip_forecast_recomposed)): - precip_forecast_recomposed = ( - probmatching.nonparam_match_empirical_cdf( - precip_forecast_recomposed, - precip_forecast_probability_matching_resampled, - nan_indices, - ) - ) - precip_forecast_probability_matching_resampled = None - elif probmatching_method == "mean": - # Use R_pm_blended as benchmark field and - mean_probabiltity_matching_forecast = np.mean( - precip_forecast_probability_matching_resampled[ - precip_forecast_probability_matching_resampled - >= precip_thr - ] - ) - no_rain_mask = precip_forecast_recomposed >= precip_thr - mean_precip_forecast = np.mean( - precip_forecast_recomposed[no_rain_mask] - ) - precip_forecast_recomposed[no_rain_mask] = ( - precip_forecast_recomposed[no_rain_mask] - - mean_precip_forecast - + mean_probabiltity_matching_forecast - ) - precip_forecast_probability_matching_resampled = None - - final_blended_forecast.append(precip_forecast_recomposed) - - precip_forecast_workers[j] = final_blended_forecast - - res = [] - - if DASK_IMPORTED and n_ens_members > 1: - for j in range(n_ens_members): - res.append(dask.delayed(worker)(j)) - dask.compute(*res, num_workers=num_ensemble_workers) - else: - for j in range(n_ens_members): - worker(j) - - res = None - - if is_nowcast_time_step: - if measure_time: - print(f"{time.time() - starttime:.2f} seconds.") - else: - print("done.") - - if callback is not None: - precip_forecast_final = np.stack(precip_forecast_workers) - if precip_forecast_final.shape[1] > 0: - callback(precip_forecast_final.squeeze()) - - if return_output: - for j in range(n_ens_members): - precip_forecast[j].extend(precip_forecast_workers[j]) - - precip_forecast_workers = None - - if measure_time: - mainloop_time = time.time() - starttime_mainloop - - if return_output: - precip_forecast_all_members_all_times = np.stack( - [np.stack(precip_forecast[j]) for j in range(n_ens_members)] - ) - if measure_time: - return precip_forecast_all_members_all_times, init_time, mainloop_time - else: - return precip_forecast_all_members_all_times - else: - return None + forecast_steps_nowcast = blended_nowcaster.compute_forecast() + blended_nowcaster.reset_states_and_params() + # Call the appropriate methods within the class + return forecast_steps_nowcast def calculate_weights_spn(correlations, covariance): From 8d8905a9624de39af087d3b01d84e2311642973f Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 5 Dec 2024 11:00:46 +0100 Subject: [PATCH 13/33] Removed old code which is no longer used --- pysteps/blending/steps.py | 532 +------------------------------------- 1 file changed, 1 insertion(+), 531 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 0e229a43..0d461bcf 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -232,9 +232,6 @@ def __init__( self.__params = StepsBlendingParams() self.__state = StepsBlendingState() - # Perform input validation - self.__check_inputs() - # Initialize nowcast components and parameters self.__initialize_nowcast_components() @@ -2804,6 +2801,7 @@ def forecast( return forecast_steps_nowcast +# TODO: add the following code to the main body def calculate_weights_spn(correlations, covariance): """Calculate SPN blending weights for STEPS blending from correlation. @@ -2951,48 +2949,6 @@ def blend_means_sigmas(means, sigmas, weights): return combined_means, combined_sigmas -def _check_inputs( - precip, precip_models, velocity, velocity_models, timesteps, ar_order -): - if precip.ndim != 3: - raise ValueError("precip must be a three-dimensional array") - if precip.shape[0] < ar_order + 1: - raise ValueError("precip.shape[0] < ar_order+1") - if precip_models.ndim != 2 and precip_models.ndim != 4: - raise ValueError( - "precip_models must be either a two-dimensional array containing dictionaries with decomposed model fields or a four-dimensional array containing the original (NWP) model forecasts" - ) - if velocity.ndim != 3: - raise ValueError("velocity must be a three-dimensional array") - if velocity_models.ndim != 5: - raise ValueError("velocity_models must be a five-dimensional array") - if velocity.shape[0] != 2 or velocity_models.shape[2] != 2: - raise ValueError( - "velocity and velocity_models must have an x- and y-component, check the shape" - ) - if precip.shape[1:3] != velocity.shape[1:3]: - raise ValueError( - "dimension mismatch between precip and velocity: shape(precip)=%s, shape(velocity)=%s" - % (str(precip.shape), str(velocity.shape)) - ) - if precip_models.shape[0] != velocity_models.shape[0]: - raise ValueError( - "precip_models and velocity_models must consist of the same number of models" - ) - if isinstance(timesteps, list) and not sorted(timesteps) == timesteps: - raise ValueError("timesteps is not in ascending order") - if isinstance(timesteps, list): - if precip_models.shape[1] != math.ceil(timesteps[-1]) + 1: - raise ValueError( - "precip_models does not contain sufficient lead times for this forecast" - ) - else: - if precip_models.shape[1] != timesteps + 1: - raise ValueError( - "precip_models does not contain sufficient lead times for this forecast" - ) - - def _compute_incremental_mask(Rbin, kr, r): # buffer the observation mask Rbin using the kernel kr # add a grayscale rim r (for smooth rain/no-rain transition) @@ -3050,107 +3006,6 @@ def f(precip, i): return precip -def _init_noise( - precip, - precip_thr, - n_cascade_levels, - bp_filter, - decompositor, - fft, - noise_method, - noise_kwargs, - noise_stddev_adj, - measure_time, - num_workers, - seed, -): - """Initialize the noise method.""" - if noise_method is None: - return None, None, None - - # get methods for perturbations - init_noise, generate_noise = noise.get_method(noise_method) - - # initialize the perturbation generator for the precipitation field - generate_perturb = init_noise(precip, fft_method=fft, **noise_kwargs) - - if noise_stddev_adj == "auto": - print("Computing noise adjustment coefficients... ", end="", flush=True) - if measure_time: - starttime = time.time() - - precip_forecast_min = np.min(precip) - noise_std_coeffs = noise.utils.compute_noise_stddev_adjs( - precip[-1, :, :], - precip_thr, - precip_forecast_min, - bp_filter, - decompositor, - generate_perturb, - generate_noise, - 20, - conditional=True, - num_workers=num_workers, - seed=seed, - ) - - if measure_time: - print(f"{time.time() - starttime:.2f} seconds.") - else: - print("done.") - elif noise_stddev_adj == "fixed": - f = lambda k: 1.0 / (0.75 + 0.09 * k) - noise_std_coeffs = [f(k) for k in range(1, n_cascade_levels + 1)] - else: - noise_std_coeffs = np.ones(n_cascade_levels) - - if noise_stddev_adj is not None: - print(f"noise std. dev. coeffs: {noise_std_coeffs}") - - return generate_perturb, generate_noise, noise_std_coeffs - - -def _compute_cascade_decomposition_radar( - precip, - ar_order, - n_cascade_levels, - n_ens_members, - MASK_thr, - domain, - bp_filter, - decompositor, - fft, -): - """Compute the cascade decompositions of the input precipitation fields.""" - precip_forecast_decomp = [] - for i in range(ar_order + 1): - precip_forecast = decompositor( - precip[i, :, :], - bp_filter, - mask=MASK_thr, - fft_method=fft, - output_domain=domain, - normalize=True, - compute_stats=True, - compact_output=True, - ) - precip_forecast_decomp.append(precip_forecast) - - # Rearrange the cascaded into a four-dimensional array of shape - # (n_cascade_levels,ar_order+1,m,n) for the autoregressive model - precip_forecast_cascades = nowcast_utils.stack_cascades( - precip_forecast_decomp, n_cascade_levels - ) - - precip_forecast_decomp = precip_forecast_decomp[-1] - mu_extrapolation = np.array(precip_forecast_decomp["means"]) - sigma_extrapolation = np.array(precip_forecast_decomp["stds"]) - precip_forecast_decomp = [ - precip_forecast_decomp.copy() for j in range(n_ens_members) - ] - return precip_forecast_cascades, mu_extrapolation, sigma_extrapolation - - def _compute_cascade_recomposition_nwp(precip_models_cascade, recompositor): """If necessary, recompose (NWP) model forecasts.""" precip_models = None @@ -3169,388 +3024,3 @@ def _compute_cascade_recomposition_nwp(precip_models_cascade, recompositor): precip_model = None return precip_models - - -def _estimate_ar_parameters_radar( - precip_forecast_cascades, ar_order, n_cascade_levels, MASK_thr, zero_precip_radar -): - """Estimate AR parameters for the radar rainfall field.""" - # If there are values in the radar fields, compute the autocorrelations - GAMMA = np.empty((n_cascade_levels, ar_order)) - if not zero_precip_radar: - # compute lag-l temporal autocorrelation coefficients for each cascade level - for i in range(n_cascade_levels): - GAMMA[i, :] = correlation.temporal_autocorrelation( - precip_forecast_cascades[i], mask=MASK_thr - ) - - # Else, use standard values for the autocorrelations - else: - # Get the climatological lag-1 and lag-2 autocorrelation values from Table 2 - # in `BPS2004`. - # Hard coded, change to own (climatological) values when present. - GAMMA = np.array( - [ - [0.99805, 0.9925, 0.9776, 0.9297, 0.796, 0.482, 0.079, 0.0006], - [0.9933, 0.9752, 0.923, 0.750, 0.367, 0.069, 0.0018, 0.0014], - ] - ) - - # Check whether the number of cascade_levels is correct - if GAMMA.shape[1] > n_cascade_levels: - GAMMA = GAMMA[:, 0:n_cascade_levels] - elif GAMMA.shape[1] < n_cascade_levels: - # Get the number of cascade levels that is missing - n_extra_lev = n_cascade_levels - GAMMA.shape[1] - # Append the array with correlation values of 10e-4 - GAMMA = np.append( - GAMMA, - [np.repeat(0.0006, n_extra_lev), np.repeat(0.0014, n_extra_lev)], - axis=1, - ) - - # Finally base GAMMA.shape[0] on the AR-level - if ar_order == 1: - GAMMA = GAMMA[0, :] - if ar_order > 2: - for repeat_index in range(ar_order - 2): - GAMMA = np.vstack((GAMMA, GAMMA[1, :])) - - # Finally, transpose GAMMA to ensure that the shape is the same as np.empty((n_cascade_levels, ar_order)) - GAMMA = GAMMA.transpose() - assert GAMMA.shape == (n_cascade_levels, ar_order) - - # Print the GAMMA value - nowcast_utils.print_corrcoefs(GAMMA) - - if ar_order == 2: - # adjust the lag-2 correlation coefficient to ensure that the AR(p) - # process is stationary - for i in range(n_cascade_levels): - GAMMA[i, 1] = autoregression.adjust_lag2_corrcoef2(GAMMA[i, 0], GAMMA[i, 1]) - - # estimate the parameters of the AR(p) model from the autocorrelation - # coefficients - PHI = np.empty((n_cascade_levels, ar_order + 1)) - for i in range(n_cascade_levels): - PHI[i, :] = autoregression.estimate_ar_params_yw(GAMMA[i, :]) - - nowcast_utils.print_ar_params(PHI) - return PHI - - -def _find_nwp_combination( - precip_models, - precip_forecast_probability_matching, - velocity_models, - mu_models, - sigma_models, - n_ens_members, - ar_order, - n_cascade_levels, - blend_nwp_members, -): - """Determine which (NWP) models will be combined with which nowcast ensemble members. - With the way it is implemented at this moment: n_ens_members of the output equals - the maximum number of (ensemble) members in the input (either the nowcasts or NWP). - """ - # Make sure the number of model members is not larger than than or equal to - # n_ens_members - n_model_members = precip_models.shape[0] - if n_model_members > n_ens_members: - raise ValueError( - "The number of NWP model members is larger than the given number of ensemble members. n_model_members <= n_ens_members." - ) - - # Check if NWP models/members should be used individually, or if all of - # them are blended together per nowcast ensemble member. - if blend_nwp_members: - n_model_indices = None - - else: - # Start with determining the maximum and mimimum number of members/models - # in both input products - n_ens_members_max = max(n_ens_members, n_model_members) - n_ens_members_min = min(n_ens_members, n_model_members) - # Also make a list of the model index numbers. These indices are needed - # for indexing the right climatological skill file when pysteps calculates - # the blended forecast in parallel. - if n_model_members > 1: - n_model_indices = np.arange(n_model_members) - else: - n_model_indices = [0] - - # Now, repeat the nowcast ensemble members or the nwp models/members until - # it has the same amount of members as n_ens_members_max. For instance, if - # you have 10 ensemble nowcasts members and 3 NWP members, the output will - # be an ensemble of 10 members. Hence, the three NWP members are blended - # with the first three members of the nowcast (member one with member one, - # two with two, etc.), subsequently, the same NWP members are blended with - # the next three members (NWP member one with member 4, NWP member 2 with - # member 5, etc.), until 10 is reached. - if n_ens_members_min != n_ens_members_max: - if n_model_members == 1: - precip_models = np.repeat(precip_models, n_ens_members_max, axis=0) - mu_models = np.repeat(mu_models, n_ens_members_max, axis=0) - sigma_models = np.repeat(sigma_models, n_ens_members_max, axis=0) - velocity_models = np.repeat(velocity_models, n_ens_members_max, axis=0) - # For the prob. matching - precip_forecast_probability_matching = np.repeat( - precip_forecast_probability_matching, n_ens_members_max, axis=0 - ) - # Finally, for the model indices - n_model_indices = np.repeat(n_model_indices, n_ens_members_max, axis=0) - - elif n_model_members == n_ens_members_min: - repeats = [ - (n_ens_members_max + i) // n_ens_members_min - for i in range(n_ens_members_min) - ] - if n_model_members == n_ens_members_min: - precip_models = np.repeat(precip_models, repeats, axis=0) - mu_models = np.repeat(mu_models, repeats, axis=0) - sigma_models = np.repeat(sigma_models, repeats, axis=0) - velocity_models = np.repeat(velocity_models, repeats, axis=0) - # For the prob. matching - precip_forecast_probability_matching = np.repeat( - precip_forecast_probability_matching, repeats, axis=0 - ) - # Finally, for the model indices - n_model_indices = np.repeat(n_model_indices, repeats, axis=0) - - return ( - precip_models, - precip_forecast_probability_matching, - velocity_models, - mu_models, - sigma_models, - n_model_indices, - ) - - -def _init_random_generators( - velocity, - noise_method, - vel_pert_method, - vp_par, - vp_perp, - seed, - n_ens_members, - kmperpixel, - timestep, -): - """Initialize all the random generators.""" - if noise_method is not None: - randgen_precip = [] - randgen_motion = [] - for j in range(n_ens_members): - rs = np.random.RandomState(seed) - randgen_precip.append(rs) - seed = rs.randint(0, high=1e9) - rs = np.random.RandomState(seed) - randgen_motion.append(rs) - seed = rs.randint(0, high=1e9) - - if vel_pert_method is not None: - init_vel_noise, generate_vel_noise = noise.get_method(vel_pert_method) - - # initialize the perturbation generators for the motion field - velocity_perturbations = [] - for j in range(n_ens_members): - kwargs = { - "randstate": randgen_motion[j], - "p_par": vp_par, - "p_perp": vp_perp, - } - vp_ = init_vel_noise(velocity, 1.0 / kmperpixel, timestep, **kwargs) - velocity_perturbations.append(vp_) - else: - velocity_perturbations, generate_vel_noise = None, None - - return randgen_precip, velocity_perturbations, generate_vel_noise - - -def _prepare_forecast_loop( - precip_forecast_cascades, - noise_method, - fft_method, - n_cascade_levels, - n_ens_members, - mask_method, - mask_kwargs, - timestep, - kmperpixel, -): - """Prepare for the forecast loop.""" - # Empty arrays for the previous displacements and the forecast cascade - previous_displacement = np.stack([None for j in range(n_ens_members)]) - previous_displacement_noise_cascade = np.stack([None for j in range(n_ens_members)]) - previous_displacement_prob_matching = np.stack([None for j in range(n_ens_members)]) - precip_forecast = [[] for j in range(n_ens_members)] - - if mask_method == "incremental": - # get mask parameters - mask_rim = mask_kwargs.get("mask_rim", 10) - mask_f = mask_kwargs.get("mask_f", 1.0) - # initialize the structuring element - struct = generate_binary_structure(2, 1) - # iterate it to expand it nxn - n = mask_f * timestep / kmperpixel - struct = iterate_structure(struct, int((n - 1) / 2.0)) - else: - mask_rim, struct = None, None - - if noise_method is None: - precip_forecast_non_perturbed = [ - precip_forecast_cascades[0][i].copy() for i in range(n_cascade_levels) - ] - else: - precip_forecast_non_perturbed = None - - fft_objs = [] - for i in range(n_ens_members): - fft_objs.append( - utils.get_method(fft_method, shape=precip_forecast_cascades.shape[-2:]) - ) - - return ( - previous_displacement, - previous_displacement_noise_cascade, - previous_displacement_prob_matching, - precip_forecast, - precip_forecast_non_perturbed, - mask_rim, - struct, - fft_objs, - ) - - -def _compute_initial_nwp_skill( - precip_forecast_cascades, - precip_models, - domain_mask, - issuetime, - outdir_path_skill, - clim_kwargs, -): - """Calculate the initial skill of the (NWP) model forecasts at t=0.""" - rho_nwp_models = [ - blending.skill_scores.spatial_correlation( - obs=precip_forecast_cascades[0, :, -1, :, :].copy(), - mod=precip_models[n_model, :, :, :].copy(), - domain_mask=domain_mask, - ) - for n_model in range(precip_models.shape[0]) - ] - rho_nwp_models = np.stack(rho_nwp_models) - - # Ensure that the model skill decreases with increasing scale level. - for n_model in range(precip_models.shape[0]): - for i in range(1, precip_models.shape[1]): - if rho_nwp_models[n_model, i] > rho_nwp_models[n_model, i - 1]: - # Set it equal to the previous scale level - rho_nwp_models[n_model, i] = rho_nwp_models[n_model, i - 1] - - # Save this in the climatological skill file - blending.clim.save_skill( - current_skill=rho_nwp_models, - validtime=issuetime, - outdir_path=outdir_path_skill, - **clim_kwargs, - ) - return rho_nwp_models - - -def _init_noise_cascade( - shape, - n_ens_members, - n_cascade_levels, - generate_noise, - decompositor, - generate_perturb, - randgen_precip, - fft_objs, - bp_filter, - domain, - noise_method, - noise_std_coeffs, - ar_order, -): - """Initialize the noise cascade with identical noise for all AR(n) steps - We also need to return the mean and standard deviations of the noise - for the recombination of the noise before advecting it. - """ - noise_cascade = np.zeros(shape) - mu_noise = np.zeros((n_ens_members, n_cascade_levels)) - sigma_noise = np.zeros((n_ens_members, n_cascade_levels)) - if noise_method: - for j in range(n_ens_members): - epsilon = generate_noise( - generate_perturb, - randstate=randgen_precip[j], - fft_method=fft_objs[j], - domain=domain, - ) - epsilon_decomposed = decompositor( - epsilon, - bp_filter, - fft_method=fft_objs[j], - input_domain=domain, - output_domain=domain, - compute_stats=True, - normalize=True, - compact_output=True, - ) - mu_noise[j] = epsilon_decomposed["means"] - sigma_noise[j] = epsilon_decomposed["stds"] - for i in range(n_cascade_levels): - epsilon_temp = epsilon_decomposed["cascade_levels"][i] - epsilon_temp *= noise_std_coeffs[i] - for n in range(ar_order): - noise_cascade[j][i][n] = epsilon_temp - epsilon_decomposed = None - epsilon_temp = None - return noise_cascade, mu_noise, sigma_noise - - -def _fill_nans_infs_nwp_cascade( - precip_models_cascade, - precip_models, - precip_cascade, - precip, - mu_models, - sigma_models, -): - """Ensure that the NWP cascade and fields do no contain any nans or infinite number""" - # Fill nans and infinite numbers with the minimum value present in precip - # (corresponding to zero rainfall in the radar observations) - min_cascade = np.nanmin(precip_cascade) - min_precip = np.nanmin(precip) - precip_models_cascade[~np.isfinite(precip_models_cascade)] = min_cascade - precip_models[~np.isfinite(precip_models)] = min_precip - # Also set any nans or infs in the mean and sigma of the cascade to - # respectively 0.0 and 1.0 - mu_models[~np.isfinite(mu_models)] = 0.0 - sigma_models[~np.isfinite(sigma_models)] = 0.0 - - return precip_models_cascade, precip_models, mu_models, sigma_models - - -def _determine_max_nr_rainy_cells_nwp(precip_models, precip_thr, n_models, timesteps): - """Initialize noise based on the NWP field time step where the fraction of rainy cells is highest""" - if precip_thr is None: - precip_thr = np.nanmin(precip_models) - - max_rain_pixels = -1 - max_rain_pixels_j = -1 - max_rain_pixels_t = -1 - for j in range(n_models): - for t in timesteps: - rain_pixels = precip_models[j][t][precip_models[j][t] > precip_thr].size - if rain_pixels > max_rain_pixels: - max_rain_pixels = rain_pixels - max_rain_pixels_j = j - max_rain_pixels_t = t - precip_noise_input = precip_models[max_rain_pixels_j][max_rain_pixels_t] - - return precip_noise_input.astype(np.float64, copy=False) From d6249f5673a9a5a38a0fcb71a6ae9b9dcd485fcb Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 5 Dec 2024 17:22:05 +0100 Subject: [PATCH 14/33] 6 more tests that fail --- .gitignore | 1 + pysteps/blending/steps.py | 166 +++++++++++++++------------ pysteps/tests/test_blending_steps.py | 2 + 3 files changed, 96 insertions(+), 73 deletions(-) diff --git a/.gitignore b/.gitignore index 4588187d..8955e65b 100644 --- a/.gitignore +++ b/.gitignore @@ -94,3 +94,4 @@ venv.bak/ # Running lcoal tests /tmp +./pysteps/tests/tmp/* diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 0d461bcf..23635336 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -42,6 +42,8 @@ calculate_weights_spn blend_means_sigmas """ +# TODO: remove sys after debugging +import sys import math import time @@ -100,7 +102,7 @@ class StepsBlendingConfig: filter_kwargs: Dict[str, Any] = field(default_factory=dict) noise_kwargs: Dict[str, Any] = field(default_factory=dict) velocity_perturbation_kwargs: Dict[str, Any] = field(default_factory=dict) - clim_kwargs: Dict[str, Any] = field(default_factory=dict) + climatology_kwargs: Dict[str, Any] = field(default_factory=dict) mask_kwargs: Dict[str, Any] = field(default_factory=dict) measure_time: bool = False callback: Optional[Any] = None @@ -110,15 +112,12 @@ class StepsBlendingConfig: # TODO: typing could be improved here @dataclass class StepsBlendingParams: - PHI: np.ndarray # AR(p) model parameters - noise_std_coeffs: np.ndarray # Noise standard deviation coefficients - mu_extrapolation: np.ndarray # Means of extrapolated cascades - sigma_extrapolation: np.ndarray # Std devs of extrapolated cascades - bandpass_filter: Any # Band-pass filter object - fft: Any # FFT method object - perturbation_generator: Callable # Perturbation generator - noise_generator: Callable # Noise generator - generate_vel_noise: Optional[Callable] # Velocity noise generator + noise_std_coeffs: np.ndarray = None # Noise standard deviation coefficients + bandpass_filter: Any = None # Band-pass filter object + fft: Any = None # FFT method object + perturbation_generator: Callable = None # Perturbation generator + noise_generator: Callable = None # Noise generator + PHI: np.ndarray = None # AR(p) model parameters extrapolation_method: Any = None decomposition_method: Any = None recomposition_method: Any = None @@ -143,7 +142,6 @@ class StepsBlendingParams: mask_threshold: Any = None zero_precip_radar: bool = False zero_precip_model_fields: bool = False - PHI: Any = None original_timesteps: Any = None num_ensemble_workers: int = None rho_nwp_models: Any = None @@ -159,15 +157,12 @@ class StepsBlendingState: mean_extrapolation: Any = None std_extrapolation: Any = None precip_models_cascades: Any = None - PHI: Any = None randgen_precip: Any = None previous_displacement: Any = None previous_displacement_noise_cascade: Any = None previous_displacement_prob_matching: Any = None precip_forecast: Any = None precip_forecast_non_perturbed: Any = None - mask_rim: Any = None - struct: Any = None fft_objs: Any = None t_prev_timestep: Any = None t_leadtime_since_start_forecast: Any = None @@ -194,7 +189,6 @@ class StepsBlendingState: precip_forecast_extrapolated_probability_matching: Any = None precip_forecast_prev_subtimestep: Any = None noise_prev_subtimestep: Any = None - final_blended_forecast_single_member: Any = None means_blended: Any = None sigmas_blended: Any = None means_blended_mod_only: Any = None @@ -232,9 +226,6 @@ def __init__( self.__params = StepsBlendingParams() self.__state = StepsBlendingState() - # Initialize nowcast components and parameters - self.__initialize_nowcast_components() - # Additional variables for time measurement self.__start_time_init = None self.__zero_precip_time = None @@ -273,7 +264,7 @@ def compute_forecast(self): "initialization", self.__start_time_init ) - self.__blended_nowcast_main() + self.__blended_nowcast_main_loop() # Stack and return the forecast output if self.__config.return_output: self.__state.precip_forecast = np.stack( @@ -293,7 +284,7 @@ def compute_forecast(self): else: return None - def __blended_nowcast_main(self): + def __blended_nowcast_main_loop(self): """ Main nowcast loop that iterates through the ensemble members and time steps to generate forecasts. @@ -333,7 +324,9 @@ def __blended_nowcast_main(self): self.__find_nowcast_NWP_combination(t) self.__determine_skill_for_current_timestep(t) # the nowcast iteration for each ensemble member - precip_forecast_workers = [None for _ in range(self.__config.n_ens_members)] + precip_ensemble_single_timestep = [ + None for _ in range(self.__config.n_ens_members) + ] def worker(j): self.__determine_skill_for_next_timestep(t, j) @@ -343,29 +336,36 @@ def worker(j): t, j ) # 8.5 Blend the cascades - self.__state.final_blended_forecast_single_member = [] + final_blended_forecast_single_member = [] for t_sub in self.__state.subtimesteps: # TODO: does it make sense to use sub time steps - check if it works? if t_sub > 0: self.__blend_cascades(t_sub, j) self.__recompose_cascade_to_rainfall_field(j) # TODO: could be I need to return and ave final_blended_forecast_single_member - self.__post_process_output(j) - precip_forecast_workers[j] = ( - self.__state.final_blended_forecast_single_member + final_blended_forecast_single_member = ( + self.__post_process_output( + j, final_blended_forecast_single_member + ) + ) + precip_ensemble_single_timestep[j] = ( + final_blended_forecast_single_member ) - result = [] + dask_worker_collection = [] if DASK_IMPORTED and self.__config.n_ens_members > 1: for j in range(self.__config.n_ens_members): - result.append(dask.delayed(worker)(j)) - dask.compute(*result, num_workers=self.__params.num_ensemble_workers) + dask_worker_collection.append(dask.delayed(worker)(j)) + dask.compute( + *dask_worker_collection, + num_workers=self.__params.num_ensemble_workers, + ) else: for j in range(self.__config.n_ens_members): worker(j) - result = None + dask_worker_collection = None if self.__state.is_nowcast_time_step: if self.__config.measure_time: @@ -374,15 +374,17 @@ def worker(j): print("done.") if self.__config.callback is not None: - precip_forecast_final = np.stack(precip_forecast_workers) + precip_forecast_final = np.stack(precip_ensemble_single_timestep) if precip_forecast_final.shape[1] > 0: self.__config.callback(precip_forecast_final.squeeze()) if self.__config.return_output: for j in range(self.__config.n_ens_members): - self.__state.precip_forecast[j].extend(precip_forecast_workers[j]) + self.__state.precip_forecast[j].extend( + precip_ensemble_single_timestep[j] + ) - precip_forecast_workers = None + precip_ensemble_single_timestep = None if self.__config.measure_time: self.__mainloop_time = time.time() - starttime_mainloop @@ -424,30 +426,33 @@ def __check_inputs(self): raise ValueError( "The number of members in the precipitation models and velocity models must match" ) - + print(self.__timesteps, file=sys.stderr) if isinstance(self.__timesteps, list): self.__params.time_steps_is_list = True - self.__params.original_timesteps = [0] + list(self.__timesteps) - self.__timesteps = nowcast_utils.binned_timesteps( - self.__params.original_timesteps - ) if not sorted(self.__timesteps) == self.__timesteps: - raise ValueError("timesteps is not in ascending order") + raise ValueError( + "timesteps is not in ascending order", self.__timesteps + ) if self.__precip_models.shape[1] != math.ceil(self.__timesteps[-1]) + 1: raise ValueError( "precip_models does not contain sufficient lead times for this forecast" ) + self.__params.original_timesteps = [0] + list(self.__timesteps) + self.__timesteps = nowcast_utils.binned_timesteps( + self.__params.original_timesteps + ) else: self.__params.time_steps_is_list = False - self.__timesteps = list(range(self.__timesteps + 1)) if self.__precip_models.shape[1] != self.__timesteps + 1: raise ValueError( "precip_models does not contain sufficient lead times for this forecast" ) + self.__timesteps = list(range(self.__timesteps + 1)) + print(self.__timesteps, file=sys.stderr) precip_nwp_dim = self.__precip_models.ndim if precip_nwp_dim == 2: - if isinstance(self.__precip_models[0], dict): + if isinstance(self.__precip_models[0][0], dict): # It's a 2D array of dictionaries with decomposed cascades self.__params.precip_models_provided_is_cascade = True else: @@ -474,15 +479,14 @@ def __check_inputs(self): if self.__config.velocity_perturbation_kwargs is None: self.__config.velocity_perturbation_kwargs = dict() - if not self.__params.precip_models_provided_is_cascade: - if self.__config.clim_kwargs is None: - # Make sure clim_kwargs at least contains the number of models - self.__config.clim_kwargs = dict( - {"n_models": self.__precip_models.shape[0]} - ) + if self.__config.climatology_kwargs is None: + # Make sure clim_kwargs at least contains the number of models + self.__config.climatology_kwargs = dict( + {"n_models": self.__precip_models.shape[0]} + ) if self.__config.mask_kwargs is None: - mask_kwargs = dict() + self.__config.mask_kwargs = dict() if np.any(~np.isfinite(self.__velocity)): raise ValueError("velocity contains non-finite values") @@ -649,10 +653,11 @@ def __initialize_nowcast_components(self): x_values, y_values = np.meshgrid(np.arange(N), np.arange(M)) self.__params.xy_coordinates = np.stack([x_values, y_values]) - precip_copy = self.__precip[-(self.__config.ar_order + 1) :, :, :].copy() + # TODO: changed precip_copy for self.__precip + self.__precip = self.__precip[-(self.__config.ar_order + 1) :, :, :].copy() # Determine the domain mask from non-finite values in the precipitation data self.__params.domain_mask = np.logical_or.reduce( - [~np.isfinite(precip_copy[i, :]) for i in range(precip_copy.shape[0])] + [~np.isfinite(self.__precip[i, :]) for i in range(self.__precip.shape[0])] ) print("Blended nowcast components initialized successfully.") @@ -692,7 +697,7 @@ def __prepare_radar_and_NWP_fields(self): """Compute the cascade decompositions of the input precipitation fields.""" precip_forecast_decomp = [] for i in range(self.__config.ar_order + 1): - precip_forecast = self.__params.extrapolation_method( + precip_forecast = self.__params.decomposition_method( self.__precip[i, :, :], self.__params.bandpass_filter, mask=self.__params.mask_threshold, @@ -754,7 +759,7 @@ def __zero_precipitation_forecast(self): # Create an empty np array with shape [n_ens_members, rows, cols] # and fill it with the minimum value from precip (corresponding to # zero precipitation) - N, M = self.__precip.shape + N, M = self.__precip.shape[1:] precip_forecast_workers = np.full( (self.__config.n_ens_members, N, M), self.__params.precip_zerovalue ) @@ -861,6 +866,9 @@ def __prepare_nowcast_for_zero_radar(self): self.__state.precip_noise_input = self.__precip_models[max_rain_pixels_j][ max_rain_pixels_t ] + self.__state.precip_noise_input = self.__state.precip_noise_input.astype( + np.float64, copy=False + ) # Make sure precip_noise_input is three-dimensional if len(self.__state.precip_noise_input.shape) != 3: @@ -878,7 +886,7 @@ def __initialize_noise(self): # initialize the perturbation generator for the precipitation field self.__params.perturbation_generator = init_noise( - self.__precip, + self.__state.precip_noise_input, fft_method=self.__params.fft, **self.__config.noise_kwargs, ) @@ -888,9 +896,9 @@ def __initialize_noise(self): if self.__config.measure_time: starttime = time.time() - precip_forecast_min = np.min(self.__precip) + precip_forecast_min = np.min(self.__state.precip_noise_input) self.__params.noise_std_coeffs = noise.utils.compute_noise_stddev_adjs( - self.__precip[-1, :, :], + self.__state.precip_noise_input[-1, :, :], self.__config.precip_threshold, precip_forecast_min, self.__params.bandpass_filter, @@ -1069,15 +1077,15 @@ def __prepare_forecast_loop(self): if self.__config.mask_method == "incremental": # get mask parameters - self.__state.mask_rim = self.__config.mask_kwargs.get("mask_rim", 10) + self.__params.mask_rim = self.__config.mask_kwargs.get("mask_rim", 10) mask_f = self.__config.mask_kwargs.get("mask_f", 1.0) # initialize the structuring element struct = generate_binary_structure(2, 1) # iterate it to expand it nxn n = mask_f * self.__config.timestep / self.__config.kmperpixel - self.__state.struct = iterate_structure(struct, int((n - 1) / 2.0)) + self.__params.struct = iterate_structure(struct, int((n - 1) / 2.0)) else: - self.__state.mask_rim, self.__state.struct = None, None + self.__params.mask_rim, self.__params.struct = None, None if self.__config.noise_method is None: self.__state.precip_forecast_non_perturbed = [ @@ -1348,33 +1356,35 @@ def __find_nowcast_NWP_combination(self, t): def __determine_skill_for_current_timestep(self, t): if t == 0: """Calculate the initial skill of the (NWP) model forecasts at t=0.""" - # TODO: n_model is not defined here, how does this work? + # TODO rewrite loop self.__params.rho_nwp_models = [ blending.skill_scores.spatial_correlation( obs=self.__state.precip_cascades[0, :, -1, :, :].copy(), mod=self.__state.precip_models_cascades_timestep[ - n_model, :, :, : + model_index, :, :, : ].copy(), domain_mask=self.__params.domain_mask, ) - for n_model in range( + for model_index in range( self.__state.precip_models_cascades_timestep.shape[0] ) ] self.__params.rho_nwp_models = np.stack(self.__params.rho_nwp_models) # Ensure that the model skill decreases with increasing scale level. - for n_model in range(self.__state.precip_models_cascades_timestep.shape[0]): + for model_index in range( + self.__state.precip_models_cascades_timestep.shape[0] + ): for i in range( 1, self.__state.precip_models_cascades_timestep.shape[1] ): if ( - self.__params.rho_nwp_models[n_model, i] - > self.__params.rho_nwp_models[n_model, i - 1] + self.__params.rho_nwp_models[model_index, i] + > self.__params.rho_nwp_models[model_index, i - 1] ): # Set it equal to the previous scale level - self.__params.rho_nwp_models[n_model, i] = ( - self.__params.rho_nwp_models[n_model, i - 1] + self.__params.rho_nwp_models[model_index, i] = ( + self.__params.rho_nwp_models[model_index, i - 1] ) # Save this in the climatological skill file @@ -1382,7 +1392,7 @@ def __determine_skill_for_current_timestep(self, t): current_skill=self.__params.rho_nwp_models, validtime=self.__issuetime, outdir_path=self.__config.outdir_path_skill, - **self.__config.clim_kwargs, + **self.__config.climatology_kwargs, ) if t > 0: # 8.1.3 Determine the skill of the components for lead time (t0 + t) @@ -1400,15 +1410,16 @@ def __determine_skill_for_next_timestep(self, t, j): # 8.1.2 Determine the skill of the nwp components for lead time (t0 + t) # Then for the model components if self.__config.blend_nwp_members: + # TODO rewrite loop rho_nwp_forecast = [ blending.skill_scores.lt_dependent_cor_nwp( lt=(t * int(self.__config.timestep)), - correlations=self.__params.rho_nwp_models[n_model], + correlations=self.__params.rho_nwp_models[model_index], outdir_path=self.__config.outdir_path_skill, - n_model=n_model, - skill_kwargs=self.__config.clim_kwargs, + n_model=model_index, + skill_kwargs=self.__config.climatology_kwargs, ) - for n_model in range(self.__params.rho_nwp_models.shape[0]) + for model_index in range(self.__params.rho_nwp_models.shape[0]) ] rho_nwp_forecast = np.stack(rho_nwp_forecast) # Concatenate rho_extrap_cascade and rho_nwp @@ -1416,12 +1427,13 @@ def __determine_skill_for_next_timestep(self, t, j): (self.__state.rho_extrap_cascade[None, :], rho_nwp_forecast), axis=0 ) else: + # TODO: check if j is the best accessor for this variable rho_nwp_forecast = blending.skill_scores.lt_dependent_cor_nwp( lt=(t * int(self.__config.timestep)), correlations=self.__params.rho_nwp_models[j], outdir_path=self.__config.outdir_path_skill, n_model=self.__state.n_model_indices[j], - skill_kwargs=self.__config.clim_kwargs, + skill_kwargs=self.__config.climatology_kwargs, ) # Concatenate rho_extrap_cascade and rho_nwp self.__state.rho_forecast = np.concatenate( @@ -2072,7 +2084,7 @@ def __recompose_cascade_to_rainfall_field(self, j): j ].irfft2(self.__state.precip_forecast_recomposed_mod_only) - def __post_process_output(self, j): + def __post_process_output(self, j, final_blended_forecast_single_member): # 8.7 Post-processing steps - use the mask and fill no data with # the blended NWP forecast. Probability matching following # Lagrangian blended probability matching which uses the @@ -2316,9 +2328,10 @@ def __post_process_output(self, j): ) precip_forecast_probability_matching_resampled = None - self.__state.final_blended_forecast_single_member.append( + final_blended_forecast_single_member.append( self.__state.precip_forecast_recomposed ) + return final_blended_forecast_single_member def __measure_time(self, label, start_time): """ @@ -2757,7 +2770,9 @@ def forecast( blending_config = StepsBlendingConfig( n_ens_members=n_ens_members, n_cascade_levels=n_cascade_levels, + blend_nwp_members=blend_nwp_members, precip_threshold=precip_thr, + norain_threshold=norain_thr, kmperpixel=kmperpixel, timestep=timestep, extrapolation_method=extrap_method, @@ -2767,17 +2782,22 @@ def forecast( noise_stddev_adj=noise_stddev_adj, ar_order=ar_order, velocity_perturbation_method=vel_pert_method, + weights_method=weights_method, conditional=conditional, probmatching_method=probmatching_method, mask_method=mask_method, + resample_distribution=resample_distribution, + smooth_radar_mask_range=smooth_radar_mask_range, seed=seed, num_workers=num_workers, fft_method=fft_method, domain=domain, + outdir_path_skill=outdir_path_skill, extrapolation_kwargs=extrap_kwargs, filter_kwargs=filter_kwargs, noise_kwargs=noise_kwargs, velocity_perturbation_kwargs=vel_pert_kwargs, + climatology_kwargs=clim_kwargs, mask_kwargs=mask_kwargs, measure_time=measure_time, callback=callback, diff --git a/pysteps/tests/test_blending_steps.py b/pysteps/tests/test_blending_steps.py index ac2d16b9..18a4e90a 100644 --- a/pysteps/tests/test_blending_steps.py +++ b/pysteps/tests/test_blending_steps.py @@ -25,6 +25,8 @@ (5, 3, 5, 8, "incremental", "cdf", False, "spn", True, 5, False, False, 0, False), (1, 10, 1, 8, "incremental", "cdf", False, "spn", True, 1, False, False, 0, False), (2, 3, 2, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False), + # TODO: make next test work! This is currently not working on the main branch + # (2, 3, 4, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False), (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False), # Test the case where the radar image contains no rain. (1, 3, 6, 8, None, None, False, "spn", True, 6, True, False, 0, False), From 38702b3bc46f22fb02561d9b190cb93da7318af9 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 5 Dec 2024 18:36:39 +0100 Subject: [PATCH 15/33] All tests pass, still need to fix TODOs --- .gitignore | 2 +- pysteps/blending/steps.py | 371 ++++++++++++++------------- pysteps/tests/test_blending_steps.py | 2 +- 3 files changed, 188 insertions(+), 187 deletions(-) diff --git a/.gitignore b/.gitignore index 8955e65b..c865918f 100644 --- a/.gitignore +++ b/.gitignore @@ -94,4 +94,4 @@ venv.bak/ # Running lcoal tests /tmp -./pysteps/tests/tmp/* +./pysteps/tests/tmp/ diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 23635336..fd110187 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -42,9 +42,6 @@ calculate_weights_spn blend_means_sigmas """ -# TODO: remove sys after debugging -import sys - import math import time from copy import deepcopy @@ -246,7 +243,7 @@ def compute_forecast(self): # Determine if rain is present in both radar and NWP fields if self.__params.zero_precip_radar and self.__params.zero_precip_model_fields: - self.__zero_precipitation_forecast() + return self.__zero_precipitation_forecast() else: # Prepare the data for the zero precipitation radar case and initialize the noise correctly if self.__params.zero_precip_radar: @@ -329,28 +326,30 @@ def __blended_nowcast_main_loop(self): ] def worker(j): - self.__determine_skill_for_next_timestep(t, j) - self.__determine_weights_per_component() - self.__regress_extrapolation_and_noise_cascades(j) + # The state needs to be copied as a dataclass is not threadsafe in python + worker_state = deepcopy(self.__state) + self.__determine_skill_for_next_timestep(t, j, worker_state) + self.__determine_weights_per_component(worker_state) + self.__regress_extrapolation_and_noise_cascades(j, worker_state) self.__perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( - t, j + t, j, worker_state ) # 8.5 Blend the cascades final_blended_forecast_single_member = [] for t_sub in self.__state.subtimesteps: # TODO: does it make sense to use sub time steps - check if it works? if t_sub > 0: - self.__blend_cascades(t_sub, j) - self.__recompose_cascade_to_rainfall_field(j) + self.__blend_cascades(t_sub, j, worker_state) + self.__recompose_cascade_to_rainfall_field(j, worker_state) # TODO: could be I need to return and ave final_blended_forecast_single_member final_blended_forecast_single_member = ( self.__post_process_output( - j, final_blended_forecast_single_member + j, final_blended_forecast_single_member, worker_state ) ) - precip_ensemble_single_timestep[j] = ( - final_blended_forecast_single_member - ) + precip_ensemble_single_timestep[j] = ( + final_blended_forecast_single_member + ) dask_worker_collection = [] @@ -426,7 +425,7 @@ def __check_inputs(self): raise ValueError( "The number of members in the precipitation models and velocity models must match" ) - print(self.__timesteps, file=sys.stderr) + if isinstance(self.__timesteps, list): self.__params.time_steps_is_list = True if not sorted(self.__timesteps) == self.__timesteps: @@ -448,7 +447,6 @@ def __check_inputs(self): "precip_models does not contain sufficient lead times for this forecast" ) self.__timesteps = list(range(self.__timesteps + 1)) - print(self.__timesteps, file=sys.stderr) precip_nwp_dim = self.__precip_models.ndim if precip_nwp_dim == 2: @@ -770,7 +768,6 @@ def __zero_precipitation_forecast(self): if self.__config.return_output: for j in range(self.__config.n_ens_members): precip_forecast[j].append(precip_forecast_workers[j]) - precip_forecast_workers = None if self.__config.measure_time: @@ -783,6 +780,7 @@ def __zero_precipitation_forecast(self): for j in range(self.__config.n_ens_members) ] ) + if self.__config.measure_time: return ( precip_forecast_all_members_all_times, @@ -1406,7 +1404,7 @@ def __determine_skill_for_current_timestep(self, t): correlations_prev=self.__state.rho_extrap_cascade_prev, ) - def __determine_skill_for_next_timestep(self, t, j): + def __determine_skill_for_next_timestep(self, t, j, worker_state): # 8.1.2 Determine the skill of the nwp components for lead time (t0 + t) # Then for the model components if self.__config.blend_nwp_members: @@ -1423,8 +1421,8 @@ def __determine_skill_for_next_timestep(self, t, j): ] rho_nwp_forecast = np.stack(rho_nwp_forecast) # Concatenate rho_extrap_cascade and rho_nwp - self.__state.rho_forecast = np.concatenate( - (self.__state.rho_extrap_cascade[None, :], rho_nwp_forecast), axis=0 + worker_state.rho_forecast = np.concatenate( + (worker_state.rho_extrap_cascade[None, :], rho_nwp_forecast), axis=0 ) else: # TODO: check if j is the best accessor for this variable @@ -1432,16 +1430,16 @@ def __determine_skill_for_next_timestep(self, t, j): lt=(t * int(self.__config.timestep)), correlations=self.__params.rho_nwp_models[j], outdir_path=self.__config.outdir_path_skill, - n_model=self.__state.n_model_indices[j], + n_model=worker_state.n_model_indices[j], skill_kwargs=self.__config.climatology_kwargs, ) # Concatenate rho_extrap_cascade and rho_nwp - self.__state.rho_forecast = np.concatenate( - (self.__state.rho_extrap_cascade[None, :], rho_nwp_forecast[None, :]), + worker_state.rho_forecast = np.concatenate( + (worker_state.rho_extrap_cascade[None, :], rho_nwp_forecast[None, :]), axis=0, ) - def __determine_weights_per_component(self): + def __determine_weights_per_component(self, worker_state): # 8.2 Determine the weights per component # Weights following the bps method. These are needed for the velocity @@ -1449,14 +1447,14 @@ def __determine_weights_per_component(self): # selected, weights will be overwritten with those weights prior to # blending step. # weight = [(extr_field, n_model_fields, noise), n_cascade_levels, ...] - self.__state.weights = calculate_weights_bps(self.__state.rho_forecast) + worker_state.weights = calculate_weights_bps(worker_state.rho_forecast) # The model only weights if self.__config.weights_method == "bps": # Determine the weights of the components without the extrapolation # cascade, in case this is no data or outside the mask. - self.__state.weights_model_only = calculate_weights_bps( - self.__state.rho_forecast[1:, :] + worker_state.weights_model_only = calculate_weights_bps( + worker_state.rho_forecast[1:, :] ) elif self.__config.weights_method == "spn": # Only the weights of the components without the extrapolation @@ -1464,11 +1462,11 @@ def __determine_weights_per_component(self): # determined after the extrapolation step in this method. if ( self.__config.blend_nwp_members - and self.__state.precip_models_cascades_timestep.shape[0] > 1 + and worker_state.precip_models_cascades_timestep.shape[0] > 1 ): - self.__state.weights_model_only = np.zeros( + worker_state.weights_model_only = np.zeros( ( - self.__state.precip_models_cascades_timestep.shape[0] + 1, + worker_state.precip_models_cascades_timestep.shape[0] + 1, self.__config.n_cascade_levels, ) ) @@ -1478,11 +1476,11 @@ def __determine_weights_per_component(self): covariance_nwp_models = np.corrcoef( np.stack( [ - self.__state.precip_models_cascades_timestep[ + worker_state.precip_models_cascades_timestep[ n_model, i, :, : ].flatten() for n_model in range( - self.__state.precip_models_cascades_timestep.shape[ + worker_state.precip_models_cascades_timestep.shape[ 0 ] ) @@ -1490,14 +1488,14 @@ def __determine_weights_per_component(self): ) ) # Determine the weights for this cascade level - self.__state.weights_model_only[:, i] = calculate_weights_spn( - correlations=self.__state.rho_forecast[1:, i], + worker_state.weights_model_only[:, i] = calculate_weights_spn( + correlations=worker_state.rho_forecast[1:, i], covariance=covariance_nwp_models, ) else: # Same as correlation and noise is 1 - correlation - self.__state.weights_model_only = calculate_weights_bps( - self.__state.rho_forecast[1:, :] + worker_state.weights_model_only = calculate_weights_bps( + worker_state.rho_forecast[1:, :] ) else: raise ValueError( @@ -1505,7 +1503,7 @@ def __determine_weights_per_component(self): % self.__config.weights_method ) - def __regress_extrapolation_and_noise_cascades(self, j): + def __regress_extrapolation_and_noise_cascades(self, j, worker_state): # 8.3 Determine the noise cascade and regress this to the subsequent # time step + regress the extrapolation component to the subsequent # time step @@ -1516,8 +1514,8 @@ def __regress_extrapolation_and_noise_cascades(self, j): # generate noise field epsilon = self.__params.noise_generator( self.__params.perturbation_generator, - randstate=self.__state.randgen_precip[j], - fft_method=self.__state.fft_objs[j], + randstate=worker_state.randgen_precip[j], + fft_method=worker_state.fft_objs[j], domain=self.__config.domain, ) @@ -1525,7 +1523,7 @@ def __regress_extrapolation_and_noise_cascades(self, j): epsilon_decomposed = self.__params.decomposition_method( epsilon, self.__params.bandpass_filter, - fft_method=self.__state.fft_objs[j], + fft_method=worker_state.fft_objs[j], input_domain=self.__config.domain, output_domain=self.__config.domain, compute_stats=True, @@ -1544,18 +1542,18 @@ def __regress_extrapolation_and_noise_cascades(self, j): epsilon_decomposed is not None or self.__config.velocity_perturbation_method is not None ): - self.__state.precip_cascades[j][i] = autoregression.iterate_ar_model( - self.__state.precip_cascades[j][i], self.__params.PHI[i, :] + worker_state.precip_cascades[j][i] = autoregression.iterate_ar_model( + worker_state.precip_cascades[j][i], self.__params.PHI[i, :] ) # Renormalize the cascade - self.__state.precip_cascades[j][i][1] /= np.std( - self.__state.precip_cascades[j][i][1] + worker_state.precip_cascades[j][i][1] /= np.std( + worker_state.precip_cascades[j][i][1] ) else: # use the deterministic AR(p) model computed above if # perturbations are disabled - self.__state.precip_cascades[j][i] = ( - self.__state.precip_forecast_non_perturbed[i] + worker_state.precip_cascades[j][i] = ( + worker_state.precip_forecast_non_perturbed[i] ) # 8.3.3 regress the noise component to the subsequent time step @@ -1569,8 +1567,8 @@ def __regress_extrapolation_and_noise_cascades(self, j): epsilon_temp = None # apply AR(p) process to noise cascade level # (Returns zero noise if epsilon_decomposed is None) - self.__state.precip_noise_cascades[j][i] = autoregression.iterate_ar_model( - self.__state.precip_noise_cascades[j][i], + worker_state.precip_noise_cascades[j][i] = autoregression.iterate_ar_model( + worker_state.precip_noise_cascades[j][i], self.__params.PHI[i, :], eps=epsilon_temp, ) @@ -1579,7 +1577,7 @@ def __regress_extrapolation_and_noise_cascades(self, j): epsilon_temp = None def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( - self, t, j + self, t, j, worker_state ): # 8.4 Perturb and blend the advection fields + advect the # extrapolation and noise cascade to the current time step @@ -1591,37 +1589,37 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( extrap_kwargs_pb = self.__config.extrapolation_kwargs.copy() velocity_perturbations_extrapolation = self.__velocity # The following should be accesseble after this function - self.__state.precip_forecast_extrapolated_decomp_done = [] - self.__state.noise_extrapolated_decomp_done = [] - self.__state.precip_forecast_extrapolated_probability_matching = [] + worker_state.precip_forecast_extrapolated_decomp_done = [] + worker_state.noise_extrapolated_decomp_done = [] + worker_state.precip_forecast_extrapolated_probability_matching = [] # Extrapolate per sub time step - for t_sub in self.__state.subtimesteps: + for t_sub in worker_state.subtimesteps: if t_sub > 0: t_diff_prev_subtimestep_int = t_sub - int(t_sub) if t_diff_prev_subtimestep_int > 0.0: precip_forecast_cascade_subtimestep = [ (1.0 - t_diff_prev_subtimestep_int) - * self.__state.precip_forecast_prev_subtimestep[j][i][-1, :] + * worker_state.precip_forecast_prev_subtimestep[j][i][-1, :] + t_diff_prev_subtimestep_int - * self.__state.precip_cascades[j][i][-1, :] + * worker_state.precip_cascades[j][i][-1, :] for i in range(self.__config.n_cascade_levels) ] noise_cascade_subtimestep = [ (1.0 - t_diff_prev_subtimestep_int) - * self.__state.noise_prev_subtimestep[j][i][-1, :] + * worker_state.noise_prev_subtimestep[j][i][-1, :] + t_diff_prev_subtimestep_int - * self.__state.precip_noise_cascades[j][i][-1, :] + * worker_state.precip_noise_cascades[j][i][-1, :] for i in range(self.__config.n_cascade_levels) ] else: precip_forecast_cascade_subtimestep = [ - self.__state.precip_forecast_prev_subtimestep[j][i][-1, :] + worker_state.precip_forecast_prev_subtimestep[j][i][-1, :] for i in range(self.__config.n_cascade_levels) ] noise_cascade_subtimestep = [ - self.__state.noise_prev_subtimestep[j][i][-1, :] + worker_state.noise_prev_subtimestep[j][i][-1, :] for i in range(self.__config.n_cascade_levels) ] @@ -1630,8 +1628,8 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( ) noise_cascade_subtimestep = np.stack(noise_cascade_subtimestep) - t_diff_prev_subtimestep = t_sub - self.__state.t_prev_timestep[j] - self.__state.t_leadtime_since_start_forecast[ + t_diff_prev_subtimestep = t_sub - worker_state.t_prev_timestep[j] + worker_state.t_leadtime_since_start_forecast[ j ] += t_diff_prev_subtimestep @@ -1644,7 +1642,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( self.__velocity + self.__params.generate_velocity_noise( self.__params.velocity_perturbations[j], - self.__state.t_leadtime_since_start_forecast[j] + worker_state.t_leadtime_since_start_forecast[j] * self.__config.timestep, ) ) @@ -1654,12 +1652,12 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( velocity_stack_all = np.concatenate( ( velocity_perturbations_extrapolation[None, :, :, :], - self.__state.velocity_models_timestep, + worker_state.velocity_models_timestep, ), axis=0, ) else: - velocity_models = self.__state.velocity_models_timestep[j] + velocity_models = worker_state.velocity_models_timestep[j] velocity_stack_all = np.concatenate( ( velocity_perturbations_extrapolation[None, :, :, :], @@ -1673,7 +1671,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( # second cascade following eq. 24 in BPS2006 velocity_blended = blending.utils.blend_optical_flows( flows=velocity_stack_all, - weights=self.__state.weights[ + weights=worker_state.weights[ :-1, 1 ], # [(extr_field, n_model_fields), cascade_level=2] ) @@ -1685,8 +1683,8 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( # A. Radar Rain precip_forecast_recomp_subtimestep = blending.utils.recompose_cascade( combined_cascade=precip_forecast_cascade_subtimestep, - combined_mean=self.__state.mean_extrapolation, - combined_sigma=self.__state.std_extrapolation, + combined_mean=worker_state.mean_extrapolation, + combined_sigma=worker_state.std_extrapolation, ) # Make sure we have values outside the mask if self.__params.zero_precip_radar: @@ -1702,11 +1700,11 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( # TODO: problem with the config here! This variable changes over time... # extrap_kwargs is in config but by adding info to it, the next run of a blended forecast will have issues! self.__config.extrapolation_kwargs["displacement_prev"] = ( - self.__state.previous_displacement[j] + worker_state.previous_displacement[j] ) ( precip_forecast_extrapolated_recomp_subtimestep_temp, - self.__state.previous_displacement[j], + worker_state.previous_displacement[j], ) = self.__params.extrapolation_method( precip_forecast_recomp_subtimestep, velocity_blended, @@ -1750,16 +1748,16 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( # B. Noise noise_cascade_subtimestep_recomp = blending.utils.recompose_cascade( combined_cascade=noise_cascade_subtimestep, - combined_mean=self.__state.precip_mean_noise[j], - combined_sigma=self.__state.precip_std_noise[j], + combined_mean=worker_state.precip_mean_noise[j], + combined_sigma=worker_state.precip_std_noise[j], ) extrap_kwargs_noise["displacement_prev"] = ( - self.__state.previous_displacement_noise_cascade[j] + worker_state.previous_displacement_noise_cascade[j] ) extrap_kwargs_noise["map_coordinates_mode"] = "wrap" ( noise_extrapolated_recomp_temp, - self.__state.previous_displacement_noise_cascade[j], + worker_state.previous_displacement_noise_cascade[j], ) = self.__params.extrapolation_method( noise_cascade_subtimestep_recomp, velocity_blended, @@ -1782,10 +1780,10 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( noise_extrapolated_decomp[i] *= self.__params.noise_std_coeffs[i] # Append the results to the output lists - self.__state.precip_forecast_extrapolated_decomp_done.append( + worker_state.precip_forecast_extrapolated_decomp_done.append( precip_forecast_extrapolated_decomp.copy() ) - self.__state.noise_extrapolated_decomp_done.append( + worker_state.noise_extrapolated_decomp_done.append( noise_extrapolated_decomp.copy() ) precip_forecast_cascade_subtimestep = None @@ -1804,7 +1802,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( # of the (NWP) model(s) for Lagrangian blended prob. matching # min_R = np.min(precip) extrap_kwargs_pb["displacement_prev"] = ( - self.__state.previous_displacement_prob_matching[j] + worker_state.previous_displacement_prob_matching[j] ) # Apply the domain mask to the extrapolation component precip_forecast_temp_for_probability_matching = self.__precip.copy() @@ -1813,7 +1811,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( ] = np.nan ( precip_forecast_extrapolated_probability_matching_temp, - self.__state.previous_displacement_prob_matching[j], + worker_state.previous_displacement_prob_matching[j], ) = self.__params.extrapolation_method( precip_forecast_temp_for_probability_matching, velocity_blended, @@ -1821,28 +1819,28 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( allow_nonfinite_values=True, **extrap_kwargs_pb, ) - self.__state.precip_forecast_extrapolated_probability_matching.append( + worker_state.precip_forecast_extrapolated_probability_matching.append( precip_forecast_extrapolated_probability_matching_temp[0] ) - self.__state.t_prev_timestep[j] = t_sub + worker_state.t_prev_timestep[j] = t_sub - if len(self.__state.precip_forecast_extrapolated_decomp_done) > 0: - self.__state.precip_forecast_extrapolated_decomp_done = np.stack( - self.__state.precip_forecast_extrapolated_decomp_done + if len(worker_state.precip_forecast_extrapolated_decomp_done) > 0: + worker_state.precip_forecast_extrapolated_decomp_done = np.stack( + worker_state.precip_forecast_extrapolated_decomp_done ) - self.__state.noise_extrapolated_decomp_done = np.stack( - self.__state.noise_extrapolated_decomp_done + worker_state.noise_extrapolated_decomp_done = np.stack( + worker_state.noise_extrapolated_decomp_done ) - self.__state.precip_forecast_extrapolated_probability_matching = np.stack( - self.__state.precip_forecast_extrapolated_probability_matching + worker_state.precip_forecast_extrapolated_probability_matching = np.stack( + worker_state.precip_forecast_extrapolated_probability_matching ) # advect the forecast field by one time step if no subtimesteps in the # current interval were found - if not self.__state.subtimesteps: - t_diff_prev_subtimestep = t + 1 - self.__state.t_prev_timestep[j] - self.__state.t_leadtime_since_start_forecast[j] += t_diff_prev_subtimestep + if not worker_state.subtimesteps: + t_diff_prev_subtimestep = t + 1 - worker_state.t_prev_timestep[j] + worker_state.t_leadtime_since_start_forecast[j] += t_diff_prev_subtimestep # compute the perturbed motion field - include the NWP # velocities and the weights @@ -1851,7 +1849,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( self.__velocity + self.__params.generate_velocity_noise( self.__params.velocity_perturbations[j], - self.__state.t_leadtime_since_start_forecast[j] + worker_state.t_leadtime_since_start_forecast[j] * self.__config.timestep, ) ) @@ -1861,12 +1859,12 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( velocity_stack_all = np.concatenate( ( velocity_perturbations_extrapolation[None, :, :, :], - self.__state.velocity_models_timestep, + worker_state.velocity_models_timestep, ), axis=0, ) else: - velocity_models = self.__state.velocity_models_timestep[j] + velocity_models = worker_state.velocity_models_timestep[j] velocity_stack_all = np.concatenate( ( velocity_perturbations_extrapolation[None, :, :, :], @@ -1880,20 +1878,20 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( # second cascade following eq. 24 in BPS2006 velocity_blended = blending.utils.blend_optical_flows( flows=velocity_stack_all, - weights=self.__state.weights[ + weights=worker_state.weights[ :-1, 1 ], # [(extr_field, n_model_fields), cascade_level=2] ) # Extrapolate the extrapolation and noise cascade - extrap_kwargs_["displacement_prev"] = self.__state.previous_displacement[j] + extrap_kwargs_["displacement_prev"] = worker_state.previous_displacement[j] extrap_kwargs_noise["displacement_prev"] = ( - self.__state.previous_displacement_noise_cascade[j] + worker_state.previous_displacement_noise_cascade[j] ) extrap_kwargs_noise["map_coordinates_mode"] = "wrap" - _, self.__state.previous_displacement[j] = ( + _, worker_state.previous_displacement[j] = ( self.__params.extrapolation_method( None, velocity_blended, @@ -1903,7 +1901,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( ) ) - _, self.__state.previous_displacement_noise_cascade[j] = ( + _, worker_state.previous_displacement_noise_cascade[j] = ( self.__params.extrapolation_method( None, velocity_blended, @@ -1916,9 +1914,9 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( # Also extrapolate the radar observation, used for the probability # matching and post-processing steps extrap_kwargs_pb["displacement_prev"] = ( - self.__state.previous_displacement_prob_matching[j] + worker_state.previous_displacement_prob_matching[j] ) - _, self.__state.previous_displacement_prob_matching[j] = ( + _, worker_state.previous_displacement_prob_matching[j] = ( self.__params.extrapolation_method( None, velocity_blended, @@ -1928,15 +1926,15 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( ) ) - self.__state.t_prev_timestep[j] = t + 1 + worker_state.t_prev_timestep[j] = t + 1 - self.__state.precip_forecast_prev_subtimestep[j] = self.__state.precip_cascades[ + worker_state.precip_forecast_prev_subtimestep[j] = worker_state.precip_cascades[ j ] - self.__state.noise_prev_subtimestep[j] = self.__state.precip_noise_cascades[j] + worker_state.noise_prev_subtimestep[j] = worker_state.precip_noise_cascades[j] - def __blend_cascades(self, t_sub, j): - self.__state.t_index = np.where(np.array(self.__state.subtimesteps) == t_sub)[ + def __blend_cascades(self, t_sub, j, worker_state): + worker_state.t_index = np.where(np.array(worker_state.subtimesteps) == t_sub)[ 0 ][0] # First concatenate the cascades and the means and sigmas @@ -1944,54 +1942,54 @@ def __blend_cascades(self, t_sub, j): if self.__config.blend_nwp_members: cascade_stack_all_components = np.concatenate( ( - self.__state.precip_forecast_extrapolated_decomp_done[ - None, self.__state.t_index + worker_state.precip_forecast_extrapolated_decomp_done[ + None, worker_state.t_index ], - self.__state.precip_models_cascades_timestep, - self.__state.noise_extrapolated_decomp_done[ - None, self.__state.t_index + worker_state.precip_models_cascades_timestep, + worker_state.noise_extrapolated_decomp_done[ + None, worker_state.t_index ], ), axis=0, ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] means_stacked = np.concatenate( ( - self.__state.mean_extrapolation[None, :], - self.__state.mean_models_timestep, + worker_state.mean_extrapolation[None, :], + worker_state.mean_models_timestep, ), axis=0, ) sigmas_stacked = np.concatenate( ( - self.__state.std_extrapolation[None, :], - self.__state.std_models_timestep, + worker_state.std_extrapolation[None, :], + worker_state.std_models_timestep, ), axis=0, ) else: cascade_stack_all_components = np.concatenate( ( - self.__state.precip_forecast_extrapolated_decomp_done[ - None, self.__state.t_index + worker_state.precip_forecast_extrapolated_decomp_done[ + None, worker_state.t_index ], - self.__state.precip_models_cascades_timestep[None, j], - self.__state.noise_extrapolated_decomp_done[ - None, self.__state.t_index + worker_state.precip_models_cascades_timestep[None, j], + worker_state.noise_extrapolated_decomp_done[ + None, worker_state.t_index ], ), axis=0, ) # [(extr_field, n_model_fields, noise), n_cascade_levels, ...] means_stacked = np.concatenate( ( - self.__state.mean_extrapolation[None, :], - self.__state.mean_models_timestep[None, j], + worker_state.mean_extrapolation[None, :], + worker_state.mean_models_timestep[None, j], ), axis=0, ) sigmas_stacked = np.concatenate( ( - self.__state.std_extrapolation[None, :], - self.__state.std_models_timestep[None, j], + worker_state.std_extrapolation[None, :], + worker_state.std_models_timestep[None, j], ), axis=0, ) @@ -2002,9 +2000,9 @@ def __blend_cascades(self, t_sub, j): # method is given? Or does this mean that in all other circumstances the weights # have been calculated in a different way? - # TODO: changed weights to self.__state.weights + # TODO: changed weights to worker_state.weights if self.__config.weights_method == "spn": - self.__state.weights = np.zeros( + worker_state.weights = np.zeros( ( cascade_stack_all_components.shape[0], self.__config.n_cascade_levels, @@ -2023,68 +2021,70 @@ def __blend_cascades(self, t_sub, j): np.ma.masked_invalid(cascade_stack_all_components_temp) ) # Determine the weights for this cascade level - self.__state.weights[:, i] = calculate_weights_spn( - correlations=self.__state.rho_forecast[:, i], + worker_state.weights[:, i] = calculate_weights_spn( + correlations=worker_state.rho_forecast[:, i], covariance=covariance_nwp_models, ) # Blend the extrapolation, (NWP) model(s) and noise cascades - self.__state.precip_forecast_blended = blending.utils.blend_cascades( - cascades_norm=cascade_stack_all_components, weights=self.__state.weights + worker_state.precip_forecast_blended = blending.utils.blend_cascades( + cascades_norm=cascade_stack_all_components, weights=worker_state.weights ) # Also blend the cascade without the extrapolation component - self.__state.precip_forecast_blended_mod_only = blending.utils.blend_cascades( + worker_state.precip_forecast_blended_mod_only = blending.utils.blend_cascades( cascades_norm=cascade_stack_all_components[1:, :], - weights=self.__state.weights_model_only, + weights=worker_state.weights_model_only, ) # Blend the means and standard deviations # Input is array of shape [number_components, scale_level, ...] - self.__state.means_blended, self.__state.sigmas_blended = blend_means_sigmas( - means=means_stacked, sigmas=sigmas_stacked, weights=self.__state.weights + worker_state.means_blended, worker_state.sigmas_blended = blend_means_sigmas( + means=means_stacked, sigmas=sigmas_stacked, weights=worker_state.weights ) # Also blend the means and sigmas for the cascade without extrapolation ( - self.__state.means_blended_mod_only, - self.__state.sigmas_blended_mod_only, + worker_state.means_blended_mod_only, + worker_state.sigmas_blended_mod_only, ) = blend_means_sigmas( means=means_stacked[1:, :], sigmas=sigmas_stacked[1:, :], - weights=self.__state.weights_model_only, + weights=worker_state.weights_model_only, ) - def __recompose_cascade_to_rainfall_field(self, j): + def __recompose_cascade_to_rainfall_field(self, j, worker_state): # 8.6 Recompose the cascade to a precipitation field # (The function first normalizes the blended cascade, precip_forecast_blended # again) - self.__state.precip_forecast_recomposed = blending.utils.recompose_cascade( - combined_cascade=self.__state.precip_forecast_blended, - combined_mean=self.__state.means_blended, - combined_sigma=self.__state.sigmas_blended, + worker_state.precip_forecast_recomposed = blending.utils.recompose_cascade( + combined_cascade=worker_state.precip_forecast_blended, + combined_mean=worker_state.means_blended, + combined_sigma=worker_state.sigmas_blended, ) # The recomposed cascade without the extrapolation (for NaN filling # outside the radar domain) - self.__state.precip_forecast_recomposed_mod_only = ( + worker_state.precip_forecast_recomposed_mod_only = ( blending.utils.recompose_cascade( - combined_cascade=self.__state.precip_forecast_blended_mod_only, - combined_mean=self.__state.means_blended_mod_only, - combined_sigma=self.__state.sigmas_blended_mod_only, + combined_cascade=worker_state.precip_forecast_blended_mod_only, + combined_mean=worker_state.means_blended_mod_only, + combined_sigma=worker_state.sigmas_blended_mod_only, ) ) if self.__config.domain == "spectral": # TODO: Check this! (Only tested with domain == 'spatial') # TODO: what needs to happen with above TODO? - self.__state.precip_forecast_recomposed = self.__state.fft_objs[j].irfft2( - self.__state.precip_forecast_recomposed + worker_state.precip_forecast_recomposed = worker_state.fft_objs[j].irfft2( + worker_state.precip_forecast_recomposed ) - self.__state.precip_forecast_recomposed_mod_only = self.__state.fft_objs[ + worker_state.precip_forecast_recomposed_mod_only = worker_state.fft_objs[ j - ].irfft2(self.__state.precip_forecast_recomposed_mod_only) + ].irfft2(worker_state.precip_forecast_recomposed_mod_only) - def __post_process_output(self, j, final_blended_forecast_single_member): + def __post_process_output( + self, j, final_blended_forecast_single_member, worker_state + ): # 8.7 Post-processing steps - use the mask and fill no data with # the blended NWP forecast. Probability matching following # Lagrangian blended probability matching which uses the @@ -2095,14 +2095,14 @@ def __post_process_output(self, j, final_blended_forecast_single_member): # that is only used for post-processing steps) with the NWP # rainfall forecast for this time step using the weights # at scale level 2. - weights_probability_matching = self.__state.weights[ + weights_probability_matching = worker_state.weights[ :-1, 1 ] # Weights without noise, level 2 weights_probability_matching_normalized = weights_probability_matching / np.sum( weights_probability_matching ) # And the weights for outside the radar domain - weights_probability_matching_mod_only = self.__state.weights_model_only[ + weights_probability_matching_mod_only = worker_state.weights_model_only[ :-1, 1 ] # Weights without noise, level 2 weights_probability_matching_normalized_mod_only = ( @@ -2113,20 +2113,20 @@ def __post_process_output(self, j, final_blended_forecast_single_member): if self.__config.blend_nwp_members: precip_forecast_probability_matching_final = np.concatenate( ( - self.__state.precip_forecast_extrapolated_probability_matching[ - None, self.__state.t_index + worker_state.precip_forecast_extrapolated_probability_matching[ + None, worker_state.t_index ], - self.__state.precip_models_timestep, + worker_state.precip_models_timestep, ), axis=0, ) else: precip_forecast_probability_matching_final = np.concatenate( ( - self.__state.precip_forecast_extrapolated_probability_matching[ - None, self.__state.t_index + worker_state.precip_forecast_extrapolated_probability_matching[ + None, worker_state.t_index ], - self.__state.precip_models_timestep[None, j], + worker_state.precip_models_timestep[None, j], ), axis=0, ) @@ -2145,12 +2145,12 @@ def __post_process_output(self, j, final_blended_forecast_single_member): 1, 1, ) - * self.__state.precip_models_timestep, + * worker_state.precip_models_timestep, axis=0, ) else: precip_forecast_probability_matching_blended_mod_only = ( - self.__state.precip_models_timestep[j] + worker_state.precip_models_timestep[j] ) # The extrapolation components are NaN outside the advected @@ -2159,7 +2159,7 @@ def __post_process_output(self, j, final_blended_forecast_single_member): # areas with the "..._mod_only" blended forecasts, consisting # of the NWP and noise components. - nan_indices = np.isnan(self.__state.precip_forecast_recomposed) + nan_indices = np.isnan(worker_state.precip_forecast_recomposed) if self.__config.smooth_radar_mask_range != 0: # Compute the smooth dilated mask new_mask = blending.utils.compute_smooth_dilated_mask( @@ -2173,10 +2173,10 @@ def __post_process_output(self, j, final_blended_forecast_single_member): # Handle NaNs in precip_forecast_new and precip_forecast_new_mod_only by setting NaNs to 0 in the blending step precip_forecast_recomposed_mod_only_no_nan = np.nan_to_num( - self.__state.precip_forecast_recomposed_mod_only, nan=0 + worker_state.precip_forecast_recomposed_mod_only, nan=0 ) precip_forecast_recomposed_no_nan = np.nan_to_num( - self.__state.precip_forecast_recomposed, nan=0 + worker_state.precip_forecast_recomposed, nan=0 ) # Perform the blending of radar and model inside the radar domain using a weighted combination @@ -2196,8 +2196,8 @@ def __post_process_output(self, j, final_blended_forecast_single_member): axis=0, ) else: - self.__state.precip_forecast_recomposed[nan_indices] = ( - self.__state.precip_forecast_recomposed_mod_only[nan_indices] + worker_state.precip_forecast_recomposed[nan_indices] = ( + worker_state.precip_forecast_recomposed_mod_only[nan_indices] ) nan_indices = np.isnan(precip_forecast_probability_matching_blended) precip_forecast_probability_matching_blended[nan_indices] = ( @@ -2206,9 +2206,9 @@ def __post_process_output(self, j, final_blended_forecast_single_member): # Finally, fill the remaining nan values, if present, with # the minimum value in the forecast - nan_indices = np.isnan(self.__state.precip_forecast_recomposed) - self.__state.precip_forecast_recomposed[nan_indices] = np.nanmin( - self.__state.precip_forecast_recomposed + nan_indices = np.isnan(worker_state.precip_forecast_recomposed) + worker_state.precip_forecast_recomposed[nan_indices] = np.nanmin( + worker_state.precip_forecast_recomposed ) nan_indices = np.isnan(precip_forecast_probability_matching_blended) precip_forecast_probability_matching_blended[nan_indices] = np.nanmin( @@ -2221,7 +2221,7 @@ def __post_process_output(self, j, final_blended_forecast_single_member): # apply the precipitation mask to prevent generation of new # precipitation into areas where it was not originally # observed - precip_forecast_min_value = self.__state.precip_forecast_recomposed.min() + precip_forecast_min_value = worker_state.precip_forecast_recomposed.min() if self.__config.mask_method == "incremental": # The incremental mask is slightly different from # the implementation in the non-blended steps.py, as @@ -2238,16 +2238,16 @@ def __post_process_output(self, j, final_blended_forecast_single_member): precip_field_mask, self.__params.struct, self.__params.mask_rim ) # Get the final mask - self.__state.precip_forecast_recomposed = ( + worker_state.precip_forecast_recomposed = ( precip_forecast_min_value + ( - self.__state.precip_forecast_recomposed + worker_state.precip_forecast_recomposed - precip_forecast_min_value ) * precip_field_mask ) precip_field_mask_temp = ( - self.__state.precip_forecast_recomposed > precip_forecast_min_value + worker_state.precip_forecast_recomposed > precip_forecast_min_value ) elif self.__config.mask_method == "obs": # The mask equals the most recent benchmark @@ -2258,7 +2258,7 @@ def __post_process_output(self, j, final_blended_forecast_single_member): ) # Set to min value outside of mask - self.__state.precip_forecast_recomposed[~precip_field_mask_temp] = ( + worker_state.precip_forecast_recomposed[~precip_field_mask_temp] = ( precip_forecast_min_value ) @@ -2269,10 +2269,10 @@ def __post_process_output(self, j, final_blended_forecast_single_member): self.__config.probmatching_method is not None and self.__config.resample_distribution ): - arr1 = self.__state.precip_forecast_extrapolated_probability_matching[ - self.__state.t_index + arr1 = worker_state.precip_forecast_extrapolated_probability_matching[ + worker_state.t_index ] - arr2 = self.__state.precip_models_timestep[j] + arr2 = worker_state.precip_models_timestep[j] # resample weights based on cascade level 2. # Areas where one of the fields is nan are not included. precip_forecast_probability_matching_resampled = ( @@ -2290,17 +2290,17 @@ def __post_process_output(self, j, final_blended_forecast_single_member): if self.__config.probmatching_method == "cdf": # nan indices in the extrapolation nowcast nan_indices = np.isnan( - self.__state.precip_forecast_extrapolated_probability_matching[ - self.__state.t_index + worker_state.precip_forecast_extrapolated_probability_matching[ + worker_state.t_index ] ) # Adjust the CDF of the forecast to match the resampled distribution combined from # extrapolation and model fields. # Rainfall outside the pure extrapolation domain is not taken into account. - if np.any(np.isfinite(self.__state.precip_forecast_recomposed)): - self.__state.precip_forecast_recomposed = ( + if np.any(np.isfinite(worker_state.precip_forecast_recomposed)): + worker_state.precip_forecast_recomposed = ( probmatching.nonparam_match_empirical_cdf( - self.__state.precip_forecast_recomposed, + worker_state.precip_forecast_recomposed, precip_forecast_probability_matching_resampled, nan_indices, ) @@ -2315,21 +2315,21 @@ def __post_process_output(self, j, final_blended_forecast_single_member): ] ) no_rain_mask = ( - self.__state.precip_forecast_recomposed + worker_state.precip_forecast_recomposed >= self.__config.precip_threshold ) mean_precip_forecast = np.mean( - self.__state.precip_forecast_recomposed[no_rain_mask] + worker_state.precip_forecast_recomposed[no_rain_mask] ) - self.__state.precip_forecast_recomposed[no_rain_mask] = ( - self.__state.precip_forecast_recomposed[no_rain_mask] + worker_state.precip_forecast_recomposed[no_rain_mask] = ( + worker_state.precip_forecast_recomposed[no_rain_mask] - mean_precip_forecast + mean_probabiltity_matching_forecast ) precip_forecast_probability_matching_resampled = None final_blended_forecast_single_member.append( - self.__state.precip_forecast_recomposed + worker_state.precip_forecast_recomposed ) return final_blended_forecast_single_member @@ -2816,6 +2816,7 @@ def forecast( ) forecast_steps_nowcast = blended_nowcaster.compute_forecast() + print(forecast_steps_nowcast) blended_nowcaster.reset_states_and_params() # Call the appropriate methods within the class return forecast_steps_nowcast diff --git a/pysteps/tests/test_blending_steps.py b/pysteps/tests/test_blending_steps.py index 18a4e90a..5840279f 100644 --- a/pysteps/tests/test_blending_steps.py +++ b/pysteps/tests/test_blending_steps.py @@ -9,6 +9,7 @@ from pysteps import blending, cascade steps_arg_values = [ + # Test the case where both the radar image and the NWP fields contain no rain. (1, 3, 4, 8, None, None, False, "spn", True, 4, False, False, 0, False), (1, 3, 4, 8, "obs", None, False, "spn", True, 4, False, False, 0, False), (1, 3, 4, 8, "incremental", None, False, "spn", True, 4, False, False, 0, False), @@ -35,7 +36,6 @@ # Test the case where the NWP fields contain no rain. (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 0, False), (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, True, 0, True), - # Test the case where both the radar image and the NWP fields contain no rain. (1, 3, 6, 8, None, None, False, "spn", True, 6, True, True, 0, False), (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, True, 0, False), (5, 3, 5, 6, "obs", "mean", True, "spn", True, 5, True, True, 0, False), From 5ff1713900495e58bcff41536cf7632df09ddf8c Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 5 Dec 2024 18:37:48 +0100 Subject: [PATCH 16/33] Updated gitignore --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index c865918f..136c1e46 100644 --- a/.gitignore +++ b/.gitignore @@ -94,4 +94,4 @@ venv.bak/ # Running lcoal tests /tmp -./pysteps/tests/tmp/ +/pysteps/tests/tmp/ From d999501c7d5269a322055b21f2294aafdfaa6764 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Fri, 6 Dec 2024 16:05:30 +0100 Subject: [PATCH 17/33] Cleanup of params and state dataclasses, next step: better typing --- pysteps/blending/steps.py | 426 +++++++++++++++++++++----------------- 1 file changed, 235 insertions(+), 191 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index fd110187..7fc42aeb 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -67,6 +67,10 @@ from dataclasses import dataclass, field from typing import Optional, List, Dict, Any, Callable +# TODO: compare old and new version of the code +# TODO: look for better typing in state and params +# TODO: GO over all other todos and check if they can be removed + @dataclass class StepsBlendingConfig: @@ -118,6 +122,9 @@ class StepsBlendingParams: extrapolation_method: Any = None decomposition_method: Any = None recomposition_method: Any = None + # TODO: check of the following two are relevant or can be replaced vel_pert_... and noise_generator + velocity_perturbations: Any = None + generate_velocity_noise: Any = None velocity_perturbations_parallel: Optional[np.ndarray] = ( None # Velocity perturbation parameters (parallel) ) @@ -129,9 +136,6 @@ class StepsBlendingParams: ) # FFT objects for ensemble members mask_rim: Optional[int] = None # Rim size for masking struct: Optional[np.ndarray] = None # Structuring element for mask - noise_method: Optional[str] = None # Noise method used - ar_order: int = 2 # Order of the AR model - seed: Optional[int] = None # Random seed for reproducibility time_steps_is_list: bool = False # Time steps is a list precip_models_provided_is_cascade: bool = False # Precip models are decomposed xy_coordinates: np.ndarray | None = None @@ -143,58 +147,75 @@ class StepsBlendingParams: num_ensemble_workers: int = None rho_nwp_models: Any = None domain_mask: Any = None - velocity_perturbations: Any = None - generate_velocity_noise: Any = None # TODO: typing could be improved here @dataclass class StepsBlendingState: + # States related to the observations precip_cascades: Any = None - mean_extrapolation: Any = None - std_extrapolation: Any = None - precip_models_cascades: Any = None - randgen_precip: Any = None - previous_displacement: Any = None - previous_displacement_noise_cascade: Any = None - previous_displacement_prob_matching: Any = None - precip_forecast: Any = None - precip_forecast_non_perturbed: Any = None - fft_objs: Any = None - t_prev_timestep: Any = None - t_leadtime_since_start_forecast: Any = None precip_noise_input: Any = None precip_noise_cascades: Any = None precip_mean_noise: Any = None precip_std_noise: Any = None + + # States related to the extrapolation + mean_extrapolation: Any = None + std_extrapolation: Any = None rho_extrap_cascade_prev: Any = None rho_extrap_cascade: Any = None - subtimesteps: Any = None - is_nowcast_time_step: bool = None - # Variables to save data over (sub)time steps + precip_cascades_prev_subtimestep: Any = None + cascade_noise_prev_subtimestep: Any = None + precip_extrapolated_after_decomp: Any = None + noise_extrapolated_after_decomp: Any = None + precip_extrapolated_probability_matching: Any = None + + # States related to the NWP models + precip_models_cascades: Any = None + # States related to NWP models for (sub)time steps precip_models_cascades_timestep: Any = None precip_models_timestep: Any = None mean_models_timestep: Any = None std_models_timestep: Any = None velocity_models_timestep: Any = None - n_model_indices: Optional[np.ndarray] = None # NWP model indices - rho_forecast: Any = None + + # State that links NWP member to final output ensemble member + mapping_list_NWP_member_to_ensemble_member: Optional[np.ndarray] = ( + None # NWP model indices + ) + + # States related to the random generation of precip and motion + randgen_precip: Any = None + randgen_motion: Any = None + + # Variables related to the (sub)timestep calculations of the final forecast + previous_displacement: Any = None + previous_displacement_noise_cascade: Any = None + previous_displacement_prob_matching: Any = None + rho_final_blended_forecast: Any = None + final_blended_forecast_means: Any = None + final_blended_forecast_stds: Any = None + final_blended_forecast_means_mod_only: Any = None + final_blended_forecast_stds_mod_only: Any = None + final_blended_forecast_cascades: Any = None + final_blended_forecast_cascades_mod_only: Any = None + final_blended_forecast_recomposed: Any = None + final_blended_forecast_recomposed_mod_only: Any = None + + # The return outputs and probability matching are stored in these states: + final_blended_forecast: Any = None + final_blended_forecast_non_perturbed: Any = None + + # Variables to keep track of the times for the forecast + time_prev_timestep: Any = None + leadtime_since_start_forecast: Any = None + subtimesteps: Any = None + is_nowcast_time_step: bool = None + subtimestep_index: Any = None + + # States related to weights weights: Any = None weights_model_only: Any = None - precip_forecast_extrapolated_decomp_done: Any = None - noise_extrapolated_decomp_done: Any = None - precip_forecast_extrapolated_probability_matching: Any = None - precip_forecast_prev_subtimestep: Any = None - noise_prev_subtimestep: Any = None - means_blended: Any = None - sigmas_blended: Any = None - means_blended_mod_only: Any = None - sigmas_blended_mod_only: Any = None - precip_forecast_blended: Any = None - precip_forecast_blended_mod_only: Any = None - precip_forecast_recomposed: Any = None - precip_forecast_recomposed_mod_only: Any = None - t_index: Any = None class StepsBlendingNowcaster: @@ -264,20 +285,20 @@ def compute_forecast(self): self.__blended_nowcast_main_loop() # Stack and return the forecast output if self.__config.return_output: - self.__state.precip_forecast = np.stack( + self.__state.final_blended_forecast = np.stack( [ - np.stack(self.__state.precip_forecast[j]) + np.stack(self.__state.final_blended_forecast[j]) for j in range(self.__config.n_ens_members) ] ) if self.__config.measure_time: return ( - self.__state.precip_forecast, + self.__state.final_blended_forecast, self.__init_time, self.__mainloop_time, ) else: - return self.__state.precip_forecast + return self.__state.final_blended_forecast else: return None @@ -300,15 +321,17 @@ def __blended_nowcast_main_loop(self): # extrap_kwargs is in config but by adding info to it, the next run of a blended forecast will have issues! self.__config.extrapolation_kwargs["return_displacement"] = True - self.__state.precip_forecast_prev_subtimestep = deepcopy( + self.__state.precip_cascades_prev_subtimestep = deepcopy( self.__state.precip_cascades ) - self.__state.noise_prev_subtimestep = deepcopy( + self.__state.cascade_noise_prev_subtimestep = deepcopy( self.__state.precip_noise_cascades ) - self.__state.t_prev_timestep = [0.0 for j in range(self.__config.n_ens_members)] - self.__state.t_leadtime_since_start_forecast = [ + self.__state.time_prev_timestep = [ + 0.0 for j in range(self.__config.n_ens_members) + ] + self.__state.leadtime_since_start_forecast = [ 0.0 for j in range(self.__config.n_ens_members) ] @@ -321,7 +344,7 @@ def __blended_nowcast_main_loop(self): self.__find_nowcast_NWP_combination(t) self.__determine_skill_for_current_timestep(t) # the nowcast iteration for each ensemble member - precip_ensemble_single_timestep = [ + final_blended_forecast_all_members_one_timestep = [ None for _ in range(self.__config.n_ens_members) ] @@ -347,7 +370,7 @@ def worker(j): j, final_blended_forecast_single_member, worker_state ) ) - precip_ensemble_single_timestep[j] = ( + final_blended_forecast_all_members_one_timestep[j] = ( final_blended_forecast_single_member ) @@ -373,17 +396,19 @@ def worker(j): print("done.") if self.__config.callback is not None: - precip_forecast_final = np.stack(precip_ensemble_single_timestep) + precip_forecast_final = np.stack( + final_blended_forecast_all_members_one_timestep + ) if precip_forecast_final.shape[1] > 0: self.__config.callback(precip_forecast_final.squeeze()) if self.__config.return_output: for j in range(self.__config.n_ens_members): - self.__state.precip_forecast[j].extend( - precip_ensemble_single_timestep[j] + self.__state.final_blended_forecast[j].extend( + final_blended_forecast_all_members_one_timestep[j] ) - precip_ensemble_single_timestep = None + final_blended_forecast_all_members_one_timestep = None if self.__config.measure_time: self.__mainloop_time = time.time() - starttime_mainloop @@ -721,7 +746,7 @@ def __prepare_radar_and_NWP_fields(self): self.__state.precip_models_cascades = None # TODO: This type of check needs to be changed when going to xarray - if self.__precip_models.ndim != 4: + if self.__params.precip_models_provided_is_cascade: self.__state.precip_models_cascades = self.__precip_models self.__precip_models = _compute_cascade_recomposition_nwp( self.__precip_models, self.__params.recomposition_method @@ -1020,16 +1045,15 @@ def __multiply_precip_cascade_to_match_ensemble_members(self): def __initialize_random_generators(self): # 6. Initialize all the random generators and prepare for the forecast loop """Initialize all the random generators.""" - # TODO: randgen_motion and randgen_precip are not defined if no noise method is given? Should we end the program in that case? if self.__config.noise_method is not None: self.__state.randgen_precip = [] - randgen_motion = [] + self.__state.randgen_motion = [] for j in range(self.__config.n_ens_members): rs = np.random.RandomState(self.__config.seed) self.__state.randgen_precip.append(rs) seed = rs.randint(0, high=1e9) rs = np.random.RandomState(seed) - randgen_motion.append(rs) + self.__state.randgen_motion.append(rs) seed = rs.randint(0, high=1e9) if self.__config.velocity_perturbation_method is not None: @@ -1042,7 +1066,7 @@ def __initialize_random_generators(self): self.__state.velocity_perturbations = [] for j in range(self.__config.n_ens_members): kwargs = { - "randstate": randgen_motion[j], + "randstate": self.__state.randgen_motion[j], "p_par": self.__params.velocity_perturbations_parallel, "p_perp": self.__params.velocity_perturbations_perpendicular, } @@ -1071,7 +1095,9 @@ def __prepare_forecast_loop(self): self.__state.previous_displacement_prob_matching = np.stack( [None for j in range(self.__config.n_ens_members)] ) - self.__state.precip_forecast = [[] for j in range(self.__config.n_ens_members)] + self.__state.final_blended_forecast = [ + [] for j in range(self.__config.n_ens_members) + ] if self.__config.mask_method == "incremental": # get mask parameters @@ -1086,16 +1112,16 @@ def __prepare_forecast_loop(self): self.__params.mask_rim, self.__params.struct = None, None if self.__config.noise_method is None: - self.__state.precip_forecast_non_perturbed = [ + self.__state.final_blended_forecast_non_perturbed = [ self.__state.precip_cascades[0][i].copy() for i in range(self.__config.n_cascade_levels) ] else: - self.__state.precip_forecast_non_perturbed = None + self.__state.final_blended_forecast_non_perturbed = None - self.__state.fft_objs = [] + self.__params.fft_objs = [] for i in range(self.__config.n_ens_members): - self.__state.fft_objs.append( + self.__params.fft_objs.append( utils.get_method( self.__config.fft_method, shape=self.__state.precip_cascades.shape[-2:], @@ -1130,13 +1156,13 @@ def __initialize_noise_cascades(self): epsilon = self.__params.noise_generator( self.__params.perturbation_generator, randstate=self.__state.randgen_precip[j], - fft_method=self.__state.fft_objs[j], + fft_method=self.__params.fft_objs[j], domain=self.__config.domain, ) epsilon_decomposed = self.__params.decomposition_method( epsilon, self.__params.bandpass_filter, - fft_method=self.__state.fft_objs[j], + fft_method=self.__params.fft_objs[j], input_domain=self.__config.domain, output_domain=self.__config.domain, compute_stats=True, @@ -1266,7 +1292,7 @@ def __find_nowcast_NWP_combination(self, t): # Check if NWP models/members should be used individually, or if all of # them are blended together per nowcast ensemble member. if self.__config.blend_nwp_members: - self.__state.n_model_indices = None + self.__state.mapping_list_NWP_member_to_ensemble_member = None else: # Start with determining the maximum and mimimum number of members/models @@ -1277,9 +1303,11 @@ def __find_nowcast_NWP_combination(self, t): # for indexing the right climatological skill file when pysteps calculates # the blended forecast in parallel. if n_model_members > 1: - self.__state.n_model_indices = np.arange(n_model_members) + self.__state.mapping_list_NWP_member_to_ensemble_member = np.arange( + n_model_members + ) else: - self.__state.n_model_indices = [0] + self.__state.mapping_list_NWP_member_to_ensemble_member = [0] # Now, repeat the nowcast ensemble members or the nwp models/members until # it has the same amount of members as n_ens_members_max. For instance, if @@ -1310,8 +1338,10 @@ def __find_nowcast_NWP_combination(self, t): self.__state.precip_models_timestep, n_ens_members_max, axis=0 ) # Finally, for the model indices - self.__state.n_model_indices = np.repeat( - self.__state.n_model_indices, n_ens_members_max, axis=0 + self.__state.mapping_list_NWP_member_to_ensemble_member = np.repeat( + self.__state.mapping_list_NWP_member_to_ensemble_member, + n_ens_members_max, + axis=0, ) elif n_model_members == n_ens_members_min: @@ -1339,8 +1369,12 @@ def __find_nowcast_NWP_combination(self, t): self.__state.precip_models_timestep, repeats, axis=0 ) # Finally, for the model indices - self.__state.n_model_indices = np.repeat( - self.__state.n_model_indices, repeats, axis=0 + self.__state.mapping_list_NWP_member_to_ensemble_member = ( + np.repeat( + self.__state.mapping_list_NWP_member_to_ensemble_member, + repeats, + axis=0, + ) ) # TODO: is this not duplicate from part 2.3.5? @@ -1421,7 +1455,7 @@ def __determine_skill_for_next_timestep(self, t, j, worker_state): ] rho_nwp_forecast = np.stack(rho_nwp_forecast) # Concatenate rho_extrap_cascade and rho_nwp - worker_state.rho_forecast = np.concatenate( + worker_state.rho_final_blended_forecast = np.concatenate( (worker_state.rho_extrap_cascade[None, :], rho_nwp_forecast), axis=0 ) else: @@ -1430,11 +1464,11 @@ def __determine_skill_for_next_timestep(self, t, j, worker_state): lt=(t * int(self.__config.timestep)), correlations=self.__params.rho_nwp_models[j], outdir_path=self.__config.outdir_path_skill, - n_model=worker_state.n_model_indices[j], + n_model=worker_state.mapping_list_NWP_member_to_ensemble_member[j], skill_kwargs=self.__config.climatology_kwargs, ) # Concatenate rho_extrap_cascade and rho_nwp - worker_state.rho_forecast = np.concatenate( + worker_state.rho_final_blended_forecast = np.concatenate( (worker_state.rho_extrap_cascade[None, :], rho_nwp_forecast[None, :]), axis=0, ) @@ -1447,14 +1481,16 @@ def __determine_weights_per_component(self, worker_state): # selected, weights will be overwritten with those weights prior to # blending step. # weight = [(extr_field, n_model_fields, noise), n_cascade_levels, ...] - worker_state.weights = calculate_weights_bps(worker_state.rho_forecast) + worker_state.weights = calculate_weights_bps( + worker_state.rho_final_blended_forecast + ) # The model only weights if self.__config.weights_method == "bps": # Determine the weights of the components without the extrapolation # cascade, in case this is no data or outside the mask. worker_state.weights_model_only = calculate_weights_bps( - worker_state.rho_forecast[1:, :] + worker_state.rho_final_blended_forecast[1:, :] ) elif self.__config.weights_method == "spn": # Only the weights of the components without the extrapolation @@ -1489,13 +1525,13 @@ def __determine_weights_per_component(self, worker_state): ) # Determine the weights for this cascade level worker_state.weights_model_only[:, i] = calculate_weights_spn( - correlations=worker_state.rho_forecast[1:, i], + correlations=worker_state.rho_final_blended_forecast[1:, i], covariance=covariance_nwp_models, ) else: # Same as correlation and noise is 1 - correlation worker_state.weights_model_only = calculate_weights_bps( - worker_state.rho_forecast[1:, :] + worker_state.rho_final_blended_forecast[1:, :] ) else: raise ValueError( @@ -1553,7 +1589,7 @@ def __regress_extrapolation_and_noise_cascades(self, j, worker_state): # use the deterministic AR(p) model computed above if # perturbations are disabled worker_state.precip_cascades[j][i] = ( - worker_state.precip_forecast_non_perturbed[i] + worker_state.final_blended_forecast_non_perturbed[i] ) # 8.3.3 regress the noise component to the subsequent time step @@ -1589,9 +1625,9 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( extrap_kwargs_pb = self.__config.extrapolation_kwargs.copy() velocity_perturbations_extrapolation = self.__velocity # The following should be accesseble after this function - worker_state.precip_forecast_extrapolated_decomp_done = [] - worker_state.noise_extrapolated_decomp_done = [] - worker_state.precip_forecast_extrapolated_probability_matching = [] + worker_state.precip_extrapolated_decomp = [] + worker_state.noise_extrapolated_decomp = [] + worker_state.precip_extrapolated_probability_matching = [] # Extrapolate per sub time step for t_sub in worker_state.subtimesteps: @@ -1600,14 +1636,14 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( if t_diff_prev_subtimestep_int > 0.0: precip_forecast_cascade_subtimestep = [ (1.0 - t_diff_prev_subtimestep_int) - * worker_state.precip_forecast_prev_subtimestep[j][i][-1, :] + * worker_state.precip_cascades_prev_subtimestep[j][i][-1, :] + t_diff_prev_subtimestep_int * worker_state.precip_cascades[j][i][-1, :] for i in range(self.__config.n_cascade_levels) ] noise_cascade_subtimestep = [ (1.0 - t_diff_prev_subtimestep_int) - * worker_state.noise_prev_subtimestep[j][i][-1, :] + * worker_state.cascade_noise_prev_subtimestep[j][i][-1, :] + t_diff_prev_subtimestep_int * worker_state.precip_noise_cascades[j][i][-1, :] for i in range(self.__config.n_cascade_levels) @@ -1615,11 +1651,11 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( else: precip_forecast_cascade_subtimestep = [ - worker_state.precip_forecast_prev_subtimestep[j][i][-1, :] + worker_state.precip_cascades_prev_subtimestep[j][i][-1, :] for i in range(self.__config.n_cascade_levels) ] noise_cascade_subtimestep = [ - worker_state.noise_prev_subtimestep[j][i][-1, :] + worker_state.cascade_noise_prev_subtimestep[j][i][-1, :] for i in range(self.__config.n_cascade_levels) ] @@ -1628,10 +1664,8 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( ) noise_cascade_subtimestep = np.stack(noise_cascade_subtimestep) - t_diff_prev_subtimestep = t_sub - worker_state.t_prev_timestep[j] - worker_state.t_leadtime_since_start_forecast[ - j - ] += t_diff_prev_subtimestep + t_diff_prev_subtimestep = t_sub - worker_state.time_prev_timestep[j] + worker_state.leadtime_since_start_forecast[j] += t_diff_prev_subtimestep # compute the perturbed motion field - include the NWP # velocities and the weights. Note that we only perturb @@ -1642,7 +1676,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( self.__velocity + self.__params.generate_velocity_noise( self.__params.velocity_perturbations[j], - worker_state.t_leadtime_since_start_forecast[j] + worker_state.leadtime_since_start_forecast[j] * self.__config.timestep, ) ) @@ -1712,39 +1746,35 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( allow_nonfinite_values=True, **self.__config.extrapolation_kwargs, ) - precip_forecast_extrapolated_recomp_subtimestep = ( + precip_extrapolated_recomp_subtimestep = ( precip_forecast_extrapolated_recomp_subtimestep_temp[0].copy() ) - temp_mask = ~np.isfinite( - precip_forecast_extrapolated_recomp_subtimestep - ) + temp_mask = ~np.isfinite(precip_extrapolated_recomp_subtimestep) # TODO: WHERE DO CAN I FIND THIS -15.0 - precip_forecast_extrapolated_recomp_subtimestep[ - ~np.isfinite(precip_forecast_extrapolated_recomp_subtimestep) + precip_extrapolated_recomp_subtimestep[ + ~np.isfinite(precip_extrapolated_recomp_subtimestep) ] = self.__params.precip_zerovalue - precip_forecast_extrapolated_decomp = ( - self.__params.decomposition_method( - precip_forecast_extrapolated_recomp_subtimestep, - self.__params.bandpass_filter, - mask=self.__params.mask_threshold, - fft_method=self.__params.fft, - output_domain=self.__config.domain, - normalize=True, - compute_stats=True, - compact_output=True, - )["cascade_levels"] - ) + precip_extrapolated_decomp = self.__params.decomposition_method( + precip_extrapolated_recomp_subtimestep, + self.__params.bandpass_filter, + mask=self.__params.mask_threshold, + fft_method=self.__params.fft, + output_domain=self.__config.domain, + normalize=True, + compute_stats=True, + compact_output=True, + )["cascade_levels"] # Make sure we have values outside the mask if self.__params.zero_precip_radar: - precip_forecast_extrapolated_decomp = np.nan_to_num( - precip_forecast_extrapolated_decomp, + precip_extrapolated_decomp = np.nan_to_num( + precip_extrapolated_decomp, copy=True, nan=np.nanmin(precip_forecast_cascade_subtimestep), posinf=np.nanmin(precip_forecast_cascade_subtimestep), neginf=np.nanmin(precip_forecast_cascade_subtimestep), ) for i in range(self.__config.n_cascade_levels): - precip_forecast_extrapolated_decomp[i][temp_mask] = np.nan + precip_extrapolated_decomp[i][temp_mask] = np.nan # B. Noise noise_cascade_subtimestep_recomp = blending.utils.recompose_cascade( combined_cascade=noise_cascade_subtimestep, @@ -1780,17 +1810,17 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( noise_extrapolated_decomp[i] *= self.__params.noise_std_coeffs[i] # Append the results to the output lists - worker_state.precip_forecast_extrapolated_decomp_done.append( - precip_forecast_extrapolated_decomp.copy() + worker_state.precip_extrapolated_decomp.append( + precip_extrapolated_decomp.copy() ) - worker_state.noise_extrapolated_decomp_done.append( + worker_state.noise_extrapolated_decomp.append( noise_extrapolated_decomp.copy() ) precip_forecast_cascade_subtimestep = None precip_forecast_recomp_subtimestep = None precip_forecast_extrapolated_recomp_subtimestep_temp = None - precip_forecast_extrapolated_recomp_subtimestep = None - precip_forecast_extrapolated_decomp = None + precip_extrapolated_recomp_subtimestep = None + precip_extrapolated_decomp = None noise_cascade_subtimestep = None noise_cascade_subtimestep_recomp = None noise_extrapolated_recomp_temp = None @@ -1819,28 +1849,28 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( allow_nonfinite_values=True, **extrap_kwargs_pb, ) - worker_state.precip_forecast_extrapolated_probability_matching.append( + worker_state.precip_extrapolated_probability_matching.append( precip_forecast_extrapolated_probability_matching_temp[0] ) - worker_state.t_prev_timestep[j] = t_sub + worker_state.time_prev_timestep[j] = t_sub - if len(worker_state.precip_forecast_extrapolated_decomp_done) > 0: - worker_state.precip_forecast_extrapolated_decomp_done = np.stack( - worker_state.precip_forecast_extrapolated_decomp_done + if len(worker_state.precip_extrapolated_decomp) > 0: + worker_state.precip_extrapolated_decomp = np.stack( + worker_state.precip_extrapolated_decomp ) - worker_state.noise_extrapolated_decomp_done = np.stack( - worker_state.noise_extrapolated_decomp_done + worker_state.noise_extrapolated_decomp = np.stack( + worker_state.noise_extrapolated_decomp ) - worker_state.precip_forecast_extrapolated_probability_matching = np.stack( - worker_state.precip_forecast_extrapolated_probability_matching + worker_state.precip_extrapolated_probability_matching = np.stack( + worker_state.precip_extrapolated_probability_matching ) # advect the forecast field by one time step if no subtimesteps in the # current interval were found if not worker_state.subtimesteps: - t_diff_prev_subtimestep = t + 1 - worker_state.t_prev_timestep[j] - worker_state.t_leadtime_since_start_forecast[j] += t_diff_prev_subtimestep + t_diff_prev_subtimestep = t + 1 - worker_state.time_prev_timestep[j] + worker_state.leadtime_since_start_forecast[j] += t_diff_prev_subtimestep # compute the perturbed motion field - include the NWP # velocities and the weights @@ -1849,7 +1879,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( self.__velocity + self.__params.generate_velocity_noise( self.__params.velocity_perturbations[j], - worker_state.t_leadtime_since_start_forecast[j] + worker_state.leadtime_since_start_forecast[j] * self.__config.timestep, ) ) @@ -1926,28 +1956,30 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( ) ) - worker_state.t_prev_timestep[j] = t + 1 + worker_state.time_prev_timestep[j] = t + 1 - worker_state.precip_forecast_prev_subtimestep[j] = worker_state.precip_cascades[ + worker_state.precip_cascades_prev_subtimestep[j] = worker_state.precip_cascades[ j ] - worker_state.noise_prev_subtimestep[j] = worker_state.precip_noise_cascades[j] + worker_state.cascade_noise_prev_subtimestep[j] = ( + worker_state.precip_noise_cascades[j] + ) def __blend_cascades(self, t_sub, j, worker_state): - worker_state.t_index = np.where(np.array(worker_state.subtimesteps) == t_sub)[ - 0 - ][0] + worker_state.subtimestep_index = np.where( + np.array(worker_state.subtimesteps) == t_sub + )[0][0] # First concatenate the cascades and the means and sigmas # precip_models = [n_models,timesteps,n_cascade_levels,m,n] if self.__config.blend_nwp_members: cascade_stack_all_components = np.concatenate( ( - worker_state.precip_forecast_extrapolated_decomp_done[ - None, worker_state.t_index + worker_state.precip_extrapolated_decomp[ + None, worker_state.subtimestep_index ], worker_state.precip_models_cascades_timestep, - worker_state.noise_extrapolated_decomp_done[ - None, worker_state.t_index + worker_state.noise_extrapolated_decomp[ + None, worker_state.subtimestep_index ], ), axis=0, @@ -1969,12 +2001,12 @@ def __blend_cascades(self, t_sub, j, worker_state): else: cascade_stack_all_components = np.concatenate( ( - worker_state.precip_forecast_extrapolated_decomp_done[ - None, worker_state.t_index + worker_state.precip_extrapolated_decomp[ + None, worker_state.subtimestep_index ], worker_state.precip_models_cascades_timestep[None, j], - worker_state.noise_extrapolated_decomp_done[ - None, worker_state.t_index + worker_state.noise_extrapolated_decomp[ + None, worker_state.subtimestep_index ], ), axis=0, @@ -2022,31 +2054,36 @@ def __blend_cascades(self, t_sub, j, worker_state): ) # Determine the weights for this cascade level worker_state.weights[:, i] = calculate_weights_spn( - correlations=worker_state.rho_forecast[:, i], + correlations=worker_state.rho_final_blended_forecast[:, i], covariance=covariance_nwp_models, ) # Blend the extrapolation, (NWP) model(s) and noise cascades - worker_state.precip_forecast_blended = blending.utils.blend_cascades( + worker_state.final_blended_forecast_cascades = blending.utils.blend_cascades( cascades_norm=cascade_stack_all_components, weights=worker_state.weights ) # Also blend the cascade without the extrapolation component - worker_state.precip_forecast_blended_mod_only = blending.utils.blend_cascades( - cascades_norm=cascade_stack_all_components[1:, :], - weights=worker_state.weights_model_only, + worker_state.final_blended_forecast_cascades_mod_only = ( + blending.utils.blend_cascades( + cascades_norm=cascade_stack_all_components[1:, :], + weights=worker_state.weights_model_only, + ) ) # Blend the means and standard deviations # Input is array of shape [number_components, scale_level, ...] - worker_state.means_blended, worker_state.sigmas_blended = blend_means_sigmas( + ( + worker_state.final_blended_forecast_means, + worker_state.final_blended_forecast_stds, + ) = blend_means_sigmas( means=means_stacked, sigmas=sigmas_stacked, weights=worker_state.weights ) # Also blend the means and sigmas for the cascade without extrapolation ( - worker_state.means_blended_mod_only, - worker_state.sigmas_blended_mod_only, + worker_state.final_blended_forecast_means_mod_only, + worker_state.final_blended_forecast_stds_mod_only, ) = blend_means_sigmas( means=means_stacked[1:, :], sigmas=sigmas_stacked[1:, :], @@ -2057,30 +2094,34 @@ def __recompose_cascade_to_rainfall_field(self, j, worker_state): # 8.6 Recompose the cascade to a precipitation field # (The function first normalizes the blended cascade, precip_forecast_blended # again) - worker_state.precip_forecast_recomposed = blending.utils.recompose_cascade( - combined_cascade=worker_state.precip_forecast_blended, - combined_mean=worker_state.means_blended, - combined_sigma=worker_state.sigmas_blended, + worker_state.final_blended_forecast_recomposed = ( + blending.utils.recompose_cascade( + combined_cascade=worker_state.final_blended_forecast_cascades, + combined_mean=worker_state.final_blended_forecast_means, + combined_sigma=worker_state.final_blended_forecast_stds, + ) ) # The recomposed cascade without the extrapolation (for NaN filling # outside the radar domain) - worker_state.precip_forecast_recomposed_mod_only = ( + worker_state.final_blended_forecast_recomposed_mod_only = ( blending.utils.recompose_cascade( - combined_cascade=worker_state.precip_forecast_blended_mod_only, - combined_mean=worker_state.means_blended_mod_only, - combined_sigma=worker_state.sigmas_blended_mod_only, + combined_cascade=worker_state.final_blended_forecast_cascades_mod_only, + combined_mean=worker_state.final_blended_forecast_means_mod_only, + combined_sigma=worker_state.final_blended_forecast_stds_mod_only, ) ) if self.__config.domain == "spectral": # TODO: Check this! (Only tested with domain == 'spatial') # TODO: what needs to happen with above TODO? - worker_state.precip_forecast_recomposed = worker_state.fft_objs[j].irfft2( - worker_state.precip_forecast_recomposed - ) - worker_state.precip_forecast_recomposed_mod_only = worker_state.fft_objs[ + worker_state.final_blended_forecast_recomposed = worker_state.fft_objs[ j - ].irfft2(worker_state.precip_forecast_recomposed_mod_only) + ].irfft2(worker_state.final_blended_forecast_recomposed) + worker_state.final_blended_forecast_recomposed_mod_only = ( + worker_state.fft_objs[j].irfft2( + worker_state.final_blended_forecast_recomposed_mod_only + ) + ) def __post_process_output( self, j, final_blended_forecast_single_member, worker_state @@ -2113,8 +2154,8 @@ def __post_process_output( if self.__config.blend_nwp_members: precip_forecast_probability_matching_final = np.concatenate( ( - worker_state.precip_forecast_extrapolated_probability_matching[ - None, worker_state.t_index + worker_state.precip_extrapolated_probability_matching[ + None, worker_state.subtimestep_index ], worker_state.precip_models_timestep, ), @@ -2123,8 +2164,8 @@ def __post_process_output( else: precip_forecast_probability_matching_final = np.concatenate( ( - worker_state.precip_forecast_extrapolated_probability_matching[ - None, worker_state.t_index + worker_state.precip_extrapolated_probability_matching[ + None, worker_state.subtimestep_index ], worker_state.precip_models_timestep[None, j], ), @@ -2159,7 +2200,7 @@ def __post_process_output( # areas with the "..._mod_only" blended forecasts, consisting # of the NWP and noise components. - nan_indices = np.isnan(worker_state.precip_forecast_recomposed) + nan_indices = np.isnan(worker_state.final_blended_forecast_recomposed) if self.__config.smooth_radar_mask_range != 0: # Compute the smooth dilated mask new_mask = blending.utils.compute_smooth_dilated_mask( @@ -2173,10 +2214,10 @@ def __post_process_output( # Handle NaNs in precip_forecast_new and precip_forecast_new_mod_only by setting NaNs to 0 in the blending step precip_forecast_recomposed_mod_only_no_nan = np.nan_to_num( - worker_state.precip_forecast_recomposed_mod_only, nan=0 + worker_state.final_blended_forecast_recomposed_mod_only, nan=0 ) precip_forecast_recomposed_no_nan = np.nan_to_num( - worker_state.precip_forecast_recomposed, nan=0 + worker_state.final_blended_forecast_recomposed, nan=0 ) # Perform the blending of radar and model inside the radar domain using a weighted combination @@ -2196,8 +2237,8 @@ def __post_process_output( axis=0, ) else: - worker_state.precip_forecast_recomposed[nan_indices] = ( - worker_state.precip_forecast_recomposed_mod_only[nan_indices] + worker_state.final_blended_forecast_recomposed[nan_indices] = ( + worker_state.final_blended_forecast_recomposed_mod_only[nan_indices] ) nan_indices = np.isnan(precip_forecast_probability_matching_blended) precip_forecast_probability_matching_blended[nan_indices] = ( @@ -2206,9 +2247,9 @@ def __post_process_output( # Finally, fill the remaining nan values, if present, with # the minimum value in the forecast - nan_indices = np.isnan(worker_state.precip_forecast_recomposed) - worker_state.precip_forecast_recomposed[nan_indices] = np.nanmin( - worker_state.precip_forecast_recomposed + nan_indices = np.isnan(worker_state.final_blended_forecast_recomposed) + worker_state.final_blended_forecast_recomposed[nan_indices] = np.nanmin( + worker_state.final_blended_forecast_recomposed ) nan_indices = np.isnan(precip_forecast_probability_matching_blended) precip_forecast_probability_matching_blended[nan_indices] = np.nanmin( @@ -2221,7 +2262,9 @@ def __post_process_output( # apply the precipitation mask to prevent generation of new # precipitation into areas where it was not originally # observed - precip_forecast_min_value = worker_state.precip_forecast_recomposed.min() + precip_forecast_min_value = ( + worker_state.final_blended_forecast_recomposed.min() + ) if self.__config.mask_method == "incremental": # The incremental mask is slightly different from # the implementation in the non-blended steps.py, as @@ -2238,16 +2281,17 @@ def __post_process_output( precip_field_mask, self.__params.struct, self.__params.mask_rim ) # Get the final mask - worker_state.precip_forecast_recomposed = ( + worker_state.final_blended_forecast_recomposed = ( precip_forecast_min_value + ( - worker_state.precip_forecast_recomposed + worker_state.final_blended_forecast_recomposed - precip_forecast_min_value ) * precip_field_mask ) precip_field_mask_temp = ( - worker_state.precip_forecast_recomposed > precip_forecast_min_value + worker_state.final_blended_forecast_recomposed + > precip_forecast_min_value ) elif self.__config.mask_method == "obs": # The mask equals the most recent benchmark @@ -2258,7 +2302,7 @@ def __post_process_output( ) # Set to min value outside of mask - worker_state.precip_forecast_recomposed[~precip_field_mask_temp] = ( + worker_state.final_blended_forecast_recomposed[~precip_field_mask_temp] = ( precip_forecast_min_value ) @@ -2269,8 +2313,8 @@ def __post_process_output( self.__config.probmatching_method is not None and self.__config.resample_distribution ): - arr1 = worker_state.precip_forecast_extrapolated_probability_matching[ - worker_state.t_index + arr1 = worker_state.precip_extrapolated_probability_matching[ + worker_state.subtimestep_index ] arr2 = worker_state.precip_models_timestep[j] # resample weights based on cascade level 2. @@ -2290,17 +2334,17 @@ def __post_process_output( if self.__config.probmatching_method == "cdf": # nan indices in the extrapolation nowcast nan_indices = np.isnan( - worker_state.precip_forecast_extrapolated_probability_matching[ - worker_state.t_index + worker_state.precip_extrapolated_probability_matching[ + worker_state.subtimestep_index ] ) # Adjust the CDF of the forecast to match the resampled distribution combined from # extrapolation and model fields. # Rainfall outside the pure extrapolation domain is not taken into account. - if np.any(np.isfinite(worker_state.precip_forecast_recomposed)): - worker_state.precip_forecast_recomposed = ( + if np.any(np.isfinite(worker_state.final_blended_forecast_recomposed)): + worker_state.final_blended_forecast_recomposed = ( probmatching.nonparam_match_empirical_cdf( - worker_state.precip_forecast_recomposed, + worker_state.final_blended_forecast_recomposed, precip_forecast_probability_matching_resampled, nan_indices, ) @@ -2315,21 +2359,21 @@ def __post_process_output( ] ) no_rain_mask = ( - worker_state.precip_forecast_recomposed + worker_state.final_blended_forecast_recomposed >= self.__config.precip_threshold ) mean_precip_forecast = np.mean( - worker_state.precip_forecast_recomposed[no_rain_mask] + worker_state.final_blended_forecast_recomposed[no_rain_mask] ) - worker_state.precip_forecast_recomposed[no_rain_mask] = ( - worker_state.precip_forecast_recomposed[no_rain_mask] + worker_state.final_blended_forecast_recomposed[no_rain_mask] = ( + worker_state.final_blended_forecast_recomposed[no_rain_mask] - mean_precip_forecast + mean_probabiltity_matching_forecast ) precip_forecast_probability_matching_resampled = None final_blended_forecast_single_member.append( - worker_state.precip_forecast_recomposed + worker_state.final_blended_forecast_recomposed ) return final_blended_forecast_single_member From ed20ecc1b7ca3d83e09f2a435d16f330d2520482 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Fri, 6 Dec 2024 16:09:55 +0100 Subject: [PATCH 18/33] Cleanup of params and state dataclasses, now all tests pass --- pysteps/blending/steps.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 7fc42aeb..4e567a10 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -1551,7 +1551,7 @@ def __regress_extrapolation_and_noise_cascades(self, j, worker_state): epsilon = self.__params.noise_generator( self.__params.perturbation_generator, randstate=worker_state.randgen_precip[j], - fft_method=worker_state.fft_objs[j], + fft_method=self.__params.fft_objs[j], domain=self.__config.domain, ) @@ -1559,7 +1559,7 @@ def __regress_extrapolation_and_noise_cascades(self, j, worker_state): epsilon_decomposed = self.__params.decomposition_method( epsilon, self.__params.bandpass_filter, - fft_method=worker_state.fft_objs[j], + fft_method=self.__params.fft_objs[j], input_domain=self.__config.domain, output_domain=self.__config.domain, compute_stats=True, @@ -2114,11 +2114,11 @@ def __recompose_cascade_to_rainfall_field(self, j, worker_state): # TODO: Check this! (Only tested with domain == 'spatial') # TODO: what needs to happen with above TODO? - worker_state.final_blended_forecast_recomposed = worker_state.fft_objs[ + worker_state.final_blended_forecast_recomposed = self.__params.fft_objs[ j ].irfft2(worker_state.final_blended_forecast_recomposed) worker_state.final_blended_forecast_recomposed_mod_only = ( - worker_state.fft_objs[j].irfft2( + self.__params.fft_objs[j].irfft2( worker_state.final_blended_forecast_recomposed_mod_only ) ) From 701e726ec02660c821af7e922796d4555545c78e Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Fri, 6 Dec 2024 16:47:10 +0100 Subject: [PATCH 19/33] Added correct typing to all parts of params and state --- pysteps/blending/steps.py | 211 ++++++++++++++++++-------------------- 1 file changed, 99 insertions(+), 112 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 4e567a10..a9f2becf 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -65,11 +65,12 @@ DASK_IMPORTED = False from dataclasses import dataclass, field -from typing import Optional, List, Dict, Any, Callable +from typing import Optional, List, Dict, Any, Callable, Union # TODO: compare old and new version of the code # TODO: look for better typing in state and params # TODO: GO over all other todos and check if they can be removed +# TODO: look at the documentation and try to improve it, lots of things are now combined together @dataclass @@ -113,109 +114,104 @@ class StepsBlendingConfig: # TODO: typing could be improved here @dataclass class StepsBlendingParams: - noise_std_coeffs: np.ndarray = None # Noise standard deviation coefficients - bandpass_filter: Any = None # Band-pass filter object - fft: Any = None # FFT method object - perturbation_generator: Callable = None # Perturbation generator - noise_generator: Callable = None # Noise generator - PHI: np.ndarray = None # AR(p) model parameters - extrapolation_method: Any = None - decomposition_method: Any = None - recomposition_method: Any = None - # TODO: check of the following two are relevant or can be replaced vel_pert_... and noise_generator - velocity_perturbations: Any = None - generate_velocity_noise: Any = None - velocity_perturbations_parallel: Optional[np.ndarray] = ( - None # Velocity perturbation parameters (parallel) + noise_std_coeffs: Optional[np.ndarray] = ( + None # Noise standard deviation coefficients ) - velocity_perturbations_perpendicular: Optional[np.ndarray] = ( - None # Velocity perturbation parameters (perpendicular) + bandpass_filter: Optional[Any] = None # Band-pass filter object + fft: Optional[Any] = None # FFT method object + perturbation_generator: Optional[Callable[..., np.ndarray]] = ( + None # Perturbation generator ) - fft_objs: List[Any] = field( - default_factory=list - ) # FFT objects for ensemble members + noise_generator: Optional[Callable[..., np.ndarray]] = None # Noise generator + PHI: Optional[np.ndarray] = None # AR(p) model parameters + extrapolation_method: Optional[Callable[..., Any]] = None + decomposition_method: Optional[Callable[..., dict]] = None + recomposition_method: Optional[Callable[..., np.ndarray]] = None + # TODO: check of the following two are relevant or can be replaced vel_pert_... and noise_generator + velocity_perturbations: Optional[Any] = None + generate_velocity_noise: Optional[Callable[[Any, float], np.ndarray]] = None + velocity_perturbations_parallel: Optional[np.ndarray] = None + velocity_perturbations_perpendicular: Optional[np.ndarray] = None + fft_objs: List[Any] = field(default_factory=list) mask_rim: Optional[int] = None # Rim size for masking struct: Optional[np.ndarray] = None # Structuring element for mask time_steps_is_list: bool = False # Time steps is a list precip_models_provided_is_cascade: bool = False # Precip models are decomposed - xy_coordinates: np.ndarray | None = None - precip_zerovalue: Any = None - mask_threshold: Any = None + xy_coordinates: Optional[np.ndarray] = None + precip_zerovalue: Optional[float] = None + mask_threshold: Optional[np.ndarray] = None zero_precip_radar: bool = False zero_precip_model_fields: bool = False - original_timesteps: Any = None - num_ensemble_workers: int = None - rho_nwp_models: Any = None - domain_mask: Any = None + original_timesteps: Optional[Union[list, np.ndarray]] = None + num_ensemble_workers: Optional[int] = None + rho_nwp_models: Optional[np.ndarray] = None + domain_mask: Optional[np.ndarray] = None # TODO: typing could be improved here @dataclass class StepsBlendingState: - # States related to the observations - precip_cascades: Any = None - precip_noise_input: Any = None - precip_noise_cascades: Any = None - precip_mean_noise: Any = None - precip_std_noise: Any = None - - # States related to the extrapolation - mean_extrapolation: Any = None - std_extrapolation: Any = None - rho_extrap_cascade_prev: Any = None - rho_extrap_cascade: Any = None - precip_cascades_prev_subtimestep: Any = None - cascade_noise_prev_subtimestep: Any = None - precip_extrapolated_after_decomp: Any = None - noise_extrapolated_after_decomp: Any = None - precip_extrapolated_probability_matching: Any = None - - # States related to the NWP models - precip_models_cascades: Any = None - # States related to NWP models for (sub)time steps - precip_models_cascades_timestep: Any = None - precip_models_timestep: Any = None - mean_models_timestep: Any = None - std_models_timestep: Any = None - velocity_models_timestep: Any = None - - # State that links NWP member to final output ensemble member - mapping_list_NWP_member_to_ensemble_member: Optional[np.ndarray] = ( - None # NWP model indices - ) - - # States related to the random generation of precip and motion - randgen_precip: Any = None - randgen_motion: Any = None - - # Variables related to the (sub)timestep calculations of the final forecast - previous_displacement: Any = None - previous_displacement_noise_cascade: Any = None - previous_displacement_prob_matching: Any = None - rho_final_blended_forecast: Any = None - final_blended_forecast_means: Any = None - final_blended_forecast_stds: Any = None - final_blended_forecast_means_mod_only: Any = None - final_blended_forecast_stds_mod_only: Any = None - final_blended_forecast_cascades: Any = None - final_blended_forecast_cascades_mod_only: Any = None - final_blended_forecast_recomposed: Any = None - final_blended_forecast_recomposed_mod_only: Any = None - - # The return outputs and probability matching are stored in these states: - final_blended_forecast: Any = None - final_blended_forecast_non_perturbed: Any = None - - # Variables to keep track of the times for the forecast - time_prev_timestep: Any = None - leadtime_since_start_forecast: Any = None - subtimesteps: Any = None - is_nowcast_time_step: bool = None - subtimestep_index: Any = None - - # States related to weights - weights: Any = None - weights_model_only: Any = None + # Radar and noise states + precip_cascades: Optional[np.ndarray] = None + precip_noise_input: Optional[np.ndarray] = None + precip_noise_cascades: Optional[np.ndarray] = None + precip_mean_noise: Optional[np.ndarray] = None + precip_std_noise: Optional[np.ndarray] = None + + # Extrapolation states + mean_extrapolation: Optional[np.ndarray] = None + std_extrapolation: Optional[np.ndarray] = None + rho_extrap_cascade_prev: Optional[np.ndarray] = None + rho_extrap_cascade: Optional[np.ndarray] = None + precip_cascades_prev_subtimestep: Optional[np.ndarray] = None + cascade_noise_prev_subtimestep: Optional[np.ndarray] = None + precip_extrapolated_after_decomp: Optional[np.ndarray] = None + noise_extrapolated_after_decomp: Optional[np.ndarray] = None + precip_extrapolated_probability_matching: Optional[np.ndarray] = None + + # NWP model states + precip_models_cascades: Optional[np.ndarray] = None + precip_models_cascades_timestep: Optional[np.ndarray] = None + precip_models_timestep: Optional[np.ndarray] = None + mean_models_timestep: Optional[np.ndarray] = None + std_models_timestep: Optional[np.ndarray] = None + velocity_models_timestep: Optional[np.ndarray] = None + + # Mapping from NWP members to ensemble members + mapping_list_NWP_member_to_ensemble_member: Optional[np.ndarray] = None + + # Random states for precipitation and motion + randgen_precip: Optional[List[np.random.RandomState]] = None + randgen_motion: Optional[List[np.random.RandomState]] = None + + # Variables for final forecast computation + previous_displacement: Optional[List[Any]] = None + previous_displacement_noise_cascade: Optional[List[Any]] = None + previous_displacement_prob_matching: Optional[List[Any]] = None + rho_final_blended_forecast: Optional[np.ndarray] = None + final_blended_forecast_means: Optional[np.ndarray] = None + final_blended_forecast_stds: Optional[np.ndarray] = None + final_blended_forecast_means_mod_only: Optional[np.ndarray] = None + final_blended_forecast_stds_mod_only: Optional[np.ndarray] = None + final_blended_forecast_cascades: Optional[np.ndarray] = None + final_blended_forecast_cascades_mod_only: Optional[np.ndarray] = None + final_blended_forecast_recomposed: Optional[np.ndarray] = None + final_blended_forecast_recomposed_mod_only: Optional[np.ndarray] = None + + # Final outputs + final_blended_forecast: Optional[np.ndarray] = None + final_blended_forecast_non_perturbed: Optional[np.ndarray] = None + + # Timing and indexing + time_prev_timestep: Optional[List[float]] = None + leadtime_since_start_forecast: Optional[List[float]] = None + subtimesteps: Optional[List[float]] = None + is_nowcast_time_step: Optional[bool] = None + subtimestep_index: Optional[int] = None + + # Weights used for blending + weights: Optional[np.ndarray] = None + weights_model_only: Optional[np.ndarray] = None class StepsBlendingNowcaster: @@ -311,7 +307,6 @@ def __blended_nowcast_main_loop(self): # 8. Start the forecasting loop ### # Isolate the last time slice of observed precipitation - # TODO: This precip was "precip = self.__precip[-1, :, :]", changed to self.__precip = self.__precip[-1, :, :]. Might need to chage again and user local variable precip in all following functions self.__precip = self.__precip[-1, :, :] print("Starting blended nowcast computation.") @@ -364,7 +359,6 @@ def worker(j): if t_sub > 0: self.__blend_cascades(t_sub, j, worker_state) self.__recompose_cascade_to_rainfall_field(j, worker_state) - # TODO: could be I need to return and ave final_blended_forecast_single_member final_blended_forecast_single_member = ( self.__post_process_output( j, final_blended_forecast_single_member, worker_state @@ -676,7 +670,6 @@ def __initialize_nowcast_components(self): x_values, y_values = np.meshgrid(np.arange(N), np.arange(M)) self.__params.xy_coordinates = np.stack([x_values, y_values]) - # TODO: changed precip_copy for self.__precip self.__precip = self.__precip[-(self.__config.ar_order + 1) :, :, :].copy() # Determine the domain mask from non-finite values in the precipitation data self.__params.domain_mask = np.logical_or.reduce( @@ -745,7 +738,6 @@ def __prepare_radar_and_NWP_fields(self): # 2.2 If necessary, recompose (NWP) model forecasts self.__state.precip_models_cascades = None - # TODO: This type of check needs to be changed when going to xarray if self.__params.precip_models_provided_is_cascade: self.__state.precip_models_cascades = self.__precip_models self.__precip_models = _compute_cascade_recomposition_nwp( @@ -1127,7 +1119,7 @@ def __prepare_forecast_loop(self): shape=self.__state.precip_cascades.shape[-2:], ) ) - # TODO: moved this from # 7 to here as it seems to fit better here. The only parameter used and needed is PHI, this is its last use untill # 7 + # initizalize the current and previous extrapolation forecast scale for the nowcasting component # phi1 / (1 - phi2), see BPS2004 self.__state.rho_extrap_cascade_prev = np.repeat( @@ -1377,7 +1369,7 @@ def __find_nowcast_NWP_combination(self, t): ) ) - # TODO: is this not duplicate from part 2.3.5? + # TODO: is this not duplicate from part 2.3.5? If so, is it still needed here? # If zero_precip_radar is True, set the velocity field equal to the NWP # velocity field for the current time step (velocity_models_temp). if self.__params.zero_precip_radar: @@ -1388,19 +1380,18 @@ def __find_nowcast_NWP_combination(self, t): def __determine_skill_for_current_timestep(self, t): if t == 0: """Calculate the initial skill of the (NWP) model forecasts at t=0.""" - # TODO rewrite loop - self.__params.rho_nwp_models = [ - blending.skill_scores.spatial_correlation( + self.__params.rho_nwp_models = [] + for model_index in range( + self.__state.precip_models_cascades_timestep.shape[0] + ): + rho_value = blending.skill_scores.spatial_correlation( obs=self.__state.precip_cascades[0, :, -1, :, :].copy(), mod=self.__state.precip_models_cascades_timestep[ model_index, :, :, : ].copy(), domain_mask=self.__params.domain_mask, ) - for model_index in range( - self.__state.precip_models_cascades_timestep.shape[0] - ) - ] + self.__params.rho_nwp_models.append(rho_value) self.__params.rho_nwp_models = np.stack(self.__params.rho_nwp_models) # Ensure that the model skill decreases with increasing scale level. @@ -1442,17 +1433,16 @@ def __determine_skill_for_next_timestep(self, t, j, worker_state): # 8.1.2 Determine the skill of the nwp components for lead time (t0 + t) # Then for the model components if self.__config.blend_nwp_members: - # TODO rewrite loop - rho_nwp_forecast = [ - blending.skill_scores.lt_dependent_cor_nwp( + rho_nwp_forecast = [] + for model_index in range(self.__params.rho_nwp_models.shape[0]): + rho_value = blending.skill_scores.lt_dependent_cor_nwp( lt=(t * int(self.__config.timestep)), correlations=self.__params.rho_nwp_models[model_index], outdir_path=self.__config.outdir_path_skill, n_model=model_index, skill_kwargs=self.__config.climatology_kwargs, ) - for model_index in range(self.__params.rho_nwp_models.shape[0]) - ] + rho_nwp_forecast.append(rho_value) rho_nwp_forecast = np.stack(rho_nwp_forecast) # Concatenate rho_extrap_cascade and rho_nwp worker_state.rho_final_blended_forecast = np.concatenate( @@ -2032,7 +2022,6 @@ def __blend_cascades(self, t_sub, j, worker_state): # method is given? Or does this mean that in all other circumstances the weights # have been calculated in a different way? - # TODO: changed weights to worker_state.weights if self.__config.weights_method == "spn": worker_state.weights = np.zeros( ( @@ -2112,8 +2101,6 @@ def __recompose_cascade_to_rainfall_field(self, j, worker_state): ) if self.__config.domain == "spectral": # TODO: Check this! (Only tested with domain == 'spatial') - - # TODO: what needs to happen with above TODO? worker_state.final_blended_forecast_recomposed = self.__params.fft_objs[ j ].irfft2(worker_state.final_blended_forecast_recomposed) From b9de5113b96b34ffa426aae8e87009d794bcd67e Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Fri, 6 Dec 2024 17:20:03 +0100 Subject: [PATCH 20/33] Ready for pull request --- pysteps/blending/steps.py | 192 +++++++++++++++++--------------------- 1 file changed, 84 insertions(+), 108 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index a9f2becf..ee53462c 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -67,9 +67,7 @@ from dataclasses import dataclass, field from typing import Optional, List, Dict, Any, Callable, Union -# TODO: compare old and new version of the code -# TODO: look for better typing in state and params -# TODO: GO over all other todos and check if they can be removed +# TODO: compare old and new version of the code, run a benchmark to compare the two # TODO: look at the documentation and try to improve it, lots of things are now combined together @@ -111,7 +109,6 @@ class StepsBlendingConfig: return_output: bool = True -# TODO: typing could be improved here @dataclass class StepsBlendingParams: noise_std_coeffs: Optional[np.ndarray] = ( @@ -127,7 +124,6 @@ class StepsBlendingParams: extrapolation_method: Optional[Callable[..., Any]] = None decomposition_method: Optional[Callable[..., dict]] = None recomposition_method: Optional[Callable[..., np.ndarray]] = None - # TODO: check of the following two are relevant or can be replaced vel_pert_... and noise_generator velocity_perturbations: Optional[Any] = None generate_velocity_noise: Optional[Callable[[Any, float], np.ndarray]] = None velocity_perturbations_parallel: Optional[np.ndarray] = None @@ -148,7 +144,6 @@ class StepsBlendingParams: domain_mask: Optional[np.ndarray] = None -# TODO: typing could be improved here @dataclass class StepsBlendingState: # Radar and noise states @@ -213,6 +208,9 @@ class StepsBlendingState: weights: Optional[np.ndarray] = None weights_model_only: Optional[np.ndarray] = None + # Is stores here aswell becasue this is changed during the forecast loop and thus not part of the config + extrapolation_kwargs: Dict[str, Any] = field(default_factory=dict) + class StepsBlendingNowcaster: def __init__( @@ -312,9 +310,8 @@ def __blended_nowcast_main_loop(self): if self.__config.measure_time: starttime_mainloop = time.time() - # TODO: problem with the config here! This variable changes over time... - # extrap_kwargs is in config but by adding info to it, the next run of a blended forecast will have issues! - self.__config.extrapolation_kwargs["return_displacement"] = True + self.__state.extrapolation_kwargs = deepcopy(self.__config.extrapolation_kwargs) + self.__state.extrapolation_kwargs["return_displacement"] = True self.__state.precip_cascades_prev_subtimestep = deepcopy( self.__state.precip_cascades @@ -697,15 +694,48 @@ def __prepare_radar_and_NWP_fields(self): # 1. Start with the radar rainfall fields. We want the fields in a # Lagrangian space - self.__precip = _transform_to_lagrangian( - self.__precip, - self.__velocity, - self.__config.ar_order, - self.__params.xy_coordinates, - self.__params.extrapolation_method, - self.__config.extrapolation_kwargs, - self.__config.num_workers, - ) + + """Advect the previous precipitation fields to the same position with the + most recent one (i.e. transform them into the Lagrangian coordinates). + """ + self.__config.extrapolation_kwargs["xy_coords"] = self.__params.xy_coordinates + res = [] + + # TODO: create beter names here fore this part, adapted from previous code which is now inlined (old function was called _transform_to_lagrangian) + def f(precip, i): + return self.__params.extrapolation_method( + precip[i, :, :], + self.__velocity, + self.__config.ar_order - i, + "min", + allow_nonfinite_values=True, + **self.__config.extrapolation_kwargs.copy(), + )[-1] + + if not DASK_IMPORTED: + # Process each earlier precipitation field directly + for i in range(self.__config.ar_order): + self.__precip[i, :, :] = f(self.__precip, i) + else: + # Use Dask delayed for parallelization if DASK_IMPORTED is True + for i in range(self.__config.ar_order): + res.append(dask.delayed(f)(self.__precip, i)) + num_workers_ = ( + len(res) + if self.__config.num_workers > len(res) + else self.__config.num_workers + ) + self.__precip = np.stack( + list(dask.compute(*res, num_workers=num_workers_)) + + [self.__precip[-1, :, :]] + ) + + # Replace non-finite values with the minimum value for each field + self.__precip = self.__precip.copy() + for i in range(self.__precip.shape[0]): + self.__precip[i, ~np.isfinite(self.__precip[i, :])] = np.nanmin( + self.__precip[i, :] + ) # 2. Perform the cascade decomposition for the input precip fields and, # if necessary, for the (NWP) model fields @@ -740,9 +770,19 @@ def __prepare_radar_and_NWP_fields(self): if self.__params.precip_models_provided_is_cascade: self.__state.precip_models_cascades = self.__precip_models - self.__precip_models = _compute_cascade_recomposition_nwp( - self.__precip_models, self.__params.recomposition_method - ) + # Inline logic of _compute_cascade_recomposition_nwp + temp_precip_models = [] + for i in range(self.__precip_models.shape[0]): + precip_model = [] + for time_step in range(self.__precip_models.shape[1]): + # Use the recomposition method to rebuild the rainfall fields + recomposed = self.__params.recomposition_method( + self.__precip_models[i, time_step] + ) + precip_model.append(recomposed) + temp_precip_models.append(precip_model) + + self.__precip_models = np.stack(temp_precip_models) # 2.3 Check for zero input fields in the radar and NWP data. self.__params.zero_precip_radar = blending.utils.check_norain( @@ -1610,9 +1650,9 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( # (or subtimesteps if non-integer time steps are given) # Settings and initialize the output - extrap_kwargs_ = self.__config.extrapolation_kwargs.copy() - extrap_kwargs_noise = self.__config.extrapolation_kwargs.copy() - extrap_kwargs_pb = self.__config.extrapolation_kwargs.copy() + extrap_kwargs_ = self.__state.extrapolation_kwargs.copy() + extrap_kwargs_noise = self.__state.extrapolation_kwargs.copy() + extrap_kwargs_pb = self.__state.extrapolation_kwargs.copy() velocity_perturbations_extrapolation = self.__velocity # The following should be accesseble after this function worker_state.precip_extrapolated_decomp = [] @@ -1721,9 +1761,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( ) # Put back the mask precip_forecast_recomp_subtimestep[self.__params.domain_mask] = np.nan - # TODO: problem with the config here! This variable changes over time... - # extrap_kwargs is in config but by adding info to it, the next run of a blended forecast will have issues! - self.__config.extrapolation_kwargs["displacement_prev"] = ( + self.__state.extrapolation_kwargs["displacement_prev"] = ( worker_state.previous_displacement[j] ) ( @@ -1734,7 +1772,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( velocity_blended, [t_diff_prev_subtimestep], allow_nonfinite_values=True, - **self.__config.extrapolation_kwargs, + **self.__state.extrapolation_kwargs, ) precip_extrapolated_recomp_subtimestep = ( precip_forecast_extrapolated_recomp_subtimestep_temp[0].copy() @@ -2264,9 +2302,23 @@ def __post_process_output( >= self.__config.precip_threshold ) # Buffer the mask - precip_field_mask = _compute_incremental_mask( - precip_field_mask, self.__params.struct, self.__params.mask_rim + # buffer the observation mask Rbin using the kernel kr + # add a grayscale rim r (for smooth rain/no-rain transition) + + # buffer observation mask + precip_field_mask_temp = np.ndarray.astype( + precip_field_mask.copy(), "uint8" ) + Rd = binary_dilation(precip_field_mask_temp, self.__params.struct) + + # add grayscale rim + kr1 = generate_binary_structure(2, 1) + mask = Rd.astype(float) + for n in range(self.__params.mask_rim): + Rd = binary_dilation(Rd, kr1) + mask += Rd + # normalize between 0 and 1 + precip_field_mask = mask / mask.max() # Get the final mask worker_state.final_blended_forecast_recomposed = ( precip_forecast_min_value @@ -2853,7 +2905,7 @@ def forecast( return forecast_steps_nowcast -# TODO: add the following code to the main body +# TODO: Where does this piece of code best fit: in utils or inside the class? def calculate_weights_spn(correlations, covariance): """Calculate SPN blending weights for STEPS blending from correlation. @@ -2935,6 +2987,7 @@ def calculate_weights_spn(correlations, covariance): return weights +# TODO: Where does this piece of code best fit: in utils or inside the class? def blend_means_sigmas(means, sigmas, weights): """Calculate the blended means and sigmas, the normalization parameters needed to recompose the cascade. This procedure uses the weights of the @@ -2999,80 +3052,3 @@ def blend_means_sigmas(means, sigmas, weights): # TODO: substract covariances to weigthed sigmas - still necessary? return combined_means, combined_sigmas - - -def _compute_incremental_mask(Rbin, kr, r): - # buffer the observation mask Rbin using the kernel kr - # add a grayscale rim r (for smooth rain/no-rain transition) - - # buffer observation mask - Rbin = np.ndarray.astype(Rbin.copy(), "uint8") - Rd = binary_dilation(Rbin, kr) - - # add grayscale rim - kr1 = generate_binary_structure(2, 1) - mask = Rd.astype(float) - for n in range(r): - Rd = binary_dilation(Rd, kr1) - mask += Rd - # normalize between 0 and 1 - return mask / mask.max() - - -def _transform_to_lagrangian( - precip, velocity, ar_order, xy_coords, extrapolator, extrap_kwargs, num_workers -): - """Advect the previous precipitation fields to the same position with the - most recent one (i.e. transform them into the Lagrangian coordinates). - """ - extrap_kwargs = extrap_kwargs.copy() - extrap_kwargs["xy_coords"] = xy_coords - res = list() - - def f(precip, i): - return extrapolator( - precip[i, :, :], - velocity, - ar_order - i, - "min", - allow_nonfinite_values=True, - **extrap_kwargs, - )[-1] - - for i in range(ar_order): - if not DASK_IMPORTED: - precip[i, :, :] = f(precip, i) - else: - res.append(dask.delayed(f)(precip, i)) - - if DASK_IMPORTED: - num_workers_ = len(res) if num_workers > len(res) else num_workers - precip = np.stack( - list(dask.compute(*res, num_workers=num_workers_)) + [precip[-1, :, :]] - ) - - # replace non-finite values with the minimum value - precip = precip.copy() - for i in range(precip.shape[0]): - precip[i, ~np.isfinite(precip[i, :])] = np.nanmin(precip[i, :]) - return precip - - -def _compute_cascade_recomposition_nwp(precip_models_cascade, recompositor): - """If necessary, recompose (NWP) model forecasts.""" - precip_models = None - - # Recompose the (NWP) model cascades to have rainfall fields per - # model and time step, which will be used in the probability matching steps. - # Recomposed cascade will have shape: [n_models, n_timesteps, m, n] - precip_models = [] - for i in range(precip_models_cascade.shape[0]): - precip_model = [] - for time_step in range(precip_models_cascade.shape[1]): - precip_model.append(recompositor(precip_models_cascade[i, time_step])) - precip_models.append(precip_model) - - precip_models = np.stack(precip_models) - precip_model = None - - return precip_models From 38ed19513d80bfd7077ec4b7a2033f3e87d2c2dd Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Fri, 6 Dec 2024 17:47:06 +0100 Subject: [PATCH 21/33] Made changes for Codacy review --- pysteps/blending/steps.py | 107 ++++++++++++++++++++------------------ 1 file changed, 56 insertions(+), 51 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index ee53462c..f889276d 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -382,7 +382,7 @@ def worker(j): if self.__state.is_nowcast_time_step: if self.__config.measure_time: - __ = self.__measure_time("subtimestep", starttime) + _ = self.__measure_time("subtimestep", starttime) else: print("done.") @@ -695,9 +695,9 @@ def __prepare_radar_and_NWP_fields(self): # 1. Start with the radar rainfall fields. We want the fields in a # Lagrangian space - """Advect the previous precipitation fields to the same position with the - most recent one (i.e. transform them into the Lagrangian coordinates). - """ + # Advect the previous precipitation fields to the same position with the + # most recent one (i.e. transform them into the Lagrangian coordinates). + self.__config.extrapolation_kwargs["xy_coords"] = self.__params.xy_coordinates res = [] @@ -740,7 +740,7 @@ def f(precip, i): # 2. Perform the cascade decomposition for the input precip fields and, # if necessary, for the (NWP) model fields # 2.1 Compute the cascade decompositions of the input precipitation fields - """Compute the cascade decompositions of the input precipitation fields.""" + # Compute the cascade decompositions of the input precipitation fields. precip_forecast_decomp = [] for i in range(self.__config.ar_order + 1): precip_forecast = self.__params.decomposition_method( @@ -902,7 +902,7 @@ def __prepare_nowcast_for_zero_radar(self): # step where the fraction of rainy cells is highest (because other lead times # might be zero as well). Else, initialize the noise with the radar # rainfall data - """Initialize noise based on the NWP field time step where the fraction of rainy cells is highest""" + # Initialize noise based on the NWP field time step where the fraction of rainy cells is highest if self.__config.precip_threshold is None: self.__config.precip_threshold = np.nanmin(self.__precip_models) @@ -967,7 +967,7 @@ def __initialize_noise(self): ) if self.__config.measure_time: - __ = self.__measure_time("Initialize noise", starttime) + _ = self.__measure_time("Initialize noise", starttime) else: print("done.") elif self.__config.noise_stddev_adj == "fixed": @@ -1028,7 +1028,7 @@ def __estimate_ar_parameters_radar(self): if self.__config.ar_order == 1: GAMMA = GAMMA[0, :] if self.__config.ar_order > 2: - for repeat_index in range(self.__config.ar_order - 2): + for _ in range(self.__config.ar_order - 2): GAMMA = np.vstack((GAMMA, GAMMA[1, :])) # Finally, transpose GAMMA to ensure that the shape is the same as np.empty((n_cascade_levels, ar_order)) @@ -1279,7 +1279,7 @@ def __decompose_nwp_if_needed_and_fill_nans_in_nwp(self, t): # fill these with the minimum value present in precip (corresponding to # zero rainfall in the radar observations) - """Ensure that the NWP cascade and fields do no contain any nans or infinite number""" + # Ensure that the NWP cascade and fields do no contain any nans or infinite number # Fill nans and infinite numbers with the minimum value present in precip self.__state.precip_models_timestep = self.__precip_models[:, t, :, :].astype( np.float64, copy=False @@ -1419,7 +1419,7 @@ def __find_nowcast_NWP_combination(self, t): def __determine_skill_for_current_timestep(self, t): if t == 0: - """Calculate the initial skill of the (NWP) model forecasts at t=0.""" + # Calculate the initial skill of the (NWP) model forecasts at t=0. self.__params.rho_nwp_models = [] for model_index in range( self.__state.precip_models_cascades_timestep.shape[0] @@ -1949,24 +1949,26 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( ) extrap_kwargs_noise["map_coordinates_mode"] = "wrap" - _, worker_state.previous_displacement[j] = ( - self.__params.extrapolation_method( - None, - velocity_blended, - [t_diff_prev_subtimestep], - allow_nonfinite_values=True, - **extrap_kwargs_, - ) + ( + _, + worker_state.previous_displacement[j], + ) = self.__params.extrapolation_method( + None, + velocity_blended, + [t_diff_prev_subtimestep], + allow_nonfinite_values=True, + **extrap_kwargs_, ) - _, worker_state.previous_displacement_noise_cascade[j] = ( - self.__params.extrapolation_method( - None, - velocity_blended, - [t_diff_prev_subtimestep], - allow_nonfinite_values=True, - **extrap_kwargs_noise, - ) + ( + _, + worker_state.previous_displacement_noise_cascade[j], + ) = self.__params.extrapolation_method( + None, + velocity_blended, + [t_diff_prev_subtimestep], + allow_nonfinite_values=True, + **extrap_kwargs_noise, ) # Also extrapolate the radar observation, used for the probability @@ -1974,14 +1976,15 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( extrap_kwargs_pb["displacement_prev"] = ( worker_state.previous_displacement_prob_matching[j] ) - _, worker_state.previous_displacement_prob_matching[j] = ( - self.__params.extrapolation_method( - None, - velocity_blended, - [t_diff_prev_subtimestep], - allow_nonfinite_values=True, - **extrap_kwargs_pb, - ) + ( + _, + worker_state.previous_displacement_prob_matching[j], + ) = self.__params.extrapolation_method( + None, + velocity_blended, + [t_diff_prev_subtimestep], + allow_nonfinite_values=True, + **extrap_kwargs_pb, ) worker_state.time_prev_timestep[j] = t + 1 @@ -2246,7 +2249,7 @@ def __post_process_output( ) # Perform the blending of radar and model inside the radar domain using a weighted combination - precip_forecast_recomposed = np.nansum( + worker_state.final_blended_forecast_recomposed = np.nansum( [ mask_model * precip_forecast_recomposed_mod_only_no_nan, mask_radar * precip_forecast_recomposed_no_nan, @@ -2302,23 +2305,25 @@ def __post_process_output( >= self.__config.precip_threshold ) # Buffer the mask - # buffer the observation mask Rbin using the kernel kr - # add a grayscale rim r (for smooth rain/no-rain transition) + # Convert the precipitation field mask into an 8-bit unsigned integer mask + obs_mask_uint8 = precip_field_mask.astype("uint8") - # buffer observation mask - precip_field_mask_temp = np.ndarray.astype( - precip_field_mask.copy(), "uint8" - ) - Rd = binary_dilation(precip_field_mask_temp, self.__params.struct) - - # add grayscale rim - kr1 = generate_binary_structure(2, 1) - mask = Rd.astype(float) - for n in range(self.__params.mask_rim): - Rd = binary_dilation(Rd, kr1) - mask += Rd - # normalize between 0 and 1 - precip_field_mask = mask / mask.max() + # Perform an initial binary dilation using the provided structuring element + dilated_mask = binary_dilation(obs_mask_uint8, self.__params.struct) + + # Create a binary structure element for incremental dilations + struct_element = generate_binary_structure(2, 1) + + # Initialize a floating-point mask to accumulate dilations for a smooth transition + accumulated_mask = dilated_mask.astype(float) + + # Iteratively dilate the mask and accumulate the results to create a grayscale rim + for _ in range(self.__params.mask_rim): + dilated_mask = binary_dilation(dilated_mask, struct_element) + accumulated_mask += dilated_mask + + # Normalize the accumulated mask values between 0 and 1 + precip_field_mask = accumulated_mask / np.max(accumulated_mask) # Get the final mask worker_state.final_blended_forecast_recomposed = ( precip_forecast_min_value From 32b656f475f12ba2d781c6c85628c42cab6f5c34 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 16 Dec 2024 11:35:36 +0100 Subject: [PATCH 22/33] Added aditional tests which currently fail in master branch --- pysteps/blending/steps.py | 2 +- pysteps/tests/test_blending_steps.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index f889276d..fcf07222 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -701,7 +701,7 @@ def __prepare_radar_and_NWP_fields(self): self.__config.extrapolation_kwargs["xy_coords"] = self.__params.xy_coordinates res = [] - # TODO: create beter names here fore this part, adapted from previous code which is now inlined (old function was called _transform_to_lagrangian) + # TODO: create beter names here for this part, adapted from previous code which is now inlined (old function was called _transform_to_lagrangian) def f(precip, i): return self.__params.extrapolation_method( precip[i, :, :], diff --git a/pysteps/tests/test_blending_steps.py b/pysteps/tests/test_blending_steps.py index 5840279f..99bd1794 100644 --- a/pysteps/tests/test_blending_steps.py +++ b/pysteps/tests/test_blending_steps.py @@ -28,6 +28,7 @@ (2, 3, 2, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False), # TODO: make next test work! This is currently not working on the main branch # (2, 3, 4, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False), + # (2, 3, 4, 8, "incremental", "cdf", False, "spn", True, 2, False, False, 0, False), (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False), # Test the case where the radar image contains no rain. (1, 3, 6, 8, None, None, False, "spn", True, 6, True, False, 0, False), From 4fe9f78cbb5bcd2db5c080e01d64317e47b43159 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 16 Dec 2024 15:03:21 +0100 Subject: [PATCH 23/33] Update .gitignore Co-authored-by: mats-knmi <145579783+mats-knmi@users.noreply.github.com> --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 136c1e46..c2006db3 100644 --- a/.gitignore +++ b/.gitignore @@ -92,6 +92,6 @@ venv.bak/ # Mac OS Stuff .DS_Store -# Running lcoal tests +# Running local tests /tmp /pysteps/tests/tmp/ From b31d55c3079f3ee0c52a0db0474c16e496747af3 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 16 Dec 2024 18:23:15 +0100 Subject: [PATCH 24/33] Used the __zero_precip_time in __zero_precipitation_forecast() --- pysteps/blending/steps.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index fcf07222..1eacfa21 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -828,7 +828,7 @@ def __zero_precipitation_forecast(self): precip_forecast_workers = None if self.__config.measure_time: - zero_precip_time = time.time() - self.__start_time_init + self.__zero_precip_time = time.time() - self.__start_time_init if self.__config.return_output: precip_forecast_all_members_all_times = np.stack( @@ -841,8 +841,8 @@ def __zero_precipitation_forecast(self): if self.__config.measure_time: return ( precip_forecast_all_members_all_times, - zero_precip_time, - zero_precip_time, + self.__zero_precip_time, + self.__zero_precip_time, ) else: return precip_forecast_all_members_all_times From cc025938b29d1ec200c33328d00d9ecb0845ac9b Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 16 Dec 2024 19:26:22 +0100 Subject: [PATCH 25/33] Changed typing hints to python 3.10+ version --- pysteps/blending/steps.py | 202 ++++++++++++++++++-------------------- 1 file changed, 94 insertions(+), 108 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 1eacfa21..589e3ddb 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -65,7 +65,7 @@ DASK_IMPORTED = False from dataclasses import dataclass, field -from typing import Optional, List, Dict, Any, Callable, Union +from typing import Any, Callable # TODO: compare old and new version of the code, run a benchmark to compare the two # TODO: look at the documentation and try to improve it, lots of things are now combined together @@ -73,7 +73,7 @@ @dataclass class StepsBlendingConfig: - precip_threshold: Optional[float] + precip_threshold: float | None norain_threshold: float kmperpixel: float timestep: float @@ -83,133 +83,119 @@ class StepsBlendingConfig: extrapolation_method: str decomposition_method: str bandpass_filter_method: str - noise_method: Optional[str] - noise_stddev_adj: Optional[str] + noise_method: str | None + noise_stddev_adj: str | None ar_order: int - velocity_perturbation_method: Optional[str] + velocity_perturbation_method: str | None weights_method: str conditional: bool - probmatching_method: Optional[str] - mask_method: Optional[str] + probmatching_method: str | None + mask_method: str | None resample_distribution: bool smooth_radar_mask_range: int - seed: Optional[int] + seed: int | None num_workers: int fft_method: str domain: str outdir_path_skill: str - extrapolation_kwargs: Dict[str, Any] = field(default_factory=dict) - filter_kwargs: Dict[str, Any] = field(default_factory=dict) - noise_kwargs: Dict[str, Any] = field(default_factory=dict) - velocity_perturbation_kwargs: Dict[str, Any] = field(default_factory=dict) - climatology_kwargs: Dict[str, Any] = field(default_factory=dict) - mask_kwargs: Dict[str, Any] = field(default_factory=dict) + extrapolation_kwargs: dict[str, Any] = field(default_factory=dict) + filter_kwargs: dict[str, Any] = field(default_factory=dict) + noise_kwargs: dict[str, Any] = field(default_factory=dict) + velocity_perturbation_kwargs: dict[str, Any] = field(default_factory=dict) + climatology_kwargs: dict[str, Any] = field(default_factory=dict) + mask_kwargs: dict[str, Any] = field(default_factory=dict) measure_time: bool = False - callback: Optional[Any] = None + callback: Any | None = None return_output: bool = True @dataclass class StepsBlendingParams: - noise_std_coeffs: Optional[np.ndarray] = ( - None # Noise standard deviation coefficients - ) - bandpass_filter: Optional[Any] = None # Band-pass filter object - fft: Optional[Any] = None # FFT method object - perturbation_generator: Optional[Callable[..., np.ndarray]] = ( - None # Perturbation generator - ) - noise_generator: Optional[Callable[..., np.ndarray]] = None # Noise generator - PHI: Optional[np.ndarray] = None # AR(p) model parameters - extrapolation_method: Optional[Callable[..., Any]] = None - decomposition_method: Optional[Callable[..., dict]] = None - recomposition_method: Optional[Callable[..., np.ndarray]] = None - velocity_perturbations: Optional[Any] = None - generate_velocity_noise: Optional[Callable[[Any, float], np.ndarray]] = None - velocity_perturbations_parallel: Optional[np.ndarray] = None - velocity_perturbations_perpendicular: Optional[np.ndarray] = None - fft_objs: List[Any] = field(default_factory=list) - mask_rim: Optional[int] = None # Rim size for masking - struct: Optional[np.ndarray] = None # Structuring element for mask - time_steps_is_list: bool = False # Time steps is a list - precip_models_provided_is_cascade: bool = False # Precip models are decomposed - xy_coordinates: Optional[np.ndarray] = None - precip_zerovalue: Optional[float] = None - mask_threshold: Optional[np.ndarray] = None + noise_std_coeffs: np.ndarray | None = None + bandpass_filter: Any | None = None + fft: Any | None = None + perturbation_generator: Callable[..., np.ndarray] | None = None + noise_generator: Callable[..., np.ndarray] | None = None + PHI: np.ndarray | None = None + extrapolation_method: Callable[..., Any] | None = None + decomposition_method: Callable[..., dict] | None = None + recomposition_method: Callable[..., np.ndarray] | None = None + velocity_perturbations: Any | None = None + generate_velocity_noise: Callable[[Any, float], np.ndarray] | None = None + velocity_perturbations_parallel: np.ndarray | None = None + velocity_perturbations_perpendicular: np.ndarray | None = None + fft_objs: list[Any] = field(default_factory=list) + mask_rim: int | None = None + struct: np.ndarray | None = None + time_steps_is_list: bool = False + precip_models_provided_is_cascade: bool = False + xy_coordinates: np.ndarray | None = None + precip_zerovalue: float | None = None + mask_threshold: np.ndarray | None = None zero_precip_radar: bool = False zero_precip_model_fields: bool = False - original_timesteps: Optional[Union[list, np.ndarray]] = None - num_ensemble_workers: Optional[int] = None - rho_nwp_models: Optional[np.ndarray] = None - domain_mask: Optional[np.ndarray] = None + original_timesteps: list | np.ndarray | None = None + num_ensemble_workers: int | None = None + rho_nwp_models: np.ndarray | None = None + domain_mask: np.ndarray | None = None @dataclass class StepsBlendingState: - # Radar and noise states - precip_cascades: Optional[np.ndarray] = None - precip_noise_input: Optional[np.ndarray] = None - precip_noise_cascades: Optional[np.ndarray] = None - precip_mean_noise: Optional[np.ndarray] = None - precip_std_noise: Optional[np.ndarray] = None - - # Extrapolation states - mean_extrapolation: Optional[np.ndarray] = None - std_extrapolation: Optional[np.ndarray] = None - rho_extrap_cascade_prev: Optional[np.ndarray] = None - rho_extrap_cascade: Optional[np.ndarray] = None - precip_cascades_prev_subtimestep: Optional[np.ndarray] = None - cascade_noise_prev_subtimestep: Optional[np.ndarray] = None - precip_extrapolated_after_decomp: Optional[np.ndarray] = None - noise_extrapolated_after_decomp: Optional[np.ndarray] = None - precip_extrapolated_probability_matching: Optional[np.ndarray] = None - - # NWP model states - precip_models_cascades: Optional[np.ndarray] = None - precip_models_cascades_timestep: Optional[np.ndarray] = None - precip_models_timestep: Optional[np.ndarray] = None - mean_models_timestep: Optional[np.ndarray] = None - std_models_timestep: Optional[np.ndarray] = None - velocity_models_timestep: Optional[np.ndarray] = None - - # Mapping from NWP members to ensemble members - mapping_list_NWP_member_to_ensemble_member: Optional[np.ndarray] = None - - # Random states for precipitation and motion - randgen_precip: Optional[List[np.random.RandomState]] = None - randgen_motion: Optional[List[np.random.RandomState]] = None - - # Variables for final forecast computation - previous_displacement: Optional[List[Any]] = None - previous_displacement_noise_cascade: Optional[List[Any]] = None - previous_displacement_prob_matching: Optional[List[Any]] = None - rho_final_blended_forecast: Optional[np.ndarray] = None - final_blended_forecast_means: Optional[np.ndarray] = None - final_blended_forecast_stds: Optional[np.ndarray] = None - final_blended_forecast_means_mod_only: Optional[np.ndarray] = None - final_blended_forecast_stds_mod_only: Optional[np.ndarray] = None - final_blended_forecast_cascades: Optional[np.ndarray] = None - final_blended_forecast_cascades_mod_only: Optional[np.ndarray] = None - final_blended_forecast_recomposed: Optional[np.ndarray] = None - final_blended_forecast_recomposed_mod_only: Optional[np.ndarray] = None - - # Final outputs - final_blended_forecast: Optional[np.ndarray] = None - final_blended_forecast_non_perturbed: Optional[np.ndarray] = None - - # Timing and indexing - time_prev_timestep: Optional[List[float]] = None - leadtime_since_start_forecast: Optional[List[float]] = None - subtimesteps: Optional[List[float]] = None - is_nowcast_time_step: Optional[bool] = None - subtimestep_index: Optional[int] = None - - # Weights used for blending - weights: Optional[np.ndarray] = None - weights_model_only: Optional[np.ndarray] = None - - # Is stores here aswell becasue this is changed during the forecast loop and thus not part of the config - extrapolation_kwargs: Dict[str, Any] = field(default_factory=dict) + precip_cascades: np.ndarray | None = None + precip_noise_input: np.ndarray | None = None + precip_noise_cascades: np.ndarray | None = None + precip_mean_noise: np.ndarray | None = None + precip_std_noise: np.ndarray | None = None + + mean_extrapolation: np.ndarray | None = None + std_extrapolation: np.ndarray | None = None + rho_extrap_cascade_prev: np.ndarray | None = None + rho_extrap_cascade: np.ndarray | None = None + precip_cascades_prev_subtimestep: np.ndarray | None = None + cascade_noise_prev_subtimestep: np.ndarray | None = None + precip_extrapolated_after_decomp: np.ndarray | None = None + noise_extrapolated_after_decomp: np.ndarray | None = None + precip_extrapolated_probability_matching: np.ndarray | None = None + + precip_models_cascades: np.ndarray | None = None + precip_models_cascades_timestep: np.ndarray | None = None + precip_models_timestep: np.ndarray | None = None + mean_models_timestep: np.ndarray | None = None + std_models_timestep: np.ndarray | None = None + velocity_models_timestep: np.ndarray | None = None + + mapping_list_NWP_member_to_ensemble_member: np.ndarray | None = None + + randgen_precip: list[np.random.RandomState] | None = None + randgen_motion: list[np.random.RandomState] | None = None + + previous_displacement: list[Any] | None = None + previous_displacement_noise_cascade: list[Any] | None = None + previous_displacement_prob_matching: list[Any] | None = None + rho_final_blended_forecast: np.ndarray | None = None + final_blended_forecast_means: np.ndarray | None = None + final_blended_forecast_stds: np.ndarray | None = None + final_blended_forecast_means_mod_only: np.ndarray | None = None + final_blended_forecast_stds_mod_only: np.ndarray | None = None + final_blended_forecast_cascades: np.ndarray | None = None + final_blended_forecast_cascades_mod_only: np.ndarray | None = None + final_blended_forecast_recomposed: np.ndarray | None = None + final_blended_forecast_recomposed_mod_only: np.ndarray | None = None + + final_blended_forecast: np.ndarray | None = None + final_blended_forecast_non_perturbed: np.ndarray | None = None + + time_prev_timestep: list[float] | None = None + leadtime_since_start_forecast: list[float] | None = None + subtimesteps: list[float] | None = None + is_nowcast_time_step: bool | None = None + subtimestep_index: int | None = None + + weights: np.ndarray | None = None + weights_model_only: np.ndarray | None = None + + extrapolation_kwargs: dict[str, Any] = field(default_factory=dict) class StepsBlendingNowcaster: From 4e4a148dd971163dcaec1b0d292beb8fef71d09f Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Mon, 16 Dec 2024 19:31:20 +0100 Subject: [PATCH 26/33] Added comments back to the State dataclass --- pysteps/blending/steps.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 589e3ddb..ba285fbb 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -142,12 +142,14 @@ class StepsBlendingParams: @dataclass class StepsBlendingState: + # Radar and noise states precip_cascades: np.ndarray | None = None precip_noise_input: np.ndarray | None = None precip_noise_cascades: np.ndarray | None = None precip_mean_noise: np.ndarray | None = None precip_std_noise: np.ndarray | None = None + # Extrapolation states mean_extrapolation: np.ndarray | None = None std_extrapolation: np.ndarray | None = None rho_extrap_cascade_prev: np.ndarray | None = None @@ -158,6 +160,7 @@ class StepsBlendingState: noise_extrapolated_after_decomp: np.ndarray | None = None precip_extrapolated_probability_matching: np.ndarray | None = None + # NWP model states precip_models_cascades: np.ndarray | None = None precip_models_cascades_timestep: np.ndarray | None = None precip_models_timestep: np.ndarray | None = None @@ -165,11 +168,14 @@ class StepsBlendingState: std_models_timestep: np.ndarray | None = None velocity_models_timestep: np.ndarray | None = None + # Mapping from NWP members to ensemble members mapping_list_NWP_member_to_ensemble_member: np.ndarray | None = None + # Random states for precipitation and motion randgen_precip: list[np.random.RandomState] | None = None randgen_motion: list[np.random.RandomState] | None = None + # Variables for final forecast computation previous_displacement: list[Any] | None = None previous_displacement_noise_cascade: list[Any] | None = None previous_displacement_prob_matching: list[Any] | None = None @@ -183,18 +189,22 @@ class StepsBlendingState: final_blended_forecast_recomposed: np.ndarray | None = None final_blended_forecast_recomposed_mod_only: np.ndarray | None = None + # Final outputs final_blended_forecast: np.ndarray | None = None final_blended_forecast_non_perturbed: np.ndarray | None = None + # Timing and indexing time_prev_timestep: list[float] | None = None leadtime_since_start_forecast: list[float] | None = None subtimesteps: list[float] | None = None is_nowcast_time_step: bool | None = None subtimestep_index: int | None = None + # Weights used for blending weights: np.ndarray | None = None weights_model_only: np.ndarray | None = None + # This is stores here as well because this is changed during the forecast loop and thus no longer part of the config extrapolation_kwargs: dict[str, Any] = field(default_factory=dict) From 0f4e037912fcddb85f52366b3c7bbfb661c87304 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Tue, 17 Dec 2024 12:46:45 +0100 Subject: [PATCH 27/33] Changed the self.__state.velocity_perturbations = [] to self.__params.velocity_perturbations = [] in __initialize_random_generators --- pysteps/blending/steps.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index ba285fbb..37976f63 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -1091,7 +1091,7 @@ def __initialize_random_generators(self): ) = noise.get_method(self.__config.velocity_perturbation_method) # initialize the perturbation generators for the motion field - self.__state.velocity_perturbations = [] + self.__params.velocity_perturbations = [] for j in range(self.__config.n_ens_members): kwargs = { "randstate": self.__state.randgen_motion[j], @@ -1104,7 +1104,7 @@ def __initialize_random_generators(self): self.__config.timestep, **kwargs, ) - self.__state.velocity_perturbations.append(vp_) + self.__params.velocity_perturbations.append(vp_) else: ( self.__params.velocity_perturbations, @@ -2902,7 +2902,6 @@ def forecast( forecast_steps_nowcast = blended_nowcaster.compute_forecast() print(forecast_steps_nowcast) blended_nowcaster.reset_states_and_params() - # Call the appropriate methods within the class return forecast_steps_nowcast From 9f413aa4431f41a53e97bb9c2ac4a70a8aaca12e Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Wed, 18 Dec 2024 23:53:22 +0100 Subject: [PATCH 28/33] Added code changes as suggested by Ruben, comments and documentation to come later --- pysteps/blending/steps.py | 238 ++++++++++++++++++-------------------- 1 file changed, 114 insertions(+), 124 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 37976f63..14866bc4 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -339,7 +339,7 @@ def __blended_nowcast_main_loop(self): def worker(j): # The state needs to be copied as a dataclass is not threadsafe in python worker_state = deepcopy(self.__state) - self.__determine_skill_for_next_timestep(t, j, worker_state) + self.__determine_NWP_skill_for_next_timestep(t, j, worker_state) self.__determine_weights_per_component(worker_state) self.__regress_extrapolation_and_noise_cascades(j, worker_state) self.__perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( @@ -697,7 +697,7 @@ def __prepare_radar_and_NWP_fields(self): self.__config.extrapolation_kwargs["xy_coords"] = self.__params.xy_coordinates res = [] - # TODO: create beter names here for this part, adapted from previous code which is now inlined (old function was called _transform_to_lagrangian) + # TODO: create better names here for this part, adapted from previous code which is now inlined (old function was called _transform_to_lagrangian) def f(precip, i): return self.__params.extrapolation_method( precip[i, :, :], @@ -1377,41 +1377,30 @@ def __find_nowcast_NWP_combination(self, t): (n_ens_members_max + i) // n_ens_members_min for i in range(n_ens_members_min) ] - if n_model_members == n_ens_members_min: - self.__state.precip_models_cascades_timestep = np.repeat( - self.__state.precip_models_cascades_timestep, - repeats, - axis=0, - ) - self.__state.mean_models_timestep = np.repeat( - self.__state.mean_models_timestep, repeats, axis=0 - ) - self.__state.std_models_timestep = np.repeat( - self.__state.std_models_timestep, repeats, axis=0 - ) - self.__state.velocity_models_timestep = np.repeat( - self.__state.velocity_models_timestep, repeats, axis=0 - ) - # For the prob. matching - self.__state.precip_models_timestep = np.repeat( - self.__state.precip_models_timestep, repeats, axis=0 - ) - # Finally, for the model indices - self.__state.mapping_list_NWP_member_to_ensemble_member = ( - np.repeat( - self.__state.mapping_list_NWP_member_to_ensemble_member, - repeats, - axis=0, - ) - ) - - # TODO: is this not duplicate from part 2.3.5? If so, is it still needed here? - # If zero_precip_radar is True, set the velocity field equal to the NWP - # velocity field for the current time step (velocity_models_temp). - if self.__params.zero_precip_radar: - # Use the velocity from velocity_models and take the average over - # n_models (axis=0) - self.__velocity = np.mean(self.__state.velocity_models_timestep, axis=0) + self.__state.precip_models_cascades_timestep = np.repeat( + self.__state.precip_models_cascades_timestep, + repeats, + axis=0, + ) + self.__state.mean_models_timestep = np.repeat( + self.__state.mean_models_timestep, repeats, axis=0 + ) + self.__state.std_models_timestep = np.repeat( + self.__state.std_models_timestep, repeats, axis=0 + ) + self.__state.velocity_models_timestep = np.repeat( + self.__state.velocity_models_timestep, repeats, axis=0 + ) + # For the prob. matching + self.__state.precip_models_timestep = np.repeat( + self.__state.precip_models_timestep, repeats, axis=0 + ) + # Finally, for the model indices + self.__state.mapping_list_NWP_member_to_ensemble_member = np.repeat( + self.__state.mapping_list_NWP_member_to_ensemble_member, + repeats, + axis=0, + ) def __determine_skill_for_current_timestep(self, t): if t == 0: @@ -1465,7 +1454,7 @@ def __determine_skill_for_current_timestep(self, t): correlations_prev=self.__state.rho_extrap_cascade_prev, ) - def __determine_skill_for_next_timestep(self, t, j, worker_state): + def __determine_NWP_skill_for_next_timestep(self, t, j, worker_state): # 8.1.2 Determine the skill of the nwp components for lead time (t0 + t) # Then for the model components if self.__config.blend_nwp_members: @@ -1650,7 +1639,7 @@ def __perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( extrap_kwargs_noise = self.__state.extrapolation_kwargs.copy() extrap_kwargs_pb = self.__state.extrapolation_kwargs.copy() velocity_perturbations_extrapolation = self.__velocity - # The following should be accesseble after this function + # The following should be accessible after this function worker_state.precip_extrapolated_decomp = [] worker_state.noise_extrapolated_decomp = [] worker_state.precip_extrapolated_probability_matching = [] @@ -2447,89 +2436,6 @@ def reset_states_and_params(self): self.__mainloop_time = None -def calculate_ratios(correlations): - """Calculate explained variance ratios from correlation. - - Parameters - ---------- - Array of shape [component, scale_level, ...] - containing correlation (skills) for each component (NWP and nowcast), - scale level, and optionally along [y, x] dimensions. - - Returns - ------- - out : numpy array - An array containing the ratios of explain variance for each - component, scale level, ... - """ - # correlations: [component, scale, ...] - square_corrs = np.square(correlations) - # Calculate the ratio of the explained variance to the unexplained - # variance of the nowcast and NWP model components - out = square_corrs / (1 - square_corrs) - # out: [component, scale, ...] - return out - - -def calculate_weights_bps(correlations): - """Calculate BPS blending weights for STEPS blending from correlation. - - Parameters - ---------- - correlations : array-like - Array of shape [component, scale_level, ...] - containing correlation (skills) for each component (NWP and nowcast), - scale level, and optionally along [y, x] dimensions. - - Returns - ------- - weights : array-like - Array of shape [component+1, scale_level, ...] - containing the weights to be used in STEPS blending for - each original component plus an addtional noise component, scale level, - and optionally along [y, x] dimensions. - - References - ---------- - :cite:`BPS2006` - - Notes - ----- - The weights in the BPS method can sum op to more than 1.0. - """ - # correlations: [component, scale, ...] - # Check if the correlations are positive, otherwise rho = 10e-5 - correlations = np.where(correlations < 10e-5, 10e-5, correlations) - - # If we merge more than one component with the noise cascade, we follow - # the weights impolementation in either :cite:`BPS2006` or :cite:`SPN2013`. - if correlations.shape[0] > 1: - # Calculate weights for each source - ratios = calculate_ratios(correlations) - # ratios: [component, scale, ...] - total_ratios = np.sum(ratios, axis=0) - # total_ratios: [scale, ...] - the denominator of eq. 11 & 12 in BPS2006 - weights = correlations * np.sqrt(ratios / total_ratios) - # weights: [component, scale, ...] - # Calculate the weight of the noise component. - # Original BPS2006 method in the following two lines (eq. 13) - total_square_weights = np.sum(np.square(weights), axis=0) - noise_weight = np.sqrt(1.0 - total_square_weights) - # Finally, add the noise_weights to the weights variable. - weights = np.concatenate((weights, noise_weight[None, ...]), axis=0) - - # Otherwise, the weight equals the correlation on that scale level and - # the noise component weight equals 1 - this weight. This only occurs for - # the weights calculation outside the radar domain where in the case of 1 - # NWP model or ensemble member, no blending of multiple models has to take - # place - else: - noise_weight = 1.0 - correlations - weights = np.concatenate((correlations, noise_weight), axis=0) - - return weights - - def forecast( precip, precip_models, @@ -2798,7 +2704,7 @@ def forecast( fmi=Finland, mch=Switzerland, fmi+mch=both pooled into the same data set - The above parameters have been fitten by using run_vel_pert_analysis.py + The above parameters have been fitted by using run_vel_pert_analysis.py and fit_vel_pert_params.py located in the scripts directory. See :py:mod:`pysteps.noise.motion` for additional documentation. @@ -2905,6 +2811,91 @@ def forecast( return forecast_steps_nowcast +# TODO: Where does this piece of code best fit: in utils or inside the class? +def calculate_ratios(correlations): + """Calculate explained variance ratios from correlation. + + Parameters + ---------- + Array of shape [component, scale_level, ...] + containing correlation (skills) for each component (NWP and nowcast), + scale level, and optionally along [y, x] dimensions. + + Returns + ------- + out : numpy array + An array containing the ratios of explain variance for each + component, scale level, ... + """ + # correlations: [component, scale, ...] + square_corrs = np.square(correlations) + # Calculate the ratio of the explained variance to the unexplained + # variance of the nowcast and NWP model components + out = square_corrs / (1 - square_corrs) + # out: [component, scale, ...] + return out + + +# TODO: Where does this piece of code best fit: in utils or inside the class? +def calculate_weights_bps(correlations): + """Calculate BPS blending weights for STEPS blending from correlation. + + Parameters + ---------- + correlations : array-like + Array of shape [component, scale_level, ...] + containing correlation (skills) for each component (NWP and nowcast), + scale level, and optionally along [y, x] dimensions. + + Returns + ------- + weights : array-like + Array of shape [component+1, scale_level, ...] + containing the weights to be used in STEPS blending for + each original component plus an addtional noise component, scale level, + and optionally along [y, x] dimensions. + + References + ---------- + :cite:`BPS2006` + + Notes + ----- + The weights in the BPS method can sum op to more than 1.0. + """ + # correlations: [component, scale, ...] + # Check if the correlations are positive, otherwise rho = 10e-5 + correlations = np.where(correlations < 10e-5, 10e-5, correlations) + + # If we merge more than one component with the noise cascade, we follow + # the weights impolementation in either :cite:`BPS2006` or :cite:`SPN2013`. + if correlations.shape[0] > 1: + # Calculate weights for each source + ratios = calculate_ratios(correlations) + # ratios: [component, scale, ...] + total_ratios = np.sum(ratios, axis=0) + # total_ratios: [scale, ...] - the denominator of eq. 11 & 12 in BPS2006 + weights = correlations * np.sqrt(ratios / total_ratios) + # weights: [component, scale, ...] + # Calculate the weight of the noise component. + # Original BPS2006 method in the following two lines (eq. 13) + total_square_weights = np.sum(np.square(weights), axis=0) + noise_weight = np.sqrt(1.0 - total_square_weights) + # Finally, add the noise_weights to the weights variable. + weights = np.concatenate((weights, noise_weight[None, ...]), axis=0) + + # Otherwise, the weight equals the correlation on that scale level and + # the noise component weight equals 1 - this weight. This only occurs for + # the weights calculation outside the radar domain where in the case of 1 + # NWP model or ensemble member, no blending of multiple models has to take + # place + else: + noise_weight = 1.0 - correlations + weights = np.concatenate((correlations, noise_weight), axis=0) + + return weights + + # TODO: Where does this piece of code best fit: in utils or inside the class? def calculate_weights_spn(correlations, covariance): """Calculate SPN blending weights for STEPS blending from correlation. @@ -3049,6 +3040,5 @@ def blend_means_sigmas(means, sigmas, weights): for i in range(weights.shape[0]): combined_means += (weights[i] / total_weight) * means[i] combined_sigmas += (weights[i] / total_weight) * sigmas[i] - # TODO: substract covariances to weigthed sigmas - still necessary? return combined_means, combined_sigmas From c72d9539871d2aede891784a681933eceda50744 Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 19 Dec 2024 11:47:43 +0100 Subject: [PATCH 29/33] Added frozen functionality to dataclasses, removed reset_state and fixed seed assingments --- pysteps/blending/steps.py | 82 ++++++++++++++++++++++----------------- 1 file changed, 46 insertions(+), 36 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 14866bc4..4a1048e5 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -71,7 +71,7 @@ # TODO: look at the documentation and try to improve it, lots of things are now combined together -@dataclass +@dataclass(frozen=True) class StepsBlendingConfig: precip_threshold: float | None norain_threshold: float @@ -138,6 +138,11 @@ class StepsBlendingParams: num_ensemble_workers: int | None = None rho_nwp_models: np.ndarray | None = None domain_mask: np.ndarray | None = None + filter_kwargs: dict | None = None + noise_kwargs: dict | None = None + velocity_perturbation_kwargs: dict | None = None + climatology_kwargs: dict | None = None + mask_kwargs: dict | None = None @dataclass @@ -306,7 +311,7 @@ def __blended_nowcast_main_loop(self): if self.__config.measure_time: starttime_mainloop = time.time() - self.__state.extrapolation_kwargs = deepcopy(self.__config.extrapolation_kwargs) + # self.__state.extrapolation_kwargs = deepcopy(self.__config.extrapolation_kwargs) self.__state.extrapolation_kwargs["return_displacement"] = True self.__state.precip_cascades_prev_subtimestep = deepcopy( @@ -478,25 +483,43 @@ def __check_inputs(self): ) if self.__config.extrapolation_kwargs is None: - self.__config.extrapolation_kwargs = dict() + self.__state.extrapolation_kwargs = dict() + else: + self.__state.extrapolation_kwargs = deepcopy( + self.__config.extrapolation_kwargs + ) if self.__config.filter_kwargs is None: - self.__config.filter_kwargs = dict() + self.__params.filter_kwargs = dict() + else: + self.__params.filter_kwargs = deepcopy(self.__config.filter_kwargs) if self.__config.noise_kwargs is None: - self.__config.noise_kwargs = dict() + self.__params.noise_kwargs = dict() + else: + self.__params.noise_kwargs = deepcopy(self.__config.noise_kwargs) if self.__config.velocity_perturbation_kwargs is None: - self.__config.velocity_perturbation_kwargs = dict() + self.__params.velocity_perturbation_kwargs = dict() + else: + self.__params.velocity_perturbation_kwargs = deepcopy( + self.__config.velocity_perturbation_kwargs + ) if self.__config.climatology_kwargs is None: # Make sure clim_kwargs at least contains the number of models - self.__config.climatology_kwargs = dict( + self.__params.climatology_kwargs = dict( {"n_models": self.__precip_models.shape[0]} ) + else: + self.__params.climatology_kwargs = deepcopy( + self.__config.climatology_kwargs + ) if self.__config.mask_kwargs is None: - self.__config.mask_kwargs = dict() + self.__params.mask_kwargs = dict() + else: + self.__params.mask_kwargs = deepcopy(self.__config.mask_kwargs) if np.any(~np.isfinite(self.__velocity)): raise ValueError("velocity contains non-finite values") @@ -598,12 +621,12 @@ def __print_forecast_info(self): print(f"order of the AR(p) model: {self.__config.ar_order}") if self.__config.velocity_perturbation_method == "bps": self.__params.velocity_perturbations_parallel = ( - self.__config.velocity_perturbation_kwargs.get( + self.__params.velocity_perturbation_kwargs.get( "p_par", noise.motion.get_default_params_bps_par() ) ) self.__params.velocity_perturbations_perpendicular = ( - self.__config.velocity_perturbation_kwargs.get( + self.__params.velocity_perturbation_kwargs.get( "p_perp", noise.motion.get_default_params_bps_perp() ) ) @@ -645,7 +668,7 @@ def __initialize_nowcast_components(self): self.__params.bandpass_filter = filter_method( (M, N), self.__config.n_cascade_levels, - **(self.__config.filter_kwargs or {}), + **(self.__params.filter_kwargs or {}), ) # Get the decomposition method (e.g., FFT) @@ -694,7 +717,7 @@ def __prepare_radar_and_NWP_fields(self): # Advect the previous precipitation fields to the same position with the # most recent one (i.e. transform them into the Lagrangian coordinates). - self.__config.extrapolation_kwargs["xy_coords"] = self.__params.xy_coordinates + self.__state.extrapolation_kwargs["xy_coords"] = self.__params.xy_coordinates res = [] # TODO: create better names here for this part, adapted from previous code which is now inlined (old function was called _transform_to_lagrangian) @@ -705,7 +728,7 @@ def f(precip, i): self.__config.ar_order - i, "min", allow_nonfinite_values=True, - **self.__config.extrapolation_kwargs.copy(), + **self.__state.extrapolation_kwargs.copy(), )[-1] if not DASK_IMPORTED: @@ -939,7 +962,7 @@ def __initialize_noise(self): self.__params.perturbation_generator = init_noise( self.__state.precip_noise_input, fft_method=self.__params.fft, - **self.__config.noise_kwargs, + **self.__params.noise_kwargs, ) if self.__config.noise_stddev_adj == "auto": @@ -1076,8 +1099,9 @@ def __initialize_random_generators(self): if self.__config.noise_method is not None: self.__state.randgen_precip = [] self.__state.randgen_motion = [] + seed = self.__config.seed for j in range(self.__config.n_ens_members): - rs = np.random.RandomState(self.__config.seed) + rs = np.random.RandomState(seed) self.__state.randgen_precip.append(rs) seed = rs.randint(0, high=1e9) rs = np.random.RandomState(seed) @@ -1129,8 +1153,8 @@ def __prepare_forecast_loop(self): if self.__config.mask_method == "incremental": # get mask parameters - self.__params.mask_rim = self.__config.mask_kwargs.get("mask_rim", 10) - mask_f = self.__config.mask_kwargs.get("mask_f", 1.0) + self.__params.mask_rim = self.__params.mask_kwargs.get("mask_rim", 10) + mask_f = self.__params.mask_kwargs.get("mask_f", 1.0) # initialize the structuring element struct = generate_binary_structure(2, 1) # iterate it to expand it nxn @@ -1440,7 +1464,7 @@ def __determine_skill_for_current_timestep(self, t): current_skill=self.__params.rho_nwp_models, validtime=self.__issuetime, outdir_path=self.__config.outdir_path_skill, - **self.__config.climatology_kwargs, + **self.__params.climatology_kwargs, ) if t > 0: # 8.1.3 Determine the skill of the components for lead time (t0 + t) @@ -1465,7 +1489,7 @@ def __determine_NWP_skill_for_next_timestep(self, t, j, worker_state): correlations=self.__params.rho_nwp_models[model_index], outdir_path=self.__config.outdir_path_skill, n_model=model_index, - skill_kwargs=self.__config.climatology_kwargs, + skill_kwargs=self.__params.climatology_kwargs, ) rho_nwp_forecast.append(rho_value) rho_nwp_forecast = np.stack(rho_nwp_forecast) @@ -1480,7 +1504,7 @@ def __determine_NWP_skill_for_next_timestep(self, t, j, worker_state): correlations=self.__params.rho_nwp_models[j], outdir_path=self.__config.outdir_path_skill, n_model=worker_state.mapping_list_NWP_member_to_ensemble_member[j], - skill_kwargs=self.__config.climatology_kwargs, + skill_kwargs=self.__params.climatology_kwargs, ) # Concatenate rho_extrap_cascade and rho_nwp worker_state.rho_final_blended_forecast = np.concatenate( @@ -2420,21 +2444,6 @@ def __measure_time(self, label, start_time): return elapsed_time return None - def reset_states_and_params(self): - """ - Reset the internal state and parameters of the nowcaster to allow multiple forecasts. - This method resets the state and params to their initial conditions without reinitializing - the inputs like precip, velocity, time_steps, or config. - """ - # Re-initialize the state and parameters - self.__state = StepsBlendingState() - self.__params = StepsBlendingParams() - - # Reset time measurement variables - self.__start_time_init = None - self.__init_time = None - self.__mainloop_time = None - def forecast( precip, @@ -2794,6 +2803,8 @@ def forecast( return_output=return_output, ) + # TODO: add comment about how this class based method is supposed to be used: for each forecast run, a new forecaster needs to be made. The config file can stay the same. + # Create an instance of the new class with all the provided arguments blended_nowcaster = StepsBlendingNowcaster( precip, @@ -2807,7 +2818,6 @@ def forecast( forecast_steps_nowcast = blended_nowcaster.compute_forecast() print(forecast_steps_nowcast) - blended_nowcaster.reset_states_and_params() return forecast_steps_nowcast From 00f057b078609e13af28a812abea269a5396e2dc Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 19 Dec 2024 12:06:11 +0100 Subject: [PATCH 30/33] Added frozen dataclass to nowcast --- pysteps/blending/steps.py | 1 - pysteps/nowcasts/steps.py | 71 +++++++++++++++++++++++++-------------- 2 files changed, 45 insertions(+), 27 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 4a1048e5..94cce042 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -311,7 +311,6 @@ def __blended_nowcast_main_loop(self): if self.__config.measure_time: starttime_mainloop = time.time() - # self.__state.extrapolation_kwargs = deepcopy(self.__config.extrapolation_kwargs) self.__state.extrapolation_kwargs["return_displacement"] = True self.__state.precip_cascades_prev_subtimestep = deepcopy( diff --git a/pysteps/nowcasts/steps.py b/pysteps/nowcasts/steps.py index b61ee8e7..818123da 100644 --- a/pysteps/nowcasts/steps.py +++ b/pysteps/nowcasts/steps.py @@ -14,6 +14,7 @@ import numpy as np from scipy.ndimage import generate_binary_structure, iterate_structure import time +from copy import deepcopy from pysteps import cascade from pysteps import extrapolation @@ -35,7 +36,7 @@ DASK_IMPORTED = False -@dataclass +@dataclass(frozen=True) class StepsNowcasterConfig: """ Parameters @@ -247,6 +248,10 @@ class StepsNowcasterParams: xy_coordinates: np.ndarray | None = None velocity_perturbation_parallel: list[float] | None = None velocity_perturbation_perpendicular: list[float] | None = None + filter_kwargs: dict | None = None + noise_kwargs: dict | None = None + velocity_perturbation_kwargs: dict | None = None + mask_kwargs: dict | None = None @dataclass @@ -268,6 +273,7 @@ class StepsNowcasterState: ) velocity_perturbations: list[Callable] | None = field(default_factory=list) fft_objects: list[Any] | None = field(default_factory=list) + extrapolation_kwargs: dict[str, Any] | None = field(default_factory=dict) class StepsNowcaster: @@ -408,7 +414,7 @@ def __nowcast_main(self): self.__time_steps, self.__config.extrapolation_method, self.__update_state, # Reference to the update function - extrap_kwargs=self.__config.extrapolation_kwargs, + extrap_kwargs=self.__state.extrapolation_kwargs, velocity_pert_gen=self.__state.velocity_perturbations, params=params, ensemble=True, @@ -483,15 +489,33 @@ def __check_inputs(self): # Handle None values for various kwargs if self.__config.extrapolation_kwargs is None: - self.__config.extrapolation_kwargs = {} + self.__state.extrapolation_kwargs = dict() + else: + self.__state.extrapolation_kwargs = deepcopy( + self.__config.extrapolation_kwargs + ) + if self.__config.filter_kwargs is None: - self.__config.filter_kwargs = {} + self.__params.filter_kwargs = dict() + else: + self.__params.filter_kwargs = deepcopy(self.__config.filter_kwargs) + if self.__config.noise_kwargs is None: - self.__config.noise_kwargs = {} + self.__params.noise_kwargs = dict() + else: + self.__params.noise_kwargs = deepcopy(self.__config.noise_kwargs) + if self.__config.velocity_perturbation_kwargs is None: - self.__config.velocity_perturbation_kwargs = {} + self.__params.velocity_perturbation_kwargs = dict() + else: + self.__params.velocity_perturbation_kwargs = deepcopy( + self.__config.velocity_perturbation_kwargs + ) + if self.__config.mask_kwargs is None: - self.__config.mask_kwargs = {} + self.__params.mask_kwargs = dict() + else: + self.__params.mask_kwargs = deepcopy(self.__config.mask_kwargs) print("Inputs validated and initialized successfully.") @@ -548,12 +572,12 @@ def __print_forecast_info(self): if self.__config.velocity_perturbation_method == "bps": self.__params.velocity_perturbation_parallel = ( - self.__config.velocity_perturbation_kwargs.get( + self.__params.velocity_perturbation_kwargs.get( "p_par", noise.motion.get_default_params_bps_par() ) ) self.__params.velocity_perturbation_perpendicular = ( - self.__config.velocity_perturbation_kwargs.get( + self.__params.velocity_perturbation_kwargs.get( "p_perp", noise.motion.get_default_params_bps_perp() ) ) @@ -588,7 +612,7 @@ def __initialize_nowcast_components(self): self.__params.bandpass_filter = filter_method( (M, N), self.__config.n_cascade_levels, - **(self.__config.filter_kwargs or {}), + **(self.__params.filter_kwargs or {}), ) # Get the decomposition method (e.g., FFT) @@ -629,7 +653,7 @@ def __perform_extrapolation(self): else: self.__state.mask_threshold = None - extrap_kwargs = self.__config.extrapolation_kwargs.copy() + extrap_kwargs = self.__state.extrapolation_kwargs.copy() extrap_kwargs["xy_coords"] = self.__params.xy_coordinates extrap_kwargs["allow_nonfinite_values"] = ( True if np.any(~np.isfinite(self.__precip)) else False @@ -691,7 +715,7 @@ def __apply_noise_and_ar_model(self): self.__params.perturbation_generator = init_noise( self.__precip, fft_method=self.__params.fft, - **self.__config.noise_kwargs, + **self.__params.noise_kwargs, ) # Handle noise standard deviation adjustments if necessary @@ -831,21 +855,16 @@ def __apply_noise_and_ar_model(self): if self.__config.noise_method is not None: self.__state.random_generator_precip = [] self.__state.random_generator_motion = [] - + seed = self.__config.seed for _ in range(self.__config.n_ens_members): # Create random state for precipitation noise generator - rs = np.random.RandomState(self.__config.seed) + rs = np.random.RandomState(seed) self.__state.random_generator_precip.append(rs) - self.__config.seed = rs.randint( - 0, high=int(1e9) - ) # Update seed after generating - + seed = rs.randint(0, high=int(1e9)) # Create random state for motion perturbations generator - rs = np.random.RandomState(self.__config.seed) + rs = np.random.RandomState(seed) self.__state.random_generator_motion.append(rs) - self.__config.seed = rs.randint( - 0, high=int(1e9) - ) # Update seed after generating + seed = rs.randint(0, high=int(1e9)) else: self.__state.random_generator_precip = None self.__state.random_generator_motion = None @@ -865,10 +884,10 @@ def __initialize_velocity_perturbations(self): for j in range(self.__config.n_ens_members): kwargs = { "randstate": self.__state.random_generator_motion[j], - "p_par": self.__config.velocity_perturbation_kwargs.get( + "p_par": self.__params.velocity_perturbation_kwargs.get( "p_par", self.__params.velocity_perturbation_parallel ), - "p_perp": self.__config.velocity_perturbation_kwargs.get( + "p_perp": self.__params.velocity_perturbation_kwargs.get( "p_perp", self.__params.velocity_perturbation_perpendicular ), } @@ -920,8 +939,8 @@ def __initialize_precipitation_mask(self): elif self.__config.mask_method == "incremental": # Get mask parameters - self.__params.mask_rim = self.__config.mask_kwargs.get("mask_rim", 10) - mask_f = self.__config.mask_kwargs.get("mask_f", 1.0) + self.__params.mask_rim = self.__params.mask_kwargs.get("mask_rim", 10) + mask_f = self.__params.mask_kwargs.get("mask_f", 1.0) # Initialize the structuring element self.__params.structuring_element = generate_binary_structure(2, 1) # Expand the structuring element based on mask factor and timestep From 1b8251283fe9d57b11df5dad348470a7c96b19df Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Thu, 19 Dec 2024 17:21:18 +0100 Subject: [PATCH 31/33] The needed checks are done for this TODO so it can be removed --- pysteps/blending/steps.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 94cce042..3c8317d8 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -2067,9 +2067,6 @@ def __blend_cascades(self, t_sub, j, worker_state): # First determine the blending weights if method is spn. The # weights for method bps have already been determined. - # TODO: no other weight method is possible, should we not al least give a user warning if a different weight - # method is given? Or does this mean that in all other circumstances the weights - # have been calculated in a different way? if self.__config.weights_method == "spn": worker_state.weights = np.zeros( From 47ab6c3f274e5e333560ae929da20c782bc8d3e6 Mon Sep 17 00:00:00 2001 From: mats-knmi <145579783+mats-knmi@users.noreply.github.com> Date: Thu, 2 Jan 2025 18:05:11 +0100 Subject: [PATCH 32/33] Use the seed in all rng in blending code (#449) * Use seed for all rng to make a test run completely deterministic * fix probmatching test and some copy paste oversights * Add test for vel_pert_method * Change the test so that it actually runs the lines that need to be covered --- pysteps/blending/steps.py | 19 ++++- pysteps/noise/utils.py | 5 +- pysteps/postprocessing/probmatching.py | 4 +- pysteps/tests/test_blending_steps.py | 78 ++++++++++--------- .../tests/test_postprocessing_probmatching.py | 25 +++--- 5 files changed, 74 insertions(+), 57 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 3c8317d8..1c86ee2e 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -176,9 +176,10 @@ class StepsBlendingState: # Mapping from NWP members to ensemble members mapping_list_NWP_member_to_ensemble_member: np.ndarray | None = None - # Random states for precipitation and motion + # Random states for precipitation, motion and probmatching randgen_precip: list[np.random.RandomState] | None = None randgen_motion: list[np.random.RandomState] | None = None + randgen_probmatching: list[np.random.RandomState] | None = None # Variables for final forecast computation previous_displacement: list[Any] | None = None @@ -1095,19 +1096,28 @@ def __multiply_precip_cascade_to_match_ensemble_members(self): def __initialize_random_generators(self): # 6. Initialize all the random generators and prepare for the forecast loop """Initialize all the random generators.""" + seed = self.__config.seed if self.__config.noise_method is not None: self.__state.randgen_precip = [] - self.__state.randgen_motion = [] - seed = self.__config.seed for j in range(self.__config.n_ens_members): rs = np.random.RandomState(seed) self.__state.randgen_precip.append(rs) seed = rs.randint(0, high=1e9) + + if self.__config.probmatching_method is not None: + self.__state.randgen_probmatching = [] + for j in range(self.__config.n_ens_members): rs = np.random.RandomState(seed) - self.__state.randgen_motion.append(rs) + self.__state.randgen_probmatching.append(rs) seed = rs.randint(0, high=1e9) if self.__config.velocity_perturbation_method is not None: + self.__state.randgen_motion = [] + for j in range(self.__config.n_ens_members): + rs = np.random.RandomState(seed) + self.__state.randgen_motion.append(rs) + seed = rs.randint(0, high=1e9) + ( init_velocity_noise, self.__params.generate_velocity_noise, @@ -2373,6 +2383,7 @@ def __post_process_output( first_array=arr1, second_array=arr2, probability_first_array=weights_probability_matching_normalized[0], + randgen=self.__state.randgen_probmatching[j], ) ) else: diff --git a/pysteps/noise/utils.py b/pysteps/noise/utils.py index 58495b8a..aaae82f9 100644 --- a/pysteps/noise/utils.py +++ b/pysteps/noise/utils.py @@ -101,8 +101,9 @@ def compute_noise_stddev_adjs( randstates = [] for k in range(num_iter): - randstates.append(np.random.RandomState(seed=seed)) - seed = np.random.randint(0, high=1e9) + rs = np.random.RandomState(seed=seed) + randstates.append(rs) + seed = rs.randint(0, high=1e9) def worker(k): # generate Gaussian white noise field, filter it using the chosen diff --git a/pysteps/postprocessing/probmatching.py b/pysteps/postprocessing/probmatching.py index afc95154..efda2b10 100644 --- a/pysteps/postprocessing/probmatching.py +++ b/pysteps/postprocessing/probmatching.py @@ -274,7 +274,7 @@ def _get_error(scale): return shift, scale, R.reshape(shape) -def resample_distributions(first_array, second_array, probability_first_array): +def resample_distributions(first_array, second_array, probability_first_array, randgen): """ Merges two distributions (e.g., from the extrapolation nowcast and NWP in the blending module) to effectively combine two distributions for probability matching without losing extremes. @@ -324,7 +324,7 @@ def resample_distributions(first_array, second_array, probability_first_array): n = asort.shape[0] # Resample the distributions - idxsamples = np.random.binomial(1, probability_first_array, n).astype(bool) + idxsamples = randgen.binomial(1, probability_first_array, n).astype(bool) csort = np.where(idxsamples, asort, bsort) csort = np.sort(csort)[::-1] diff --git a/pysteps/tests/test_blending_steps.py b/pysteps/tests/test_blending_steps.py index 99bd1794..4752ea32 100644 --- a/pysteps/tests/test_blending_steps.py +++ b/pysteps/tests/test_blending_steps.py @@ -8,48 +8,48 @@ import pysteps from pysteps import blending, cascade +# fmt:off steps_arg_values = [ - # Test the case where both the radar image and the NWP fields contain no rain. - (1, 3, 4, 8, None, None, False, "spn", True, 4, False, False, 0, False), - (1, 3, 4, 8, "obs", None, False, "spn", True, 4, False, False, 0, False), - (1, 3, 4, 8, "incremental", None, False, "spn", True, 4, False, False, 0, False), - (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False, 0, False), - (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False, 0, True), - (1, 3, 4, 8, None, "cdf", False, "spn", True, 4, False, False, 0, False), - (1, [1, 2, 3], 4, 8, None, "cdf", False, "spn", True, 4, False, False, 0, False), - (1, 3, 4, 8, "incremental", "cdf", False, "spn", True, 4, False, False, 0, False), - (1, 3, 4, 6, "incremental", "cdf", False, "bps", True, 4, False, False, 0, False), - (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False, 0, False), - (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False, 0, True), - (1, 3, 4, 9, "incremental", "cdf", False, "spn", True, 4, False, False, 0, False), - (2, 3, 10, 8, "incremental", "cdf", False, "spn", True, 10, False, False, 0, False), - (5, 3, 5, 8, "incremental", "cdf", False, "spn", True, 5, False, False, 0, False), - (1, 10, 1, 8, "incremental", "cdf", False, "spn", True, 1, False, False, 0, False), - (2, 3, 2, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False), - # TODO: make next test work! This is currently not working on the main branch - # (2, 3, 4, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False), - # (2, 3, 4, 8, "incremental", "cdf", False, "spn", True, 2, False, False, 0, False), - (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False), + (1, 3, 4, 8, None, None, False, "spn", True, 4, False, False, 0, False, None), + (1, 3, 4, 8, "obs", None, False, "spn", True, 4, False, False, 0, False, None), + (1, 3, 4, 8, "incremental", None, False, "spn", True, 4, False, False, 0, False, None), + (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False, 0, False, None), + (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False, 0, True, None), + (1, 3, 4, 8, None, "cdf", False, "spn", True, 4, False, False, 0, False, None), + (1, [1, 2, 3], 4, 8, None, "cdf", False, "spn", True, 4, False, False, 0, False, None), + (1, 3, 4, 8, "incremental", "cdf", False, "spn", True, 4, False, False, 0, False, None), + (1, 3, 4, 6, "incremental", "cdf", False, "bps", True, 4, False, False, 0, False, None), + (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False, 0, False, None), + (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False, 0, True, None), + (1, 3, 4, 9, "incremental", "cdf", False, "spn", True, 4, False, False, 0, False, None), + (2, 3, 10, 8, "incremental", "cdf", False, "spn", True, 10, False, False, 0, False, None), + (5, 3, 5, 8, "incremental", "cdf", False, "spn", True, 5, False, False, 0, False, None), + (1, 10, 1, 8, "incremental", "cdf", False, "spn", True, 1, False, False, 0, False, None), + (2, 3, 2, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False, None), + (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False, None), + (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False, "bps"), # Test the case where the radar image contains no rain. - (1, 3, 6, 8, None, None, False, "spn", True, 6, True, False, 0, False), - (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0, False), - (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0, True), + (1, 3, 6, 8, None, None, False, "spn", True, 6, True, False, 0, False, None), + (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0, False, None), + (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0, True, None), # Test the case where the NWP fields contain no rain. - (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 0, False), - (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, True, 0, True), - (1, 3, 6, 8, None, None, False, "spn", True, 6, True, True, 0, False), - (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, True, 0, False), - (5, 3, 5, 6, "obs", "mean", True, "spn", True, 5, True, True, 0, False), + (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 0, False, None), + (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, True, 0, True, None), + # Test the case where both the radar image and the NWP fields contain no rain. + (1, 3, 6, 8, None, None, False, "spn", True, 6, True, True, 0, False, None), + (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, True, 0, False, None), + (5, 3, 5, 6, "obs", "mean", True, "spn", True, 5, True, True, 0, False, None), # Test for smooth radar mask - (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 80, False), - (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, False, 80, False), - (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, False, False, 80, False), - (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 80, False), - (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 80, True), - (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False), - (5, [1, 2, 3], 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False), - (5, [1, 3], 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False), + (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 80, False, None), + (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, False, 80, False, None), + (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, False, False, 80, False, None), + (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 80, False, None), + (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 80, True, None), + (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False, None), + (5, [1, 2, 3], 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False, None), + (5, [1, 3], 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False, None), ] +# fmt:on steps_arg_names = ( "n_models", @@ -66,6 +66,7 @@ "zero_nwp", "smooth_radar_mask_range", "resample_distribution", + "vel_pert_method", ) @@ -85,6 +86,7 @@ def test_steps_blending( zero_nwp, smooth_radar_mask_range, resample_distribution, + vel_pert_method, ): pytest.importorskip("cv2") @@ -278,7 +280,7 @@ def test_steps_blending( noise_method="nonparametric", noise_stddev_adj="auto", ar_order=2, - vel_pert_method=None, + vel_pert_method=vel_pert_method, weights_method=weights_method, conditional=False, probmatching_method=probmatching_method, diff --git a/pysteps/tests/test_postprocessing_probmatching.py b/pysteps/tests/test_postprocessing_probmatching.py index 8c7c12f4..b42a95e7 100644 --- a/pysteps/tests/test_postprocessing_probmatching.py +++ b/pysteps/tests/test_postprocessing_probmatching.py @@ -1,7 +1,10 @@ import numpy as np import pytest -from pysteps.postprocessing.probmatching import resample_distributions -from pysteps.postprocessing.probmatching import nonparam_match_empirical_cdf + +from pysteps.postprocessing.probmatching import ( + nonparam_match_empirical_cdf, + resample_distributions, +) class TestResampleDistributions: @@ -16,7 +19,7 @@ def test_valid_inputs(self): second_array = np.array([2, 4, 6, 8, 10]) probability_first_array = 0.6 result = resample_distributions( - first_array, second_array, probability_first_array + first_array, second_array, probability_first_array, np.random ) expected_result = np.array([9, 8, 6, 3, 1]) # Expected result based on the seed assert result.shape == first_array.shape @@ -27,7 +30,7 @@ def test_probability_zero(self): second_array = np.array([2, 4, 6, 8, 10]) probability_first_array = 0.0 result = resample_distributions( - first_array, second_array, probability_first_array + first_array, second_array, probability_first_array, np.random ) assert np.array_equal(result, np.sort(second_array)[::-1]) @@ -36,7 +39,7 @@ def test_probability_one(self): second_array = np.array([2, 4, 6, 8, 10]) probability_first_array = 1.0 result = resample_distributions( - first_array, second_array, probability_first_array + first_array, second_array, probability_first_array, np.random ) assert np.array_equal(result, np.sort(first_array)[::-1]) @@ -45,7 +48,7 @@ def test_nan_in_arr1_prob_1(self): array_without_nan = np.array([2.0, 4, 6, 8, 10]) probability_first_array = 1.0 result = resample_distributions( - array_with_nan, array_without_nan, probability_first_array + array_with_nan, array_without_nan, probability_first_array, np.random ) expected_result = np.array([np.nan, 9, 7, 3, 1], dtype=float) assert np.allclose(result, expected_result, equal_nan=True) @@ -55,7 +58,7 @@ def test_nan_in_arr1_prob_0(self): array_without_nan = np.array([2, 4, 6, 8, 10]) probability_first_array = 0.0 result = resample_distributions( - array_with_nan, array_without_nan, probability_first_array + array_with_nan, array_without_nan, probability_first_array, np.random ) expected_result = np.array([np.nan, 10, 8, 4, 2], dtype=float) assert np.allclose(result, expected_result, equal_nan=True) @@ -65,7 +68,7 @@ def test_nan_in_arr2_prob_1(self): array_with_nan = np.array([2.0, 4, 6, np.nan, 10]) probability_first_array = 1.0 result = resample_distributions( - array_without_nan, array_with_nan, probability_first_array + array_without_nan, array_with_nan, probability_first_array, np.random ) expected_result = np.array([np.nan, 9, 5, 3, 1], dtype=float) assert np.allclose(result, expected_result, equal_nan=True) @@ -75,7 +78,7 @@ def test_nan_in_arr2_prob_0(self): array_with_nan = np.array([2, 4, 6, np.nan, 10]) probability_first_array = 0.0 result = resample_distributions( - array_without_nan, array_with_nan, probability_first_array + array_without_nan, array_with_nan, probability_first_array, np.random ) expected_result = np.array([np.nan, 10, 6, 4, 2], dtype=float) assert np.allclose(result, expected_result, equal_nan=True) @@ -85,7 +88,7 @@ def test_nan_in_both_prob_1(self): array2_with_nan = np.array([2.0, 4, np.nan, np.nan, 10]) probability_first_array = 1.0 result = resample_distributions( - array1_with_nan, array2_with_nan, probability_first_array + array1_with_nan, array2_with_nan, probability_first_array, np.random ) expected_result = np.array([np.nan, np.nan, np.nan, 9, 1], dtype=float) assert np.allclose(result, expected_result, equal_nan=True) @@ -95,7 +98,7 @@ def test_nan_in_both_prob_0(self): array2_with_nan = np.array([2.0, 4, np.nan, np.nan, 10]) probability_first_array = 0.0 result = resample_distributions( - array1_with_nan, array2_with_nan, probability_first_array + array1_with_nan, array2_with_nan, probability_first_array, np.random ) expected_result = np.array([np.nan, np.nan, np.nan, 10, 2], dtype=float) assert np.allclose(result, expected_result, equal_nan=True) From 48187c4a55f03844e78a332e97584b46180efb6d Mon Sep 17 00:00:00 2001 From: Simon De Kock Date: Fri, 3 Jan 2025 15:49:22 +0100 Subject: [PATCH 33/33] Removed deepcopy of worker_state. The state is now accessable to all workers at the same time --- pysteps/blending/steps.py | 16 +++++++--------- pysteps/tests/test_blending_steps.py | 3 +++ 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 1c86ee2e..5ec243ea 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -342,24 +342,22 @@ def __blended_nowcast_main_loop(self): ] def worker(j): - # The state needs to be copied as a dataclass is not threadsafe in python - worker_state = deepcopy(self.__state) - self.__determine_NWP_skill_for_next_timestep(t, j, worker_state) - self.__determine_weights_per_component(worker_state) - self.__regress_extrapolation_and_noise_cascades(j, worker_state) + self.__determine_NWP_skill_for_next_timestep(t, j, self.__state) + self.__determine_weights_per_component(self.__state) + self.__regress_extrapolation_and_noise_cascades(j, self.__state) self.__perturb_blend_and_advect_extrapolation_and_noise_to_current_timestep( - t, j, worker_state + t, j, self.__state ) # 8.5 Blend the cascades final_blended_forecast_single_member = [] for t_sub in self.__state.subtimesteps: # TODO: does it make sense to use sub time steps - check if it works? if t_sub > 0: - self.__blend_cascades(t_sub, j, worker_state) - self.__recompose_cascade_to_rainfall_field(j, worker_state) + self.__blend_cascades(t_sub, j, self.__state) + self.__recompose_cascade_to_rainfall_field(j, self.__state) final_blended_forecast_single_member = ( self.__post_process_output( - j, final_blended_forecast_single_member, worker_state + j, final_blended_forecast_single_member, self.__state ) ) final_blended_forecast_all_members_one_timestep[j] = ( diff --git a/pysteps/tests/test_blending_steps.py b/pysteps/tests/test_blending_steps.py index 4752ea32..77f975c7 100644 --- a/pysteps/tests/test_blending_steps.py +++ b/pysteps/tests/test_blending_steps.py @@ -28,6 +28,9 @@ (2, 3, 2, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False, None), (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False, None), (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False, "bps"), + # TODO: make next test work! This is currently not working on the main branch + # (2, 3, 4, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False), + # (2, 3, 4, 8, "incremental", "cdf", False, "spn", True, 2, False, False, 0, False), # Test the case where the radar image contains no rain. (1, 3, 6, 8, None, None, False, "spn", True, 6, True, False, 0, False, None), (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0, False, None),