Add ecomip runscript and outputs

preprocessing_scripts/vertical_remap.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+
+import xarray, numpy, argparse
+
+def main():
+    parser = argparse.ArgumentParser(description="Do vertical remap from one sigma hybrid grid to another")
+    parser.add_argument("input_file"                , help="Input data file")
+    parser.add_argument("output_vertical_grid_file" , help="File with vertical coordinate information for target grid")
+    parser.add_argument("output_file"               , help="Output data file with remapped fields")
+    parser.add_argument("--input_vertical_grid_file", help="File with vertical coordinate information for source grid, if not contained in input_file", default=None)
+    parser.add_argument("--vertical_dimension"      , help="Name of vertical dimension", default='lev')
+    parser.add_argument("--use_log_surface_pressure", help="Use lnsp for hybrid calculation (ECMWF IFS)", default=True)
+    parser.add_argument("--fields"                  , help="Names of fields in input_file to remap", nargs="+", type=str, default=None)
+    args = parser.parse_args()
+    _main(args.input_file, args.input_vertical_grid_file,
+          args.output_vertical_grid_file, args.output_file,
+          args.vertical_dimension, args.fields, args.use_log_surface_pressure)
+
+
+# Do linear interpolation with extrapolation at the endpoints
+def linear_interp(x, xp, fp):
+
+    # Start with linear interpolation within domain of xp
+    y = numpy.interp(x, xp, fp)
+
+    # Extrapolate left
+    mask_left = x < xp[0]
+    if numpy.any(mask_left):
+        slope_left = (fp[1] - fp[0]) / (xp[1] - xp[0])
+        y[mask_left] = fp[0] + (x[mask_left] - xp[0]) * slope_left
+
+    # Extrapolate right
+    mask_right = x > xp[-1]
+    if numpy.any(mask_right):
+        slope_right = (fp[-1] - fp[-2]) / (xp[-1] - xp[-2])
+        y[mask_right] = fp[-1] + (x[mask_right] - xp[-1]) * slope_right
+
+    return y
+
+
+def _main(input_file, input_vertical_grid_file, output_vertical_grid_file,
+          output_file, vertical_dimension, fields, use_log_surface_pressure):
+
+    # Open input dataset
+    # TODO: figure out how to abstract out the chunking for large datasets, or
+    # allow passing this on the command line
+    # Make sure we do NOT chunk on vertical dimension
+    ds_in = xarray.open_dataset(input_file, chunks={vertical_dimension: -1}) #chunks={'ncells': 10000, vertical_dimension: -1, 'time': 1})
+
+    # Open input_vertical_grid_file to get grid information, if separate from input file
+    if input_vertical_grid_file is not None:
+        with xarray.open_dataset(input_vertical_grid_file) as ds_vert_in:
+            for v in ('hyam', 'hybm', 'hyai', 'hybi'):
+                ds_in[v] = ds_vert_in[v]
+
+    # Get list of fields to remap
+    if fields is None: fields = [f for f in ds_in.variables.keys() if vert_dim in ds_in[f].dims]
+
+    # Compute pressure on input vertical grid
+    if use_log_surface_pressure and all([v in ds_in.variables.keys() for v in ('lnsp', 'hyam', 'hybm', 'hyai', 'hybi')]):
+        # ECMWF IFS vertical grid
+        ps = numpy.exp(ds_in['lnsp'].isel({vertical_dimension: 0}))
+        p_in = (ds_in['hyam'] + ds_in['hybm'] * ps).rename({'nhym': vertical_dimension})
+    elif all([v in ds_in.variables.keys() for v in ('P0', 'PS', 'hyam', 'hybm', 'hyai', 'hybi')]):
+        # EAM/EAMxx vertical grid
+        ps = ds_in['PS']
+        p_in = ds_in['hyam'] * ds_in['P0'] + ds_in['hybm'] * ps
+    else:
+        raise RuntimeError('Cannot infer vertical grid type of input_file.')
+
+    # Compute pressure on output vertical grid and figure out name of output vertical coordinate dimension
+    # Also copy this information to output dataset
+    with xarray.open_dataset(output_vertical_grid_file) as ds_vert:
+        output_vertical_dimension = ds_vert['hyam'].dims[0]
+        p_out = ds_vert['P0'] * ds_vert['hyam'] + ps * ds_vert['hybm']
+        ds_out = xarray.Dataset({v: ds_vert[v] for v in ('hyam', 'hybm', 'hyai', 'hybi')})
+        ds_out['PS'] = ps  # Make sure surface pressure is copied as well
+        ds_out['P0'] = ds_vert['P0']
+
+    # Remap fields to target vertical grid, linear in pressure
+    for f in fields:
+        # Sanity check: make sure input does not contain missing data
+        if ds_in[f].isnull().any():
+            raise RuntimeError('Cannot handle missing data.')
+
+        # Do vertical interpolation vectorized over non-level dimensions
+        print(p_out.dims, p_in.dims, ds_in[f].dims)
+        ds_out[f] = xarray.apply_ufunc(
+                numpy.interp,
+                p_out, p_in, ds_in[f],
+                input_core_dims=[[output_vertical_dimension], [vertical_dimension], [vertical_dimension]],
+                output_core_dims=[[output_vertical_dimension]],
+                vectorize=True, dask='parallelized',
+                output_dtypes=[float],
+            )
+
+        # Copy attributes from input data
+        ds_out[f].attrs = ds_in[f].attrs
+
+    print('Write to file...', end='')
+    ds_out.to_netcdf(output_file)
+    print('done.')
+    ds_in.close()
+    ds_out.close()
+
+
+if __name__ == '__main__':
+    main()

run_scripts/ecomip/eamxx_ecomip_output.yaml

@@ -0,0 +1,57 @@
+%YAML 1.2
+---
+averaging_type: instant
+max_snapshots_per_file: 1
+fields:
+  physics_pg2:
+    field_names:
+    # For Earthcare simulators
+    - ps  # For computing pressure from hybrid levels
+    - z_mid
+    - p_int
+    - p_mid
+    - T_mid
+    #- horiz_winds
+    - U
+    - V
+    #- tracers
+    - qv
+    - qc
+    - qi
+    - qr
+    - qm
+    - nc
+    - ni
+    - nr
+    - bm
+    - cldfrac_tot
+    - precip_total_surf_mass_flux
+    - precip_liq_surf_mass_flux
+    - precip_ice_surf_mass_flux
+    - omega
+    #- convective_mass_flux # Conditional diagnostic
+    - surf_radiative_T #skin_temperature
+    - wind_speed_10m
+    #- U_at_10m
+    #- V_at_10m
+    - SeaLevelPressure
+    - phis
+    - pseudo_density
+    # For diagnosis
+    - SW_flux_up_at_model_top
+    - SW_flux_dn_at_model_top
+    - LW_flux_up_at_model_top
+    - LW_flux_dn_at_model_top
+    - SW_clrsky_flux_up_at_model_top
+    - LW_clrsky_flux_up_at_model_top
+    - IceWaterPath
+    - LiqWaterPath
+    - VapWaterPath
+    #- precip_accumulated_over_30min
+    - diag_equiv_reflectivity
+filename_prefix: ecomip.eamxx.h
+output_control:
+  frequency: 30 #${HIST_N}
+  frequency_units: nminutes #${HIST_OPTION}
+  skip_t0_output: 'false'
+...

run_scripts/ecomip/run.ecomip.sh

@@ -0,0 +1,119 @@
+#!/bin/bash
+set -e
+umask 022
+
+script_root=${PWD}
+branch="master" #"decadal-production-run4" #"fix-nanobug-in-horiz-remap" #"decadal-production-run4" #"scorpio-update"
+code_root=${HOME}/codes/e3sm/branches/${branch}
+res=ne1024pg2_ne1024pg2 #ne30pg2_ne30pg2 #ne1024pg2_ne1024pg2  #ne1024pg2_ne1024pg2 #ne1024pg2_ne1024pg2 #ne30pg2_EC30to60E2r2
+compset=F2010-SCREAMv1
+machine="frontier" #"pm-gpu" #chrysalis
+compiler="craygnu-mphipcc" #"gnugpu" #intel
+project="cli115" #"e3sm"
+walltime="02:00:00"
+datestring=20260422 #`date +'%Y%m%d'` #"20251031"
+casename=ecomip.${res}.${compset}.${datestring}
+caseroot=/lustre/orion/cli115/proj-shared/brhillman/e3sm/cases/${casename}
+#pecount=32x1
+
+if [ "${res}" == "ne1024pg2_ne1024pg2" ]; then
+    #pecount="2560x6" # 320 nodes
+    pecount="16384x6" # 2048 nodes
+elif [ "${res}" == "ne256pg2_ne256pg2" ]; then
+    pecount="768x6" # 96 nodes
+elif [ "${res}" == "ne30pg2_EC30to60E2r2" ]; then
+    pecount="16x6"
+    #pecount=5540x1 #128x1
+fi
+
+mkdir -p `dirname ${caseroot}`
+${code_root}/cime/scripts/create_newcase \
+    --case=${caseroot} \
+    --res=${res} \
+    --compset=${compset} \
+    --machine=${machine} \
+    --compiler=${compiler} \
+    --pecount=${pecount} \
+    --project=${project} \
+    --walltime=${walltime}
+
+cd ${caseroot}
+
+# Extract input_data_dir for user edits to the namelist
+DIN_LOC_ROOT=`./xmlquery DIN_LOC_ROOT --value`
+
+# Change run length
+./xmlchange STOP_OPTION=ndays,STOP_N=2
+./xmlchange REST_OPTION=ndays,REST_N=2
+./xmlchange RESUBMIT=0
+./xmlchange RUN_STARTDATE="2024-09-02"
+#./xmlchange ATM_NCPL=48
+
+# Run setup before configuring components
+./case.setup
+
+# Namelist options for EAMxx
+if [ "${res}" == "ne1024pg2_ne1024pg2" ]; then
+    ./atmchange initial_conditions::filename="/lustre/orion/cli115/proj-shared/brhillman/scream/data/ecomip_ifs_od-0001_analysis20240902T0000Z_0.05deg_ml.ne1024np4L128.eamxx.nc"
+fi
+
+# Turn on cosp and set frequency
+if [ 1 -eq 0 ]; then
+./atmchange physics::atm_procs_list="mac_aero_mic,rrtmgp,cosp"
+./atmchange physics::cosp::cosp_frequency_units="hours"
+./atmchange physics::cosp::cosp_frequency=1
+fi
+
+output_yaml_files=(`ls ${script_root}/*.yaml`)
+for file in ${output_yaml_files[@]}; do
+    cp -v ${file} ./
+    ./atmchange output_yaml_files+=${file}
+done
+
+# Namelist options for ELM
+# Specify land initial condition and surface datasets
+if [ "${res}" == "ne1024pg2_ne1024pg2" ]; then
+cat << EOF >> user_nl_elm
+  ! Set input data paths
+  finidat = "${DIN_LOC_ROOT}/lnd/clm2/initdata/20231226.I2010CRUELM.ne1024pg2_ICOS10.elm.r.1994-10-01-00000.nc"
+EOF
+elif [ "${res}" == "ne256pg2_ne256pg2" ]; then
+cat << EOF >> user_nl_elm
+  ! Set input data paths
+  finidat = "${DIN_LOC_ROOT}/lnd/clm2/initdata/20240104.I2010CRUELM.ne256pg2.elm.r.1994-10-01-00000.nc"
+EOF
+elif [ "${res}" == "ne30pg2_EC30to60E2r2" ]; then
+cat << EOF >> user_nl_elm
+  finidat = "$DIN_LOC_ROOT/lnd/clm2/initdata/20210802.ICRUELM-1950.ne30pg2_EC30to60E2r2.elm.r.0051-01-01-00000.nc"
+EOF
+fi
+
+# Set MOSART initial condition
+if [ "${res}" == "ne1024pg2_ne1024pg2" ]; then
+cat << EOF >> user_nl_mosart
+  finidat_rtm = "${DIN_LOC_ROOT}/rof/mosart/initdata/20231226.I2010CRUELM.ne1024pg2_ICOS10.mosart.r.1994-10-01-00000.nc"
+EOF
+elif [ "${res}" == "ne256pg2_ne256pg2" ]; then
+cat << EOF >> user_nl_mosart
+  finidat_rtm = "${DIN_LOC_ROOT}/rof/mosart/initdata/20240104.I2010CRUELM.ne256pg2.mosart.r.1994-10-01-00000.nc"
+EOF
+fi
+
+# Point to new SST forcing
+./xmlchange --file env_run.xml --id SSTICE_DATA_FILENAME --val "/lustre/orion/cli115/proj-shared/brhillman/scream/data/sst_ice.ecomip_ifs_od-0001_analysis20240902T0000Z_0.05deg_sfc.3600x7200wst.cdocon.nc"
+./xmlchange --file env_run.xml --id SSTICE_GRID_FILENAME --val "/lustre/orion/cli115/proj-shared/brhillman/scream/data/domain.ecomip_ifs_od_0.05deg.3600x7200wst.nc"
+./xmlchange --file env_run.xml --id SSTICE_YEAR_ALIGN --val 2024
+./xmlchange --file env_run.xml --id SSTICE_YEAR_START --val 2024
+./xmlchange --file env_run.xml --id SSTICE_YEAR_END --val 2025
+
+# Link to rundir
+ln -s `./xmlquery --value RUNDIR` run
+
+# Copy runscript to case dir
+cp ${script_root}/`basename $0` ./
+
+# Setup, build, run
+./case.setup
+./case.build
+./case.submit -a="--job-name=${casename} -t ${walltime} --mail-type=ALL --mail-user=bhillma@sandia.gov"
+echo "caseroot = ${caseroot}"