diff --git a/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF.H b/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF.H index 40bc5cec2d..27a4550194 100644 --- a/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF.H +++ b/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF.H @@ -129,12 +129,14 @@ read_plot_file(amrex::PlotFileData& pf, amrex::MultiFab& mf); amrex::MultiFab -read_member_multifab(int n, +read_member_multifab(const int da_iter, + const int n, const std::string& pf_name, const amrex::Vector& varnames); amrex::MultiFab -compute_ensemble_mean(int Nens, +compute_ensemble_mean(const int da_iter, + const int Nens, const std::string& pf_name, const amrex::Vector& varnames); @@ -142,7 +144,7 @@ void Apply_H(const amrex::MultiFab& x_mf, amrex::MultiFab& y_mf); void read_in_observations(const int& da_iter, amrex::MultiFab& y_obs); -std::vector get_plotfile_list(); +std::vector get_plotfile_list(const int da_iter); void read_in_observations(const int& da_iter, const amrex::Vector& varnames, amrex::MultiFab& y_obs); @@ -156,12 +158,14 @@ void compute_d_prime_vec(amrex::MultiFab& d_prime_vec, const amrex::MultiFab& d_vec, const amrex::Vector& R_diag); -void compute_mean_H_xf(amrex::MultiFab& mean_H_xf, +void compute_mean_H_xf(const int da_iter, + amrex::MultiFab& mean_H_xf, const int Nens, const std::string& last_pf_name, const amrex::Vector& varnames); -void compute_r_vec (int Nens, +void compute_r_vec (const int da_iter, + int Nens, const std::string& last_pf_name, const amrex::Vector& varnames, const amrex::MultiFab& mean_H_xf, @@ -173,14 +177,16 @@ void compute_alpha_vec (const int& Nens, const amrex::Vector& r_vec, amrex::Vector& alpha_vec); -void compute_Xf_prime_times_vector (const int Nens, - const std::string& last_pf_name, - const amrex::Vector& varnames, - const amrex::MultiFab& xf_bar, - const amrex::Vector& vec_in, - amrex::MultiFab& result); +void compute_Xf_prime_times_vector (const int da_iter, + const int Nens, + const std::string& last_pf_name, + const amrex::Vector& varnames, + const amrex::MultiFab& xf_bar, + const amrex::Vector& vec_in, + amrex::MultiFab& result); -void update_ensemble (const int Nens, +void update_ensemble (const int da_iter, + const int Nens, const std::string& last_pf_name, const amrex::Vector& varnames, const amrex::MultiFab& xf_bar, @@ -200,7 +206,8 @@ Matrix matrix_multiply(const Matrix& A, const Matrix& B); void matrix_print(const Matrix& A); void lapack_testing(); -void compute_S_matrix(Matrix& S, +void compute_S_matrix(const int da_iter, + Matrix& S, const int& Nens, const amrex::MultiFab& mean_H_xf, const amrex::Vector& R_diag, @@ -208,4 +215,7 @@ void compute_S_matrix(Matrix& S, const amrex::Vector& varnames); std::string MakeEnsembleCheckpointName (int da_iter, int ens_no); + +std::string +MakeFullPath (const std::string& filename, const int da_iter); #endif diff --git a/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF.cpp b/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF.cpp index a5934f8655..394fa0ebe3 100644 --- a/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF.cpp +++ b/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF.cpp @@ -140,10 +140,10 @@ WriteUpdatedEnsembleToERFClassData (const MultiFab& mf_cc_fine, } void -ERF::ComputeAndWriteEnsemblePerturbations() +ERF::ComputeAndWriteEnsemblePerturbations(const int da_iter) { - auto pltfiles = get_plotfile_list(); + auto pltfiles = get_plotfile_list(da_iter); // Step 2: loop over all plotfiles (timestamps at which the plotfiles are written) // ie.find the ensemble mean at iteration 100, loop over the plt00100 file in each of the @@ -153,13 +153,13 @@ ERF::ComputeAndWriteEnsemblePerturbations() const std::string member_prefix = "member_"; for (const auto& pf_name : pltfiles) { - MultiFab mf_mean = compute_ensemble_mean(Nens, pf_name, varnames); + MultiFab mf_mean = compute_ensemble_mean(da_iter, Nens, pf_name, varnames); // ------------------------------- // Step 3 & 4: Compute perturbations and write plotfiles // ------------------------------- for (int n = 0; n < Nens; ++n) { - MultiFab mf_pert = read_member_multifab(n, pf_name, varnames); + MultiFab mf_pert = read_member_multifab(da_iter, n, pf_name, varnames); MultiFab::Subtract(mf_pert, mf_mean, 0, 0, mf_mean.nComp(), mf_mean.nGrow()); // Create output directory @@ -187,7 +187,7 @@ ERF::PerformDataAssimilation(int da_iter) { //lapack_testing(); - auto pltfiles = get_plotfile_list(); + auto pltfiles = get_plotfile_list(da_iter); std::string last_pf_name; if (!pltfiles.empty()) { last_pf_name = pltfiles.back(); @@ -205,12 +205,12 @@ ERF::PerformDataAssimilation(int da_iter) Vector varnames = {"density","theta", "x_velocity","y_velocity","z_velocity", "qv", "qc", "qrain"}; // Compute the ensemble mean - MultiFab xf_bar = compute_ensemble_mean(Nens, last_pf_name, varnames); + MultiFab xf_bar = compute_ensemble_mean(da_iter, Nens, last_pf_name, varnames); Print() << "Computing ensemble mean complete" << std::endl; // Construct perturbation plotfile name - std::string pltname = "plt_ens_mean"; + std::string pltname = MakeFullPath("plt_ens_mean", da_iter); WriteSingleLevelPlotfile(pltname, xf_bar, varnames, @@ -220,13 +220,13 @@ ERF::PerformDataAssimilation(int da_iter) // Compute the mean of forecast observations yf_bar = Hx_f MultiFab mean_H_xf; - compute_mean_H_xf(mean_H_xf, Nens, last_pf_name, varnames); + compute_mean_H_xf(da_iter, mean_H_xf, Nens, last_pf_name, varnames); Print() << "compute_mean_H_xf complete" << std::endl; Vector varnames1 = {"density","theta", "x_velocity","y_velocity","z_velocity", "qv", "qc", "qrain"}; // Construct perturbation plotfile name - std::string pltname1 = "plt_mean_H_xf"; + std::string pltname1 = MakeFullPath("plt_mean_H_xf", da_iter); WriteSingleLevelPlotfile(pltname1, mean_H_xf, varnames1, @@ -240,7 +240,7 @@ ERF::PerformDataAssimilation(int da_iter) Print() << "Observation reading complete" << std::endl; - /*std::string pltname2 = "plt_y_obs"; + /*std::string pltname2 = WriteSingleLevelPlotfile("plt_y_obs", da_iter); Vector varnames2 = {"x_velocity", "y_velocity"}; WriteSingleLevelPlotfile(pltname2, y_obs, @@ -279,13 +279,13 @@ ERF::PerformDataAssimilation(int da_iter) // Compute r = Y'^Td' Vector r_vec; - compute_r_vec(Nens, last_pf_name, varnames, mean_H_xf, d_prime_vec, r_vec); + compute_r_vec(da_iter, Nens, last_pf_name, varnames, mean_H_xf, d_prime_vec, r_vec); Print() << "compute_r_vec complete" << std::endl; // Compute the S matrix Matrix S_mat(Nens); - compute_S_matrix(S_mat, Nens, mean_H_xf, R_diag, last_pf_name, varnames); + compute_S_matrix(da_iter, S_mat, Nens, mean_H_xf, R_diag, last_pf_name, varnames); Print() << "compute_S_matrix complete" << std::endl; @@ -295,11 +295,11 @@ ERF::PerformDataAssimilation(int da_iter) Print() << "compute_alpha_vec complete" << std::endl; MultiFab Xf_prime_alpha; - compute_Xf_prime_times_vector(Nens, last_pf_name, varnames, xf_bar, alpha_vec, Xf_prime_alpha); + compute_Xf_prime_times_vector(da_iter, Nens, last_pf_name, varnames, xf_bar, alpha_vec, Xf_prime_alpha); Print() << "compute_Xf_prime_times_vector complete" << std::endl; - /*std::string pltname3 = "plt_Xf_prime_alpha"; + /*std::string pltname3 = WriteSingleLevelPlotfile("plt_Xf_prime_alpha", da_iter); Vector varnames3 = {"density", "theta", "x_velocity", "y_velocity", "z_velocity"}; WriteSingleLevelPlotfile(pltname3, Xf_prime_alpha, @@ -314,11 +314,10 @@ ERF::PerformDataAssimilation(int da_iter) Print() << "add_multifabs(xf_bar, Xf_prime_alpha, xf_bar_updated) complete" << std::endl; - /*std::string pltname4 = "plt_xf_bar_updated"; - Vector varnames4 = {"density", "theta", "x_velocity", "y_velocity", "z_velocity"}; + std::string pltname4 = MakeFullPath("plt_xf_bar_updated", da_iter); WriteSingleLevelPlotfile(pltname4, xf_bar_updated, - varnames4, + varnames, geom[0], 0.0, // time 0); // level*/ @@ -332,7 +331,7 @@ ERF::PerformDataAssimilation(int da_iter) for(int n=0; n< Nens; n++) { Print() << "Updating for ensemble " << n << std::endl; MultiFab mf_ens_pert; - update_ensemble(Nens, last_pf_name, varnames, xf_bar, T_mat, n, mf_ens_pert); + update_ensemble(da_iter, Nens, last_pf_name, varnames, xf_bar, T_mat, n, mf_ens_pert); // Update the ensemble MultiFab mf_ens_updated; @@ -348,6 +347,16 @@ ERF::PerformDataAssimilation(int da_iter) m_plot3d_int_1 = -1; m_check_int = -1; + // InitData() handles all the initialization including when doing a + // restart. But when doing data assimilation, in the inputs file, we do not specify a + // erf.restart. So, the ERF class does not know it is a restart. + // So, whenever InitData() is called without assigning a restart_chkfile, + // it just does the from scratch initialization. This function is called by the + // temporararily created erf class - tmp_erf, from main.cpp. We simply want to fill all the data + // structures here by calling InitData(), and then the WriteUpdatedEnsembleToERFClassData will write + // all the updated ensemble data into the lev_new data structure, which will then be + // written into the checkpoint files by MakeEnsembleCheckpointName. + InitData(); auto& lev_new = vars_new[0]; @@ -363,31 +372,3 @@ ERF::PerformDataAssimilation(int da_iter) WriteCheckpointFile(); } } - - -void -ERF::SetDirsForPlotfilesAndCheckpointsForDA (const int ens_no) -{ - // -------------------------------------------------------- - // Set up member-specific output directories - // -------------------------------------------------------- - std::string member_dir; - - std::stringstream ss; - ss << "member_" << std::setw(2) << std::setfill('0') << ens_no; - member_dir = ss.str(); - - if (ParallelDescriptor::IOProcessor()) - { - fs::create_directories(member_dir + "/plotfiles"); - fs::create_directories(member_dir + "/chkfiles"); - fs::create_directories(member_dir + "/pertfiles"); - } - - ParallelDescriptor::Barrier(); - - // Only change the basename/prefix. - // ERF will handle the step numbering as before. - check_file = member_dir + "/chkfiles/chk"; - plot3d_file_1 = member_dir + "/plotfiles/plt"; -} diff --git a/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF_SquallLine2D.cpp b/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF_SquallLine2D.cpp index d621b013b5..4e2bac6699 100644 --- a/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF_SquallLine2D.cpp +++ b/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF_SquallLine2D.cpp @@ -96,6 +96,12 @@ void compute_R_diag_vals(Vector& R_diag) // Set the size to 2 R_diag.resize(8,0.01); - R_diag[0] = 0.01; - R_diag[1] = 0.01; + R_diag[0] = 4e-4; //density + R_diag[1] = 1.0; // theta + R_diag[2] = 1.0; // x_velocity + R_diag[3] = 1.0; // y_velocity + R_diag[4] = 0.25; // z_velocity + R_diag[5] = 2.5e-7; // qv + R_diag[6] = 1e-8; // qc + R_diag[7] = 1e-8; // qrain } diff --git a/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF_Utils.cpp b/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF_Utils.cpp index 6b86485207..944896fa5c 100644 --- a/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF_Utils.cpp +++ b/.Exec_dev/DataAssimilation/SquallLine2D/ERF_DA_EnKFSRF_Utils.cpp @@ -82,12 +82,13 @@ lapack_testing() } } - std::vector -get_plotfile_list() +get_plotfile_list(const int da_iter) { + const std::string da_dir = "DA_Cycle_" + std::to_string(da_iter) + "/"; + std::vector pltfiles; - const std::string member_prefix = "member_"; + const std::string member_prefix = da_dir + "member_"; std::string pf_dir = member_prefix + "00/plotfiles"; @@ -159,11 +160,15 @@ read_plot_file(PlotFileData& pf, } MultiFab -read_member_multifab(int n, +read_member_multifab(const int da_iter, + const int n, const std::string& pf_name, const Vector& varnames) { - const std::string member_prefix = "member_"; + + const std::string da_dir = "DA_Cycle_" + std::to_string(da_iter) + "/"; + + const std::string member_prefix = da_dir + "member_"; std::string member_dir = member_prefix + amrex::Concatenate("", n, 2); std::string pf_path = member_dir + "/plotfiles/" + pf_name; @@ -181,7 +186,8 @@ read_member_multifab(int n, } MultiFab -compute_ensemble_mean(int Nens, +compute_ensemble_mean(const int da_iter, + const int Nens, const std::string& pf_name, const Vector& varnames) { @@ -191,7 +197,7 @@ compute_ensemble_mean(int Nens, for (int n = 0; n < Nens; ++n) { std::cout << "Reading ensembles for compute_ensemble_mean = " << n << std::endl; - MultiFab mf_tmp = read_member_multifab(n, pf_name, varnames); + MultiFab mf_tmp = read_member_multifab(da_iter, n, pf_name, varnames); if (!initialized) { mf_mean.define(mf_tmp.boxArray(), @@ -270,7 +276,8 @@ compute_d_prime_vec(MultiFab& d_prime_vec, } void -compute_mean_H_xf(MultiFab& mean_H_xf, +compute_mean_H_xf(const int da_iter, + MultiFab& mean_H_xf, const int Nens, const std::string& last_pf_name, const Vector& varnames) @@ -278,7 +285,7 @@ compute_mean_H_xf(MultiFab& mean_H_xf, for (int n = 0; n < Nens; ++n) { - MultiFab xf_i = read_member_multifab(n, last_pf_name, varnames); + MultiFab xf_i = read_member_multifab(da_iter, n, last_pf_name, varnames); if (n==0) { mean_H_xf.define(xf_i.boxArray(), xf_i.DistributionMap(), @@ -346,14 +353,15 @@ Compute_yf_prime_i_T_Rinv_yf_prime_j(const MultiFab& Yi, } void -compute_yf_prime (int i, - const std::string& last_pf_name, - const Vector& varnames, - const MultiFab& mean_H_xf, - MultiFab& yf_prime_i) +compute_yf_prime (const int da_iter, + const int i, + const std::string& last_pf_name, + const Vector& varnames, + const MultiFab& mean_H_xf, + MultiFab& yf_prime_i) { // Read ensemble member - MultiFab xf_i = read_member_multifab(i, last_pf_name, varnames); + MultiFab xf_i = read_member_multifab(da_iter, i, last_pf_name, varnames); // Apply observation operator MultiFab H_xf_i; @@ -380,7 +388,8 @@ compute_yf_prime (int i, } void -compute_S_matrix(Matrix& S, +compute_S_matrix(const int da_iter, + Matrix& S, const int& Nens, const MultiFab& mean_H_xf, const Vector& R_diag, @@ -389,11 +398,11 @@ compute_S_matrix(Matrix& S, { for (int i = 0; i < Nens; ++i) { MultiFab yf_prime_i; - compute_yf_prime(i, last_pf_name, varnames, mean_H_xf, yf_prime_i); + compute_yf_prime(da_iter, i, last_pf_name, varnames, mean_H_xf, yf_prime_i); for (int j = 0; j < Nens; ++j) { MultiFab yf_prime_j; - compute_yf_prime(j, last_pf_name, varnames, mean_H_xf, yf_prime_j); + compute_yf_prime(da_iter, j, last_pf_name, varnames, mean_H_xf, yf_prime_j); Real val = Compute_yf_prime_i_T_Rinv_yf_prime_j(yf_prime_i, yf_prime_j, R_diag); S(i,j) = val/(Nens-1); if(i==j) { @@ -459,7 +468,8 @@ compute_yf_prime_T_d_prime_vec (const MultiFab& yf_prime, } void -compute_r_vec (int Nens, +compute_r_vec (const int da_iter, + int Nens, const std::string& last_pf_name, const Vector& varnames, const MultiFab& mean_H_xf, @@ -472,7 +482,7 @@ compute_r_vec (int Nens, { MultiFab yf_prime_i; - compute_yf_prime(i, last_pf_name, varnames, mean_H_xf, yf_prime_i); + compute_yf_prime(da_iter, i, last_pf_name, varnames, mean_H_xf, yf_prime_i); r_vec[i] = compute_yf_prime_T_d_prime_vec(yf_prime_i, d_prime_vec); } @@ -510,18 +520,20 @@ compute_alpha_vec (const int& Nens, // column-vector element multiply void -compute_Xf_prime_times_vector (const int Nens, - const std::string& last_pf_name, - const Vector& varnames, - const MultiFab& xf_bar, - const Vector& vec_in, - MultiFab& result) +compute_Xf_prime_times_vector (const int da_iter, + const int Nens, + const std::string& last_pf_name, + const Vector& varnames, + const MultiFab& xf_bar, + const Vector& vec_in, + MultiFab& result) { AMREX_ALWAYS_ASSERT(vec_in.size() == Nens); // Read first member to define result MultiFab xf_0 = - read_member_multifab(0, + read_member_multifab(da_iter, + 0, last_pf_name, varnames); @@ -534,7 +546,7 @@ compute_Xf_prime_times_vector (const int Nens, for (int n = 0; n < Nens; ++n) { - MultiFab xf_n = read_member_multifab(n, last_pf_name, varnames); + MultiFab xf_n = read_member_multifab(da_iter, n, last_pf_name, varnames); AMREX_ALWAYS_ASSERT(xf_n.boxArray() == result.boxArray()); AMREX_ALWAYS_ASSERT(xf_n.DistributionMap() == result.DistributionMap()); @@ -601,7 +613,8 @@ compute_T_matrix (const Matrix& S_mat, void -update_ensemble (const int Nens, +update_ensemble (const int da_iter, + const int Nens, const std::string& last_pf_name, const Vector& varnames, const MultiFab& xf_bar, @@ -613,9 +626,37 @@ update_ensemble (const int Nens, for (int i = 0; i < Nens; ++i) { T_colvec[i] = T(i,n); } - compute_Xf_prime_times_vector(Nens, last_pf_name, varnames, xf_bar, T_colvec, result); + compute_Xf_prime_times_vector(da_iter, Nens, last_pf_name, varnames, xf_bar, T_colvec, result); } +void +ERF::SetDirsForPlotfilesAndCheckpointsForDA (const int da_iter, + const int ens_no) +{ + // -------------------------------------------------------- + // Set up cycle- and member-specific output directories + // -------------------------------------------------------- + std::stringstream cycle_ss; + cycle_ss << "DA_Cycle_" << da_iter; + std::string cycle_dir = cycle_ss.str(); + + std::stringstream member_ss; + member_ss << "member_" << std::setw(2) << std::setfill('0') << ens_no; + std::string member_dir = cycle_dir + "/" + member_ss.str(); + + if (ParallelDescriptor::IOProcessor()) + { + fs::create_directories(member_dir + "/plotfiles"); + fs::create_directories(member_dir + "/chkfiles"); + fs::create_directories(member_dir + "/pertfiles"); + } + + ParallelDescriptor::Barrier(); + + // Set file prefixes under the specific cycle and member path + check_file = member_dir + "/chkfiles/chk"; + plot3d_file_1 = member_dir + "/plotfiles/plt"; +} // da_iter : DA iteration number // ens_no : Ensemble number @@ -623,19 +664,47 @@ std::string MakeEnsembleCheckpointName (int da_iter, int ens_no) { // Create CheckpointAfterDA if needed - const std::string base_dir = "CheckpointAfterDA"; - fs::create_directories(base_dir); - - // Create CheckpointAfterDA/DACycle_ - const std::string da_dir = - base_dir + "/DACycle_" + std::to_string(da_iter); + const std::string da_dir = "DA_Cycle_" + std::to_string(da_iter) + "/"; fs::create_directories(da_dir); + const std::string chk_after_da_dir = da_dir + "CheckpointAfterDA"; + + // Create DA_Cycle_/CheckpointAfterDA + fs::create_directories(chk_after_da_dir); + std::ostringstream chk_name; - chk_name << da_dir + chk_name << chk_after_da_dir << "/chk_" << std::setw(2) << std::setfill('0') << ens_no << "_"; return chk_name.str(); } + +void +ERF::GetEnsembleCheckpointName (int da_iter, int ens_no) +{ + // Create CheckpointAfterDA if needed + const std::string da_dir = "DA_Cycle_" + std::to_string(da_iter) + "/"; + fs::create_directories(da_dir); + + // Create CheckpointAfterDA/DA_Cycle_ + const std::string chk_after_da_dir = da_dir + "CheckpointAfterDA"; + fs::create_directories(chk_after_da_dir); + + std::ostringstream chk_name; + chk_name << chk_after_da_dir + << "/chk_" + << std::setw(2) << std::setfill('0') << ens_no + << "_" << "00000"; + + restart_chkfile = chk_name.str(); +} + +std::string +MakeFullPath (const std::string& filename, const int da_iter) +{ + std::stringstream ss; + ss << "DA_Cycle_" << da_iter << "/" << filename; + return ss.str(); +} diff --git a/.Exec_dev/DataAssimilation/SquallLine2D/main.cpp b/.Exec_dev/DataAssimilation/SquallLine2D/main.cpp index 818779f33f..1f37f0c2ad 100644 --- a/.Exec_dev/DataAssimilation/SquallLine2D/main.cpp +++ b/.Exec_dev/DataAssimilation/SquallLine2D/main.cpp @@ -96,51 +96,81 @@ int main (int argc, char* argv[]) ParmParse pp_ens("ensemble"); - int n_ens = 1; - pp_ens.query("n_members", n_ens); + int num_da_cycles = -1; + pp_erf.query("num_da_cycles", num_da_cycles); - // Ensemble run loop - for (int ens_no = 0; ens_no < n_ens; ++ens_no) - { - // -------------------------------------------------------- - // Fresh ERF instance per ensemble - // -------------------------------------------------------- - ERF erf; - - if(is_init_for_ensemble) { - erf.SetDirsForPlotfilesAndCheckpointsForDA(ens_no); - } - erf.InitData(); - erf.Evolve(); - - Real end_total = amrex::second() - strt_total; - ParallelDescriptor::ReduceRealMax( - end_total, ParallelDescriptor::IOProcessorNumber()); - - if (erf.Verbose()) { - amrex::Print() << "Ensemble " << ens_no - << " wallclock time: " - << end_total << '\n'; - } - - // -------------------------------------------------------- - // MPI barrier to ensure all ranks finish Evolve - // -------------------------------------------------------- - ParallelDescriptor::Barrier(); - - } // Ensemble run loop complete + if(num_da_cycles != -1 and !is_init_for_ensemble) { + Abort("You are trying to run data assimilation by using the erf.num_da_cycles option. For this, you also " + "have to specify erf.is_init_for_ensemble=true, as ensembles have to be run for data assimilation"); + } - ERF tmp_erf; - // This is only a post-processing step for visualization - //tmp_erf.ComputeAndWriteEnsemblePerturbations(); + int n_ens = 1; + pp_ens.query("n_members", n_ens); - // Perform data assimilation - int da_iter = 0; - tmp_erf.PerformDataAssimilation(da_iter); + // Data assimilation cycle loop + for (int da_iter = 0; da_iter < num_da_cycles; ++da_iter) { + Print() << "Doing data assimilation cycle " << da_iter <=1, + // there has to be a erf.restart_chkfile specified (for each ensemble), + // which is the checkpoint file which contains the latest updated ensemble after + // data assimilation for that ensemble. The InitData then reads that checkpoint file + // and continues the simulation from that. + + // This call will fill the restart_chkfile string + if(da_iter > 0){ + erf.GetEnsembleCheckpointName(da_iter-1, ens_no); + } + erf.InitData(); + erf.Evolve(); + + Real end_total = amrex::second() - strt_total; + ParallelDescriptor::ReduceRealMax( + end_total, ParallelDescriptor::IOProcessorNumber()); + + if (erf.Verbose()) { + amrex::Print() << "Ensemble " << ens_no + << " wallclock time: " + << end_total << '\n'; + } + + // -------------------------------------------------------- + // MPI barrier to ensure all ranks finish Evolve + // -------------------------------------------------------- + ParallelDescriptor::Barrier(); + } // Ensemble run loop complete + + ERF tmp_erf; + // This is only a post-processing step for visualization + //tmp_erf.ComputeAndWriteEnsemblePerturbations(da_iter); + + // After all ensemble runs are complete, perform data assimilation + tmp_erf.PerformDataAssimilation(da_iter); + } - BL_PROFILE_VAR_STOP(pmain); + BL_PROFILE_VAR_STOP(pmain); - amrex::Finalize(); + amrex::Finalize(); #ifdef AMREX_USE_MPI MPI_Finalize(); diff --git a/Source/ERF.H b/Source/ERF.H index 3eb0ded76b..71fb110208 100644 --- a/Source/ERF.H +++ b/Source/ERF.H @@ -732,11 +732,13 @@ public: void apply_gaussian_smoothing_to_perturbations (const int lev, amrex::MultiFab& mf_cc_pert); - void ComputeAndWriteEnsemblePerturbations (); + void ComputeAndWriteEnsemblePerturbations (const int da_iter); void PerformDataAssimilation (int da_iter); - void SetDirsForPlotfilesAndCheckpointsForDA (const int ens_no); + void SetDirsForPlotfilesAndCheckpointsForDA (const int da_iter, const int ens_no); + + void GetEnsembleCheckpointName (int da_iter, int ens_no); void init_custom (int lev); diff --git a/Source/Initialization/ERF_InitCustomPertState.cpp b/Source/Initialization/ERF_InitCustomPertState.cpp index 25bd5ebc73..3601e34a27 100644 --- a/Source/Initialization/ERF_InitCustomPertState.cpp +++ b/Source/Initialization/ERF_InitCustomPertState.cpp @@ -156,9 +156,12 @@ ERF::init_custom (int lev) // add the perturbations to the background state and then populate the "pert variables if(solverChoice.is_init_for_ensemble) { + Vector varnames = {"density","theta", "x_velocity","y_velocity","z_velocity", "qv", "qc", "qrain"}; MultiFab mf_cc_pert; create_random_perturbations(lev, mf_cc_pert); + WriteSingleLevelPlotfile("plt_pert_correlated_gaussian", mf_cc_pert, varnames, geom[lev], zero, 0); apply_gaussian_smoothing_to_perturbations(lev, mf_cc_pert); + //WriteSingleLevelPlotfile("plt_pert_correlated_gaussian", mf_cc_pert, varnames, geom[lev], zero, 0); create_background_state_for_ensemble(lev, mf_cc_pert, lev_new[Vars::cons], lev_new[Vars::xvel], lev_new[Vars::yvel], lev_new[Vars::zvel]); } } diff --git a/Source/Initialization/ERF_InitForEnsemble.cpp b/Source/Initialization/ERF_InitForEnsemble.cpp index bb9d3a72ab..6ea72012c8 100644 --- a/Source/Initialization/ERF_InitForEnsemble.cpp +++ b/Source/Initialization/ERF_InitForEnsemble.cpp @@ -10,6 +10,19 @@ using namespace amrex; namespace fs = std::filesystem; +Vector +perturb_scale = +{ + 0.01, // density [kg/m3] + 1.0, // theta [K] + 1.0, // u [m/s] + 1.0, // v [m/s] + 0.5, // w [m/s] + 5.0e-4, // qv [kg/kg] + 1.0e-4, // qc [kg/kg] + 1.0e-4 // qrain [kg/kg] +}; + /** * Create a cell-centered MultiFab of random perturbations for one AMR level. * @@ -94,85 +107,103 @@ void NormalizeMultiFabRMS_PerComponent(MultiFab& mf_cc_pert) } /** - * Apply horizontal Gaussian smoothing to cell-centered perturbations. + * Apply 3D isotropic/anisotropic Gaussian smoothing to cell-centered perturbations. * * @param lev Integer specifying the current level * @param mf_cc_pert MultiFab containing perturbations to smooth in place */ + void -ERF::apply_gaussian_smoothing_to_perturbations(const int lev, - MultiFab& mf_cc_pert) +ERF::apply_gaussian_smoothing_to_perturbations(const int lev, MultiFab& mf_cc_pert) { const Geometry& gm = geom[lev]; const Real dx = gm.CellSize(0); const Real dy = gm.CellSize(1); + const Real dz = gm.CellSize(2); - const Real dmesh = std::min(dx, dy); - - // ---- User choice ---- const Real sigma = solverChoice.ens_pert_correlated_radius; - const int r = static_cast(3.0 * sigma / dmesh); + + const int rx = static_cast(std::ceil(3.0 * sigma / dx)); + const int ry = static_cast(std::ceil(3.0 * sigma / dy)); + const int rz = static_cast(std::ceil(3.0 * sigma / dz)); const int ncomp = mf_cc_pert.nComp(); - // ---- Precompute Gaussian weights ---- - const int wsize = 2*r + 1; - Vector w_host(wsize * wsize); - - Real Z = zero; - for (int m = -r; m <= r; ++m) { - for (int n = -r; n <= r; ++n) { - Real val = std::exp(-(m*m*dx*dx + n*n*dy*dy) - /(two*sigma*sigma)); - w_host[(m+r)*wsize + (n+r)] = val; - Z += val; + // 1. Build 1D Gaussian weight vectors on host + auto compute_1d_weights = [](int r, Real d, Real sig) { + Vector w(2*r + 1); + Real sum = 0.0; + for (int m = -r; m <= r; ++m) { + Real val = std::exp(-(m * m * d * d) / (2.0 * sig * sig)); + w[m + r] = val; + sum += val; } + for (auto& v : w) v /= sum; + return w; + }; + + Vector wx_h = compute_1d_weights(rx, dx, sigma); + Vector wy_h = compute_1d_weights(ry, dy, sigma); + Vector wz_h = compute_1d_weights(rz, dz, sigma); + + Gpu::DeviceVector wx_d(wx_h.size()), wy_d(wy_h.size()), wz_d(wz_h.size()); + Gpu::copy(Gpu::hostToDevice, wx_h.begin(), wx_h.end(), wx_d.begin()); + Gpu::copy(Gpu::hostToDevice, wy_h.begin(), wy_h.end(), wy_d.begin()); + Gpu::copy(Gpu::hostToDevice, wz_h.begin(), wz_h.end(), wz_d.begin()); + + Real const* wx = wx_d.data(); + Real const* wy = wy_d.data(); + Real const* wz = wz_d.data(); + + // Temp MultiFab for separable intermediate states + IntVect ngrow(AMREX_D_DECL(rx, ry, rz)); + MultiFab mf_tmp(mf_cc_pert.boxArray(), mf_cc_pert.DistributionMap(), ncomp, ngrow); + + // --- Pass 1: X Direction --- + mf_tmp.ParallelCopy(mf_cc_pert, 0, 0, ncomp, IntVect(0), ngrow, gm.periodicity()); + for (MFIter mfi(mf_cc_pert, TilingIfNotGPU()); mfi.isValid(); ++mfi) { + const Box& bx = mfi.tilebox(); + auto const& in = mf_tmp.const_array(mfi); + auto const& out = mf_cc_pert.array(mfi); + ParallelFor(bx, ncomp, [=] AMREX_GPU_DEVICE (int i, int j, int k, int n) noexcept { + Real sum = 0.0; + for (int m = -rx; m <= rx; ++m) { + sum += wx[m + rx] * in(i + m, j, k, n); + } + out(i, j, k, n) = sum; + }); } - for (auto& v : w_host) { - v /= Z; + // --- Pass 2: Y Direction --- + mf_tmp.ParallelCopy(mf_cc_pert, 0, 0, ncomp, IntVect(0), ngrow, gm.periodicity()); + for (MFIter mfi(mf_cc_pert, TilingIfNotGPU()); mfi.isValid(); ++mfi) { + const Box& bx = mfi.tilebox(); + auto const& in = mf_tmp.const_array(mfi); + auto const& out = mf_cc_pert.array(mfi); + ParallelFor(bx, ncomp, [=] AMREX_GPU_DEVICE (int i, int j, int k, int n) noexcept { + Real sum = 0.0; + for (int nn = -ry; nn <= ry; ++nn) { + sum += wy[nn + ry] * in(i, j + nn, k, n); + } + out(i, j, k, n) = sum; + }); } - Gpu::DeviceVector w_dev(w_host.size()); - Gpu::copy(Gpu::hostToDevice, w_host.begin(), w_host.end(), w_dev.begin()); - - Real const* w = w_dev.data(); - - // ---- Create a grown copy (for stencil access) ---- - IntVect ngrow_big(AMREX_D_DECL(r, r, 0)); - - MultiFab mf_copy(mf_cc_pert.boxArray(), - mf_cc_pert.DistributionMap(), - ncomp, ngrow_big); - - mf_copy.ParallelCopy(mf_cc_pert, - 0, 0, ncomp, - IntVect(0), ngrow_big, - gm.periodicity()); - - // ---- Apply smoothing ---- - for (MFIter mfi(mf_cc_pert, TilingIfNotGPU()); mfi.isValid(); ++mfi) - { + // --- Pass 3: Z Direction --- + mf_tmp.ParallelCopy(mf_cc_pert, 0, 0, ncomp, IntVect(0), ngrow, gm.periodicity()); + for (MFIter mfi(mf_cc_pert, TilingIfNotGPU()); mfi.isValid(); ++mfi) { const Box& bx = mfi.tilebox(); - - auto const& in = mf_copy.const_array(mfi); + auto const& in = mf_tmp.const_array(mfi); auto const& out = mf_cc_pert.array(mfi); - - ParallelFor(bx, ncomp, - [=] AMREX_GPU_DEVICE (int i, int j, int k, int n) noexcept - { - Real sum = zero; - - for (int m = -r; m <= r; ++m) { - for (int nn = -r; nn <= r; ++nn) { - Real wij = w[(m+r)*wsize + (nn+r)]; - sum += wij * in(i+m, j+nn, k, n); - } + ParallelFor(bx, ncomp, [=] AMREX_GPU_DEVICE (int i, int j, int k, int n) noexcept { + Real sum = 0.0; + for (int p = -rz; p <= rz; ++p) { + sum += wz[p + rz] * in(i, j, k + p, n); } - - out(i,j,k,n) = sum; + out(i, j, k, n) = sum; }); } + NormalizeMultiFabRMS_PerComponent(mf_cc_pert); } @@ -602,9 +633,9 @@ MakeFinalMultiFabs (const MultiFab& mf_cc_fine, Real tmp_qrain = mf_cc_fine_arr(i,j,k,7); cons_pert_arr(i,j,k,Rho_comp) = tmp_rho; cons_pert_arr(i,j,k,RhoTheta_comp) = tmp_rho*tmp_theta; - if (n_qstate_moist > 0) cons_pert_arr(i,j,k,RhoQ1_comp) = tmp_rho*tmp_qv; - if (n_qstate_moist > 1) cons_pert_arr(i,j,k,RhoQ2_comp) = tmp_rho*tmp_qc; - if (n_qstate_moist > 2) cons_pert_arr(i,j,k,RhoQ3_comp) = tmp_rho*tmp_qrain; + if (n_qstate_moist > 0) cons_pert_arr(i,j,k,RhoQ1_comp) = std::max(tmp_rho*tmp_qv,0.0_rt); + if (n_qstate_moist > 1) cons_pert_arr(i,j,k,RhoQ2_comp) = std::max(tmp_rho*tmp_qc,0.0_rt); + if (n_qstate_moist > 2) cons_pert_arr(i,j,k,RhoQ3_comp) = std::max(tmp_rho*tmp_qrain,0.0_rt); }); } @@ -665,6 +696,12 @@ AddPertToBckgnd(MultiFab& mf_cc_fine, // Optional safety check (recommended) AMREX_ALWAYS_ASSERT(mf_cc_pert.nComp() == ncomp); + // Create a GPU-accessible array and copy the host vector into it + GpuArray scale_gpu; // Use the maximum possible ncomp, e.g., 8 + for (int n = 0; n < ncomp; ++n) { + scale_gpu[n] = perturb_scale[n]; + } + for (MFIter mfi(mf_cc_fine, TilingIfNotGPU()); mfi.isValid(); ++mfi) { const Box& bx = mfi.tilebox(); @@ -675,7 +712,7 @@ AddPertToBckgnd(MultiFab& mf_cc_fine, amrex::ParallelFor(bx, ncomp, [=] AMREX_GPU_DEVICE (int i, int j, int k, int n) noexcept { - Real ens_amp = ens_pert_amplitude*std::abs(bg(i,j,k,n)); + Real ens_amp = ens_pert_amplitude*scale_gpu[n]*std::abs(bg(i,j,k,n)); bg(i,j,k,n) += ens_amp*pert(i,j,k,n); }); }