Add ability to cycle GOCART aerosols

jobs/rocoto/init_chem.sh

@@ -0,0 +1,56 @@
+#!/bin/bash
+
+set -x
+
+###############################################################
+## Abstract:
+## Create FV3 initial conditions from GFS intitial conditions
+## RUN_ENVIR : runtime environment (emc | nco)
+## HOMEgfs   : /full/path/to/workflow
+## EXPDIR : /full/path/to/config/files
+## CDATE  : current date (YYYYMMDDHH)
+## CDUMP  : cycle name (gdas / gfs)
+## PDY    : current date (YYYYMMDD)
+## cyc    : current cycle (HH)
+###############################################################
+
+###############################################################
+# Source FV3GFS workflow modules
+. $USHgfs/load_fv3gfs_modules.sh
+status=$?
+[[ $status -ne 0 ]] && exit $status
+
+###############################################################
+# Source relevant configs
+configs="base"
+for config in $configs; do
+    . $EXPDIR/config.${config}
+    status=$?
+    [[ $status -ne 0 ]] && exit $status
+done
+
+###############################################################
+# Source machine runtime environment
+. $BASE_ENV/${machine}.env init_aerosol
+status=$?
+[[ $status -ne 0 ]] && exit $status
+
+# export DATAROOT="$RUNDIR/$CDATE/$CDUMP"
+# [[ ! -d $DATAROOT ]] && mkdir -p $DATAROOT
+# export DATA=${DATA:-${DATAROOT}/${jobid:?}}
+# mkdir -p $DATA
+# cd $DATA
+
+$SCRgfs/exgfs_chem_init_aerosol.py
+
+status=$?
+if [[ $status -ne 0 ]]; then
+    echo "FATAL ERROR: exgfs_chem_init_aerosol.py failed with error code $status"
+    exit $status
+fi
+
+##############################################################
+# Exit cleanly
+
+set +x
+exit 0

modulefiles/module_base.hera

@@ -4,7 +4,8 @@
 ##
 
 # Modules availble through hpc-stack
-module use /scratch2/NCEPDEV/nwprod/hpc-stack/libs/hpc-stack/modulefiles/stack
+# module use /scratch2/NCEPDEV/nwprod/hpc-stack/libs/hpc-stack/modulefiles/stack
+module use /scratch1/NCEPDEV/nems/Raffaele.Montuoro/dev/nasa/libs/hpc-stack/1.1.0-nasa/modulefiles/stack
 
 module load hpc/1.1.0
 module load hpc-intel/18.0.5.274
@@ -18,6 +19,7 @@ module load crtm 2.3.0
 module load prod_util/1.2.2
 module load grib_util/1.2.2
 module load ip/3.3.3
+
 module load sp/2.3.3
 module load w3nco/2.4.1
 module load bacio/2.4.1
@@ -32,11 +34,14 @@ module load hdf5/1.10.6
 module load esmf/8_1_0_beta_snapshot_27
 module load w3emc/2.7.3
 module load wgrib2/2.0.8
-setenv WGRIB2 wgrib2
+setenv "WGRIB2" "wgrib2"
 
 # Modules not under hpc-stack
 module load hpss/hpss
 module load nco/4.9.1
 module load gempak/7.4.2
 module load ncl/6.5.0
 module load cdo/1.9.5
+module load intelpython/3.6.8
+# netCDF4 is not in the python installation on Hera
+append-path "PYTHONPATH" "/scratch2/NCEPDEV/ensemble/save/Walter.Kolczynski/python_packages"

parm/chem/gocart_tracer.list

@@ -0,0 +1,20 @@
+so2
+sulf
+dms
+msa
+pp25
+pp10
+bc1
+bc2
+oc1
+oc2
+dust1
+dust2
+dust3
+dust4
+dust5
+seas1
+seas2
+seas3
+seas4
+seas5

scripts/exgfs_chem_init_aerosol.py

@@ -0,0 +1,148 @@
+#! /usr/bin/env python3
+
+"""
+"""
+
+import os
+import subprocess
+import typing
+from datetime import datetime, timedelta
+from functools import partial
+
+# Constants
+atm_base_pattern = "{ics_dir}/%Y%m%d%H/atmos/{case}/INPUT"                    # Location of atmosphere ICs
+atm_file_pattern = "{path}/gfs_data.{tile}.nc"                                # Atm IC file names
+tracer_base_pattern = "{rot_dir}/{cdump}.%Y%m%d/%H/atmos/RERUN_RESTART"       # Location of restart files (time of previous run)
+tracer_file_pattern = "{tracer_base}/{timestamp}fv_tracer.res.{tile}.nc"      # Name of restart files (time when restart is valid)
+tracer_list_file_pattern = "{parm_gfs}/chem/gocart_tracer.list"               # Text list of tracer names to copy
+merge_script_pattern = "{ush_gfs}/merge_fv3_chem_tile.py"
+n_tiles = 6
+max_lookback = 4                                                              # Maximum number of past cycles to look for for tracer data
+debug = True
+
+# End configurable settings
+
+# Make sure print statements are flushed immediately, otherwise
+#   print statments may be out-of-order with subprocess output
+print = partial(print, flush=True)
+
+tiles = list(map(lambda t: "tile{t}".format(t=t), range(1, n_tiles + 1)))
+
+
+def main() -> None:
+	# Read in environment variables and make sure they exist
+	cdate = get_env_var("CDATE")
+	incr = int(get_env_var('STEP_GFS'))
+	fcst_length = int(get_env_var('FHMAX_GFS'))
+	cdump = get_env_var("CDUMP")
+	rot_dir = get_env_var("ROTDIR")
+	ics_dir = get_env_var("ICSDIR")
+	case = get_env_var("CASE")
+	ush_gfs = get_env_var("USHgfs")
+	parm_gfs = get_env_var("PARMgfs")
+	# data = get_env_var('DATA')
+
+	# os.chdir(data)
+
+	merge_script = merge_script_pattern.format(ush_gfs=ush_gfs)
+	tracer_list_file = tracer_list_file_pattern.format(parm_gfs=parm_gfs)
+
+	time = datetime.strptime(cdate, "%Y%m%d%H")
+	atm_source_path = time.strftime(atm_base_pattern.format(**locals()))
+
+	if(debug):
+		for var in ['merge_script', 'tracer_list_file', 'atm_source_path']:
+			print(f'{var} = {f"{var}"}')
+
+	atm_files = get_atm_files(atm_source_path)
+	tracer_files = get_tracer_files(time, incr, max_lookback, fcst_length, rot_dir, cdump)
+
+	if (tracer_files is not None):
+		merge_tracers(merge_script, atm_files, tracer_files, tracer_list_file)
+
+	return
+
+
+# Retrieve environment variable and exit or print warning if not defined
+def get_env_var(varname: str, fail_on_missing: bool = True) -> str:
+	if(debug):
+		print(f'Trying to read envvar {varname}')
+
+	var = os.environ.get(varname)
+	if(var is None):
+		if(fail_on_missing is True):
+			print("FATAL: Environment variable {varname} not set".format(varname=varname))
+			exit(100)
+		else:
+			print("WARNING: Environment variable {varname} not set, continuing using None".format(varname=varname))
+	if(debug):
+		print(f'\tValue: {var}')
+	return(var)
+
+
+# Check if atm files exist
+def get_atm_files(path: str) -> typing.List[str]:
+	print(f'Checking for atm files in {path}')
+
+	files = list(map(lambda tile: atm_file_pattern.format(tile=tile, path=path), tiles))
+	for file_name in files:
+		if(debug):
+			print(f"\tChecking for {file_name}")
+		if(not os.path.isfile(file_name)):
+			print("FATAL: Atmosphere file {file_name} not found".format(file_name=file_name))
+			exit(101)
+		elif(debug):
+			print(f"\t\tFound {file_name}")
+	return files
+
+
+# Find last cycle with tracer data available via restart files
+def get_tracer_files(time: datetime, incr: int, max_lookback: int, fcst_length: int, rot_dir: str, cdump: str) -> typing.List[str]:
+	print(f"Looking for restart tracer files in {rot_dir}")
+	for lookback in map(lambda i: incr * (i + 1), range(max_lookback)):
+		if(lookback > fcst_length):
+			# Trying to look back farther than the length of a forecast
+			break
+		elif(lookback == fcst_length):
+			# Restart files at the end of the cycle don't have a timestamp
+			timestamp = ""
+		else:
+			timestamp = time.strftime("%Y%m%d.%H0000.")
+
+		last_time = time - timedelta(hours=lookback)
+
+		if(debug):
+			print(f"\tChecking {last_time}")
+		tracer_base = last_time.strftime(tracer_base_pattern.format(**locals()))
+		files = list(map(lambda tile: tracer_file_pattern.format(timestamp=timestamp, tracer_base=tracer_base, tile=tile), tiles))
+		if(debug):
+			print(f"\t\tLooking for files {files} in directory {tracer_base}")
+		found = [file for file in files if os.path.isfile(file)]
+		if(found):
+			if(debug):
+				print(f"\t\tAll files found!")
+			return files
+		else:
+			print(last_time.strftime("Tracer files not found for %Y%m%d_%H"))
+
+	if(not found):
+		print("WARNING: Unable to find tracer files, will use zero fields")
+		return None
+
+
+# Merge tracer data into atmospheric data
+def merge_tracers(merge_script: str, atm_files: typing.List[str], tracer_files: typing.List[str], tracer_list_file: str) -> None:
+	print(f"Merging tracers")
+	if(len(atm_files) != len(tracer_files)):
+		print("FATAL: atmosphere file list and tracer file list are not the same length")
+		exit(102)
+
+	for atm_file, tracer_file in zip(atm_files, tracer_files):
+		if(debug):
+			print(f"\tMerging tracers from {tracer_file} into {atm_file}")
+		subprocess.run([merge_script, atm_file, tracer_file, tracer_list_file], check=True)
+
+
+if __name__ == "__main__":
+	main()
+	exit(0)

ush/merge_fv3_chem_tile.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+"""
+Appends tracer data from one NetCDF file to another and updates the tracer
+count.
+
+usage: merge_fv3_chem_tile.py [-h] atm_file chem_file variable_file [out_file]
+
+Appends tracer data from one NetCDF file to another and updates the tracer
+count.
+
+positional arguments:
+  atm_file       File containing the atmospheric data
+  chem_file      File containing the chemistry tracer data to be added
+  variable_file  File containing list of tracer variable_names in the
+                  chem_file to add to the atm_file, one tracer per line
+  out_file       Name of file to create. If none is specified, the atm_file
+                  will be edited in place. New file will be a copy of atm_file
+                  with the specificed tracers listed in variable_file appended
+                  from chem_file and ntracers updated.
+
+optional arguments:
+  -h, --help     show this help message and exit
+
+"""
+
+import os, sys, subprocess
+from typing import List
+from functools import partial
+from shutil import copyfile
+import argparse
+
+try:
+	import netCDF4
+except ImportError:
+	print("FATAL ERROR: Failed to import netCDF4 in " + __file__ + "!")
+	print("Make sure you are using a version of python 3 with netCDF4 installed")
+	sys.exit(100)
+
+# Make sure print statements are flushed immediately, otherwise
+#   print statments may be out-of-order with subprocess output
+print = partial(print, flush=True)
+
+def merge_tile(base_file_name: str, append_file_name: str, tracers_to_append: List[str]) -> None:
+	if not os.path.isfile(base_file_name):
+		print("FATAL ERROR: Atmosphere file " + base_file_name + " does not exist!")
+		sys.exit(102)
+
+	if not os.path.isfile(append_file_name):
+		print("FATAL ERROR: Chemistry file " + append_file_name + " does not exist!")
+		sys.exit(103)
+
+	append_file = netCDF4.Dataset(append_file_name, "r")
+	base_file = netCDF4.Dataset(base_file_name, "r+")
+
+	# print(base_file)
+	# print(base_file.dimensions["ntracer"])
+
+	old_ntracer = base_file.dimensions["ntracer"].size
+	new_ntracer = old_ntracer + len(tracers_to_append)
+
+	# Copy over chemistry dimensions
+	for dim_name in append_file.dimensions:
+		base_file.createDimension(dim_name, append_file.dimensions[dim_name].size)
+
+	for variable_name in tracers_to_append:
+		print("Adding variable " + variable_name + " to file")
+		variable = append_file[variable_name]
+		base_file.createVariable(variable_name, variable.datatype, variable.dimensions)
+		base_file[variable_name][:] = variable[:]
+		base_file[variable_name].setncatts(variable.__dict__)
+		# print("Done adding " + variable_name)
+
+	print("Updating ntracer")
+
+	# Use ncks to rewrite file without ntracer so we can define it anew
+	subprocess.run(["ncks", "-x", "-v", "ntracer", "-O", base_file_name, base_file_name], check=True)
+
+	base_file = netCDF4.Dataset(base_file_name, "r+")
+	base_file.createDimension("ntracer", new_ntracer)
+	# print(base_file.dimensions["ntracer"])
+
+
+def main() -> None:
+	parser = argparse.ArgumentParser(
+		description="Appends tracer data from one NetCDF file to another and updates the tracer count.")
+	parser.add_argument('atm_file', type=str, help="File containing the atmospheric data")
+	parser.add_argument('chem_file', type=str, help="File containing the chemistry tracer data to be added")
+	parser.add_argument('variable_file', type=str, help="File containing list of tracer variable_names in the chem_file to add to the atm_file, one tracer per line")
+	parser.add_argument('out_file', type=str, nargs="?", help="Name of file to create. If none is specified, the atm_file will be edited in place. New file will be a copy of atm_file with the specificed tracers listed in variable_file appended from chem_file and ntracers updated.")
+
+	args = parser.parse_args()
+
+	atm_file_name = args.atm_file
+	chem_file_name = args.chem_file
+	variable_file = args.variable_file
+	out_file_name = args.out_file
+
+	if out_file_name is None:
+		print("INFO: No out_file specified, will edit atm_file in-place")
+		out_file_name = atm_file_name
+	else:
+		if os.path.isfile(out_file_name):
+			print("WARNING: Specified out file " + out_file_name + " exists and will be overwritten")
+		copyfile(atm_file_name, out_file_name)
+
+	variable_file = open(variable_file)
+	variable_names = variable_file.read().splitlines()
+	variable_file.close()
+
+	merge_tile(out_file_name, chem_file_name, variable_names)
+
+	# print(variable_names)
+
+
+if __name__ == "__main__":
+	main()

workflow/cases/aerosol_firex_forecast.yaml

@@ -11,10 +11,10 @@ case:
   settings: 
     SDATE: 2019-07-01t00:00:00
     EDATE: 2019-07-15t00:00:00
+    STEP_GFS: !timedelta "24:00:00"
+    STEP_DA: !timedelta "00:00:00"
     max_job_tries: 1
 
-    gfs_cyc: 1
-
     cplgocart: .true.
     print_esmf: .true.
     nems_temp: 'atm_aer'