EAMxx: MAM4xx+ZM [DRAFT]

components/eamxx/cime_config/namelist_defaults_eamxx.xml

@@ -463,7 +463,11 @@ be lost if SCREAM_HACK_XML is not enabled.
     </mam4_optics>
 
     <!-- MAM4xx-Wetscav -->
-    <mam4_wetscav inherit="mam4_atm_proc_base" />
+    <mam4_wetscav inherit="mam4_atm_proc_base">
+      <!-- Convective transport -->
+      <convproc_do_aer type="logical" doc="Flag to enable convective transport of aerosols">false</convproc_do_aer>
+      <convproc_do_gas type="logical" doc="Flag to enable convective transport of gases">false</convproc_do_gas>
+    </mam4_wetscav>
 
     <!-- MAM4xx-Surface-Emissions -->
     <mam4_srf_online_emiss inherit="mam4_atm_proc_base">

components/eamxx/cime_config/testdefs/testmods_dirs/eamxx/mam4xx/zm/shell_commands

@@ -0,0 +1,14 @@
+alias ATMCHANGE='$CIMEROOT/../components/eamxx/scripts/atmchange'
+# Add ZM deep convection and wetscav processes
+#------------------------------------------------------
+# Add zm before mac_aero_mic, and wetscav after it
+ATMCHANGE physics::atm_procs_list="zm,mam4_constituent_fluxes,mac_aero_mic,mam4_wetscav,mam4_optics,rrtmgp,mam4_srf_online_emiss,mam4_aero_microphys,mam4_drydep" -b
+
+#------------------------------------------------------
+# Enable convective processing in wetscav
+#------------------------------------------------------
+# This enables wetscav to use ZM outputs for aerosol processing
+# in deep convective clouds
+ATMCHANGE mam4_wetscav::convproc_do_aer=true -b
+ATMCHANGE mam4_wetscav::convproc_do_gas=false -b
+

components/eamxx/src/physics/mam/eamxx_mam_wetscav_process_interface.hpp

@@ -26,7 +26,11 @@ class MAMWetscav : public MAMGenericInterface
   using KT           = ekat::KokkosTypes<DefaultDevice>;
   using view_2d      = typename KT::template view_2d<Real>;
   using view_2d_host = typename KT::template view_2d<Real>::host_mirror_type;
+  using int_view_1d  = typename KT::template view_1d<int>;
   using int_view_2d  = typename KT::template view_2d<int>;
+  using const_view_1d = typename KT::template view_1d<const Real>;
+  using const_view_1d_int = typename KT::template view_1d<const int>;
+  using const_view_2d = typename KT::template view_2d<const Real>;
 
   // a thread team dispatched to a single vertical column
   using ThreadTeam = mam4::ThreadTeam;
@@ -74,10 +78,15 @@ class MAMWetscav : public MAMGenericInterface
 
   // Work arrays
   view_2d work_;
+  view_2d work_convproc_;  // Separate work array for convection processing scratch1Dviews
   int_view_2d isprx_;
 
   // TODO: Following variables are from convective parameterization (not
-  // implemented yet in EAMxx), so should be zero for now
+  // implemented yet in EAMxx), so should be zero for now.
+  // NOTE: rprddp, icwmrdp, and evapcdp are no longer member variables;
+  // when do_convproc_ is true they alias FM views directly to avoid
+  // redundant copies.  When false, the zero-initialized member views below
+  // serve as placeholders.
 
   view_2d sh_frac_;
 
@@ -87,23 +96,25 @@ class MAMWetscav : public MAMGenericInterface
   // Evaporation rate of shallow convective precipitation >=0. [kg/kg/s]
   view_2d evapcsh_;
 
-  view_2d evapcdp_;
-
   // Rain production, shallow convection [kg/kg/s]
   view_2d rprdsh_;
 
-  // Rain production, deep convection [kg/kg/s]
-  view_2d rprddp_;
-
-  // In cloud water mixing ratio, deep convection
-  view_2d icwmrdp_;
-
   // In cloud water mixing ratio, shallow convection
   view_2d icwmrsh_;
 
+  // Deep convection zero-placeholder views (only allocated when !do_convproc_).
+  // When do_convproc_ is true, run_impl aliases FM views directly instead.
+  view_2d rprddp_zero_;
+  view_2d icwmrdp_zero_;
+  view_2d evapcdp_zero_;
+
   // Detraining cld H20 from deep convection [kg/kg/s]
   view_2d dlf_;
 
+  // Temporary view for pressure thickness in mb (converted from Pa)
+  view_2d dp_tmp_;
+  view_2d dpdry_tmp_;
+
   int num_2d_scratch_ = 39;
   //
   view_2d scavimptblnum_;
@@ -139,6 +150,22 @@ class MAMWetscav : public MAMGenericInterface
   // activation_fraction_in_cloud_conv_ is convection activation fraction
   Real activation_fraction_in_cloud_conv_ = 0.40;
 
+  // convproc_do_aer: Enable aerosol processing in convection
+  bool convproc_do_aer_ = false;
+
+  // convproc_do_gas: Enable gas processing in convection
+  bool convproc_do_gas_ = false;
+
+  // do_convproc: Combined flag for any convection processing
+  bool do_convproc_ = false;
+
+  // ConvProc configuration for species classification and resuspension mapping
+  mam4::ConvProc::Config convproc_config_;
+
+  // Device views for species classification and resuspension mapping (for GPU kernels)
+  int_view_1d species_class_device_;
+  int_view_1d mmtoo_prevap_resusp_device_;
+
 }; // class MAMWetscav
 
 } // namespace scream

components/eamxx/src/physics/zm/eamxx_zm_process_interface.cpp

@@ -89,6 +89,8 @@ void ZMDeepConvection::create_requests ()
   add_field<Computed>("zm_cape",              scalar2d,     J/kg,   grid_name);
   add_field<Computed>("zm_dcape",             scalar2d,   J/kg/s,   grid_name);
   add_field<Computed>("zm_activity",          scalar2d,     none,   grid_name);
+  add_field<Computed>("zm_jt",                scalar2d,     none,   grid_name);
+  add_field<Computed>("zm_msemax_klev",       scalar2d,     none,   grid_name);
 
   add_field<Computed>("zm_detr_qc",           scalar3d_mid, kg/kg/s,grid_name, pack_size);
   add_field<Computed>("zm_detr_qi",           scalar3d_mid, kg/kg/s,grid_name, pack_size);
@@ -102,6 +104,10 @@ void ZMDeepConvection::create_requests ()
   add_field<Computed>("zm_mflx_dn",           scalar3d_mid, kg/m2/s,grid_name, pack_size);
   add_field<Computed>("zm_entr_dn",           scalar3d_mid, 1/s,    grid_name, pack_size);
 
+  // Fields needed by MAM wetscav for aerosol convective processing
+  add_field<Computed>("zm_rain_prod",          scalar3d_mid, kg/kg/s,grid_name, pack_size);  // rain production rate (rprddp)
+  add_field<Computed>("zm_icw_liq",            scalar3d_mid, kg/kg,  grid_name, pack_size);  // in-cloud liquid water (icwmrdp)
+
   add_field<Computed>("mcsp_ds_out",          scalar3d_mid, K/s,    grid_name, pack_size);
   add_field<Computed>("mcsp_dq_out",          scalar3d_mid, kg/kg/s,grid_name, pack_size);
   add_field<Computed>("mcsp_du_out",          scalar3d_mid, m/s/s,  grid_name, pack_size);
@@ -416,6 +422,20 @@ void ZMDeepConvection::run_impl (const double dt)
   // initialize intermediate output tendencies for zm_conv_evap()
   zm_output.init_tmp(m_ncol, nlev_mid);
 
+  // Zero evap_ds/dq_out for all columns before zm_conv_evap.
+  // zm_conv_evap only writes these for active columns (early-return guard),
+  // so inactive columns would otherwise retain stale values from a previous
+  // timestep when convection was active.  init_tmp cannot be used because it
+  // is also called before zm_transport_momentum, which would wipe the values
+  // that zm_conv_evap just wrote.
+  Kokkos::parallel_for("zm_zero_evap_out",
+    KT::RangePolicy(0, m_ncol*nlev_mid), KOKKOS_LAMBDA (const int idx) {
+    const int i = idx/nlev_mid;
+    const int k = idx%nlev_mid;
+    loc_zm_output_evap_ds_out(i,k) = 0;
+    loc_zm_output_evap_dq_out(i,k) = 0;
+  });
+
   // perform the convective evaporation calculations
   Kokkos::parallel_for(team_policy, KOKKOS_LAMBDA(const KT::MemberType& team) {
     const Int i = team.league_rank();
@@ -582,13 +602,17 @@ void ZMDeepConvection::run_impl (const double dt)
   const auto& zm_cape       = get_field_out("zm_cape")    .get_view<Real*>();
   const auto& zm_dcape      = get_field_out("zm_dcape")   .get_view<Real*>();
   const auto& zm_activity   = get_field_out("zm_activity").get_view<Real*>();
+  const auto& zm_jt         = get_field_out("zm_jt")      .get_view<Real*>();
+  const auto& zm_msemax_klev      = get_field_out("zm_msemax_klev")   .get_view<Real*>();
   Kokkos::parallel_for("zm_diag_outputs_2D",m_ncol, KOKKOS_LAMBDA (const int i) {
     zm_prec    (i) = loc_zm_output_prec    (i);
     zm_snow    (i) = loc_zm_output_snow    (i);
     zm_cape    (i) = loc_zm_output_cape    (i);
     zm_dcape   (i) = loc_zm_output_dcape   (i);
     zm_activity(i) = loc_zm_output_activity(i);
-  });
+    zm_jt   (i) = loc_zm_output_jt   (i);
+    zm_msemax_klev   (i) = loc_zm_output_msemax_klev(i);
+ });
 
   // 3D mid-level output
   const auto& mflx_up = get_field_out("zm_mflx_up").get_view<Real**>();
@@ -607,6 +631,19 @@ void ZMDeepConvection::run_impl (const double dt)
     entr_dn(i,k) = loc_zm_output_entr_dn(i,k);
   });
 
+  // 3D mid-level output for MAM wetscav coupling
+  const auto& zm_rain_prod_out = get_field_out("zm_rain_prod").get_view<Real**>();
+  const auto& zm_icw_liq        = get_field_out("zm_icw_liq").get_view<Real**>();
+  Kokkos::parallel_for("zm_aerosol_coupling_outputs",
+    KT::RangePolicy(0, m_ncol*nlev_mid), KOKKOS_LAMBDA (const int idx) {
+    const int i = idx/nlev_mid;
+    const int k = idx%nlev_mid;
+    // rain production rate (rprddp in EAM Fortran)
+    zm_rain_prod_out(i,k) = loc_zm_output_rain_prod(i,k);
+    // in-cloud liquid water mixing ratio (icwmrdp in EAM Fortran)
+    zm_icw_liq(i,k)        = loc_zm_output_ql(i,k);
+  });
+
   // 3D interface output
   const auto loc_zm_output_mass_flux = zm_output.mass_flux;
   const auto& mass_flux_int = get_field_out("zm_mass_flux_int").get_view<Real**>();

components/eamxx/src/physics/zm/impl/zm_conv_main_impl.hpp

@@ -161,8 +161,9 @@ void Functions<S,D>::zm_conv_main(
       prec(i)                = 0;
       rliq(i)                = 0;
       dsubcld(i)             = 0;
-      jctop(i)               = pver - 1;
+      jctop(i)               = pver-1;
       jcbot(i)               = 0;
+      jt(i)                  = -1;  // Initialize for inactive columns (no convection)
       cape(i)                = 0;
       dcape(i)               = 0;
       msemax_klev(i)         = 0;

components/eamxx/tests/multi-process/dynamics_physics/mam/CMakeLists.txt

@@ -1,3 +1,4 @@
 add_subdirectory(homme_shoc_cld_p3_mam_optics_rrtmgp)
 add_subdirectory(homme_shoc_cld_mam_aci_p3_mam_optics_rrtmgp_mam_drydep)
 add_subdirectory(homme_shoc_cld_spa_p3_rrtmgp_mam4_wetscav)
+add_subdirectory(homme_shoc_cld_spa_p3_rrtmgp_zm_mam4_wetscav)

components/eamxx/tests/multi-process/dynamics_physics/mam/homme_shoc_cld_spa_p3_rrtmgp_zm_mam4_wetscav/CMakeLists.txt

@@ -0,0 +1,95 @@
+include (ScreamUtils)
+
+set (TEST_BASE_NAME homme_shoc_cld_spa_p3_rrtmgp_zm_mam4_wetscav)
+set (FIXTURES_BASE_NAME ${TEST_BASE_NAME}_generate_output_nc_files)
+
+# Get or create the dynamics lib
+#                 HOMME_TARGET   NP PLEV QSIZE_D
+CreateDynamicsLib("theta-l_kokkos"  4   72   41)
+
+# Create the test
+CreateADUnitTest(${TEST_BASE_NAME}
+  LIBS cld_fraction shoc spa p3 scream_rrtmgp zm mam ${dynLibName}
+  LABELS dynamics shoc cld p3 rrtmgp zm physics mam4_wetscav
+  MPI_RANKS ${TEST_RANK_START} ${TEST_RANK_END}
+  FIXTURES_SETUP_INDIVIDUAL ${FIXTURES_BASE_NAME}
+)
+
+# Set AD configurable options
+set (ATM_TIME_STEP 1800)
+SetVarDependingOnTestSize(NUM_STEPS 2 4 48)  # 1h 2h 24h
+set (RUN_T0 2021-10-12-45000)
+
+# Determine num subcycles needed to keep shoc dt<=300s
+set (SHOC_MAX_DT 300)
+math (EXPR MAC_MIC_SUBCYCLES "(${ATM_TIME_STEP} + ${SHOC_MAX_DT} - 1) / ${SHOC_MAX_DT}")
+
+## Copy (and configure) yaml files needed by tests
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/input.yaml
+               ${CMAKE_CURRENT_BINARY_DIR}/input.yaml)
+configure_file(${CMAKE_CURRENT_SOURCE_DIR}/output.yaml
+               ${CMAKE_CURRENT_BINARY_DIR}/output.yaml)
+
+# Set homme's test options, so that we can configure the namelist correctly
+# Discretization/algorithm settings
+set (HOMME_TEST_NE 2)
+set (HOMME_TEST_LIM 9)
+set (HOMME_TEST_REMAP_FACTOR 3)
+set (HOMME_TEST_TRACERS_FACTOR 1)
+set (HOMME_TEST_TIME_STEP 300)
+set (HOMME_THETA_FORM 1)
+set (HOMME_TTYPE 5)
+set (HOMME_SE_FTYPE 0)
+set (HOMME_TEST_TRANSPORT_ALG 0)
+set (HOMME_TEST_CUBED_SPHERE_MAP 0)
+
+# Hyperviscosity settings
+set (HOMME_TEST_HVSCALING 0)
+set (HOMME_TEST_HVS 1)
+set (HOMME_TEST_HVS_TOM 0)
+set (HOMME_TEST_HVS_Q 1)
+
+set (HOMME_TEST_NU 7e15)
+set (HOMME_TEST_NUDIV 1e15)
+set (HOMME_TEST_NUTOP 2.5e5)
+
+# Testcase settings
+set (HOMME_TEST_MOISTURE notdry)
+set (HOMME_THETA_HY_MODE true)
+
+# Vert coord settings
+set (HOMME_TEST_VCOORD_INT_FILE acme-72i.ascii)
+set (HOMME_TEST_VCOORD_MID_FILE acme-72m.ascii)
+
+# Configure the namelist into the test directory
+configure_file(${SCREAM_SRC_DIR}/dynamics/homme/tests/theta.nl
+               ${CMAKE_CURRENT_BINARY_DIR}/namelist.nl)
+
+# Ensure test input files are present in the data dir
+set (TEST_INPUT_FILES
+  scream/init/spa_init_ne2np4.nc
+  scream/init/spa_from_v3LRamip_CMIP6_2010_clim_ne2_c20260930.nc
+  scream/init/${EAMxx_tests_IC_FILE_MAM4xx_72lev}
+  cam/topo/${EAMxx_tests_TOPO_FILE}
+)
+foreach (file IN ITEMS ${TEST_INPUT_FILES})
+  GetInputFile(${file})
+endforeach()
+
+# Compare output files produced by npX tests, to ensure they are bfb
+include (CompareNCFiles)
+
+CompareNCFilesFamilyMpi (
+  TEST_BASE_NAME ${TEST_BASE_NAME}
+  FILE_META_NAME ${TEST_BASE_NAME}_output.INSTANT.nsteps_x${NUM_STEPS}.npMPIRANKS.${RUN_T0}.nc
+  MPI_RANKS ${TEST_RANK_START} ${TEST_RANK_END}
+  LABELS dynamics physics shoc cld p3 rrtmgp zm spa mam4_wetscav
+  META_FIXTURES_REQUIRED ${FIXTURES_BASE_NAME}_npMPIRANKS_omp1
+)
+
+if (SCREAM_ENABLE_BASELINE_TESTS)
+  # Compare one of the output files with the baselines.
+  # Note: one is enough, since we already check that np1 is BFB with npX
+  set (OUT_FILE ${TEST_BASE_NAME}_output.INSTANT.nsteps_x${NUM_STEPS}.np${TEST_RANK_END}.${RUN_T0}.nc)
+  CreateBaselineTest(${TEST_BASE_NAME} ${TEST_RANK_END} ${OUT_FILE} ${FIXTURES_BASE_NAME})
+endif()