-
Notifications
You must be signed in to change notification settings - Fork 15
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
5 changed files
with
231 additions
and
65 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,115 @@ | ||
| import os | ||
| import pprint | ||
| import re | ||
|
|
||
| import jinja2 | ||
|
|
||
| from zppy.bundle import handle_bundles | ||
| from zppy.utils import ( | ||
| checkStatus, | ||
| getComponent, | ||
| getTasks, | ||
| getYears, | ||
| makeExecutable, | ||
| setMappingFile, | ||
| submitScript, | ||
| ) | ||
|
|
||
|
|
||
| # ----------------------------------------------------------------------------- | ||
| def e3sm_to_cmip(config, scriptDir, existing_bundles, job_ids_file): | ||
|
|
||
| # --- Initialize jinja2 template engine --- | ||
| templateLoader = jinja2.FileSystemLoader( | ||
| searchpath=config["default"]["templateDir"] | ||
| ) | ||
| templateEnv = jinja2.Environment(loader=templateLoader) | ||
| template = templateEnv.get_template("e3sm_to_cmip.bash") | ||
|
|
||
| # --- List of tasks --- | ||
| tasks = getTasks(config, "ts") | ||
| if len(tasks) == 0: | ||
| return existing_bundles | ||
|
|
||
| # --- Generate and submit ts scripts --- | ||
| for c in tasks: | ||
|
|
||
| setMappingFile(c) | ||
|
|
||
| # Grid name (if not explicitly defined) | ||
| # 'native' if no remapping | ||
| # or extracted from mapping filename | ||
| if c["grid"] == "": | ||
| if c["mapping_file"] == "": | ||
| c["grid"] = "native" | ||
| elif c["mapping_file"] == "glb": | ||
| c["grid"] = "glb" | ||
| else: | ||
| tmp = os.path.basename(c["mapping_file"]) | ||
| # FIXME: W605 invalid escape sequence '\.' | ||
| tmp = re.sub("\.[^.]*\.nc$", "", tmp) # noqa: W605 | ||
| tmp = tmp.split("_") | ||
| if tmp[0] == "map": | ||
| c["grid"] = "%s_%s" % (tmp[-2], tmp[-1]) | ||
| else: | ||
| raise ValueError( | ||
| "Cannot extract target grid name from mapping file %s" | ||
| % (c["mapping_file"]) | ||
| ) | ||
|
|
||
| # Component | ||
| c["component"] = getComponent(c["input_files"]) | ||
|
|
||
| c["cmor_tables_prefix"] = c["diagnostics_base_path"] | ||
|
|
||
| # Loop over year sets | ||
| year_sets = getYears(c["years"]) | ||
| for s in year_sets: | ||
|
|
||
| c["yr_start"] = s[0] | ||
| c["yr_end"] = s[1] | ||
| if ("last_year" in c.keys()) and (c["yr_end"] > c["last_year"]): | ||
| continue # Skip this year set | ||
| c["ypf"] = s[1] - s[0] + 1 | ||
| c["scriptDir"] = scriptDir | ||
| if c["subsection"]: | ||
| sub = c["subsection"] | ||
| else: | ||
| sub = c["grid"] | ||
| prefix = "ts_%s_%04d-%04d-%04d" % ( | ||
| sub, | ||
| c["yr_start"], | ||
| c["yr_end"], | ||
| c["ypf"], | ||
| ) | ||
| print(prefix) | ||
| c["prefix"] = prefix | ||
| scriptFile = os.path.join(scriptDir, "%s.bash" % (prefix)) | ||
| statusFile = os.path.join(scriptDir, "%s.status" % (prefix)) | ||
| settingsFile = os.path.join(scriptDir, "%s.settings" % (prefix)) | ||
| skip = checkStatus(statusFile) | ||
| if skip: | ||
| continue | ||
|
|
||
| # Create script | ||
| with open(scriptFile, "w") as f: | ||
| f.write(template.render(**c)) | ||
| makeExecutable(scriptFile) | ||
|
|
||
| with open(settingsFile, "w") as sf: | ||
| p = pprint.PrettyPrinter(indent=2, stream=sf) | ||
| p.pprint(c) | ||
| p.pprint(s) | ||
|
|
||
| export = "ALL" | ||
| existing_bundles = handle_bundles( | ||
| c, scriptFile, export, existing_bundles=existing_bundles | ||
| ) | ||
| if not c["dry_run"]: | ||
| if c["bundle"] == "": | ||
| # Submit job | ||
| submitScript(scriptFile, statusFile, export, job_ids_file) | ||
| else: | ||
| print("...adding to bundle '%s'" % (c["bundle"])) | ||
|
|
||
| return existing_bundles |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,109 @@ | ||
| #!/bin/bash | ||
| {% include 'slurm_header.sh' %} | ||
| {{ environment_commands }} | ||
|
|
||
| # Turn on debug output if needed | ||
| debug={{ debug }} | ||
| if [[ "${debug,,}" == "true" ]]; then | ||
| set -x | ||
| fi | ||
|
|
||
| # Script dir | ||
| cd {{ scriptDir }} | ||
|
|
||
| # Get jobid | ||
| id=${SLURM_JOBID} | ||
|
|
||
| # Update status file | ||
| STARTTIME=$(date +%s) | ||
| echo "RUNNING ${id}" > {{ prefix }}.status | ||
|
|
||
| # Create temporary workdir | ||
| workdir=`mktemp -d tmp.${id}.XXXX` | ||
| cd ${workdir} | ||
|
|
||
| # Create symbolic links to input files | ||
| input={{ input }}/{{ input_subdir }} | ||
| for (( year={{ yr_start }}; year<={{ yr_end }}; year++ )) | ||
| do | ||
| YYYY=`printf "%04d" ${year}` | ||
| for file in ${input}/{{ case }}.{{ input_files }}.${YYYY}-*.nc | ||
| do | ||
| ln -s ${file} . | ||
| done | ||
| done | ||
|
|
||
|
|
||
|
|
||
|
|
||
|
|
||
| tmp_dir=tmp_{{ prefix }} | ||
|
|
||
| # Generate CMIP ts | ||
| cat > default_metadata.json << EOF | ||
| {% include cmip_metadata %} | ||
| EOF | ||
| { | ||
| export cmortables_dir={{ cmor_tables_prefix }}/cmip6-cmor-tables/Tables | ||
| input_dir={{ output }}/post/{{ component }}/{{ grid }}/ts/{{ frequency }}/{{ '%dyr' % (ypf) }} | ||
| dest_cmip={{ output }}/post/{{ component }}/{{ grid }}/cmip_ts/{{ frequency }} | ||
| mkdir -p ${dest_cmip} | ||
| srun -N 1 e3sm_to_cmip \ | ||
| --output-path \ | ||
| ${dest_cmip}/${tmp_dir} \ | ||
| {% if input_files.split(".")[0] == 'clm2' or input_files.split(".")[0] == 'elm' -%} | ||
| --var-list \ | ||
| 'mrsos, mrso, mrfso, mrros, mrro, prveg, evspsblveg, evspsblsoi, tran, tsl, lai, cLitter, cProduct, cSoilFast, cSoilMedium, cSoilSlow, fFire, fHarvest, cVeg, nbp, gpp, ra, rh' \ | ||
| --realm \ | ||
| lnd \ | ||
| {% endif -%} | ||
| {% if input_files.split(".")[0] == 'cam' or input_files.split(".")[0] == 'eam' -%} | ||
| --var-list \ | ||
| 'pr, tas, rsds, rlds, rsus' \ | ||
| --realm \ | ||
| atm \ | ||
| {% endif -%} | ||
| --input-path \ | ||
| ${input_dir}\ | ||
| --user-metadata \ | ||
| {{ scriptDir }}/${workdir}/default_metadata.json \ | ||
| --num-proc \ | ||
| 12 \ | ||
| --tables-path \ | ||
| ${cmortables_dir} | ||
|
|
||
| if [ $? != 0 ]; then | ||
| cd {{ scriptDir }} | ||
| echo 'ERROR (1)' > {{ prefix }}.status | ||
| exit 1 | ||
| fi | ||
|
|
||
| # Move output ts files to final destination | ||
| mv ${dest_cmip}/${tmp_dir}/CMIP6/CMIP/*/*/*/*/*/*/*/*/*.nc ${dest_cmip} | ||
| if [ $? != 0 ]; then | ||
| cd {{ scriptDir }} | ||
| echo 'ERROR (2)' > {{ prefix }}.status | ||
| exit 2 | ||
| fi | ||
|
|
||
| rm -r ${dest_cmip}/${tmp_dir} | ||
|
|
||
| } | ||
| if [ $? != 0 ]; then | ||
| cd {{ scriptDir }} | ||
| echo 'ERROR (3)' > {{ prefix }}.status | ||
| exit 3 | ||
| fi | ||
|
|
||
|
|
||
| # Update status file and exit | ||
| {% raw %} | ||
| ENDTIME=$(date +%s) | ||
| ELAPSEDTIME=$(($ENDTIME - $STARTTIME)) | ||
| {% endraw %} | ||
| echo ============================================== | ||
| echo "Elapsed time: $ELAPSEDTIME seconds" | ||
| echo ============================================== | ||
| rm -f {{ prefix }}.status | ||
| echo 'OK' > {{ prefix }}.status | ||
| exit 0 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters