[Bug]: Re-testing land variables for unit conversion issue on main

[chengzhuzhang]

What happened?

In zppy tests, it is found that some land variables units conversion is broken, which may be related due to the recent refactor example:

Traceback (most recent call last):
  File "/gpfs/fs1/home/ac.forsyth2/miniforge3/envs/e3sm_diags_20250115/lib/python3.12/site-packages/e3sm_diags/parameter/core_parameter.py", line 340, in _run_diag
    single_result = module.run_diag(self)
                    ^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/fs1/home/ac.forsyth2/miniforge3/envs/e3sm_diags_20250115/lib/python3.12/site-packages/e3sm_diags/driver/lat_lon_land_driver.py", line 13, in run_diag
    return base_run_diag(parameter)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/fs1/home/ac.forsyth2/miniforge3/envs/e3sm_diags_20250115/lib/python3.12/site-packages/e3sm_diags/driver/lat_lon_driver.py", line 75, in run_diag
    ds_test = test_ds.get_climo_dataset(var_key, season)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/fs1/home/ac.forsyth2/miniforge3/envs/e3sm_diags_20250115/lib/python3.12/site-packages/e3sm_diags/driver/utils/dataset_xr.py", line 399, in get_climo_dataset
    ds = self._get_climo_dataset(season)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/fs1/home/ac.forsyth2/miniforge3/envs/e3sm_diags_20250115/lib/python3.12/site-packages/e3sm_diags/driver/utils/dataset_xr.py", line 428, in _get_climo_dataset
    ds = self._get_dataset_with_derived_climo_var(ds)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/fs1/home/ac.forsyth2/miniforge3/envs/e3sm_diags_20250115/lib/python3.12/site-packages/e3sm_diags/driver/utils/dataset_xr.py", line 718, in _get_dataset_with_derived_climo_var
    ds_derived = self._get_dataset_with_derivation_func(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/fs1/home/ac.forsyth2/miniforge3/envs/e3sm_diags_20250115/lib/python3.12/site-packages/e3sm_diags/driver/utils/dataset_xr.py", line 932, in _get_dataset_with_derivation_func
    derived_var = func(*func_args)  # pragma: nocover
                  ^^^^^^^^^^^^^^^^
  File "/gpfs/fs1/home/ac.forsyth2/miniforge3/envs/e3sm_diags_20250115/lib/python3.12/site-packages/e3sm_diags/derivations/derivations.py", line 1059, in <lambda>
    lambda v: convert_units(v, target_units="mg*/m^2/day"),
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/fs1/home/ac.forsyth2/miniforge3/envs/e3sm_diags_20250115/lib/python3.12/site-packages/e3sm_diags/derivations/utils.py", line 81, in convert_units
    var.attrs["units"] = "m" + var.attrs["units"][0:7] + "day"
                               ~~~~~~~~~^^^^^^^^^
KeyError: 'units'

What did you expect to happen? Are there are possible answers you came across?

We should test over all the land variables to fix unit conversion.

Minimal Complete Verifiable Example (MVCE)

import os
import numpy
from e3sm_diags.parameter.core_parameter import CoreParameter
from e3sm_diags.parameter.lat_lon_land_parameter import LatLonLandParameter


from e3sm_diags.run import runner

short_name = 'v3.LR.historical_0051'

param = CoreParameter()

# Model
param.test_name = 'v3.LR.historical_0051'


# Ref

# Output dir
param.results_dir = '/lcrc/group/e3sm/public_html/ac.zhang40/zppy_weekly_comprehensive_v3_output/test_zppy_pr651_20250115/v3.LR.historical_0051/post/'

# Additional settings
param.run_type = 'model_vs_model'
param.diff_title = 'Difference'
param.multiprocessing = True
param.num_workers = 8
param.seasons = ['ANN']
#param.fail_on_incomplete = True
params = [param]

# Model land
land_param = LatLonLandParameter()
land_param.test_data_path = '/lcrc/group/e3sm/ac.zhang40/zppy_weekly_comprehensive_v3_output/test_zppy_pr651_20250115/v3.LR.historical_0051/post/lnd/180x360_aave/clim/2yr'

# Reference
land_param.reference_data_path = '/lcrc/group/e3sm/ac.zhang40/zppy_weekly_comprehensive_v3_output/test_zppy_pr651_20250115/v3.LR.historical_0051/post/lnd/180x360_aave/clim/2yr'
land_param.ref_name = 'v3.LR.historical_0051'
land_param.short_ref_name = 'same simulation'
# Optionally, swap test and reference model
if False:
   land_param.test_data_path, param.reference_data_path = param.reference_data_path, param.test_data_path
   land_param.test_name, param.ref_name = param.ref_name, param.test_name
   land_param.short_test_name, param.short_ref_name = param.short_ref_name, param.short_test_name
params.append(land_param)
# Run
runner.sets_to_run = ['lat_lon_land']
runner.run_diags(params)

Relevant log output

Anything else we need to know?

No response

Environment

main branch after #907 merged.

[tomvothecoder]

Based on the KeyError: 'units', it looks like the attributes on the variable are being dropped/not preserved at some point before convert_units(). I will analyze #907 to see where we need to set with xr.set_options(keep_attrs=True).

Can you update the example script to fill in the Jinja templating (e.g., short_name = '${short}')? That way I can step through the code with the VS Code debugger.

[chengzhuzhang]

@tomvothecoder I just update the reproducer to fill in the parameters needed. Thank you for taking a look!