diff --git a/constraints.txt b/constraints.txt index 282e5988..dd122d08 100644 --- a/constraints.txt +++ b/constraints.txt @@ -87,7 +87,7 @@ cytoolz==0.11.2 # via # gt4py # gt4py (external/gt4py/setup.cfg) -dace==0.14.1 +dace==0.14.0 # via # -r requirements_dev.txt # pace-dsl diff --git a/doc_primer_orchestration.md b/doc_primer_orchestration.md index 16fadd4d..a10baf0c 100644 --- a/doc_primer_orchestration.md +++ b/doc_primer_orchestration.md @@ -104,6 +104,10 @@ _Parsing errors_ DaCe cannot parse _any_ dynamic Python and any code that allocates memory on the fly (think list creation). It will also complain about any arguments it can't memory describe (remember `dace_compiletime_args` ). +_GT_CACHE_DIR_NAME_ + +We do not honor the `GT_CACHE_DIR_NAME` with orchestration. `GT_CACHE_ROOT` is respected. + Conclusion ---------- diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index c4c40c19..4817869c 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -215,5 +215,5 @@ def _transform_horizontal_grid( grid.data[:, :, 0] = lon_transform[:] grid.data[:, :, 1] = lat_transform[:] - metric_terms._grid.data[:] = grid.data[:] + metric_terms._grid.data[:] = grid.data[:] # type: ignore[attr-defined] metric_terms._init_agrid() diff --git a/driver/pace/driver/performance/collector.py b/driver/pace/driver/performance/collector.py index 272e72be..cbc6c62a 100644 --- a/driver/pace/driver/performance/collector.py +++ b/driver/pace/driver/performance/collector.py @@ -66,6 +66,10 @@ def __init__(self, experiment_name: str, comm: pace.util.Comm): self.experiment_name = experiment_name self.comm = comm + def clear(self): + self.times_per_step = [] + self.hits_per_step = [] + def collect_performance(self): """ Take the accumulated timings and flush them into a new entry diff --git a/driver/pace/driver/tools.py b/driver/pace/driver/tools.py index c1dc5181..c58eb6fa 100644 --- a/driver/pace/driver/tools.py +++ b/driver/pace/driver/tools.py @@ -27,14 +27,48 @@ type=click.STRING, ) @click.option("--report_detail", is_flag=True, type=click.BOOL, default=False) -def command_line(action: str, sdfg_path: Optional[str], report_detail: Optional[bool]): +@click.option( + "--hardware_bw_in_gb_s", + required=False, + type=click.FLOAT, + default=0.0, +) +@click.option( + "--output_format", + required=False, + type=click.STRING, + default=None, +) +@click.option( + "--backend", + required=False, + type=click.STRING, + default="dace:gpu", +) +def command_line( + action: str, + sdfg_path: Optional[str], + report_detail: Optional[bool], + hardware_bw_in_gb_s: Optional[float], + output_format: Optional[str], + backend: Optional[str], +): """ Run tooling. """ if action == ACTION_SDFG_MEMORY_STATIC_ANALYSIS: print(memory_static_analysis_from_path(sdfg_path, detail_report=report_detail)) elif action == ACTION_SDFG_KERNEL_THEORETICAL_TIMING: - print(kernel_theoretical_timing_from_path(sdfg_path)) + print( + kernel_theoretical_timing_from_path( + sdfg_path, + hardware_bw_in_GB_s=( + None if hardware_bw_in_gb_s == 0 else hardware_bw_in_gb_s + ), + backend=backend, + output_format=output_format, + ) + ) if __name__ == "__main__": diff --git a/dsl/pace/dsl/dace/utils.py b/dsl/pace/dsl/dace/utils.py index cb268d3b..47eb61ad 100644 --- a/dsl/pace/dsl/dace/utils.py +++ b/dsl/pace/dsl/dace/utils.py @@ -1,15 +1,23 @@ +import json import time from dataclasses import dataclass, field -from typing import Dict, List +from typing import Dict, List, Optional import dace +import numpy as np from dace.transformation.helpers import get_parent_map +from gt4py.cartesian.gtscript import PARALLEL, computation, interval from pace.dsl.dace.dace_config import DaceConfig +from pace.dsl.stencil import CompilationConfig, FrozenStencil, StencilConfig +from pace.dsl.typing import Float, FloatField +from pace.util._optional_imports import cupy as cp from pace.util.logging import pace_log +# ---------------------------------------------------------- # Rough timer & log for major operations of DaCe build stack +# ---------------------------------------------------------- class DaCeProgress: """Timer and log to track build progress""" @@ -48,6 +56,9 @@ def _is_ref(sd: dace.sdfg.SDFG, aname: str): return found +# ---------------------------------------------------------- +# Memory analyser from SDFG +# ---------------------------------------------------------- @dataclass class ArrayReport: name: str = "" @@ -175,19 +186,38 @@ def memory_static_analysis_from_path(sdfg_path: str, detail_report=False) -> str ) -# TODO (floriand): in order for the timing analysis to be realistic the reference -# bandwidth of the hardware should be measured with GT4Py & simple in/out copy -# stencils. This allows to both measure the _actual_ deployed hardware and -# size it against the current GT4Py version. -# Below we bypass this needed automation by writing the P100 bw on Piz Daint -# measured with the above strategy. -# A better tool would allow this measure with a simple command and allow -# a one command that measure bw & report kernel analysis in one command -_HARDWARE_BW_GB_S = {"P100": 492.0} +# ---------------------------------------------------------- +# Theoritical bandwith from SDFG +# ---------------------------------------------------------- +def copy_defn(q_in: FloatField, q_out: FloatField): + with computation(PARALLEL), interval(...): + q_in = q_out + + +class MaxBandwithBenchmarkProgram: + def __init__(self, size, backend) -> None: + from pace.dsl.dace.orchestration import DaCeOrchestration, orchestrate + + dconfig = DaceConfig(None, backend, orchestration=DaCeOrchestration.BuildAndRun) + c = CompilationConfig(backend=backend) + s = StencilConfig(dace_config=dconfig, compilation_config=c) + self.copy_stencil = FrozenStencil( + func=copy_defn, + origin=(0, 0, 0), + domain=size, + stencil_config=s, + ) + orchestrate(obj=self, config=dconfig) + + def __call__(self, A, B, n: int): + for i in dace.nounroll(range(n)): + self.copy_stencil(A, B) def kernel_theoretical_timing( - sdfg: dace.sdfg.SDFG, hardware="P100", hardware_bw_in_Gb_s=None + sdfg: dace.sdfg.SDFG, + hardware_bw_in_GB_s=None, + backend=None, ) -> Dict[str, float]: """Compute a lower timing bound for kernels with the following hypothesis: @@ -197,6 +227,39 @@ def kernel_theoretical_timing( - Memory pressure is mostly in read/write from global memory, inner scalar & shared memory is not counted towards memory movement. """ + if not hardware_bw_in_GB_s: + size = np.array(sdfg.arrays["__g_self__w"].shape) + print( + f"Calculating experimental hardware bandwith on {size}" + f" arrays at {Float} precision..." + ) + bench = MaxBandwithBenchmarkProgram(size, backend) + if backend == "dace:gpu": + A = cp.ones(size, dtype=Float) + B = cp.ones(size, dtype=Float) + else: + A = np.ones(size, dtype=Float) + B = np.ones(size, dtype=Float) + n = 1000 + m = 4 + dt = [] + # Warm up run (build, allocation) + # to remove from timing the common runtime + bench(A, B, n) + # Time + for _ in range(m): + s = time.time() + bench(A, B, n) + dt.append((time.time() - s) / n) + memory_size_in_b = np.prod(size) * np.dtype(Float).itemsize * 8 + bandwidth_in_bytes_s = memory_size_in_b / np.median(dt) + print( + f"Hardware bandwith computed: {bandwidth_in_bytes_s/(1024*1024*1024)} GB/s" + ) + else: + bandwidth_in_bytes_s = hardware_bw_in_GB_s * 1024 * 1024 * 1024 + print(f"Given hardware bandwith: {bandwidth_in_bytes_s/(1024*1024*1024)} GB/s") + allmaps = [ (me, state) for me, state in sdfg.all_nodes_recursive() @@ -228,19 +291,6 @@ def kernel_theoretical_timing( ] ) - # Compute hardware memory bandwidth in bytes/us - if hardware_bw_in_Gb_s and hardware in _HARDWARE_BW_GB_S.keys(): - raise NotImplementedError("can't specify hardware bandwidth and hardware") - if hardware_bw_in_Gb_s: - bandwidth_in_bytes_s = hardware_bw_in_Gb_s * 1024 * 1024 * 1024 - elif hardware in _HARDWARE_BW_GB_S.keys(): - # Time it has to take (at least): bytes / bandwidth_in_bytes_s - bandwidth_in_bytes_s = _HARDWARE_BW_GB_S[hardware] * 1024 * 1024 * 1024 - else: - print( - f"Timing analysis: hardware {hardware} unknown and no bandwidth given" - ) - in_us = 1000 * 1000 # Theoretical fastest timing @@ -249,9 +299,12 @@ def kernel_theoretical_timing( except TypeError: newresult_in_us = (alldata_in_bytes / bandwidth_in_bytes_s) * in_us - if node.label in result: - import sympy + # We keep sympy import here because sympy is known to be a problematic + # import and an heavy module which should be avoided if possible. + # TODO: refactor it out by shadow-coding the sympy.Max/Eval functions + import sympy + if node.label in result: newresult_in_us = sympy.Max(result[node.label], newresult_in_us).expand() try: newresult_in_us = float(newresult_in_us) @@ -259,29 +312,54 @@ def kernel_theoretical_timing( pass # Bad expansion - if not isinstance(newresult_in_us, float): + if not isinstance(newresult_in_us, sympy.core.numbers.Float) and not isinstance( + newresult_in_us, float + ): continue - result[node.label] = newresult_in_us + result[node.label] = float(newresult_in_us) return result def report_kernel_theoretical_timing( - timings: Dict[str, float], human_readable: bool = True, csv: bool = False + timings: Dict[str, float], + human_readable: bool = True, + out_format: Optional[str] = None, ) -> str: """Produce a human readable or CSV of the kernel timings""" result_string = f"Maps processed: {len(timings)}.\n" if human_readable: result_string += "Timing in microseconds Map name:\n" result_string += "\n".join(f"{v:.2f}\t{k}," for k, v in sorted(timings.items())) - if csv: - result_string += "#Map name,timing in microseconds\n" - result_string += "\n".join(f"{k},{v}," for k, v in sorted(timings.items())) + if out_format == "csv": + csv_string = "" + csv_string += "#Map name,timing in microseconds\n" + csv_string += "\n".join(f"{k},{v}," for k, v in sorted(timings.items())) + with open("kernel_theoretical_timing.csv", "w") as f: + f.write(csv_string) + elif out_format == "json": + with open("kernel_theoretical_timing.json", "w") as f: + json.dump(timings, f, indent=2) + return result_string -def kernel_theoretical_timing_from_path(sdfg_path: str) -> str: +def kernel_theoretical_timing_from_path( + sdfg_path: str, + hardware_bw_in_GB_s: Optional[float] = None, + backend: Optional[str] = None, + output_format: Optional[str] = None, +) -> str: """Load an SDFG and report the theoretical kernel timings""" - timings = kernel_theoretical_timing(dace.SDFG.from_file(sdfg_path)) - return report_kernel_theoretical_timing(timings, human_readable=True, csv=False) + print(f"Running kernel_theoretical_timing for {sdfg_path}") + timings = kernel_theoretical_timing( + dace.SDFG.from_file(sdfg_path), + hardware_bw_in_GB_s=hardware_bw_in_GB_s, + backend=backend, + ) + return report_kernel_theoretical_timing( + timings, + human_readable=True, + out_format=output_format, + ) diff --git a/dsl/pace/dsl/stencil.py b/dsl/pace/dsl/stencil.py index 3b0bd781..26454ef8 100644 --- a/dsl/pace/dsl/stencil.py +++ b/dsl/pace/dsl/stencil.py @@ -325,8 +325,9 @@ def __init__( if "dace" in self.stencil_config.compilation_config.backend: dace.Config.set( "default_build_folder", - value="{gt_cache}/dacecache".format( - gt_cache=gt4py.cartesian.config.cache_settings["dir_name"] + value="{gt_root}/{gt_cache}/dacecache".format( + gt_root=gt4py.cartesian.config.cache_settings["root_path"], + gt_cache=gt4py.cartesian.config.cache_settings["dir_name"], ), ) diff --git a/fv3core/pace/fv3core/initialization/geos_wrapper.py b/fv3core/pace/fv3core/initialization/geos_wrapper.py index 00b6a1b7..4fc34052 100644 --- a/fv3core/pace/fv3core/initialization/geos_wrapper.py +++ b/fv3core/pace/fv3core/initialization/geos_wrapper.py @@ -1,6 +1,7 @@ +import enum import os from datetime import timedelta -from typing import Dict +from typing import Dict, List, Tuple import f90nml import numpy as np @@ -9,6 +10,14 @@ from pace import fv3core from pace.driver.performance.collector import PerformanceCollector from pace.dsl.dace import DaceConfig, orchestrate +from pace.dsl.gt4py_utils import is_gpu_backend +from pace.util.logging import pace_log + + +@enum.unique +class MemorySpace(enum.Enum): + HOST = 0 + DEVICE = 1 class GeosDycoreWrapper: @@ -23,6 +32,7 @@ def __init__( bdt: int, comm: pace.util.Comm, backend: str, + fortran_mem_space: MemorySpace = MemorySpace.HOST, ): # Look for an override to run on a single node gtfv3_single_rank_override = int(os.getenv("GTFV3_SINGLE_RANK_OVERRIDE", -1)) @@ -95,6 +105,7 @@ def __init__( self.dycore_state = fv3core.DycoreState.init_zeros( quantity_factory=quantity_factory ) + self.dycore_state.bdt = self.dycore_config.dt_atmos damping_coefficients = pace.util.grid.DampingCoefficients.new_from_metric_terms( metric_terms @@ -107,17 +118,28 @@ def __init__( quantity_factory=quantity_factory, damping_coefficients=damping_coefficients, config=self.dycore_config, - timestep=timedelta(seconds=self.dycore_config.dt_atmos), + timestep=timedelta(seconds=self.dycore_state.bdt), phis=self.dycore_state.phis, state=self.dycore_state, ) + self._fortran_mem_space = fortran_mem_space + self._pace_mem_space = ( + MemorySpace.DEVICE if is_gpu_backend(backend) else MemorySpace.HOST + ) + self.output_dict: Dict[str, np.ndarray] = {} self._allocate_output_dir() + pace_log.info( + "GEOS-Wrapper with: \n" + f" dt : {self.dycore_state.bdt}\n" + f" bridge : {self._fortran_mem_space} > {self._pace_mem_space}\n" + ) + def _critical_path(self): """Top-level orchestration function""" - with self.perf_collector.timestep_timer.clock("dycore"): + with self.perf_collector.timestep_timer.clock("step_dynamics"): self.dynamical_core.step_dynamics( state=self.dycore_state, timer=self.perf_collector.timestep_timer, @@ -125,6 +147,7 @@ def _critical_path(self): def __call__( self, + timings: Dict[str, List[float]], u: np.ndarray, v: np.ndarray, w: np.ndarray, @@ -149,9 +172,9 @@ def __call__( cxd: np.ndarray, cyd: np.ndarray, diss_estd: np.ndarray, - ) -> Dict[str, np.ndarray]: + ) -> Tuple[Dict[str, np.ndarray], Dict[str, List[float]]]: - with self.perf_collector.timestep_timer.clock("move_to_pace"): + with self.perf_collector.timestep_timer.clock("numpy-to-dycore"): self.dycore_state = self._put_fortran_data_in_dycore( u, v, @@ -182,20 +205,19 @@ def __call__( # Enter orchestrated code - if applicable self._critical_path() - with self.perf_collector.timestep_timer.clock("move_to_fortran"): + with self.perf_collector.timestep_timer.clock("dycore-to-numpy"): self.output_dict = self._prep_outputs_for_geos() - # Collect performance of the timestep and write - # a json file for rank 0 + # Collect performance of the timestep and write a json file for rank 0 self.perf_collector.collect_performance() - self.perf_collector.write_out_rank_0( - backend=self.backend, - is_orchestrated=self._is_orchestrated, - dt_atmos=self.dycore_config.dt_atmos, - sim_status="Ongoing", - ) + for k, v in self.perf_collector.times_per_step[0].items(): + if k not in timings.keys(): + timings[k] = [v] + else: + timings[k].append(v) + self.perf_collector.clear() - return self.output_dict + return self.output_dict, timings def _put_fortran_data_in_dycore( self, @@ -300,162 +322,232 @@ def _prep_outputs_for_geos(self) -> Dict[str, np.ndarray]: iec = self._grid_indexing.iec + 1 jec = self._grid_indexing.jec + 1 - pace.util.utils.safe_assign_array( - output_dict["u"], self.dycore_state.u.data[:-1, :, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["v"], self.dycore_state.v.data[:, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["w"], self.dycore_state.w.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["ua"], self.dycore_state.ua.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["va"], self.dycore_state.va.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["uc"], self.dycore_state.uc.data[:, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["vc"], self.dycore_state.vc.data[:-1, :, :-1] - ) + if self._fortran_mem_space != self._pace_mem_space: + pace.util.utils.safe_assign_array( + output_dict["u"], self.dycore_state.u.data[:-1, :, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["v"], self.dycore_state.v.data[:, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["w"], self.dycore_state.w.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["ua"], self.dycore_state.ua.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["va"], self.dycore_state.va.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["uc"], self.dycore_state.uc.data[:, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["vc"], self.dycore_state.vc.data[:-1, :, :-1] + ) - pace.util.utils.safe_assign_array( - output_dict["delz"], self.dycore_state.delz.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["pt"], self.dycore_state.pt.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["delp"], self.dycore_state.delp.data[:-1, :-1, :-1] - ) + pace.util.utils.safe_assign_array( + output_dict["delz"], self.dycore_state.delz.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["pt"], self.dycore_state.pt.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["delp"], self.dycore_state.delp.data[:-1, :-1, :-1] + ) - pace.util.utils.safe_assign_array( - output_dict["mfxd"], - self.dycore_state.mfxd.data[isc : iec + 1, jsc:jec, :-1], - ) - pace.util.utils.safe_assign_array( - output_dict["mfyd"], - self.dycore_state.mfyd.data[isc:iec, jsc : jec + 1, :-1], - ) - pace.util.utils.safe_assign_array( - output_dict["cxd"], self.dycore_state.cxd.data[isc : iec + 1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["cyd"], self.dycore_state.cyd.data[:-1, jsc : jec + 1, :-1] - ) + pace.util.utils.safe_assign_array( + output_dict["mfxd"], + self.dycore_state.mfxd.data[isc : iec + 1, jsc:jec, :-1], + ) + pace.util.utils.safe_assign_array( + output_dict["mfyd"], + self.dycore_state.mfyd.data[isc:iec, jsc : jec + 1, :-1], + ) + pace.util.utils.safe_assign_array( + output_dict["cxd"], self.dycore_state.cxd.data[isc : iec + 1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["cyd"], self.dycore_state.cyd.data[:-1, jsc : jec + 1, :-1] + ) - pace.util.utils.safe_assign_array( - output_dict["ps"], self.dycore_state.ps.data[:-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["pe"], - self.dycore_state.pe.data[isc - 1 : iec + 1, jsc - 1 : jec + 1, :], - ) - pace.util.utils.safe_assign_array( - output_dict["pk"], self.dycore_state.pk.data[isc:iec, jsc:jec, :] - ) - pace.util.utils.safe_assign_array( - output_dict["peln"], self.dycore_state.peln.data[isc:iec, jsc:jec, :] - ) - pace.util.utils.safe_assign_array( - output_dict["pkz"], self.dycore_state.pkz.data[isc:iec, jsc:jec, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["phis"], self.dycore_state.phis.data[:-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["q_con"], self.dycore_state.q_con.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["omga"], self.dycore_state.omga.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["diss_estd"], self.dycore_state.diss_estd.data[:-1, :-1, :-1] - ) + pace.util.utils.safe_assign_array( + output_dict["ps"], self.dycore_state.ps.data[:-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["pe"], + self.dycore_state.pe.data[isc - 1 : iec + 1, jsc - 1 : jec + 1, :], + ) + pace.util.utils.safe_assign_array( + output_dict["pk"], self.dycore_state.pk.data[isc:iec, jsc:jec, :] + ) + pace.util.utils.safe_assign_array( + output_dict["peln"], self.dycore_state.peln.data[isc:iec, jsc:jec, :] + ) + pace.util.utils.safe_assign_array( + output_dict["pkz"], self.dycore_state.pkz.data[isc:iec, jsc:jec, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["phis"], self.dycore_state.phis.data[:-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["q_con"], self.dycore_state.q_con.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["omga"], self.dycore_state.omga.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["diss_estd"], + self.dycore_state.diss_estd.data[:-1, :-1, :-1], + ) - pace.util.utils.safe_assign_array( - output_dict["qvapor"], self.dycore_state.qvapor.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["qliquid"], self.dycore_state.qliquid.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["qice"], self.dycore_state.qice.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["qrain"], self.dycore_state.qrain.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["qsnow"], self.dycore_state.qsnow.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["qgraupel"], self.dycore_state.qgraupel.data[:-1, :-1, :-1] - ) - pace.util.utils.safe_assign_array( - output_dict["qcld"], self.dycore_state.qcld.data[:-1, :-1, :-1] - ) + pace.util.utils.safe_assign_array( + output_dict["qvapor"], self.dycore_state.qvapor.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qliquid"], self.dycore_state.qliquid.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qice"], self.dycore_state.qice.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qrain"], self.dycore_state.qrain.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qsnow"], self.dycore_state.qsnow.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qgraupel"], self.dycore_state.qgraupel.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qcld"], self.dycore_state.qcld.data[:-1, :-1, :-1] + ) + else: + output_dict["u"] = self.dycore_state.u.data[:-1, :, :-1] + output_dict["v"] = self.dycore_state.v.data[:, :-1, :-1] + output_dict["w"] = self.dycore_state.w.data[:-1, :-1, :-1] + output_dict["ua"] = self.dycore_state.ua.data[:-1, :-1, :-1] + output_dict["va"] = self.dycore_state.va.data[:-1, :-1, :-1] + output_dict["uc"] = self.dycore_state.uc.data[:, :-1, :-1] + output_dict["vc"] = self.dycore_state.vc.data[:-1, :, :-1] + output_dict["delz"] = self.dycore_state.delz.data[:-1, :-1, :-1] + output_dict["pt"] = self.dycore_state.pt.data[:-1, :-1, :-1] + output_dict["delp"] = self.dycore_state.delp.data[:-1, :-1, :-1] + output_dict["mfxd"] = self.dycore_state.mfxd.data[ + isc : iec + 1, jsc:jec, :-1 + ] + output_dict["mfyd"] = self.dycore_state.mfyd.data[ + isc:iec, jsc : jec + 1, :-1 + ] + output_dict["cxd"] = self.dycore_state.cxd.data[isc : iec + 1, :-1, :-1] + output_dict["cyd"] = self.dycore_state.cyd.data[:-1, jsc : jec + 1, :-1] + output_dict["ps"] = self.dycore_state.ps.data[:-1, :-1] + output_dict["pe"] = self.dycore_state.pe.data[ + isc - 1 : iec + 1, jsc - 1 : jec + 1, : + ] + output_dict["pk"] = self.dycore_state.pk.data[isc:iec, jsc:jec, :] + output_dict["peln"] = self.dycore_state.peln.data[isc:iec, jsc:jec, :] + output_dict["pkz"] = self.dycore_state.pkz.data[isc:iec, jsc:jec, :-1] + output_dict["phis"] = self.dycore_state.phis.data[:-1, :-1] + output_dict["q_con"] = self.dycore_state.q_con.data[:-1, :-1, :-1] + output_dict["omga"] = self.dycore_state.omga.data[:-1, :-1, :-1] + output_dict["diss_estd"] = self.dycore_state.diss_estd.data[:-1, :-1, :-1] + output_dict["qvapor"] = self.dycore_state.qvapor.data[:-1, :-1, :-1] + output_dict["qliquid"] = self.dycore_state.qliquid.data[:-1, :-1, :-1] + output_dict["qice"] = self.dycore_state.qice.data[:-1, :-1, :-1] + output_dict["qrain"] = self.dycore_state.qrain.data[:-1, :-1, :-1] + output_dict["qsnow"] = self.dycore_state.qsnow.data[:-1, :-1, :-1] + output_dict["qgraupel"] = self.dycore_state.qgraupel.data[:-1, :-1, :-1] + output_dict["qcld"] = self.dycore_state.qcld.data[:-1, :-1, :-1] return output_dict def _allocate_output_dir(self): + if self._fortran_mem_space != self._pace_mem_space: + nhalo = self._grid_indexing.n_halo + shape_centered = self._grid_indexing.domain_full(add=(0, 0, 0)) + shape_x_interface = self._grid_indexing.domain_full(add=(1, 0, 0)) + shape_y_interface = self._grid_indexing.domain_full(add=(0, 1, 0)) + shape_z_interface = self._grid_indexing.domain_full(add=(0, 0, 1)) + shape_2d = shape_centered[:-1] + + self.output_dict["u"] = np.empty((shape_y_interface)) + self.output_dict["v"] = np.empty((shape_x_interface)) + self.output_dict["w"] = np.empty((shape_centered)) + self.output_dict["ua"] = np.empty((shape_centered)) + self.output_dict["va"] = np.empty((shape_centered)) + self.output_dict["uc"] = np.empty((shape_x_interface)) + self.output_dict["vc"] = np.empty((shape_y_interface)) + + self.output_dict["delz"] = np.empty((shape_centered)) + self.output_dict["pt"] = np.empty((shape_centered)) + self.output_dict["delp"] = np.empty((shape_centered)) + + self.output_dict["mfxd"] = np.empty( + (self._grid_indexing.domain_full(add=(1 - 2 * nhalo, -2 * nhalo, 0))) + ) + self.output_dict["mfyd"] = np.empty( + (self._grid_indexing.domain_full(add=(-2 * nhalo, 1 - 2 * nhalo, 0))) + ) + self.output_dict["cxd"] = np.empty( + (self._grid_indexing.domain_full(add=(1 - 2 * nhalo, 0, 0))) + ) + self.output_dict["cyd"] = np.empty( + (self._grid_indexing.domain_full(add=(0, 1 - 2 * nhalo, 0))) + ) - nhalo = self._grid_indexing.n_halo - shape_centered = self._grid_indexing.domain_full(add=(0, 0, 0)) - shape_x_interface = self._grid_indexing.domain_full(add=(1, 0, 0)) - shape_y_interface = self._grid_indexing.domain_full(add=(0, 1, 0)) - shape_z_interface = self._grid_indexing.domain_full(add=(0, 0, 1)) - shape_2d = shape_centered[:-1] - - self.output_dict["u"] = np.empty((shape_y_interface)) - self.output_dict["v"] = np.empty((shape_x_interface)) - self.output_dict["w"] = np.empty((shape_centered)) - self.output_dict["ua"] = np.empty((shape_centered)) - self.output_dict["va"] = np.empty((shape_centered)) - self.output_dict["uc"] = np.empty((shape_x_interface)) - self.output_dict["vc"] = np.empty((shape_y_interface)) - - self.output_dict["delz"] = np.empty((shape_centered)) - self.output_dict["pt"] = np.empty((shape_centered)) - self.output_dict["delp"] = np.empty((shape_centered)) - - self.output_dict["mfxd"] = np.empty( - (self._grid_indexing.domain_full(add=(1 - 2 * nhalo, -2 * nhalo, 0))) - ) - self.output_dict["mfyd"] = np.empty( - (self._grid_indexing.domain_full(add=(-2 * nhalo, 1 - 2 * nhalo, 0))) - ) - self.output_dict["cxd"] = np.empty( - (self._grid_indexing.domain_full(add=(1 - 2 * nhalo, 0, 0))) - ) - self.output_dict["cyd"] = np.empty( - (self._grid_indexing.domain_full(add=(0, 1 - 2 * nhalo, 0))) - ) - - self.output_dict["ps"] = np.empty((shape_2d)) - self.output_dict["pe"] = np.empty( - (self._grid_indexing.domain_full(add=(2 - 2 * nhalo, 2 - 2 * nhalo, 1))) - ) - self.output_dict["pk"] = np.empty( - (self._grid_indexing.domain_full(add=(-2 * nhalo, -2 * nhalo, 1))) - ) - self.output_dict["peln"] = np.empty( - (self._grid_indexing.domain_full(add=(-2 * nhalo, -2 * nhalo, 1))) - ) - self.output_dict["pkz"] = np.empty( - (self._grid_indexing.domain_full(add=(-2 * nhalo, -2 * nhalo, 0))) - ) - self.output_dict["phis"] = np.empty((shape_2d)) - self.output_dict["q_con"] = np.empty((shape_centered)) - self.output_dict["omga"] = np.empty((shape_centered)) - self.output_dict["diss_estd"] = np.empty((shape_centered)) - - self.output_dict["qvapor"] = np.empty((shape_centered)) - self.output_dict["qliquid"] = np.empty((shape_centered)) - self.output_dict["qice"] = np.empty((shape_centered)) - self.output_dict["qrain"] = np.empty((shape_centered)) - self.output_dict["qsnow"] = np.empty((shape_centered)) - self.output_dict["qgraupel"] = np.empty((shape_centered)) - self.output_dict["qcld"] = np.empty((shape_centered)) + self.output_dict["ps"] = np.empty((shape_2d)) + self.output_dict["pe"] = np.empty( + (self._grid_indexing.domain_full(add=(2 - 2 * nhalo, 2 - 2 * nhalo, 1))) + ) + self.output_dict["pk"] = np.empty( + (self._grid_indexing.domain_full(add=(-2 * nhalo, -2 * nhalo, 1))) + ) + self.output_dict["peln"] = np.empty( + (self._grid_indexing.domain_full(add=(-2 * nhalo, -2 * nhalo, 1))) + ) + self.output_dict["pkz"] = np.empty( + (self._grid_indexing.domain_full(add=(-2 * nhalo, -2 * nhalo, 0))) + ) + self.output_dict["phis"] = np.empty((shape_2d)) + self.output_dict["q_con"] = np.empty((shape_centered)) + self.output_dict["omga"] = np.empty((shape_centered)) + self.output_dict["diss_estd"] = np.empty((shape_centered)) + + self.output_dict["qvapor"] = np.empty((shape_centered)) + self.output_dict["qliquid"] = np.empty((shape_centered)) + self.output_dict["qice"] = np.empty((shape_centered)) + self.output_dict["qrain"] = np.empty((shape_centered)) + self.output_dict["qsnow"] = np.empty((shape_centered)) + self.output_dict["qgraupel"] = np.empty((shape_centered)) + self.output_dict["qcld"] = np.empty((shape_centered)) + else: + self.output_dict["u"] = None + self.output_dict["v"] = None + self.output_dict["w"] = None + self.output_dict["ua"] = None + self.output_dict["va"] = None + self.output_dict["uc"] = None + self.output_dict["vc"] = None + self.output_dict["delz"] = None + self.output_dict["pt"] = None + self.output_dict["delp"] = None + self.output_dict["mfxd"] = None + self.output_dict["mfyd"] = None + self.output_dict["cxd"] = None + self.output_dict["cyd"] = None + self.output_dict["ps"] = None + self.output_dict["pe"] = None + self.output_dict["pk"] = None + self.output_dict["peln"] = None + self.output_dict["pkz"] = None + self.output_dict["phis"] = None + self.output_dict["q_con"] = None + self.output_dict["omga"] = None + self.output_dict["diss_estd"] = None + self.output_dict["qvapor"] = None + self.output_dict["qliquid"] = None + self.output_dict["qice"] = None + self.output_dict["qrain"] = None + self.output_dict["qsnow"] = None + self.output_dict["qgraupel"] = None + self.output_dict["qcld"] = None diff --git a/fv3core/pace/fv3core/stencils/d_sw.py b/fv3core/pace/fv3core/stencils/d_sw.py index c5aa45d2..cc602f4e 100644 --- a/fv3core/pace/fv3core/stencils/d_sw.py +++ b/fv3core/pace/fv3core/stencils/d_sw.py @@ -503,7 +503,6 @@ def heat_source_from_vorticity_damping( rdx: FloatFieldIJ, rdy: FloatFieldIJ, heat_source: FloatField, - heat_source_total: FloatField, dissipation_estimate: FloatField, kinetic_energy_fraction_to_damp: FloatFieldK, ): @@ -526,14 +525,13 @@ def heat_source_from_vorticity_damping( rdy (in): 1 / dy heat_source (inout): heat source from vorticity damping implied by energy conservation - heat_source_total (inout): accumulated heat source - dissipation_estimate (out): dissipation estimate, only calculated if + dissipation_estimate (inout): dissipation estimate, only calculated if calculate_dissipation_estimate is 1. Used for stochastic kinetic energy backscatter (skeb) routine. kinetic_energy_fraction_to_damp (in): the fraction of kinetic energy to explicitly damp and convert into heat. """ - from __externals__ import ( + from __externals__ import ( # noqa (see below) d_con, do_stochastic_ke_backscatter, local_ie, @@ -570,11 +568,20 @@ def heat_source_from_vorticity_damping( heat_source - kinetic_energy_fraction_to_damp * dampterm ) - if __INLINED((d_con > dcon_threshold) or do_stochastic_ke_backscatter): + if __INLINED(do_stochastic_ke_backscatter): with horizontal(region[local_is : local_ie + 1, local_js : local_je + 1]): - heat_source_total = heat_source_total + heat_source - if __INLINED(do_stochastic_ke_backscatter): - dissipation_estimate -= dampterm + dissipation_estimate -= dampterm + + +def accumulate_heat_source_and_dissipation_estimate( + heat_source: FloatField, + heat_source_total: FloatField, + diss_est: FloatField, + diss_est_total: FloatField, +): + with computation(PARALLEL), interval(...): + heat_source_total += heat_source + diss_est_total += diss_est # TODO(eddied): Had to split this into a separate stencil to get this to validate @@ -746,6 +753,8 @@ def __init__( orchestrate(obj=self, config=stencil_factory.config.dace_config) self.grid_data = grid_data self._f0 = self.grid_data.fC_agrid + self._d_con = config.d_con + self._do_stochastic_ke_backscatter = config.do_skeb self.grid_indexing = stencil_factory.grid_indexing assert config.grid_type < 3, "ubke and vbke only implemented for grid_type < 3" @@ -771,6 +780,7 @@ def make_quantity(): ) self._tmp_heat_s = make_quantity() + self._tmp_diss_e = make_quantity() self._vort_x_delta = make_quantity() self._vort_y_delta = make_quantity() self._dt_kinetic_energy_on_cell_corners = make_quantity() @@ -921,6 +931,15 @@ def make_quantity(): }, ) ) + + if (self._d_con > 1.e-5) or (self._do_stochastic_ke_backscatter): + self._accumulate_heat_source_and_dissipation_estimate_stencil = ( + stencil_factory.from_dims_halo( + func=accumulate_heat_source_and_dissipation_estimate, + compute_dims=[X_DIM, Y_DIM, Z_DIM], + ) + ) + self._compute_vorticity_stencil = stencil_factory.from_dims_halo( compute_vorticity, compute_dims=[X_DIM, Y_DIM, Z_DIM], @@ -1055,7 +1074,7 @@ def __call__( w, self.grid_data.rarea, self._tmp_heat_s, - diss_est, + self._tmp_diss_e, self._tmp_dw, self._column_namelist["damp_w"], self._column_namelist["ke_bg"], @@ -1231,11 +1250,15 @@ def __call__( self.grid_data.rdx, self.grid_data.rdy, self._tmp_heat_s, - heat_source, - diss_est, + self._tmp_diss_e, self._column_namelist["d_con"], ) + if (self._d_con > 1.e-5) or (self._do_stochastic_ke_backscatter): + self._accumulate_heat_source_and_dissipation_estimate_stencil( + self._tmp_heat_s, heat_source, self._tmp_diss_e, diss_est + ) + self._update_u_and_v_stencil( self._tmp_ut, self._tmp_vt, diff --git a/fv3core/pace/fv3core/stencils/fv_dynamics.py b/fv3core/pace/fv3core/stencils/fv_dynamics.py index 0c418daf..3299e501 100644 --- a/fv3core/pace/fv3core/stencils/fv_dynamics.py +++ b/fv3core/pace/fv3core/stencils/fv_dynamics.py @@ -27,13 +27,6 @@ from pace.util.mpi import MPI -# nq is actually given by ncnst - pnats, where those are given in atmosphere.F90 by: -# ncnst = Atm(mytile)%ncnst -# pnats = Atm(mytile)%flagstruct%pnats -# here we hard-coded it because 8 is the only supported value, refactor this later! -NQ = 8 # state.nq_tot - spec.namelist.dnats - - def pt_to_potential_density_pt( pkz: FloatField, dp_initial: FloatField, q_con: FloatField, pt: FloatField ): @@ -179,6 +172,11 @@ def __init__( method_to_orchestrate="_checkpoint_tracer_advection_out", dace_compiletime_args=["state"], ) + if timestep == timedelta(seconds=0): + raise RuntimeError( + "Bad dynamical core configuration:" + " the atmospheric timestep is 0 seconds!" + ) # nested and stretched_grid are options in the Fortran code which we # have not implemented, so they are hard-coded here. self.call_checkpointer = checkpointer is not None @@ -207,7 +205,7 @@ def __init__( ) self.tracers = {} - for name in utils.tracer_variables[0:NQ]: + for name in utils.tracer_variables[0 : constants.NQ]: self.tracers[name] = state.__dict__[name] temporaries = fvdyn_temporaries(quantity_factory) @@ -282,7 +280,7 @@ def __init__( ) self._cappa = self.acoustic_dynamics.cappa - if not (not self.config.inline_q and NQ != 0): + if not (not self.config.inline_q and constants.NQ != 0): raise NotImplementedError("tracer_2d not implemented, turn on z_tracer") self._adjust_tracer_mixing_ratio = AdjustNegativeTracerMixingRatio( stencil_factory, @@ -296,7 +294,7 @@ def __init__( quantity_factory=quantity_factory, config=config.remapping, area_64=grid_data.area_64, - nq=NQ, + nq=constants.NQ, pfull=self._pfull, tracers=self.tracers, checkpointer=checkpointer, @@ -546,6 +544,12 @@ def _compute(self, state: DycoreState, timer: pace.util.Timer): log_on_rank_0("Remapping") with timer.clock("Remapping"): self._checkpoint_remapping_in(state) + + # TODO: When NQ=9, we shouldn't need to pass qcld explicitly + # since it's in self.tracers. It should not be an issue since + # we don't have self.tracers & qcld computation at the same + # time + # When NQ=8, we do need qcld passed explicitely self._lagrangian_to_eulerian_obj( self.tracers, state.pt, diff --git a/fv3core/pace/fv3core/stencils/saturation_adjustment.py b/fv3core/pace/fv3core/stencils/saturation_adjustment.py index 124f7734..41f3f39a 100644 --- a/fv3core/pace/fv3core/stencils/saturation_adjustment.py +++ b/fv3core/pace/fv3core/stencils/saturation_adjustment.py @@ -909,14 +909,14 @@ def satadjust( # icloud_f = 0: bug - fixed # icloud_f = 1: old fvgfs gfdl) mp implementation # icloud_f = 2: binary cloud scheme (0 / 1) - if rh > 0.75 and qpz > 1.0e-8: + if rh > 0.75 and qpz > constants.SAT_ADJUST_THRESHOLD: dq = hvar * qpz q_plus = qpz + dq q_minus = qpz - dq if icloud_f == 2: # TODO untested if qpz > qstar: qa = 1.0 - elif (qstar < q_plus) and (q_cond > 1.0e-8): + elif (qstar < q_plus) and (q_cond > constants.SAT_ADJUST_THRESHOLD): qa = min(1.0, ((q_plus - qstar) / dq) ** 2) else: qa = 0.0 @@ -932,7 +932,7 @@ def satadjust( else: qa = 0.0 # impose minimum cloudiness if substantial q_cond exist - if q_cond > 1.0e-8: + if q_cond > constants.SAT_ADJUST_THRESHOLD: qa = max(cld_min, qa) qa = min(1, qa) else: diff --git a/fv3core/pace/fv3core/stencils/temperature_adjust.py b/fv3core/pace/fv3core/stencils/temperature_adjust.py index 33feb484..0226df38 100644 --- a/fv3core/pace/fv3core/stencils/temperature_adjust.py +++ b/fv3core/pace/fv3core/stencils/temperature_adjust.py @@ -29,7 +29,6 @@ def apply_diffusive_heating( """ with computation(PARALLEL), interval(...): pkz = exp(cappa / (1.0 - cappa) * log(constants.RDG * delp / delz * pt)) - pkz = (constants.RDG * delp / delz * pt) ** (cappa / (1.0 - cappa)) dtmp = heat_source / (constants.CV_AIR * delp) with computation(PARALLEL): with interval(0, 1): diff --git a/fv3core/tests/savepoint/translate/translate_init_case.py b/fv3core/tests/savepoint/translate/translate_init_case.py index 9716aca3..4a9fafc4 100644 --- a/fv3core/tests/savepoint/translate/translate_init_case.py +++ b/fv3core/tests/savepoint/translate/translate_init_case.py @@ -7,7 +7,6 @@ import pace.dsl.gt4py_utils as utils import pace.fv3core.initialization.baroclinic as baroclinic_init import pace.fv3core.initialization.baroclinic_jablonowski_williamson as jablo_init -import pace.fv3core.stencils.fv_dynamics as fv_dynamics import pace.util import pace.util as fv3util from pace.fv3core.testing import TranslateDycoreFortranData2Py @@ -17,7 +16,6 @@ class TranslateInitCase(ParallelTranslateBaseSlicing): - outputs: Dict[str, Any] = { "u": { "name": "x_wind", @@ -175,7 +173,6 @@ def outputs_from_state(self, state: dict): state[name][tracer] = state[name][tracer].data arrays[name] = state[name] elif len(self.outputs[name]["dims"]) > 0: - arrays[name] = state[name].data else: outputs[name] = state[name] # scalar @@ -184,7 +181,10 @@ def outputs_from_state(self, state: dict): def compute_parallel(self, inputs, communicator): state = {} - full_shape = (*self.grid.domain_shape_full(add=(1, 1, 1)), fv_dynamics.NQ) + full_shape = ( + *self.grid.domain_shape_full(add=(1, 1, 1)), + pace.util.constants.NQ, + ) for variable, properties in self.outputs.items(): dims = properties["dims"] state[variable] = fv3util.Quantity( diff --git a/requirements_dev.txt b/requirements_dev.txt index 8318706a..052bf5c3 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -11,7 +11,7 @@ dask>=2021.10.0 netCDF4 cftime fv3config>=0.9.0 -dace>=0.14.1 +dace==0.14.0 f90nml>=1.1.0 numpy>=1.15 -e external/gt4py diff --git a/util/pace/util/constants.py b/util/pace/util/constants.py index a7fe5a53..ef57ed18 100644 --- a/util/pace/util/constants.py +++ b/util/pace/util/constants.py @@ -1,3 +1,30 @@ +import os +from enum import Enum + +from pace.util.logging import pace_log + + +# The FV3GFS model ships with two sets of constants, one used in the GFS physics +# package and the other used for the Dycore. Their difference are small but significant +# In addition the GSFC's GEOS model as its own variables +class ConstantVersions(Enum): + FV3DYCORE = "FV3DYCORE" # NOAA's FV3 dynamical core constants (original port) + GFS = "GFS" # Constant as defined in NOAA GFS + GEOS = "GEOS" # Constant as defined in GEOS v13 + + +CONST_VERSION_AS_STR = os.environ.get("PACE_CONSTANTS", "FV3DYCORE") + +try: + CONST_VERSION = ConstantVersions[CONST_VERSION_AS_STR] + pace_log.info(f"Constant selected: {CONST_VERSION}") +except KeyError as e: + raise RuntimeError(f"Constants {CONST_VERSION_AS_STR} is not implemented, abort.") + +##################### +# Common constants +##################### + ROOT_RANK = 0 X_DIM = "x" X_INTERFACE_DIM = "x_interface" @@ -28,15 +55,42 @@ BOUNDARY_TYPES = EDGE_BOUNDARY_TYPES + CORNER_BOUNDARY_TYPES N_HALO_DEFAULT = 3 +####################### +# Tracers configuration +####################### + +# nq is actually given by ncnst - pnats, where those are given in atmosphere.F90 by: +# ncnst = Atm(mytile)%ncnst +# pnats = Atm(mytile)%flagstruct%pnats +# here we hard-coded it because 8 is the only supported value, refactor this later! +if CONST_VERSION == ConstantVersions.GEOS: + # 'qlcd' is exchanged in GEOS + NQ = 9 +elif ( + CONST_VERSION == ConstantVersions.GFS or CONST_VERSION == ConstantVersions.FV3DYCORE +): + NQ = 8 +else: + raise RuntimeError("Constant selector failed, bad code.") + ##################### # Physical constants ##################### - -# The FV3GFS model ships with two sets of constants, one used in the GFS physics -# package and the other used for the Dycore. Their difference are small but significant -# Our Fortran executable on GCE has GFS_PHYS=True -GFS_PHYS = True -if GFS_PHYS: +if CONST_VERSION == ConstantVersions.GEOS: + RADIUS = 6.371e6 + PI = 3.14159265358979323846 + OMEGA = 2.0 * PI / 86164.0 + GRAV = 9.80665 + RGRAV = 1.0 / GRAV + RDGAS = 8314.47 / 28.965 + RVGAS = 8314.47 / 18.015 + HLV = 2.4665e6 + HLF = 3.3370e5 + KAPPA = RDGAS / (3.5 * RDGAS) + CP_AIR = RDGAS / KAPPA + TFREEZE = 273.15 + SAT_ADJUST_THRESHOLD = 1.0e-6 +elif CONST_VERSION == ConstantVersions.GFS: RADIUS = 6.3712e6 # Radius of the Earth [m] PI = 3.1415926535897931 OMEGA = 7.2921e-5 # Rotation of the earth @@ -49,7 +103,8 @@ CP_AIR = 1004.6 KAPPA = RDGAS / CP_AIR # Specific heat capacity of dry air at TFREEZE = 273.15 -else: + SAT_ADJUST_THRESHOLD = 1.0e-8 +elif CONST_VERSION == ConstantVersions.FV3DYCORE: RADIUS = 6371.0e3 # Radius of the Earth [m] #6371.0e3 PI = 3.14159265358979323846 # 3.14159265358979323846 OMEGA = 7.292e-5 # Rotation of the earth # 7.292e-5 @@ -62,6 +117,9 @@ KAPPA = 2.0 / 7.0 CP_AIR = RDGAS / KAPPA # Specific heat capacity of dry air at TFREEZE = 273.16 # Freezing temperature of fresh water [K] + SAT_ADJUST_THRESHOLD = 1.0e-8 +else: + raise RuntimeError("Constant selector failed, bad code.") DZ_MIN = 2.0 CV_AIR = CP_AIR - RDGAS # Heat capacity of dry air at constant volume