EAMxx: histogram diagnostic

components/eamxx/docs/user/diags/field_contraction.md

@@ -3,7 +3,7 @@
 In EAMxx, we can automatically calculate field reductions
 across the horizontal columns and across the model vertical levels.
 We call these horizontal and vertical reductions.
-We can also automatically calculate zonal averages.
+We can also automatically calculate zonal averages and histograms.
 We have *experimental* support for composing diagnostics; below,
 the horizontal and vertical reductions can be composed
 sequentially, but if using dp- or dz-weighted reductions,
@@ -85,6 +85,19 @@ And so on...
 | --------- | ------ | ----------- |
 | `X_zonal_avg_Y_bins` | Area fraction | Average across the zonal direction |
 
+## Histograms
+
+We currently have a utility to calculate histograms online.
+To select the histogram, you need to suffix a field name `X` with
+`_histogram_` followed by the values specifying the edges of the bins,
+separated by `_`. For example, the histogram specified by
+`T_mid_histogram_250_270_290_310` has 5 bins effectively defined by
+[$-\infty$, 250), [250, 270), [270, 290), [290, 310), and [310, $\infty$).
+
+| Reduction | Weight | Description |
+| --------- | ------ | ----------- |
+| `X_histogram_V0_V1_..._VN` | 1 or 0 | Count of field values within each range |
+
 ## Example
 
 ```yaml
@@ -99,7 +112,7 @@ fields:
       # in this example, we use T_mid in units of K
       - T_mid_horiz_avg  # K
       - T_mid_vert_avg_dp_weighted  # K
-      - T_mid_vert_sum_dp_weighted  # K * Pa * s / (m * m) 
+      - T_mid_vert_sum_dp_weighted  # K * Pa * s / (m * m)
       - T_mid_vert_avg_dz_weighted  # K
       - T_mid_vert_sum_dz_weighted  # K * m
       - T_mid_vert_avg  # K
@@ -110,6 +123,7 @@ fields:
       - T_mid_vert_avg_dp_weighted_horiz_avg  # K
       - T_mid_zonal_avg_180_bins  # K
       - T_mid_zonal_avg_90_bins  # K
+      - T_mid_histogram_250_270_290_310 # count
 output_control:
   frequency: 1
   frequency_units: nmonths

components/eamxx/src/diagnostics/CMakeLists.txt

@@ -9,6 +9,7 @@ set(DIAGNOSTIC_SRCS
   field_at_height.cpp
   field_at_level.cpp
   field_at_pressure_level.cpp
+  histogram.cpp
   horiz_avg.cpp
   longwave_cloud_forcing.cpp
   number_path.cpp

components/eamxx/src/diagnostics/histogram.cpp

@@ -0,0 +1,164 @@
+#include "diagnostics/histogram.hpp"
+
+#include <ekat_team_policy_utils.hpp>
+#include <ekat_string_utils.hpp>
+
+
+namespace scream {
+
+HistogramDiag::HistogramDiag(const ekat::Comm &comm, const ekat::ParameterList &params)
+    : AtmosphereDiagnostic(comm, params) {
+  const auto &field_name       = m_params.get<std::string>("field_name");
+  const std::string bin_config = m_params.get<std::string>("bin_configuration");
+  m_diag_name                  = field_name + "_histogram_" + bin_config;
+
+  // extract bin values from configuration, append end values, and check
+  const std::vector<std::string> bin_strings = ekat::split(bin_config, "_");
+  m_bin_reals.resize(bin_strings.size()+2);
+  m_bin_reals[0] = std::numeric_limits<Real>::lowest();
+  for (long unsigned int i=1; i < m_bin_reals.size()-1; i++)
+  {
+    m_bin_reals[i] = std::stod(bin_strings[i-1]);
+    EKAT_REQUIRE_MSG(m_bin_reals[i] > m_bin_reals[i-1],
+                     "Error! HistogramDiag bin values must be monotonically "
+                     "increasing.\n"
+                     " - bin configuration: " + bin_config + "\n");
+  }
+  m_bin_reals.back() = std::numeric_limits<Real>::max();
+}
+
+void HistogramDiag::set_grids(const std::shared_ptr<const GridsManager> grids_manager) {
+  const auto &field_name = m_params.get<std::string>("field_name");
+  const auto &grid_name  = m_params.get<std::string>("grid_name");
+  add_field<Required>(field_name, grid_name);
+}
+
+void HistogramDiag::initialize_impl(const RunType /*run_type*/) {
+  using FieldIdentifier = FieldHeader::identifier_type;
+  using FieldLayout     = FieldIdentifier::layout_type;
+  using ShortFieldTagsNames::CMP;
+  const Field &field              = get_fields_in().front();
+  const FieldIdentifier &field_id = field.get_header().get_identifier();
+  const FieldLayout &field_layout = field_id.get_layout();
+  EKAT_REQUIRE_MSG(field_layout.rank() >= 1 && field_layout.rank() <= 3,
+                   "Error! Field rank not supported by HistogramDiag.\n"
+                   " - field name: " +
+                       field_id.name() +
+                       "\n"
+                       " - field layout: " +
+                       field_layout.to_string() + "\n");
+
+  // allocate histogram field
+  const int num_bins = m_bin_reals.size()-1;
+  FieldLayout diagnostic_layout({CMP}, {num_bins}, {"bin"});
+  FieldIdentifier diagnostic_id(m_diag_name, diagnostic_layout,
+                                FieldIdentifier::Units::nondimensional(),
+                                field_id.get_grid_name());
+  m_diagnostic_output = Field(diagnostic_id);
+  m_diagnostic_output.allocate_view();
+
+  // allocate field for bin values
+  FieldLayout bin_values_layout({CMP}, {num_bins+1}, {"bin"});
+  FieldIdentifier bin_values_id(m_diag_name + "_bin_values", bin_values_layout,
+                                field_id.get_units(), field_id.get_grid_name());
+  m_bin_values = Field(bin_values_id);
+  m_bin_values.allocate_view();
+
+  // copy bin values into field
+  auto bin_values_view = m_bin_values.get_view<Real *>();
+  auto bin_values_view_host = Kokkos::create_mirror_view(bin_values_view);
+  for (long unsigned int i=0; i < bin_values_layout.dim(0); i++)
+    bin_values_view_host(i) = m_bin_reals[i];
+  Kokkos::deep_copy(bin_values_view,bin_values_view_host);
+}
+
+void HistogramDiag::compute_diagnostic_impl() {
+  const auto &field = get_fields_in().front();
+  auto field_layout = field.get_header().get_identifier().get_layout();
+  auto histogram_layout = m_diagnostic_output.get_header().get_identifier().get_layout();
+  const int num_bins = histogram_layout.dim(0);
+
+  auto bin_values_view = m_bin_values.get_view<Real *>();
+  auto histogram_view = m_diagnostic_output.get_view<Real *>();
+  using KT         = ekat::KokkosTypes<DefaultDevice>;
+  using TeamPolicy = Kokkos::TeamPolicy<Field::device_t::execution_space>;
+  using TeamMember = typename TeamPolicy::member_type;
+  using TPF        = ekat::TeamPolicyFactory<typename KT::ExeSpace>;
+  switch (field_layout.rank())
+  {
+    case 1: {
+      const int d1 = field_layout.dim(0);
+      auto field_view = field.get_view<const Real *>();
+      TeamPolicy team_policy = TPF::get_default_team_policy(num_bins, d1);
+      Kokkos::parallel_for("compute_histogram_" + field.name(), team_policy,
+          KOKKOS_LAMBDA(const TeamMember &tm) {
+            const int bin_i = tm.league_rank();
+            const Real bin_lower = bin_values_view(bin_i);
+            const Real bin_upper = bin_values_view(bin_i+1);
+            Kokkos::parallel_reduce(Kokkos::TeamVectorRange(tm, d1),
+                [&](int i, Real &val) {
+                  if ((bin_lower <= field_view(i)) && (field_view(i) < bin_upper))
+                    val += sp(1.0);
+                },
+                histogram_view(bin_i));
+          });
+    } break;
+    case 2: {
+      const int d1 = field_layout.dim(0);
+      const int d2 = field_layout.dim(1);
+      auto field_view = field.get_view<const Real **>();
+      TeamPolicy team_policy = TPF::get_default_team_policy(num_bins, d1*d2);
+      Kokkos::parallel_for("compute_histogram_" + field.name(), team_policy,
+          KOKKOS_LAMBDA(const TeamMember &tm) {
+            const int bin_i = tm.league_rank();
+            const Real bin_lower = bin_values_view(bin_i);
+            const Real bin_upper = bin_values_view(bin_i+1);
+            Kokkos::parallel_reduce(Kokkos::TeamVectorRange(tm, d1*d2),
+                [&](int ind, Real &val) {
+                  const int i1 = ind / d2;
+                  const int i2 = ind % d2;
+                  if ((bin_lower <= field_view(i1,i2)) && (field_view(i1,i2) < bin_upper))
+                    val += sp(1.0);
+                },
+                histogram_view(bin_i));
+          });
+    } break;
+    case 3: {
+      const int d1 = field_layout.dim(0);
+      const int d2 = field_layout.dim(1);
+      const int d3 = field_layout.dim(2);
+      auto field_view = field.get_view<const Real ***>();
+      TeamPolicy team_policy = TPF::get_default_team_policy(num_bins, d1*d2*d3);
+      Kokkos::parallel_for("compute_histogram_" + field.name(), team_policy,
+          KOKKOS_LAMBDA(const TeamMember &tm) {
+            const int bin_i = tm.league_rank();
+            const Real bin_lower = bin_values_view(bin_i);
+            const Real bin_upper = bin_values_view(bin_i+1);
+            Kokkos::parallel_reduce(Kokkos::TeamVectorRange(tm, d1*d2*d3),
+                [&](int ind, Real &val) {
+                  const int i1 = ind / (d2*d3);
+                  const int ind2 = ind % (d2*d3);
+                  const int i2 = ind2 / d3;
+                  const int i3 = ind2 % d3;
+                  if ((bin_lower <= field_view(i1,i2,i3)) && (field_view(i1,i2,i3) < bin_upper))
+                    val += sp(1.0);
+                },
+                histogram_view(bin_i));
+          });
+  } break;
+
+  default:
+    EKAT_ERROR_MSG("Error! Unsupported field rank for histogram.\n");
+  }
+
+  // TODO: use device-side MPI calls
+  // TODO: the dev ptr causes problems; revisit this later
+  // TODO: doing cuda-aware MPI allreduce would be ~10% faster
+  Kokkos::fence();
+  m_diagnostic_output.sync_to_host();
+  m_comm.all_reduce(m_diagnostic_output.template get_internal_view_data<Real, Host>(),
+                      histogram_layout.size(), MPI_SUM);
+  m_diagnostic_output.sync_to_dev();
+}
+
+} // namespace scream

components/eamxx/src/diagnostics/histogram.hpp

@@ -0,0 +1,40 @@
+#ifndef EAMXX_HISTOTGRAM_HPP
+#define EAMXX_HISTOGRAM_HPP
+
+#include "share/atm_process/atmosphere_diagnostic.hpp"
+
+namespace scream {
+/*
+ * This diagnostic will calculate a histogram of a field across all dimensions
+ * producing a one dimensional field, with CMP tag dimension named "bin", that
+ * indicates how many times a field value in the specified range occurred.
+ */
+
+class HistogramDiag : public AtmosphereDiagnostic {
+
+public:
+  // Constructors
+  HistogramDiag(const ekat::Comm &comm, const ekat::ParameterList &params);
+
+  // The name of the diagnostic
+  std::string name() const { return m_diag_name; }
+
+  // Set the grid
+  void set_grids(const std::shared_ptr<const GridsManager> grids_manager);
+
+protected:
+#ifdef KOKKOS_ENABLE_CUDA
+public:
+#endif
+  void initialize_impl(const RunType /*run_type*/);
+  void compute_diagnostic_impl();
+
+protected:
+  std::string m_diag_name;
+  std::vector<Real> m_bin_reals;
+  Field m_bin_values;
+};
+
+} // namespace scream
+
+#endif // EAMXX_HISTOGRAM_HPP

components/eamxx/src/diagnostics/register_diagnostics.hpp

@@ -30,6 +30,7 @@
 #include "diagnostics/zonal_avg.hpp"
 #include "diagnostics/conditional_sampling.hpp"
 #include "diagnostics/binary_ops.hpp"
+#include "diagnostics/histogram.hpp"
 
 namespace scream {
 
@@ -63,6 +64,7 @@ inline void register_diagnostics () {
   diag_factory.register_product("ZonalAvgDiag",&create_atmosphere_diagnostic<ZonalAvgDiag>);
   diag_factory.register_product("ConditionalSampling",&create_atmosphere_diagnostic<ConditionalSampling>);
   diag_factory.register_product("BinaryOpsDiag", &create_atmosphere_diagnostic<BinaryOpsDiag>);
+  diag_factory.register_product("HistogramDiag",&create_atmosphere_diagnostic<HistogramDiag>);
 }
 
 } // namespace scream

components/eamxx/src/diagnostics/tests/CMakeLists.txt

@@ -82,10 +82,13 @@ CreateDiagTest(vert_contract "vert_contract_test.cpp" MPI_RANKS 1 ${SCREAM_TEST_
 CreateDiagTest(vert_derivative "vert_derivative_test.cpp")
 
 # Test zonal averaging
-CreateDiagTest(zonal_avg zonal_avg_test.cpp MPI_RANKS 1 ${SCREAM_TEST_MAX_RANKS})
+CreateDiagTest(zonal_avg "zonal_avg_test.cpp" MPI_RANKS 1 ${SCREAM_TEST_MAX_RANKS})
 
 # Test conditional sampling
 CreateDiagTest(conditional_sampling "conditional_sampling_test.cpp")
 
 # Test binary ops
 CreateDiagTest(binary_ops "binary_ops_test.cpp")
+
+# Test histogram
+CreateDiagTest(histogram "histogram_test.cpp" MPI_RANKS 1 ${SCREAM_TEST_MAX_RANKS})