diff --git a/create_ensemble_namelists.py b/create_ensemble_namelists.py new file mode 100644 index 0000000..50e295e --- /dev/null +++ b/create_ensemble_namelists.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Create per-ensemble-member namelist files for eCLM-PDAF runs. + +For each ensemble member 1..N, writes ensemble-specific copies of: + - _modelio.nml_NNNN (atm, esp, glc, ice, lnd, ocn, rof, wav, cpl) + - datm_in_NNNN + - lnd_in_NNNN + - mosart_in_NNNN + - datm.streams_NNNN.* (XML stream description files) + +Usage +----- + python create_ensemble_namelists.py [OPTIONS] + +Options +------- + -r, --rundir DIR Run directory containing the base namelists (default: .) + -n, --num_ensemble INT Number of ensemble members (default: 96) + -b, --backend BACKEND Namelist backend: 're' (default) or 'f90nml' + --suffix-fsurdat Add ensemble suffix to clm_inparm:fsurdat in lnd_in +""" +import argparse +import os +import sys + +import f90nml +from lxml import etree +from io import BytesIO + +from namelist_utils import _re_add_ens_suffix, _re_get_streams_filenames + + +def adapt_modelio(mod, iens, backend="re"): + """Write _modelio.nml_NNNN from _modelio.nml. + + Changes: logfile gets an ensemble suffix (_NNNN) inserted before the + file extension, e.g. 'lnd.log' -> 'lnd_0003.log'. + """ + if backend == "f90nml": + nml = f90nml.read(mod + "_modelio.nml") + nml.indent = 1 + + nml["modelio"]["logfile"] = nml["modelio"]["logfile"].replace( + ".log", "_" + str(iens).zfill(4) + ".log") + + nml.write(mod + "_modelio.nml_" + str(iens).zfill(4)) + + elif backend == "re": + with open(mod + "_modelio.nml", "r") as fh: + content = fh.read() + + content = _re_add_ens_suffix(content, "logfile", iens) + + with open(mod + "_modelio.nml_" + str(iens).zfill(4), "w") as fh: + fh.write(content) + + +def adapt_mosart_in(iens, backend="re"): + """Write mosart_in_NNNN from mosart_in (unchanged copy).""" + if backend == "f90nml": + nml = f90nml.read("mosart_in") + nml.indent = 1 + nml.write("mosart_in_" + str(iens).zfill(4)) + + elif backend == "re": + with open("mosart_in", "r") as fh: + content = fh.read() + with open("mosart_in_" + str(iens).zfill(4), "w") as fh: + fh.write(content) + + +def adapt_datm_in(iens, backend="re"): + """Write datm_in_NNNN from datm_in. + + Changes: + + Each entry in shr_strdata_nml:streams gets an ensemble suffix + inserted before the first '.txt' in the filename, e.g. + 'datm.streams.txt.CLMCRUNCEPv7.Solar' -> + 'datm.streams_0003.txt.CLMCRUNCEPv7.Solar'. + + """ + if backend == "f90nml": + nml = f90nml.read("datm_in") + nml.indent = 1 + + for i, streamfile in enumerate(nml["shr_strdata_nml"]["streams"]): + nml["shr_strdata_nml"]["streams"][i] = streamfile.replace( + ".txt", "_" + str(iens).zfill(4) + ".txt") + + nml.write("datm_in_" + str(iens).zfill(4)) + + elif backend == "re": + with open("datm_in", "r") as fh: + content = fh.read() + + for sf in _re_get_streams_filenames("datm_in"): + ens_sf = sf.replace(".txt", "_" + str(iens).zfill(4) + ".txt", 1) + content = content.replace(sf, ens_sf) + + with open("datm_in_" + str(iens).zfill(4), "w") as fh: + fh.write(content) + + +def adapt_lnd_in(iens, suffix_fsurdat=False, backend="re"): + """Write lnd_in_NNNN from lnd_in. + + If suffix_fsurdat is True, clm_inparm:fsurdat gets an ensemble suffix + inserted before the file extension, e.g. 'surfdata.nc' -> 'surfdata_00003.nc' + (zero-padded to 5 digits). Otherwise lnd_in_NNNN is an unchanged copy. + """ + if backend == "f90nml": + nml = f90nml.read("lnd_in") + nml.indent = 1 + + if suffix_fsurdat: + nml["clm_inparm"]["fsurdat"] = nml["clm_inparm"]["fsurdat"].replace( + ".nc", "_" + str(iens).zfill(5) + ".nc") + + nml.write("lnd_in_" + str(iens).zfill(4)) + + elif backend == "re": + with open("lnd_in", "r") as fh: + content = fh.read() + if suffix_fsurdat: + content = _re_add_ens_suffix(content, "fsurdat", iens, pad=5) + with open("lnd_in_" + str(iens).zfill(4), "w") as fh: + fh.write(content) + + +def adapt_stream_files(iens, backend="re"): + """Write per-ensemble XML stream description files. + + Must be called after adapt_datm_in, which determines the ensemble + stream filenames. + + For each stream file listed in datm_in, reads the original XML and + writes an ensemble-specific copy (named as in datm_in_NNNN), with + the fieldInfo/filePath entries under './forcings' redirected to + './forcings/real_NNNNN'. + + """ + # Must be called after datm update + + if backend == "f90nml": + # Read original stream files and new stream files + nml_datm = f90nml.read("datm_in") + nml_datm_ens = f90nml.read("datm_in_" + str(iens).zfill(4)) + + orig_streamfiles = nml_datm["shr_strdata_nml"]["streams"] + ens_streamfiles = nml_datm_ens["shr_strdata_nml"]["streams"] + + elif backend == "re": + orig_streamfiles = _re_get_streams_filenames("datm_in") + ens_streamfiles = _re_get_streams_filenames("datm_in_" + str(iens).zfill(4)) + + for orig, ens_sf in zip(orig_streamfiles, ens_streamfiles): + orig_path = orig.split()[0] if backend == "f90nml" else orig + ens_path = ens_sf.split()[0] if backend == "f90nml" else ens_sf + + tree = etree.parse(orig_path) + root = tree.getroot() + + # Extract element "filePath" with parent element "fieldInfo" + for filePath in root.iter("filePath"): + # Check parent + if filePath.getparent().tag == "fieldInfo": + # print(filePath.getparent().tag, filePath.tag) + if filePath.text.find("./forcings") > -1: + filePath.text = filePath.text.replace( + "./forcings", "./forcings/real_" + str(iens).zfill(5) + ) + + # Write XML streamfile + # -------------------- + + # Use buffer first + fbuffer = BytesIO() + tree.write(fbuffer, xml_declaration=True, encoding="ASCII") + fstr = fbuffer.getvalue().decode("ASCII") + + # Replace XML declaration to match original + fstr = fstr.replace("version='1.0'", 'version="1.0"') + fstr = fstr.replace(" encoding='ASCII'", "") + + # Write to file + with open(ens_path, "w+") as f: + f.write(fstr) + + +def create_ensemble_namelists(iens, suffix_fsurdat=False, backend="re"): + """Create all per-ensemble namelist files for ensemble member iens. + + Calls: + + - adapt_modelio for all nine CIME components (atm, esp, glc, ice, + lnd, ocn, rof, wav, cpl) + - adapt_datm_in + - adapt_lnd_in + - adapt_mosart_in + - adapt_stream_files + """ + for m in ["atm", "esp", "glc", "ice", "lnd", "ocn", "rof", "wav", "cpl"]: + adapt_modelio(m, iens, backend=backend) + + adapt_datm_in(iens, backend=backend) + + adapt_lnd_in(iens, suffix_fsurdat=suffix_fsurdat, backend=backend) + + adapt_mosart_in(iens, backend=backend) + + adapt_stream_files(iens, backend=backend) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Create ensemble namelists for eCLM-PDAF runs" + ) + parser.add_argument( + "-r", "--rundir", + default=".", + help="Run directory containing the base namelists (default: current directory)", + ) + parser.add_argument( + "-n", "--num_ensemble", + type=int, + default=96, + help="Number of ensemble members (default: 96)", + ) + parser.add_argument( + "-b", "--backend", + choices=["f90nml", "re"], + default="re", + help="Backend for reading/writing namelists: 're' (default, regexp) or 'f90nml'", + ) + parser.add_argument( + "--suffix-fsurdat", + action="store_true", + default=False, + help="Add ensemble suffix to clm_inparm:fsurdat in lnd_in (default: off)", + ) + + args = parser.parse_args() + + os.chdir(args.rundir) + + for ens in range(1, args.num_ensemble + 1): + create_ensemble_namelists(ens, suffix_fsurdat=args.suffix_fsurdat, backend=args.backend) + sys.stdout.write("\r[%s] " % ("Done with ensemble member: " + str(ens))) + sys.stdout.flush() + + sys.stdout.write("\n") + sys.stdout.flush() diff --git a/docs/_toc.yml b/docs/_toc.yml index 227f113..344967d 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -13,3 +13,4 @@ parts: - file: users_guide/normalize_stream_files - file: users_guide/modify_stream_files - file: users_guide/modify_case_namelists + - file: users_guide/create_ensemble_namelists diff --git a/docs/users_guide/README.md b/docs/users_guide/README.md index 98e137c..5a30838 100644 --- a/docs/users_guide/README.md +++ b/docs/users_guide/README.md @@ -15,3 +15,7 @@ generation and manipulation. - [modify_case_namelists.py](modify_case_namelists.md): Modifies individual key-value entries across eCLM case namelist files. + +- [create_ensemble_namelists.py](create_ensemble_namelists.md): + Creates per-ensemble-member namelist files for eCLM-PDAF Data + Assimilation experiments. diff --git a/docs/users_guide/create_ensemble_namelists.md b/docs/users_guide/create_ensemble_namelists.md new file mode 100644 index 0000000..c013bbd --- /dev/null +++ b/docs/users_guide/create_ensemble_namelists.md @@ -0,0 +1,94 @@ +# create_ensemble_namelists.py + +> **Data Assimilation only**: this script is exclusively needed for +> eCLM-PDAF experiments. It produces the per-member namelist files +> that PDAF expects to find for each ensemble member. It is not +> required for standard (single-instance) eCLM runs. + +Creates per-ensemble-member copies of all eCLM namelist and stream +files from a single set of base namelists. + +## What it does + +Starting from the base namelist files in the run directory, the script +writes one set of ensemble-specific files for every member `1 .. N`. +Each output filename carries a four-digit zero-padded member index as +a suffix (`_NNNN`). + +The following modifications are applied per file type: + +| Output file | Source | Modification | +|--------------------------|---------------------|---------------------------------------------------------------------------------------| +| `{mod}_modelio.nml_NNNN` | `{mod}_modelio.nml` | `logfile` gets `_NNNN` inserted before `.log` | +| `datm_in_NNNN` | `datm_in` | Each stream filename in `shr_strdata_nml:streams` gets `_NNNN` inserted before `.txt` | +| `lnd_in_NNNN` | `lnd_in` | Unchanged copy; optionally `clm_inparm:fsurdat` gets a five-digit suffix | +| `mosart_in_NNNN` | `mosart_in` | Unchanged copy | +| `datm.streams_NNNN.*` | `datm.streams.*` | `fieldInfo/filePath` is redirected from `./forcings` to `./forcings/real_NNNNN` | + +The `{mod}` components processed for modelio files are: +`atm`, `esp`, `glc`, `ice`, `lnd`, `ocn`, `rof`, `wav`, `cpl`. + +## Usage + +```bash +# Create namelists for 96 ensemble members (default) in the current directory +python3 create_ensemble_namelists.py + +# Specify run directory and ensemble size +python3 create_ensemble_namelists.py \ + --rundir /path/to/rundir \ + --num_ensemble 50 + +# Also add per-member suffix to the surface data file path +python3 create_ensemble_namelists.py \ + --rundir /path/to/rundir \ + --num_ensemble 50 \ + --suffix-fsurdat +``` + +### Options + +| Option | Description | +|------------------------|------------------------------------------------------------------------------| +| `-r`, `--rundir` | Directory containing the base namelist files. Defaults to `.` | +| `-n`, `--num_ensemble` | Number of ensemble members to generate. Defaults to `96` | +| `-b`, `--backend` | Namelist backend: `re` (default, regexp-based) or `f90nml` | +| `--suffix-fsurdat` | Also add a five-digit member suffix to `clm_inparm:fsurdat` in `lnd_in_NNNN` | + +## Expected run directory layout + +The script reads the following base files from `--rundir` and must be +called **after** the base namelists have been set up (e.g. with +`modify_case_namelists.py`): + +``` +rundir/ + atm_modelio.nml cpl_modelio.nml esp_modelio.nml + glc_modelio.nml ice_modelio.nml lnd_modelio.nml + ocn_modelio.nml rof_modelio.nml wav_modelio.nml + datm_in + lnd_in + mosart_in + datm.streams.txt.* (one file per DATM stream) +``` + +## Output + +The script prints a running status line and exits after all members are +written: + +``` +[Done with ensemble member: 50] +``` + +Output files are written into the same directory as the base namelists: + +``` +rundir/ + lnd_modelio.nml_0001 lnd_modelio.nml_0002 ... lnd_modelio.nml_0050 + datm_in_0001 datm_in_0002 ... datm_in_0050 + lnd_in_0001 lnd_in_0002 ... lnd_in_0050 + mosart_in_0001 mosart_in_0002 ... mosart_in_0050 + datm.streams_0001.txt.* datm.streams_0002.txt.* ... + ... +```