From 13d16384be7e2289b3d09e6016b92fe21511cae2 Mon Sep 17 00:00:00 2001 From: "W. Andre Perkins" Date: Tue, 14 Sep 2021 16:00:11 -0700 Subject: [PATCH] Online Zhao-Carr microphysics emulation (#218) This PR adds the online emulation hooks for the ZC microphysics scheme and a barebones module to test it. Downstream users should override the emulation package with two top-level functions, store and microphysics. During runtime, two new namelist parameters under gfs_physics_nml are added which can be used to turn on/off the emulation components: emulate_zc_microphysics: calls emulation.microphysics save_zc_microphysics: calls emulation.store * Add gcsfs to requirements * Add prognostic hooks and module * Add emulation namelist parameters * Add conditionals to physics driver * Add specific configuration for training creation vs emulation * remove unused global variable * Add namelist flags for emulation * emulation.emulate debugging * Ignore egg-info * Add dummy model for testing * Fixed tensorflow import for emulation hooks * Change data saving namelist param name * Emulation integration test finished * Fix dummy model to correct vertical levels * Remove storage option from emulator, use monitor instead * Undo accidental import shift * Move emulation back into fv3net, remove from here * Reduce to basic call_py_fort interface for tests * Change to basic interface for testing call_py_fort usage * Update test config to run both call_py_fort functions * Update integration test * Fix integration test * Cleanup * Update emulation/README.md Co-authored-by: Noah D. Brenowitz * Remove subprocess directory removal from callpyfort integration * Move emulation package to tests subdirectory * Update readme with more information about emulation * Add gotcha notes * Update language for emulation outputs Co-authored-by: Noah D. Brenowitz --- .gitignore | 1 + .../GFS_layer/GFS_physics_driver.F90 | 23 +- FV3/gfsphysics/GFS_layer/GFS_typedefs.F90 | 16 +- docker/Dockerfile | 2 +- emulation/README.md | 27 -- emulation/emulation/__init__.py | 0 emulation/emulation/_filesystem.py | 43 --- emulation/emulation/debug.py | 53 --- emulation/emulation/monitor.py | 220 ------------ .../microphysics_parameter_metadata.yaml | 320 ------------------ tests/emulation/README.md | 44 +++ tests/emulation/emulation/__init__.py | 2 + tests/emulation/emulation/emulate.py | 15 + tests/emulation/emulation/monitor.py | 15 + {emulation => tests/emulation}/setup.py | 13 +- tests/pytest/config/emulation.yml | 2 + tests/pytest/test_regression.py | 18 +- 17 files changed, 128 insertions(+), 686 deletions(-) delete mode 100644 emulation/README.md delete mode 100644 emulation/emulation/__init__.py delete mode 100644 emulation/emulation/_filesystem.py delete mode 100644 emulation/emulation/debug.py delete mode 100644 emulation/emulation/monitor.py delete mode 100644 emulation/microphysics_parameter_metadata.yaml create mode 100644 tests/emulation/README.md create mode 100644 tests/emulation/emulation/__init__.py create mode 100644 tests/emulation/emulation/emulate.py create mode 100644 tests/emulation/emulation/monitor.py rename {emulation => tests/emulation}/setup.py (82%) diff --git a/.gitignore b/.gitignore index 4ee9f62c3..74f6ecab4 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,4 @@ depend **/.vscode fv3.exe *.tmp.f90 +**/*.egg-info diff --git a/FV3/gfsphysics/GFS_layer/GFS_physics_driver.F90 b/FV3/gfsphysics/GFS_layer/GFS_physics_driver.F90 index 90c143ee9..0620dccf8 100644 --- a/FV3/gfsphysics/GFS_layer/GFS_physics_driver.F90 +++ b/FV3/gfsphysics/GFS_layer/GFS_physics_driver.F90 @@ -4566,8 +4566,27 @@ subroutine GFS_physics_driver & call set_state("total_precipitation", rain1) call set_state("ratio_of_snowfall_to_rainfall", Diag%sr) call set_state("tendency_of_rain_water_mixing_ratio_due_to_microphysics", rainp) - call call_function("emulation.monitor", "store_zarr") - call call_function("emulation.monitor", "store_netcdf") + + if (Model%emulate_zc_microphysics) then + ! apply microphysics emulator + call call_function("emulation", "microphysics") + endif + + if (Model%save_zc_microphysics) then + call call_function("emulation", "store") + endif + + call get_state("air_temperature_output", Stateout%gt0) + call get_state("specific_humidity_output", qv_post_precpd) + call get_state("cloud_water_mixing_ratio_output", qc_post_precpd) + call get_state("total_precipitation", rain1) + + do k=1,levs + do i=1,im + Stateout%gq0(i,k,1) = qv_post_precpd(i,k) + Stateout%gq0(i,k,ntcw) = qc_post_precpd(i,k) + enddo + enddo #endif endif diff --git a/FV3/gfsphysics/GFS_layer/GFS_typedefs.F90 b/FV3/gfsphysics/GFS_layer/GFS_typedefs.F90 index 54e6c10ce..08c46f966 100644 --- a/FV3/gfsphysics/GFS_layer/GFS_typedefs.F90 +++ b/FV3/gfsphysics/GFS_layer/GFS_typedefs.F90 @@ -1085,6 +1085,8 @@ module GFS_typedefs real(kind=kind_phys) :: sst_perturbation ! Sea surface temperature perturbation to climatology or nudging SST (default 0.0 K) logical :: override_surface_radiative_fluxes ! Whether to use Statein to override the surface radiative fluxes logical :: use_climatological_sst ! Whether to allow the Python wrapper to override the sea surface temperature + logical :: emulate_zc_microphysics ! Use an emulator in place of ZC microphysics + logical :: save_zc_microphysics ! Save ZC microphysics state #ifdef CCPP ! From physcons.F90, updated/set in control_initialize real(kind=kind_phys) :: dxinv ! inverse scaling factor for critical relative humidity, replaces dxinv in physcons.F90 @@ -3130,6 +3132,8 @@ subroutine control_initialize (Model, nlunit, fn_nml, me, master, & real(kind=kind_phys) :: sst_perturbation = 0.0 ! Sea surface temperature perturbation [K] logical :: override_surface_radiative_fluxes = .false. logical :: use_climatological_sst = .true. + logical :: emulate_zc_microphysics = .false. + logical :: save_zc_microphysics = .false. !--- END NAMELIST VARIABLES NAMELIST /gfs_physics_nml/ & @@ -3221,7 +3225,8 @@ subroutine control_initialize (Model, nlunit, fn_nml, me, master, & !--- aerosol scavenging factors ('name:value' string array) fscav_aero, & sst_perturbation, & - override_surface_radiative_fluxes, use_climatological_sst + override_surface_radiative_fluxes, use_climatological_sst, & + emulate_zc_microphysics, save_zc_microphysics !--- other parameters integer :: nctp = 0 !< number of cloud types in CS scheme @@ -3691,6 +3696,11 @@ subroutine control_initialize (Model, nlunit, fn_nml, me, master, & Model%sst_perturbation = sst_perturbation Model%override_surface_radiative_fluxes = override_surface_radiative_fluxes Model%use_climatological_sst = use_climatological_sst + + !--- emulation parameters + Model%emulate_zc_microphysics = emulate_zc_microphysics + Model%save_zc_microphysics = save_zc_microphysics + !--- tracer handling Model%ntrac = size(tracer_names) #ifdef CCPP @@ -4462,6 +4472,10 @@ subroutine control_print(Model) print *, ' evpco : ', Model%evpco print *, ' wminco : ', Model%wminco print *, ' ' + print *, ' Z-C Emulation parameters' + print *, ' use ZC emulator : ', Model%emulate_zc_microphysics + print *, ' save ZC state : ', Model%save_zc_microphysics + print *, ' ' endif if (Model%imp_physics == Model%imp_physics_wsm6 .or. Model%imp_physics == Model%imp_physics_thompson) then print *, ' Thompson microphysical parameters' diff --git a/docker/Dockerfile b/docker/Dockerfile index 6fe6b3d78..2571e929c 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -136,7 +136,7 @@ RUN cd ${CALLPY} && make && make install && ldconfig RUN python3 -m pip install git+https://github.com/VulcanClimateModeling/fv3gfs-util.git@1d7c302b836befe905d776b0a972f464bfd3a255 # copy emulation target -COPY emulation/ /emulation +COPY tests/emulation/ /emulation RUN cd /emulation && python3 -m pip install . ENV VAR_META_PATH=/emulation/microphysics_parameter_metadata.yaml \ diff --git a/emulation/README.md b/emulation/README.md deleted file mode 100644 index 441e793a1..000000000 --- a/emulation/README.md +++ /dev/null @@ -1,27 +0,0 @@ -emulation -========= - -This is a stripped down set of modules for adding into a prognostic run using `call_py_fort`. It's currently used to create training datasets directly from data from the microphysics parameterization. - -The interaction points with python are located in `GFS_physics_driver.f90` as the `set_state` or `get_state` commands. These are enabled for the emulation image (compiled using `make build_emulation`) when `-DENABLE_CALLPYFORT` is passed to the compiler. - -### Example snippet - -``` -#ifdef ENABLE_CALLPYFORT - do k=1,levs - do i=1,im - qv_post_precpd(i,k) = Stateout%gq0(i,k,1) - qc_post_precpd(i,k) = Stateout%gq0(i,k,ntcw) - enddo - enddo - - call set_state("air_temperature_output", Stateout%gt0) - call set_state("specific_humidity_output", qv_post_precpd) - call set_state("cloud_water_mixing_ratio_output", qc_post_precpd) -``` - -## Training Data -By default, the training data is saved out to the current working directory with a zarr monitor to state_output.zarr (time, tile, x, y), or individual netCDF files for each time and tile under $(cwd)/netcdf_output. - -To change the frequency for which data are saved (defaults to 5 hours [18,000 s]), prescribe the `OUTPUT_FREQ_SEC` environment variable in the runtime image. diff --git a/emulation/emulation/__init__.py b/emulation/emulation/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/emulation/emulation/_filesystem.py b/emulation/emulation/_filesystem.py deleted file mode 100644 index 87f74f8aa..000000000 --- a/emulation/emulation/_filesystem.py +++ /dev/null @@ -1,43 +0,0 @@ -# keras has several routines which interact with file paths directly as opposed to -# filesystem objects, which means we need these wrappers so we can allow remote paths - -import contextlib -import tempfile -import fsspec -import os - - -@contextlib.contextmanager -def put_dir(path: str): - with tempfile.TemporaryDirectory() as tmpdir: - yield tmpdir - fs, _, _ = fsspec.get_fs_token_paths(path) - fs.makedirs(os.path.dirname(path), exist_ok=True) - # cannot use fs.put as it cannot merge directories - _put_directory(tmpdir, path) - - -@contextlib.contextmanager -def get_dir(path: str): - with tempfile.TemporaryDirectory() as tmpdir: - fs, _, _ = fsspec.get_fs_token_paths(path) - # fsspec places the directory inside the tmpdir, as a subdirectory - fs.get(path, tmpdir, recursive=True) - yield tmpdir - - -def _put_directory( - local_source_dir: str, dest_dir: str, fs: fsspec.AbstractFileSystem = None, -): - """Copy the contents of a local directory to a local or remote directory. - """ - if fs is None: - fs, _, _ = fsspec.get_fs_token_paths(dest_dir) - fs.makedirs(dest_dir, exist_ok=True) - for token in os.listdir(local_source_dir): - source = os.path.join(os.path.abspath(local_source_dir), token) - dest = os.path.join(dest_dir, token) - if os.path.isdir(source): - _put_directory(source, dest, fs=fs) - else: - fs.put(source, dest) diff --git a/emulation/emulation/debug.py b/emulation/emulation/debug.py deleted file mode 100644 index 6550304d9..000000000 --- a/emulation/emulation/debug.py +++ /dev/null @@ -1,53 +0,0 @@ -""" -A module for testing/debugging call routines -""" -import logging -import os -import traceback -import numpy as np -from datetime import datetime - -logger = logging.getLogger(__name__) - - -def print_errors(func): - def new_func(*args, **kwargs): - try: - return func(*args, **kwargs) - except Exception as e: - logger.error(traceback.print_exc()) - raise e - - return new_func - - -def print_arr_info(state): - - logger = logging.getLogger(__name__) - for varname, arr in state.items(): - logger.info(f"{varname}: shape[{arr.shape}] isfortran[{np.isfortran(arr)}]") - - -def print_location_ping(state): - - logger = logging.getLogger(__name__) - logger.info("Ping reached!") - - -def dump_state(state): - - DUMP_PATH = str(os.environ.get("STATE_DUMP_PATH")) - - logger = logging.getLogger(__name__) - - try: - rank = state.get("rank") - except KeyError: - logger.info("Could not save state. No rank included in state.") - return - - time_str = datetime.now().strftime("%Y%m%d.%H%M%S") - filename = f"state_dump.{time_str}.tile{int(rank.squeeze()[0])}.npz" - outfile = os.path.join(DUMP_PATH, filename) - logger.info(f"Dumping state to {outfile}") - np.savez(outfile, **state) diff --git a/emulation/emulation/monitor.py b/emulation/emulation/monitor.py deleted file mode 100644 index b687c1937..000000000 --- a/emulation/emulation/monitor.py +++ /dev/null @@ -1,220 +0,0 @@ -import logging -import os -import json -import cftime -import f90nml -import yaml -import numpy as np -import xarray as xr -from datetime import timedelta -from mpi4py import MPI -from pathlib import Path - -from fv3gfs.util import ZarrMonitor, CubedSpherePartitioner, Quantity -from .debug import print_errors - - -logger = logging.getLogger(__name__) - -TIME_FMT = "%Y%m%d.%H%M%S" -DIMS_MAP = { - 1: ["sample"], - 2: ["sample", "z"], -} - -# Initialized later from within functions to print errors -VAR_META_PATH = None -NML_PATH = None -DUMP_PATH = None -NC_DUMP_PATH = None -OUTPUT_FREQ_SEC = None -DT_SEC = None -INITIAL_TIME = None - - -@print_errors -def _load_environment_vars_into_global(): - global VAR_META_PATH - global NML_PATH - global DUMP_PATH - global NC_DUMP_PATH - global OUTPUT_FREQ_SEC - - cwd = os.getcwd() - logger.debug(f"Current working directory: {cwd}") - - VAR_META_PATH = os.environ["VAR_META_PATH"] - NML_PATH = os.path.join(cwd, "input.nml") - DUMP_PATH = cwd - NC_DUMP_PATH = os.path.join(cwd, "netcdf_output") - OUTPUT_FREQ_SEC = int(os.environ["OUTPUT_FREQ_SEC"]) - - -@print_errors -def _load_nml(): - namelist = f90nml.read(NML_PATH) - logger.info(f"Loaded namelist for ZarrMonitor from {NML_PATH}") - - global DT_SEC - DT_SEC = int(namelist["coupler_nml"]["dt_atmos"]) - - return namelist - - -@print_errors -def _load_metadata(): - try: - with open(str(VAR_META_PATH), "r") as f: - variable_metadata = yaml.safe_load(f) - logger.info(f"Loaded variable metadata from: {VAR_META_PATH}") - except FileNotFoundError: - variable_metadata = {} - logger.info(f"No metadata found at: {VAR_META_PATH}") - return variable_metadata - - -@print_errors -def _load_monitor(namelist): - partitioner = CubedSpherePartitioner.from_namelist(namelist) - output_zarr = os.path.join(str(DUMP_PATH), "state_output.zarr") - output_monitor = ZarrMonitor( - output_zarr, - partitioner, - mpi_comm=MPI.COMM_WORLD - ) - logger.info(f"Initialized zarr monitor at: {output_zarr}") - return output_monitor - - -@print_errors -def _make_output_paths(): - zarr_path = Path(DUMP_PATH) - zarr_path.mkdir(exist_ok=True) - - netcdf_path = Path(NC_DUMP_PATH) - netcdf_path.mkdir(exist_ok=True) - - -_load_environment_vars_into_global() -_namelist = _load_nml() -_variable_metadata = _load_metadata() -_make_output_paths() -_output_monitor = _load_monitor(_namelist) - - -def print_rank(state): - logger.info(MPI.COMM_WORLD.Get_rank()) - - -def _remove_io_suffix(key: str): - if key.endswith("_input"): - var_key = key[:-6] - logger.debug(f"Removed _input with result {var_key} for metadata mapping") - elif key.endswith("_output"): - var_key = key[:-7] - logger.debug(f"Removed _output with result {var_key} for metadata mapping") - else: - var_key = key - - return var_key - - -def _get_attrs(key: str): - key = _remove_io_suffix(key) - if key in _variable_metadata: - meta = dict(**_variable_metadata[key]) - meta = {k: json.dumps(v) for k,v in meta.items()} - else: - logger.debug(f"No metadata found for {key}... skipping") - meta = {} - - return meta - - -def _convert_to_quantities(state): - - quantities = {} - for key, data in state.items(): - data = np.squeeze(data.astype(np.float32)) - data_t = data.T - dims = DIMS_MAP[data.ndim] - attrs = _get_attrs(key) - units = attrs.pop("units", "unknown") - quantities[key] = Quantity(data_t, dims, units) - # Access to private member could break TODO: Quantity kwarg for attrs? - quantities[key]._attrs.update(attrs) - - return quantities - - -def _convert_to_xr_dataset(state): - - dataset = {} - for key, data in state.items(): - data = np.squeeze(data.astype(np.float32)) - data_t = data.T - dims = DIMS_MAP[data.ndim] - attrs = _get_attrs(key) - attrs["units"] = attrs.pop("units", "unknown") - dataset[key] = xr.DataArray(data_t, dims=dims, attrs=attrs) - - return xr.Dataset(dataset) - - -def _translate_time(time): - year = time[0] - month = time[1] - day = time[2] - hour = time[4] - min = time[5] - datetime = cftime.DatetimeJulian(year, month, day, hour, min) - logger.debug(f"Translated input time: {datetime}") - - return datetime - - -def _store_interval_check(time): - - global INITIAL_TIME - - if INITIAL_TIME is None: - INITIAL_TIME = time - - # add increment since we are in the middle of timestep - increment = timedelta(seconds=DT_SEC) - elapsed = (time + increment) - INITIAL_TIME - - logger.debug(f"Time elapsed after increment: {elapsed}") - logger.debug(f"Output frequency modulus: {elapsed.seconds % OUTPUT_FREQ_SEC}") - - return elapsed.seconds % OUTPUT_FREQ_SEC == 0 - - -@print_errors -def store_netcdf(state): - state = dict(**state) - time = _translate_time(state.pop("model_time")) - - if _store_interval_check(time): - logger.debug(f"Model fields: {list(state.keys())}") - logger.info(f"Storing state to netcdf on rank {MPI.COMM_WORLD.Get_rank()}") - ds = _convert_to_xr_dataset(state) - rank = MPI.COMM_WORLD.Get_rank() - coords = {"time": time, "tile": rank} - ds = ds.assign_coords(coords) - filename = f"state_{time.strftime(TIME_FMT)}_{rank}.nc" - out_path = os.path.join(NC_DUMP_PATH, filename) - ds.to_netcdf(out_path) - - -@print_errors -def store_zarr(state): - state = dict(**state) - time = _translate_time(state.pop("model_time")) - - if _store_interval_check(time): - logger.info(f"Storing model state on rank {MPI.COMM_WORLD.Get_rank()}") - logger.debug(f"Model fields: {list(state.keys())}") - state = _convert_to_quantities(state) - state["time"] = time - _output_monitor.store(state) diff --git a/emulation/microphysics_parameter_metadata.yaml b/emulation/microphysics_parameter_metadata.yaml deleted file mode 100644 index 897bc550f..000000000 --- a/emulation/microphysics_parameter_metadata.yaml +++ /dev/null @@ -1,320 +0,0 @@ -air_pressure: - fortran_name: prsl - standard_name: air_pressure - long_name: layer mean air pressure - units: Pa - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: in - optional: F -pressure_thickness_of_atmospheric_layer: - standard_name: air_pressure_difference_between_midlayers - fortran_name: del - long_name: pressure level thickness - units: Pa - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: in - optional: F -specific_humidity: - fortran_name: q - standard_name: specific_humidity - long_name: water vapor specific humidity - units: kg kg-1 - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: inout - optional: F -cloud_water_mixing_ratio: - fortran_name: cwm - standard_name: cloud_water_mixing_ratio - long_name: moist cloud condensed water mixing ratio - units: kg kg-1 - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: out - optional: F -air_temperature: - fortran_name: t - standard_name: air_temperature - long_name: layer mean air temperature - units: K - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: inout - optional: F -latitude: - fortran_name: xlat - standard_name: latitude - long_name: grid latitude in radians - units: radians - field_dims: (horizontal_dimension) - type: real - kind: kind_phys - intent: in - optional: F -longitude: - fortran_name: xlon - standard_name: longitude - long_name: grid longitude in radians - units: radians - field_dims: (horizontal_dimension) - type: real - kind: kind_phys - intent: in - optional: F - - -# gscond specific - -time_step_for_dynamics: - fortran_name: dtf - standard_name: time_step_for_dynamics - long_name: dynamics time step - units: s - field_dims: [] - type: real - kind: kind_phys - intent: in - optional: F - - -surface_air_pressure: - fortran_name: ps - standard_name: surface_air_pressure - long_name: surface pressure - units: Pa - field_dims: [horizontal_dimension] - type: real - kind: kind_phys - intent: in - optional: F -ice_water_mixing_ratio_convective_transport_tracer: - fortran_name: clw1 - standard_name: ice_water_mixing_ratio_convective_transport_tracer - long_name: moist [dry+vapor, no]] mixing ratio of ice water in the convectively transported tracer array - units: kg kg-1 - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: in - optional: F -cloud_condensed_water_mixing_ratio_convective_transport_tracer: - fortran_name: clw2 - standard_name: cloud_condensed_water_mixing_ratio_convective_transport_tracer - long_name: moist [dry+vapor, no]] mixing ratio of cloud water [condensate] in the convectively transported tracer array - units: kg kg-1 - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: in - optional: F -air_temperature_two_time_steps_back: - fortran_name: tp - standard_name: air_temperature_two_time_steps_back - long_name: air temperature two time steps back - units: K - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: inout - optional: F -specific_humidity_two_time_steps_back: - fortran_name: qp - standard_name: specific_humidity_two_time_steps_back - long_name: water vapor specific humidity two time steps back - units: kg kg-1 - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: inout - optional: F -surface_air_pressure_two_time_steps_back: - fortran_name: psp - standard_name: surface_air_pressure_two_time_steps_back - long_name: surface air pressure two time steps back - units: Pa - field_dims: [horizontal_dimension] - type: real - kind: kind_phys - intent: inout - optional: F -air_temperature_at_previous_time_step: - fortran_name: tp1 - standard_name: air_temperature_at_previous_time_step - long_name: air temperature at previous time step - units: K - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: inout - optional: F -specific_humidity_at_previous_time_step: - fortran_name: qp1 - standard_name: specific_humidity_at_previous_time_step - long_name: water vapor specific humidity at previous time step - units: kg kg-1 - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: inout - optional: F -surface_air_pressure_at_previous_time_step: - fortran_name: psp1 - standard_name: surface_air_pressure_at_previous_time_step - long_name: surface air surface pressure at previous time step - units: Pa - field_dims: [horizontal_dimension] - type: real - kind: kind_phys - intent: inout - optional: F -critical_relative_humidity: - fortran_name: u - standard_name: critical_relative_humidity - long_name: critical relative humidity - units: frac - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: in - optional: F -lprnt: - standard_name: flag_print - long_name: flag for printing diagnostics to output - units: flag - field_dims: [] - type: logical - intent: in - optional: F -ipr: - standard_name: horizontal_index_of_printed_column - long_name: horizontal index of printed column - units: index - field_dims: [] - type: integer - intent: in - optional: F -errmsg: - standard_name: ccpp_error_message - long_name: error message for error handling in CCPP - units: none - field_dims: [] - type: character - kind: le:* - intent: out - optional: F -errflg: - standard_name: ccpp_error_flag - long_name: error flag for error handling in CCPP - units: flag - field_dims: [] - type: integer - intent: out - optional: F - -# precpd specific -total_precipitation: - fortran_name: rn - standard_name: total_precipitation - long_name: explicit precipitation amount on physics timestep - units: m - field_dims: [horizontal_dimension] - type: real - kind: kind_phys - intent: out - optional: F -ratio_of_snowfall_to_rainfall: - fortran_name: sr - standard_name: ratio_of_snowfall_to_rainfall - long_name: ratio of snowfall to large-scale rainfall - units: frac - field_dims: [horizontal_dimension] - type: real - kind: kind_phys - intent: out - optional: F -tendency_of_rain_water_mixing_ratio_due_to_microphysics: - fortran_name: rainp - standard_name: tendency_of_rain_water_mixing_ratio_due_to_microphysics - long_name: tendency of rain water mixing ratio due to microphysics - units: kg kg-1 s-1 - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: out - optional: F -critical_relative_humidity: - fortran_name: u00k - standard_name: critical_relative_humidity - long_name: critical relative humidity - units: frac - field_dims: [horizontal_dimension, vertical_dimension] - type: real - kind: kind_phys - intent: in - optional: F -coefficient_from_cloud_ice_to_snow: - fortran_name: psautco - standard_name: coefficient_from_cloud_ice_to_snow - long_name: conversion coefficient from cloud ice to snow - units: none - field_dims: [2] - type: real - kind: kind_phys - intent: in - optional: F -coefficient_from_cloud_water_to_rain: - fortran_name: prautco - standard_name: coefficient_from_cloud_water_to_rain - long_name: conversion coefficient from cloud water to rain - units: none - field_dims: [2] - type: real - kind: kind_phys - intent: in - optional: F -coefficient_for_evaporation_of_rainfall: - fortran_name: evpco - standard_name: coefficient_for_evaporation_of_rainfall - long_name: coefficient for evaporation of rainfall - units: none - field_dims: [] - type: real - kind: kind_phys - intent: in - optional: F -cloud_condensed_water_conversion_threshold: - fortran_name: wminco - standard_name: cloud_condensed_water_conversion_threshold - long_name: conversion coefficient from cloud liquid and ice to precipitation - units: none - field_dims: [2] - type: real - kind: kind_phys - intent: in - optional: F -grid_size_related_coefficient_used_in_scale_sensitive_schemes: - fortran_name: wk1 - standard_name: grid_size_related_coefficient_used_in_scale_sensitive_schemes - long_name: grid size related coefficient used in scale-sensitive schemes - units: none - field_dims: [horizontal_dimension] - type: real - kind: kind_phys - intent: in - optional: F -jpr: - standard_name: horizontal_index_of_printed_column - long_name: horizontal index of printed column - units: index - field_dims: [] - type: integer - intent: in - optional: F - diff --git a/tests/emulation/README.md b/tests/emulation/README.md new file mode 100644 index 000000000..8f34648d7 --- /dev/null +++ b/tests/emulation/README.md @@ -0,0 +1,44 @@ +emulation +========= + +A basic implementation of `emulation` is provided to test the call_py_fort calls from fv3gfs. Downstream projects should provide their own definitions `emulation` package by defining two functions the two functions `store` and `microphysics`. + +The interaction points with python are located in `GFS_physics_driver.f90` as the `set_state`, `get_state`, and `call_function` commands. Using `make build_emulation` from the parent directory will provide a compiled image with call_py_fort functionality enabled. Two namelist flags can be used to turn off storage (`gfs_physics_nml.save_zc_microphysics`) and emulation (`gfs_physics_nml.emulate_zc_microphysics`) while call_py_fort is enabled. + +The emulation function `microphysics` should overwrite the fields `air_temperature_output`, `specific_humidity_output`, `cloud_water_mixing_ratio_output`, `total_precipitation`, to control the microphysical updates during runtime. Otherwise, the default parameterization will be in control. + +Available input state fields for emulation +------------------------------------------ +- model_time +- latitude +- longitude +- pressure_thickness_of_atmospheric_layer +- air_pressure +- surface_air_pressure +- air_temperature_input +- specific_humidity_input +- cloud_water_mixing_ratio_input +- air_temperature_two_time_steps_back +- specific_humidity_two_time_steps_back +- surface_air_pressure_two_time_steps_back +- air_temperature_at_previous_time_step +- specific_humidity_at_previous_time_step +- surface_air_pressure_at_previous_time_step +- ratio_of_snowfall_to_rainfall +- tendency_of_rain_water_mixing_ratio_due_to_microphysics + +Notes +----- + +The usage of call_py_fort results in some subtle differences with python initialization that can lead to some hard-to-debug issues. + +- the current directory is not added to `PYTHONPATH` +- `sys.argv` is not initialized and will lead to errors for anything that expects to use it (this was an issue for tensorflow). Remedied by instantiating in the module prior to imports that need it. E.g., + +``` +import sys +if not hasattr(sys, "argv"): + sys.argv = [""] + +import tensorflow +``` diff --git a/tests/emulation/emulation/__init__.py b/tests/emulation/emulation/__init__.py new file mode 100644 index 000000000..82f23e846 --- /dev/null +++ b/tests/emulation/emulation/__init__.py @@ -0,0 +1,2 @@ +from .emulate import microphysics +from .monitor import store diff --git a/tests/emulation/emulation/emulate.py b/tests/emulation/emulation/emulate.py new file mode 100644 index 000000000..a5460e1f9 --- /dev/null +++ b/tests/emulation/emulation/emulate.py @@ -0,0 +1,15 @@ +import logging + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + +def microphysics(state): + + try: + with open(f"microphysics_success.txt", "w") as f: + f.write("SUCCESS") + except Exception as e: + logger.error(e) + raise e + + logger.info("Successful call to microphysics in emulate.py") \ No newline at end of file diff --git a/tests/emulation/emulation/monitor.py b/tests/emulation/emulation/monitor.py new file mode 100644 index 000000000..264f29144 --- /dev/null +++ b/tests/emulation/emulation/monitor.py @@ -0,0 +1,15 @@ +import logging + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + +def store(state): + + try: + with open(f"store_success.txt", "w") as f: + f.write("SUCCESS") + except Exception as e: + logger.error(e) + raise e + + logging.info("Successful call to store in monitor.py") \ No newline at end of file diff --git a/emulation/setup.py b/tests/emulation/setup.py similarity index 82% rename from emulation/setup.py rename to tests/emulation/setup.py index 474c40de1..d305facc3 100644 --- a/emulation/setup.py +++ b/tests/emulation/setup.py @@ -5,18 +5,7 @@ from setuptools import setup, find_packages -requirements = [ - "cftime==1.2.1", - "f90nml==1.1.2", - "fsspec==2021.6.0", - "mpi4py==3.0.3", - "numpy==1.19.4", - "pyyaml==5.3", - "tensorflow==2.4.0", - "xarray==0.16.2", - "zarr==2.7.0", - "scipy", -] +requirements = [] setup_requirements = [] diff --git a/tests/pytest/config/emulation.yml b/tests/pytest/config/emulation.yml index 776b922a0..d1d2cd8c6 100644 --- a/tests/pytest/config/emulation.yml +++ b/tests/pytest/config/emulation.yml @@ -247,6 +247,8 @@ namelist: swhtr: true trans_trac: true use_ufo: true + emulate_zc_microphysics: true + save_zc_microphysics: true interpolator_nml: interp_method: conserve_great_circle nam_stochy: diff --git a/tests/pytest/test_regression.py b/tests/pytest/test_regression.py index b6de0d739..a34746f92 100644 --- a/tests/pytest/test_regression.py +++ b/tests/pytest/test_regression.py @@ -38,6 +38,11 @@ def reference_dir(request): return request.config.getoption("--refdir") +@pytest.fixture +def code_root(request): + return request.config.getoption("--code_root") + + @pytest.fixture def image_runner(request): return request.config.getoption("--image_runner") @@ -96,19 +101,18 @@ def test_regression( shutil.rmtree(run_dir) -def test_run_emulation(image, image_version, monkeypatch): +def test_callpyfort_integration(image, image_version): config = get_config("emulation.yml") model_image_tag = "{version}-emulation".format(version=image_version) model_image = f"{image}:{model_image_tag}" run_dir = get_run_dir(model_image_tag, config) - monkeypatch.setenv("OUTPUT_FREQ_SEC", str(900*2)) - env_vars = ["--env", "OUTPUT_FREQ_SEC"] - run_model(config, run_dir, model_image, "docker", additional_env_vars=env_vars) - nc_files = os.listdir(os.path.join(run_dir, "netcdf_output")) - assert os.path.exists(os.path.join(run_dir, "state_output.zarr")) - assert len(nc_files) > 0 + run_model(config, run_dir, model_image, "docker") + + assert os.path.exists(join(run_dir, "microphysics_success.txt")) + assert os.path.exists(join(run_dir, "store_success.txt")) + shutil.rmtree(run_dir) def check_rundir_md5sum(run_dir, md5sum_filename):