filter aodvis by day indices

components/eamxx/src/diagnostics/aodvis.cpp

@@ -2,6 +2,8 @@
 
 #include <ekat/kokkos/ekat_kokkos_utils.hpp>
 
+#include "share/util/scream_universal_constants.hpp"
+
 namespace scream {
 
 AODVis::AODVis(const ekat::Comm &comm, const ekat::ParameterList &params)
@@ -23,11 +25,13 @@ void AODVis::set_grids(
   m_nlevs = grid->get_num_vertical_levels();
 
   // Define layouts we need (both inputs and outputs)
-  FieldLayout scalar3d_swband_layout = grid->get_3d_vector_layout(true,m_swbands,"swband");
-  FieldLayout scalar1d_layout        = grid->get_2d_scalar_layout();
+  FieldLayout scalar3d_swband_layout =
+      grid->get_3d_vector_layout(true, m_swbands, "swband");
+  FieldLayout scalar1d_layout = grid->get_2d_scalar_layout();
 
   // The fields required for this diagnostic to be computed
   add_field<Required>("aero_tau_sw", scalar3d_swband_layout, nondim, grid_name);
+  add_field<Required>("sunlit", scalar1d_layout, nondim, grid_name);
 
   // Construct and allocate the aodvis field
   FieldIdentifier fid("AerosolOpticalDepth550nm", scalar1d_layout, nondim,
@@ -41,18 +45,25 @@ void AODVis::compute_diagnostic_impl() {
   using MT  = typename KT::MemberType;
   using ESU = ekat::ExeSpaceUtils<typename KT::ExeSpace>;
 
+  Real var_fill_value = constants::DefaultFillValue<Real>().value;
+
   const auto aod     = m_diagnostic_output.get_view<Real *>();
   const auto tau_vis = get_field_in("aero_tau_sw")
                            .subfield(1, m_vis_bnd)
                            .get_view<const Real **>();
+  const auto sunlit = get_field_in("sunlit").get_view<const Real *>();
 
   const auto num_levs = m_nlevs;
   const auto policy   = ESU::get_default_team_policy(m_ncols, m_nlevs);
   Kokkos::parallel_for(
       "Compute " + name(), policy, KOKKOS_LAMBDA(const MT &team) {
         const int icol = team.league_rank();
-        auto tau_icol  = ekat::subview(tau_vis, icol);
-        aod(icol)      = ESU::view_reduction(team, 0, num_levs, tau_icol);
+        if(sunlit(icol) == 0.0) {
+          aod(icol) = var_fill_value;
+        } else {
+          auto tau_icol = ekat::subview(tau_vis, icol);
+          aod(icol)     = ESU::view_reduction(team, 0, num_levs, tau_icol);
+        }
       });
 }
 

components/eamxx/src/diagnostics/tests/aodvis_test.cpp

@@ -1,10 +1,11 @@
+#include <share/field/field_utils.hpp>
+
 #include "catch2/catch.hpp"
 #include "diagnostics/register_diagnostics.hpp"
 #include "share/field/field_utils.hpp"
 #include "share/grid/mesh_free_grids_manager.hpp"
 #include "share/util/scream_setup_random_test.hpp"
-#include "share/util/scream_utils.hpp"
-
+#include "share/util/scream_universal_constants.hpp"
 namespace scream {
 
 std::shared_ptr<GridsManager> create_gm(const ekat::Comm &comm, const int ncols,
@@ -30,6 +31,10 @@ TEST_CASE("aodvis") {
   using namespace ShortFieldTagsNames;
   using namespace ekat::units;
 
+  Real var_fill_value = constants::DefaultFillValue<Real>().value;
+
+  Real some_limit = 0.0025;
+
   // A world comm
   ekat::Comm comm(MPI_COMM_WORLD);
 
@@ -40,7 +45,7 @@ TEST_CASE("aodvis") {
 
   // Create a grids manager - single column for these tests
   constexpr int nlevs = 33;
-  const int ngcols    = 2 * comm.size();
+  const int ngcols    = 10 * comm.size();
 
   int nbnds = eamxx_swbands();
   int swvis = eamxx_vis_swband_idx();
@@ -49,16 +54,23 @@ TEST_CASE("aodvis") {
   auto grid = gm->get_grid("Physics");
 
   // Input (randomized) tau
-  FieldLayout scalar3d_swband_layout = grid->get_3d_vector_layout(true,nbnds,"swband");
+  FieldLayout scalar3d_swband_layout =
+      grid->get_3d_vector_layout(true, nbnds, "swband");
   FieldIdentifier tau_fid("aero_tau_sw", scalar3d_swband_layout, nondim,
                           grid->name());
   Field tau(tau_fid);
   tau.allocate_view();
   tau.get_header().get_tracking().update_time_stamp(t0);
+  // Input (randomized) sunlit
+  FieldLayout scalar2d_layout = grid->get_2d_scalar_layout();
+  FieldIdentifier sunlit_fid("sunlit", scalar2d_layout, nondim, grid->name());
+  Field sunlit(sunlit_fid);
+  sunlit.allocate_view();
+  sunlit.get_header().get_tracking().update_time_stamp(t0);
 
   // Construct random number generator stuff
   using RPDF = std::uniform_real_distribution<Real>;
-  RPDF pdf(0, 0.05);
+  RPDF pdf(0, 0.005);
   auto engine = scream::setup_random_test();
 
   // Construct the Diagnostics
@@ -71,13 +83,27 @@ TEST_CASE("aodvis") {
     // Randomize tau
     randomize(tau, engine, pdf);
 
+    // Randomize sunlit
+    randomize(sunlit, engine, pdf);
+
     // Create and set up the diagnostic
     ekat::ParameterList params;
     auto diag = diag_factory.create("AerosolOpticalDepth550nm", comm, params);
     diag->set_grids(gm);
     diag->set_required_field(tau);
+    diag->set_required_field(sunlit);
     diag->initialize(t0, RunType::Initial);
 
+    auto sun_h = sunlit.get_view<Real *, Host>();
+    for(int icol = 0; icol < grid->get_num_local_dofs(); ++icol) {
+      // zero out all sun_h if below 0.05
+      if(sun_h(icol) < some_limit) {
+        sun_h(icol) = 0.0;
+      }
+    }
+    // sync to device
+    sunlit.sync_to_dev();
+
     // Run diag
     diag->compute_diagnostic();
 
@@ -93,12 +119,17 @@ TEST_CASE("aodvis") {
     auto aod_t = aod_tf.get_view<Real *, Host>();
 
     for(int icol = 0; icol < grid->get_num_local_dofs(); ++icol) {
-      for(int ilev = 0; ilev < nlevs; ++ilev) {
-        aod_t(icol) += tau_h(icol, swvis, ilev);
+      if(sun_h(icol) < some_limit) {
+        aod_t(icol) = var_fill_value;
+      } else {
+        for(int ilev = 0; ilev < nlevs; ++ilev) {
+          aod_t(icol) += tau_h(icol, swvis, ilev);
+        }
       }
     }
     aod_hf.sync_to_dev();
     aod_tf.sync_to_dev();
+
     // Workaround for non-bfb behavior of view_reduction() in release builds
     if(SCREAM_BFB_TESTING) {
       REQUIRE(views_are_equal(aod_hf, aod_tf));