TMP: fix aggr, add test, will be revised

Author: mahf708

components/eamxx/cime_config/testdefs/testmods_dirs/eamxx/rad_cosp_async/postrun_check.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+
+
+import os
+import sys
+
+# add eamxx scripts dir to get stuff from there
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../../../../scripts'))
+# get the pylib stuff and ensure we have xarray
+from utils import _ensure_pylib_impl
+_ensure_pylib_impl('xarray')
+_ensure_pylib_impl('numpy')
+
+
+# get the input arg which is the case dir
+case_dir = sys.argv[1]
+# parse the case name, which is the string after the last "/"
+case_name = case_dir.split('/')[-1]
+
+import xarray as xr
+
+# before the restart
+ds_avg = xr.open_dataset(case_dir + '/run/' + case_name + '.scream.average_5_steps.h.AVERAGE.nsteps_x5.0001-01-01-00000.nc')
+ds_ins = xr.open_dataset(case_dir + '/run/' + case_name + '.scream.instant_every_step.h.INSTANT.nsteps_x1.0001-01-01-03600.nc')
+
+import numpy as np
+assert np.allclose(ds_avg.isccp_cldtot.values, ds_ins.isccp_cldtot.mean('time', skipna=True).values, equal_nan=True)
+
+print("PASS")
+sys.exit(0)

components/eamxx/cime_config/testdefs/testmods_dirs/eamxx/rad_cosp_async/shell_commands

@@ -0,0 +1,68 @@
+#!/bin/bash
+
+cime_root=$(./xmlquery --value CIMEROOT)
+input_data_dir=$(./xmlquery --value DIN_LOC_ROOT)
+atmchange=$cime_root/../components/eamxx/scripts/atmchange
+case_name=$(./xmlquery --value CASE)
+
+# Turn on cosp and set frequency
+$atmchange -b physics::atm_procs_list="mac_aero_mic,rrtmgp,cosp"
+$atmchange -b physics::cosp::cosp_frequency_units="steps"
+$atmchange -b physics::cosp::cosp_frequency=4
+
+# set rrtmgp frequency to be half of that of cosp
+$atmchange -b physics::rrtmgp::rad_frequency=1
+
+# Need to explicitly turn on computing tendencies
+$atmchange -b physics::mac_aero_mic::shoc::compute_tendencies=T_mid,qv
+$atmchange -b physics::mac_aero_mic::p3::compute_tendencies=T_mid,qv
+$atmchange -b physics::rrtmgp::compute_tendencies=T_mid
+$atmchange -b homme::compute_tendencies=T_mid,qv
+
+# determine grid and set remap files
+atm_grid=$(./xmlquery --value ATM_GRID)
+if [[ "${atm_grid}" = "ne30np4.pg2" ]]; then
+    hmapfile="${input_data_dir}/atm/scream/maps/map_ne30pg2_to_ne4pg2_20231201.nc"
+    armmapfile="${input_data_dir}/atm/scream/maps/map_ne30pg2_to_DecadalSites_c20240130.nc"
+    # Run with bugfixed SPA file
+    $atmchange -b spa_data_file="${input_data_dir}/atm/scream/init/spa_v3.LR.F2010.2011-2025.c_20240405.nc"
+elif [[ "${atm_grid}" = "ne4np4.pg2" ]]; then
+    hmapfile="${input_data_dir}/atm/scream/maps/map_ne4pg2_to_ne2pg2_c20240902.nc"
+    echo "Note: arm remap only works for ne30pg2 atm grids for now"
+    armmapfile="not-supported-yet"
+    # Keep default SPA file
+    # ... (do nothing)
+else
+    echo "Note: horiz/arm remaps only work for ne30pg2 and ne4pg2 atm grids for now"
+    hmapfile="not-supported-yet"
+    armmapfile="not-supported-yet"
+fi
+
+# set the output yaml files
+output_yaml_files=$(find ${cime_root}/../components/eamxx/cime_config/testdefs/testmods_dirs/eamxx/rad_cosp_async/yaml_outs/ -maxdepth 1 -type f)
+for file in ${output_yaml_files[@]}; do
+    # if the word "coarse" is in the file name, do nothing
+    if [[ "${file}" == *"_coarse.yaml" && "${hmapfile}" == "not-supported-yet" ]]; then
+        continue
+    elif [[ "${file}" == *"_arm.yaml" && "${armmapfile}" == "not-supported-yet" ]]; then
+        continue
+    else
+        # TODO: add remap file replacement for different grids
+        cp -v ${file} ./
+        if [ "${file}" == "${output_yaml_files[0]}" ]; then
+            # First file, reset output list
+            $atmchange -b output_yaml_files="./$(basename ${file})"
+        else
+            # Append to output list
+            $atmchange -b output_yaml_files+="./$(basename ${file})"
+        fi
+        # Replace remap files
+        sed -i "s|horiz_remap_file:.*_to_ne30.*|horiz_remap_file: ${hmapfile}|" ./$(basename ${file})
+        sed -i "s|horiz_remap_file:.*_to_DecadalSites.*|horiz_remap_file: ${armmapfile}|" ./$(basename ${file})
+    fi
+    # replace all filename prefixes so that st_archive works...
+    sed -i "s|eamxx_output.decadal|${case_name}.scream|" ./$(basename ${file})
+done
+
+chmod +x postrun_check.py
+./xmlchange POSTRUN_SCRIPT=${cime_root}/../components/eamxx/cime_config/testdefs/testmods_dirs/eamxx/rad_cosp_async/postrun_check.py

components/eamxx/cime_config/testdefs/testmods_dirs/eamxx/rad_cosp_async/yaml_outs/average_every_5_steps.yaml

@@ -0,0 +1,16 @@
+%YAML 1.1
+---
+filename_prefix: eamxx_output.decadal.average_5_steps.h
+iotype: pnetcdf
+averaging_type: average
+max_snapshots_per_file: 1
+fields:
+  physics_pg2:
+    field_names:
+    - isccp_ctptau
+    - modis_ctptau
+    - misr_cthtau
+    - isccp_cldtot
+output_control:
+  frequency: 5
+  frequency_units: nsteps

components/eamxx/cime_config/testdefs/testmods_dirs/eamxx/rad_cosp_async/yaml_outs/instant_every_step.yaml

@@ -0,0 +1,16 @@
+%YAML 1.1
+---
+filename_prefix: eamxx_output.decadal.instant_every_step.h
+iotype: pnetcdf
+averaging_type: instant
+max_snapshots_per_file: 5
+fields:
+  physics_pg2:
+    field_names:
+    - isccp_ctptau
+    - modis_ctptau
+    - misr_cthtau
+    - isccp_cldtot
+output_control:
+  frequency: 1
+  frequency_units: nsteps

components/eamxx/src/physics/cosp/eamxx_cosp.cpp

@@ -113,6 +113,7 @@ void Cosp::initialize_impl (const RunType /* run_type */)
   for (const auto& field_name : vnames) {
     // the mask here is just the sunlit mask, so set it
     get_field_out(field_name).get_header().set_extra_data("mask_field", get_field_in("sunlit_mask"));
+    get_field_out(field_name).get_header().set_may_be_filled(true);
   }
 }
 

components/eamxx/src/share/io/scorpio_output.cpp

@@ -146,7 +146,7 @@ AtmosphereOutput (const ekat::Comm& comm, const ekat::ParameterList& params,
   EKAT_REQUIRE_MSG (not has_duplicates(m_fields_names),
       "[AtmosphereOutput] Error! One of the output yaml files has duplicate field entries.\n"
       " - yaml file: " + params.name() + "\n"
-      " - fields names; " + ekat::join(m_fields_names,",") + "\n");
+      " - fields names: " + ekat::join(m_fields_names,",") + "\n");
 
   // Check if remapping and if so create the appropriate remapper
   // Note: We currently support three remappers
@@ -174,6 +174,7 @@ AtmosphereOutput (const ekat::Comm& comm, const ekat::ParameterList& params,
   }
 
   // ... then add diagnostic fields
+  // ... and we also use this to init track avg cnt for regular fields
   init_diagnostics ();
 
   // Avg count only makes sense if we have
@@ -462,8 +463,8 @@ run (const std::string& filename,
       auto field = fm_after_hr->get_field(fname);
       auto mask  = count.get_header().get_extra_data<Field>("mask");
 
-      // Find where the field is NOT equal to fill_value
-      compute_mask<Comparison::NE>(field,constants::fill_value<Real>,mask);
+        // Find where the field is NOT equal to fill_value
+        compute_mask<Comparison::NE>(field,constants::fill_value<Real>,mask);
 
       // mask=1 for "good" entries, and mask=0 otherwise.
       count.update(mask,1,1);
@@ -508,6 +509,16 @@ run (const std::string& filename,
     const auto& f_in  = fm_after_hr->get_field(field_name);
           auto& f_out = fm_scorpio->get_field(field_name);
 
+    // Safety check: if a field may contain fill values and we are computing an Average,
+    // we must have created an avg-count tracking field; otherwise division by the raw
+    // number of steps would bias the result wherever fill values occurred.
+    if (m_avg_type==OutputAvgType::Average && f_in.get_header().may_be_filled()) {
+      EKAT_REQUIRE_MSG(m_field_to_avg_count.count(field_name),
+        "[AtmosphereOutput::run] Error! Averaging a fill-aware field without avg-count tracking.\n"
+        " - field name : " + field_name + "\n"
+        "This indicates the field was marked may_be_filled after output initialization or tracking logic missed it." );
+    }
+
     switch (m_avg_type) {
       case OutputAvgType::Instant:
         f_out.deep_copy(f_in);  break; // Note: if f_in aliases f_out, this is a no-op
@@ -1003,6 +1014,26 @@ init_diagnostics ()
   for (const auto& fname : m_fields_names) {
     if (not m_field_mgrs[FromModel]->has_field(fname)) {
       create_diag(fname);
+    } else {
+      // This is a regular field, not a diagnostic.
+      // Still, we might need to do some extra setup, like for avg_count.
+      const auto& f = m_field_mgrs[FromModel]->get_field(fname);
+      // We need avg-count tracking for any averaged (non-instant) field that:
+      //  - supplies explicit mask info (mask_data or mask_field), OR
+      //  - is marked as potentially containing fill values (may_be_filled()).
+      // Without this, fill-aware updates skip fill_value during accumulation (good)
+      // but we would still divide by the raw nsteps, biasing the result low.
+      if (m_avg_type!=OutputAvgType::Instant) {
+        const bool has_mask = f.get_header().has_extra_data("mask_data") || f.get_header().has_extra_data("mask_field");
+        const bool may_be_filled = f.get_header().may_be_filled();
+        if (has_mask || may_be_filled) {
+          m_track_avg_cnt = true;
+          // Avoid duplicate insertion if already present (e.g., mask + filled both true)
+          if (m_field_to_avg_cnt_suffix.count(f.name())==0) {
+            m_field_to_avg_cnt_suffix.emplace(f.name(), "_" + f.name());
+          }
+        }
+      }
     }
   }
 }