diff --git a/e3sm_to_cmip/cmor_handlers/_formulas.py b/e3sm_to_cmip/cmor_handlers/_formulas.py index 5230a04f..ae825589 100644 --- a/e3sm_to_cmip/cmor_handlers/_formulas.py +++ b/e3sm_to_cmip/cmor_handlers/_formulas.py @@ -466,20 +466,51 @@ def pr(ds: xr.Dataset) -> xr.DataArray: High frequency version: pr = PRECT * 1000.0 + + EAMxx version: + pr = (precip_liq_surf_mass_flux + precip_ice_surf_mass_flux) * 1000.0 + or: + pr = precip_total_surf_mass_flux * 1000.0 """ if all(key in ds.data_vars for key in ["PRECC", "PRECL"]): result = (ds["PRECC"] + ds["PRECL"]) * 1000.0 elif "PRECT" in ds: result = ds["PRECT"] * 1000.0 + elif all( + key in ds.data_vars + for key in ["precip_liq_surf_mass_flux", "precip_ice_surf_mass_flux"] + ): + result = ( + ds["precip_liq_surf_mass_flux"] + ds["precip_ice_surf_mass_flux"] + ) * 1000.0 + elif "precip_total_surf_mass_flux" in ds: + result = ds["precip_total_surf_mass_flux"] * 1000.0 else: raise KeyError( "No formula could be applied for 'pr'. Check the handler entry for 'pr' " - "and input file(s) contain either 'PRECC' and 'PRECL', or 'PRECT." + "and input file(s) contain either 'PRECC' and 'PRECL', 'PRECT', " + "'precip_liq_surf_mass_flux' and 'precip_ice_surf_mass_flux', " + "or 'precip_total_surf_mass_flux'." ) return result +def clwvi(ds: xr.Dataset) -> xr.DataArray: + """ + CMIP6 clwvi = condensed water path (liquid + ice). + + EAM version: + clwvi = TGCLDCWP (total condensed = liquid + ice, single variable) + + EAMxx version: + clwvi = LiqWaterPath + IceWaterPath + In EAMxx, LiqWaterPath is liquid-only and IceWaterPath is ice-only; + they must be summed to match the CMIP6 definition. + """ + return ds["LiqWaterPath"] + ds["IceWaterPath"] + + def prsn(ds: xr.Dataset) -> xr.DataArray: """ prsn = (PRECSC + PRECSL) * 1000.0 diff --git a/e3sm_to_cmip/cmor_handlers/handlers.yaml b/e3sm_to_cmip/cmor_handlers/handlers.yaml index beeb67c2..1c9bcd08 100644 --- a/e3sm_to_cmip/cmor_handlers/handlers.yaml +++ b/e3sm_to_cmip/cmor_handlers/handlers.yaml @@ -1040,3 +1040,288 @@ formula: null positive: null levels: { name: plev19, units: Pa, e3sm_axis_name: plev } +# ============================================================ +# EAMxx variable handlers (issue #339) +# Maps EAMxx output variable names to CMIP6 variables. +# These entries coexist with EAM entries above; the handler +# system selects whichever raw_variables are present in input. +# ============================================================ +# --- Precipitation and column water --- +- name: pr + units: kg m-2 s-1 + raw_variables: [precip_liq_surf_mass_flux, precip_ice_surf_mass_flux] + table: CMIP6_Amon.json + unit_conversion: null + formula: (precip_liq_surf_mass_flux + precip_ice_surf_mass_flux) * 1000.0 + positive: null + levels: null +- name: pr + units: kg m-2 s-1 + raw_variables: [precip_total_surf_mass_flux] + table: CMIP6_Amon.json + unit_conversion: null + formula: precip_total_surf_mass_flux * 1000.0 + positive: null + levels: null +- name: prw + units: kg m-2 + raw_variables: [VapWaterPath] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: null +# --- Surface temperature and humidity --- +- name: ts + units: K + raw_variables: [surf_radiative_T] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: null +- name: tas + units: K + raw_variables: [T_2m] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: null +- name: huss + units: "1" + raw_variables: [qv_2m] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: null +# --- Surface energy and momentum fluxes --- +- name: hfls + units: W m-2 + raw_variables: [surface_upward_latent_heat_flux] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: up + levels: null +- name: hfss + units: W m-2 + raw_variables: [surf_sens_flux] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: up + levels: null +- name: evspsbl + units: kg m-2 s-1 + raw_variables: [surf_evap] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: null +- name: tauu + units: Pa + raw_variables: [surf_mom_flux_U] + table: CMIP6_Amon.json + unit_conversion: "-1" + formula: null + positive: down + levels: null +- name: tauv + units: Pa + raw_variables: [surf_mom_flux_V] + table: CMIP6_Amon.json + unit_conversion: "-1" + formula: null + positive: down + levels: null +# --- 10m wind speed --- +- name: sfcWind + units: m s-1 + raw_variables: [wind_speed_10m] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: null +# --- Sea-level pressure --- +- name: psl + units: Pa + raw_variables: [SeaLevelPressure] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: null +# --- TOA radiation --- +- name: rsdt + units: W m-2 + raw_variables: [SW_flux_dn_at_model_top] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: down + levels: null +- name: rsut + units: W m-2 + raw_variables: [SW_flux_up_at_model_top] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: up + levels: null +- name: rsutcs + units: W m-2 + raw_variables: [SW_clrsky_flux_up_at_model_top] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: up + levels: null +- name: rlut + units: W m-2 + raw_variables: [LW_flux_up_at_model_top] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: up + levels: null +- name: rlutcs + units: W m-2 + raw_variables: [LW_clrsky_flux_up_at_model_top] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: up + levels: null +# --- Surface radiation --- +- name: rsds + units: W m-2 + raw_variables: [SW_flux_dn_at_model_bot] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: down + levels: null +- name: rsus + units: W m-2 + raw_variables: [SW_flux_up_at_model_bot] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: up + levels: null +- name: rsdscs + units: W m-2 + raw_variables: [SW_clrsky_flux_dn_at_model_bot] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: down + levels: null +- name: rsuscs + units: W m-2 + raw_variables: [SW_clrsky_flux_up_at_model_bot] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: up + levels: null +- name: rlds + units: W m-2 + raw_variables: [LW_flux_dn_at_model_bot] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: down + levels: null +- name: rlus + units: W m-2 + raw_variables: [LW_flux_up_at_model_bot] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: up + levels: null +- name: rldscs + units: W m-2 + raw_variables: [LW_clrsky_flux_dn_at_model_bot] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: down + levels: null +# --- Cloud water paths --- +# CMIP6 clwvi = "condensed water path" = liquid + ice (per CMIP6_Amon.json comment). +# In EAMxx, LiqWaterPath is liquid-only (equiv. EAM TGCLDLWP), and +# IceWaterPath is ice-only (equiv. EAM TGCLDIWP). Their sum gives the +# total condensed water path required by CMIP6 clwvi. +# By contrast, EAM uses TGCLDCWP (total = liquid + ice) directly for clwvi. +- name: clwvi + units: kg m-2 + raw_variables: [LiqWaterPath, IceWaterPath] + table: CMIP6_Amon.json + unit_conversion: null + formula: LiqWaterPath + IceWaterPath + positive: null + levels: null +- name: clivi + units: kg m-2 + raw_variables: [IceWaterPath] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: null +# --- Aerosol optical depth --- +- name: od550aer + units: "1" + raw_variables: [AerosolOpticalDepth550nm] + table: CMIP6_AERmon.json + unit_conversion: null + formula: null + positive: null + levels: null +# --- 3D pressure-level variables (requires data interpolated to plev) --- +- name: ta + units: K + raw_variables: [T_mid] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: { name: plev19, units: Pa, e3sm_axis_name: plev } +- name: hus + units: "1" + raw_variables: [qv] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: { name: plev19, units: Pa, e3sm_axis_name: plev } +- name: hur + units: "%" + raw_variables: [RelativeHumidity] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: { name: plev19, units: Pa, e3sm_axis_name: plev } +- name: wap + units: Pa s-1 + raw_variables: [omega] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: { name: plev19, units: Pa, e3sm_axis_name: plev } +- name: zg + units: m + raw_variables: [z_mid] + table: CMIP6_Amon.json + unit_conversion: null + formula: null + positive: null + levels: { name: plev19, units: Pa, e3sm_axis_name: plev } diff --git a/e3sm_to_cmip/cmor_handlers/utils.py b/e3sm_to_cmip/cmor_handlers/utils.py index 4f1f73f6..6c7ce747 100644 --- a/e3sm_to_cmip/cmor_handlers/utils.py +++ b/e3sm_to_cmip/cmor_handlers/utils.py @@ -4,7 +4,6 @@ from importlib.machinery import SourceFileLoader from typing import Any, Literal, get_args -import pandas as pd import yaml from e3sm_to_cmip import ( @@ -256,23 +255,21 @@ def _get_handlers_from_yaml() -> dict[str, list[dict[str, Any]]]: with open(HANDLER_DEFINITIONS_PATH, "r") as infile: handlers_file = yaml.load(infile, yaml.SafeLoader) - df_in = pd.DataFrame.from_dict(handlers_file) - handlers = defaultdict(list) - for row in df_in.itertuples(): + for entry in handlers_file: var_handler = VarHandler( - name=row.name, # type: ignore - units=row.units, # type: ignore - raw_variables=row.raw_variables, # type: ignore - table=row.table, # type: ignore - formula=row.formula, # type: ignore - unit_conversion=row.unit_conversion, # type: ignore - positive=row.positive, # type: ignore - levels=row.levels, # type: ignore + name=entry["name"], + units=entry["units"], + raw_variables=entry["raw_variables"], + table=entry["table"], + formula=entry.get("formula"), + unit_conversion=entry.get("unit_conversion"), + positive=entry.get("positive"), + levels=entry.get("levels"), ).to_dict() - handlers[row.name].append(var_handler) + handlers[entry["name"]].append(var_handler) - return dict(handlers) # type: ignore + return dict(handlers) def _get_handlers_from_modules(path: str) -> dict[str, list[dict[str, Any]]]: diff --git a/scripts/debug/339-eamxx-handlers/prep_vert_remap.sh b/scripts/debug/339-eamxx-handlers/prep_vert_remap.sh new file mode 100644 index 00000000..5754ed10 --- /dev/null +++ b/scripts/debug/339-eamxx-handlers/prep_vert_remap.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# prep_vert_remap.sh +# +# Purpose: Build vert_L128.nc – a SCREAM L128 vertical coordinate file – +# by extracting hybrid coordinate variables from a DecadalSCREAM +# output file and injecting the missing P0 scalar. +# +# The resulting vert_L128.nc is saved alongside vrt_remap_plev19.nc in +# the shared e3sm_to_cmip grids directory. +# +# NOTE: This script can only be executed on Perlmutter. +# +# Run with: +# source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh +# bash scripts/debug/339-eamxx-handlers/prep_vert_remap.sh + +set -euo pipefail + +# ------------------------------------------------------------------ +# Configuration +# ------------------------------------------------------------------ + +# Source file containing L128 hybrid coordinate variables (P0 is absent). +VERT_SRC="/global/cfs/cdirs/e3sm/mldubey/DecadalSCREAM/\ +output.scream.decadal.6hourlyINST_ne1024pg2/\ +output.scream.decadal.6hourlyINST_ne1024pg2.INSTANT.nhours_x6.2006-01-20-21600.nc" + +# Output file – same directory as vrt_remap_plev19.nc (from mvce.py). +GRIDS_DIR="/global/cfs/cdirs/e3sm/diagnostics/e3sm_to_cmip_data/grids" +VERT_L128="${GRIDS_DIR}/vert_L128.nc" + +# ------------------------------------------------------------------ +# Helper +# ------------------------------------------------------------------ +run() { echo "+ $*"; "$@"; } + +# ------------------------------------------------------------------ +# Step 1: Extract hybrid vertical coordinate variables. +# P0 is absent from VERT_SRC and is injected in Step 2. +# ------------------------------------------------------------------ +echo "=== Step 1: Extract vertical coordinates ===" +run ncks -v hyai,hybi,hyam,hybm \ + "${VERT_SRC}" \ + "${VERT_L128}" + +# ------------------------------------------------------------------ +# Step 2: Add P0 = 100000.0 Pa (missing from source). +# ------------------------------------------------------------------ +echo "=== Step 2: Add P0 to ${VERT_L128} ===" +run ncap2 -A -v -s 'P0=100000.0' "${VERT_L128}" + +echo "Done. Output: ${VERT_L128}" diff --git a/scripts/debug/339-eamxx-handlers/run_eamxx_vars.py b/scripts/debug/339-eamxx-handlers/run_eamxx_vars.py new file mode 100644 index 00000000..cc1e0816 --- /dev/null +++ b/scripts/debug/339-eamxx-handlers/run_eamxx_vars.py @@ -0,0 +1,232 @@ +"""A run script for testing newly added EAMxx cmor handlers (issue #339). + +This script is useful for testing code changes to `e3sm_to_cmip` without having +to install the package. It imports the `main` function from the `__main__` +module and passes a list of arguments to it. The arguments are the same as those +passed to the command line interface (CLI) of `e3sm_to_cmip`. +The script can be run from the command line or an IDE. + +NOTE: This script can only be executed on Perlmutter. + +Run with: + conda activate e3sm_to_cmip_dev + python scripts/debug/339-eamxx-handlers/run_eamxx_vars.py + +Workflow: + 1. Generate monthly time-series from native EAMxx run output using ncclimo + (-P eamxx --split) for 2D and 3D variables separately. + 2. Vertically remap 3D files from model levels to pressure levels (plev19). + 3. CMORize 2D variables from the horizontally regridded ncclimo output. + 4. CMORize 3D variables from the vertically remapped directory. + +Variables excluded because their raw inputs are absent in the run directory: + prw -> needs VapWaterPath + tauu -> needs surf_mom_flux_U + tauv -> needs surf_mom_flux_V + clwvi -> needs LiqWaterPath + od550aer -> needs AerosolOpticalDepth550nm +""" + +import os +import subprocess + +from e3sm_to_cmip.main import main + +# ------------------------------------------------------------------ +# Configuration +# ------------------------------------------------------------------ +CASE = "ne256pg2_ne256pg2.F20TR-SCREAMv1.July-1.spanc800.2xauto.acc150.n0032.test2.1" +FILE_NAME = "1ma_ne30pg2.AVERAGE.nmonths_x1" + +START = 1995 +END = 2004 +YPF = 5 # years per file for time-series output + +MAP_FILE = ( + "/global/cfs/cdirs/e3sm/diagnostics/maps/" + "map_ne30pg2_to_cmip6_180x360_aave.20200201.nc" +) + +# Native EAMxx run directory containing 1ma_ne30pg2.AVERAGE.nmonths_x1.*.nc files. +NATIVE_PATH = f"/pscratch/sd/z/zhan391/EAMxx/{CASE}/run" + +# Post-processing root (write to our own scratch). +POST_ROOT = "/pscratch/sd/c/chengzhu/e3sm_to_cmip_test/339-eamxx-handlers/post/atm/180x360_aave" + +# ncclimo output directories: regridded time-series and native staging (trash). +DRC_TS_RGR = os.path.join(POST_ROOT, f"ts_eamxx/monthly/{YPF}yr") +DRC_TS_NATIVE = os.path.join(DRC_TS_RGR, "trash") + +# CMORized output. +OUTPUT_PATH = "/pscratch/sd/c/chengzhu/e3sm_to_cmip_test/339-eamxx-handlers" + +# Vertically remapped 3D files (pressure levels). +INPUT_VERT_PLEV = os.path.join(OUTPUT_PATH, "rgr_vert_plev") + +# Vertical remap grid file (plev19 pressure levels). +VRT_REMAP_FILE = ( + "/global/cfs/cdirs/e3sm/diagnostics/e3sm_to_cmip_data/grids/vrt_remap_plev19.nc" +) + +# SCREAM L128 vertical coordinate file (hybrid levels → pressure levels). +VRT_IN_FILE = ( + "/global/cfs/cdirs/e3sm/diagnostics/e3sm_to_cmip_data/grids/vert_L128.nc" +) + +# CMIP tables and metadata. +TABLES_PATH = ( + "/global/cfs/cdirs/e3sm/diagnostics/e3sm_to_cmip_data/cmip6-cmor-tables/Tables" +) +USER_METADATA = ( + "/global/cfs/cdirs/e3sm/diagnostics/e3sm_to_cmip_data/default_metadata.json" +) + +# E3SM Unified environment for NCO tools (ncclimo, ncremap, etc.). +E3SM_UNIFIED_ENV = ( + "/global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh" +) + +# 2D variables for ncclimo (EAMxx native names). +VARS_2D = ( + "ps,surf_radiative_T,SeaLevelPressure,IceWaterPath,qv_2m," + "precip_liq_surf_mass_flux,precip_ice_surf_mass_flux," + "omega_at_500hPa,omega_at_700hPa,omega_at_850hPa,T_mid_at_700hPa,T_2m," + "surface_upward_latent_heat_flux,surf_sens_flux,z_mid_at_700hPa," + "wind_speed_10m,surf_evap,U_at_10m_above_surface," + "LW_clrsky_flux_dn_at_model_bot,LW_clrsky_flux_up_at_model_top," + "LW_flux_dn_at_model_bot,LW_flux_up_at_model_bot,LW_flux_up_at_model_top," + "SW_clrsky_flux_dn_at_model_bot,SW_clrsky_flux_dn_at_model_top," + "SW_clrsky_flux_up_at_model_bot,SW_clrsky_flux_up_at_model_top," + "SW_flux_dn_at_model_bot,SW_flux_dn_at_model_top," + "SW_flux_up_at_model_bot,SW_flux_up_at_model_top," + "ShortwaveCloudForcing,LongwaveCloudForcing" +) +VAR_XTR_2D = "area,landfrac,ocnfrac" + +# 3D variables for ncclimo (EAMxx native names). +VARS_3D = "U,V,T_mid,z_mid,omega,RelativeHumidity,p_mid,qv" +VAR_XTR_3D = "ps,hyai,hyam,hybi,hybm,area,landfrac,ocnfrac" + +# EAMxx native variable name prefixes for 3D files that need vertical regridding. +# Corresponds to CMIP variables: ta, hus, hur, wap, zg. +RAW_VARS_3D = ["T_mid", "qv", "RelativeHumidity", "omega", "z_mid"] + +# 2D CMIP variables (no vertical remapping needed). +VAR_LIST_2D = ( + "pr, ts, tas, huss, hfls, hfss, evspsbl, sfcWind, psl, " + "rsdt, rsut, rsutcs, rlut, rlutcs, " + "rsds, rsus, rsdscs, rsuscs, rlds, rlus, rldscs, " + "clivi" +) + +# 3D CMIP variables (require vertical remapping from model levels to plev19). +VAR_LIST_3D = "ta, hus, hur, wap, zg" + + +def run(cmd: str) -> None: + """Run a shell command after sourcing the E3SM Unified environment.""" + full = f"source {E3SM_UNIFIED_ENV} && {cmd}" + subprocess.run(full, shell=True, executable="/bin/bash", check=True) + + +# ------------------------------------------------------------------ +# Step 1a: Generate 2D monthly time-series with ncclimo. +# cd to the run dir so the eval ls brace expansion finds the files. +# ------------------------------------------------------------------ +os.makedirs(DRC_TS_NATIVE, exist_ok=True) +os.makedirs(DRC_TS_RGR, exist_ok=True) + +print("Step 1a: ncclimo 2D time-series ...") +run( + f"cd {NATIVE_PATH} && " + f"eval ls {FILE_NAME}.*{{{START}..{END}}}*.nc " + f"| ncclimo -P eamxx -c {CASE} -v {VARS_2D} --split " + f"--var_xtr={VAR_XTR_2D} " + f"--yr_srt={START} --yr_end={END} --ypf={YPF} " + f"--map={MAP_FILE} " + f"-o {DRC_TS_NATIVE} -O {DRC_TS_RGR}" +) + +# ------------------------------------------------------------------ +# Step 1b: Generate 3D monthly time-series with ncclimo. +# ------------------------------------------------------------------ +print("Step 1b: ncclimo 3D time-series ...") +run( + f"cd {NATIVE_PATH} && " + f"eval ls {FILE_NAME}.*{{{START}..{END}}}*.nc " + f"| ncclimo -P eamxx -c {CASE} -v {VARS_3D} --split " + f"--var_xtr={VAR_XTR_3D} " + f"--yr_srt={START} --yr_end={END} --ypf={YPF} " + f"--map={MAP_FILE} " + f"-o {DRC_TS_NATIVE} -O {DRC_TS_RGR}" +) + +# ------------------------------------------------------------------ +# Step 2: Vertically remap 3D files from model levels to plev19. +# Files in DRC_TS_RGR use EAMxx native names (T_mid, qv, etc.) +# and contain ps (lowercase) as the surface pressure variable. +# ------------------------------------------------------------------ +os.makedirs(INPUT_VERT_PLEV, exist_ok=True) + +import re + +print("Step 2: vertical regridding to plev19 ...") + +for fname in sorted(os.listdir(DRC_TS_RGR)): + if not fname.endswith(".nc"): + continue + + for v in RAW_VARS_3D: + # Match only: v_YYYYMM_YYYYMM.nc + if re.match(rf"^{v}_[0-9]{{6}}_[0-9]{{6}}\.nc$", fname): + src = os.path.join(DRC_TS_RGR, fname) + dst = os.path.join(INPUT_VERT_PLEV, fname) + + print(f" Vertically remapping: {fname}") + run( + f"ncremap --ps_nm={src}/ps " + f"--vrt_ntp=log --vrt_xtr=mss_val " + f"--vrt_in={VRT_IN_FILE} " + f"--vrt_out={VRT_REMAP_FILE} {src} {dst}" + ) + break + +# ncks --rgr xtr_mth=mss_val --vrt_fl=${e2c_path}/grids/vrt_remap_plev19.nc ${rgr_dir_vert}/$file ${rgr_dir}/$file + +# ------------------------------------------------------------------ +# Step 3: CMORize 2D atmosphere variables from the regridded directory. +# ------------------------------------------------------------------ +print("Step 3: CMORize 2D variables ...") +main([ + "--var-list", VAR_LIST_2D, + "--input-path", DRC_TS_RGR, + "--output-path", OUTPUT_PATH, + "--tables-path", TABLES_PATH, + "--user-metadata", USER_METADATA, + "--freq", "mon", + "--serial", + "--debug", +]) +# +# ------------------------------------------------------------------ +# Step 4: CMORize 3D atmosphere variables from vertically remapped files. +# ------------------------------------------------------------------ +print("Step 4: CMORize 3D variables ...") +main([ + "--var-list", VAR_LIST_3D, + "--input-path", INPUT_VERT_PLEV, + "--output-path", OUTPUT_PATH, + "--tables-path", TABLES_PATH, + "--user-metadata", USER_METADATA, + "--freq", "mon", + "--serial", + "--debug", +]) + +# Ensure the output and its contents have the correct permissions. +for root, dirs, files in os.walk(OUTPUT_PATH): + os.chmod(root, 0o755) + for d in dirs: + os.chmod(os.path.join(root, d), 0o755) + for f in files: + os.chmod(os.path.join(root, f), 0o644) diff --git a/tests/cmor_handlers/test__formulas.py b/tests/cmor_handlers/test__formulas.py index 8666560a..b0fd0bf5 100644 --- a/tests/cmor_handlers/test__formulas.py +++ b/tests/cmor_handlers/test__formulas.py @@ -259,6 +259,29 @@ def test_pr(): np.testing.assert_array_equal(result, expected) + ds = xr.Dataset( + data_vars={ + "precip_liq_surf_mass_flux": _dummy_dataarray(), + "precip_ice_surf_mass_flux": _dummy_dataarray(), + } + ) + result = pr(ds) + expected = xr.DataArray( + dims=["lat", "lon"], + data=np.array([[0, 2000, 4000], [0, 2000, 4000], [0, 2000, 4000]]), + ) + + np.testing.assert_array_equal(result, expected) + + ds = xr.Dataset(data_vars={"precip_total_surf_mass_flux": _dummy_dataarray()}) + result = pr(ds) + expected = xr.DataArray( + dims=["lat", "lon"], + data=np.array([[0, 1000, 2000], [0, 1000, 2000], [0, 1000, 2000]]), + ) + + np.testing.assert_array_equal(result, expected) + # Test when required variable keys are NOT in the data dictionary. with pytest.raises(KeyError): pr(xr.Dataset()) diff --git a/tests/cmor_handlers/test_utils.py b/tests/cmor_handlers/test_utils.py index 99ee9881..71f63689 100644 --- a/tests/cmor_handlers/test_utils.py +++ b/tests/cmor_handlers/test_utils.py @@ -86,6 +86,35 @@ def test_updates_CMIP_table_for_variable_based_on_freq_param(self): output_data=None, method=VarHandler.cmorize, ), + dict( + name="pr", + units="kg m-2 s-1", + raw_variables=[ + "precip_liq_surf_mass_flux", + "precip_ice_surf_mass_flux", + ], + table="CMIP6_Amon.json", + unit_conversion=None, + formula="(precip_liq_surf_mass_flux + precip_ice_surf_mass_flux) * 1000.0", + formula_method=_formulas.pr, + positive=None, + levels=None, + output_data=None, + method=VarHandler.cmorize, + ), + dict( + name="pr", + units="kg m-2 s-1", + raw_variables=["precip_total_surf_mass_flux"], + table="CMIP6_Amon.json", + unit_conversion=None, + formula="precip_total_surf_mass_flux * 1000.0", + formula_method=_formulas.pr, + positive=None, + levels=None, + output_data=None, + method=VarHandler.cmorize, + ), ] # Update each handler objects' BoundMethod to the underlying function.