diff --git a/preprocessing_scripts/vertical_remap.py b/preprocessing_scripts/vertical_remap.py new file mode 100755 index 00000000..41a4ec63 --- /dev/null +++ b/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() diff --git a/run_scripts/ecomip/eamxx_ecomip_output.yaml b/run_scripts/ecomip/eamxx_ecomip_output.yaml new file mode 100644 index 00000000..85918346 --- /dev/null +++ b/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' +... diff --git a/run_scripts/ecomip/run.ecomip.sh b/run_scripts/ecomip/run.ecomip.sh new file mode 100755 index 00000000..ccf34bd1 --- /dev/null +++ b/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}"