Skip to content

DiffusionConfig as dataclass #565

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
halungge wants to merge 4 commits into from
Closed

DiffusionConfig as dataclass #565

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
066f7ac
convert DiffusionConfig to dataclass
Sep 27, 2024
0c5e549
trying how to load config (WIP)
Sep 30, 2024
2e60cf3
merge main
Oct 9, 2024
44a066c
merge main
Dec 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
229 changes: 103 additions & 126 deletions model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,184 +78,164 @@
log = logging.getLogger(__name__)


class DiffusionType(int, enum.Enum):
class DiffusionType(enum.Enum):
"""
Order of nabla operator for diffusion.

Note: Called `hdiff_order` in `mo_diffusion_nml.f90`.
Note: We currently only support type 5.
"""

NO_DIFFUSION = -1 #: no diffusion
LINEAR_2ND_ORDER = 2 #: 2nd order linear diffusion on all vertical levels
SMAGORINSKY_NO_BACKGROUND = 3 #: Smagorinsky diffusion without background diffusion
LINEAR_4TH_ORDER = 4 #: 4th order linear diffusion on all vertical levels
SMAGORINSKY_4TH_ORDER = 5 #: Smagorinsky diffusion with fourth-order background diffusion
#: no diffusion
NO_DIFFUSION = -1

#: 2nd order linear diffusion on all vertical levels
LINEAR_2ND_ORDER = 2

class TurbulenceShearForcingType(int, enum.Enum):
#: Smagorinsky diffusion without background diffusion
SMAGORINSKY_NO_BACKGROUND = 3

#: 4th order linear diffusion on all vertical levels
LINEAR_4TH_ORDER = 4

#: Smagorinsky diffusion with fourth-order background diffusion
SMAGORINSKY_4TH_ORDER = 5


class TurbulenceShearForcingType(enum.Enum):
"""
Type of shear forcing used in turbulance.

Note: called `itype_sher` in `mo_turbdiff_nml.f90`
"""

VERTICAL_OF_HORIZONTAL_WIND = 0 #: only vertical shear of horizontal wind
VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND = (
1 #: as `VERTICAL_ONLY` plus horizontal shar correction
)
VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND = (
2 #: as `VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND` plus shear form vertical velocity
)
VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND_LTHESH = 3 #: same as `VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND` but scaling of coarse-grid horizontal shear production term with 1/sqrt(Ri) (if LTKESH = TRUE)
#: only vertical shear of horizontal wind
VERTICAL_OF_HORIZONTAL_WIND = 0

#: as `VERTICAL_ONLY` plus horizontal shar correction
VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND = 1

#: as `VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND` plus shear form vertical velocity
VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND = 2

#: same as `VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND` but scaling of coarse-grid horizontal shear production term with 1/sqrt(Ri) (if LTKESH = TRUE)
VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND_LTHESH = 3


@dataclasses.dataclass(frozen=True)
class DiffusionConfig:
"""
Contains necessary parameter to configure a diffusion run.

Encapsulates namelist parameters and derived parameters.
Values should be read from configuration.
Default values are taken from the defaults in the corresponding ICON Fortran namelist files.
"""

# TODO(Magdalena): to be read from config
# TODO(Magdalena): handle dependencies on other namelists (see below...)

def __init__(
self,
diffusion_type: DiffusionType = DiffusionType.SMAGORINSKY_4TH_ORDER,
hdiff_w=True,
hdiff_vn=True,
hdiff_temp=True,
type_vn_diffu: int = 1,
smag_3d: bool = False,
type_t_diffu: int = 2,
hdiff_efdt_ratio: float = 36.0,
hdiff_w_efdt_ratio: float = 15.0,
smagorinski_scaling_factor: float = 0.015,
n_substeps: int = 5,
zdiffu_t: bool = True,
thslp_zdiffu: float = 0.025,
thhgtd_zdiffu: float = 200.0,
velocity_boundary_diffusion_denom: float = 200.0,
temperature_boundary_diffusion_denom: float = 135.0,
max_nudging_coeff: float = 0.02,
nudging_decay_rate: float = 2.0,
shear_type: TurbulenceShearForcingType = TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND,
):
"""Set the diffusion configuration parameters with the ICON default values."""
# parameters from namelist diffusion_nml

self.diffusion_type: int = diffusion_type

#: If True, apply diffusion on the vertical wind field
#: Called `lhdiff_w` in mo_diffusion_nml.f90
self.apply_to_vertical_wind: bool = hdiff_w

#: True apply diffusion on the horizontal wind field, is ONLY used in mo_nh_stepping.f90
#: Called `lhdiff_vn` in mo_diffusion_nml.f90
self.apply_to_horizontal_wind = hdiff_vn
# TODO(Magdalena): handle dependencies on other namelists (see below...)
"""

#: If True, apply horizontal diffusion to temperature field
#: Called `lhdiff_temp` in mo_diffusion_nml.f90
self.apply_to_temperature: bool = hdiff_temp
diffusion_type: DiffusionType = DiffusionType.SMAGORINSKY_4TH_ORDER

#: If True, compute 3D Smagorinsky diffusion coefficient
#: Called `lsmag_3d` in mo_diffusion_nml.f90
self.compute_3d_smag_coeff: bool = smag_3d
#: If True, apply diffusion on the vertical wind field
#: Called `lhdiff_w` in mo_diffusion_nml.f90
apply_to_vertical_wind: bool = True

#: Options for discretizing the Smagorinsky momentum diffusion
#: Called `itype_vn_diffu` in mo_diffusion_nml.f90
self.type_vn_diffu: int = type_vn_diffu
#: True apply diffusion on the horizontal wind field, is ONLY used in mo_nh_stepping.f90
#: Called `lhdiff_vn` in mo_diffusion_nml.f90
apply_to_horizontal_wind: bool = True

#: Options for discretizing the Smagorinsky temperature diffusion
#: Called `itype_t_diffu` inmo_diffusion_nml.f90
self.type_t_diffu = type_t_diffu
#: If True, apply horizontal diffusion to temperature field
#: Called `lhdiff_temp` in mo_diffusion_nml.f90
apply_to_temperature: bool = True

#: Ratio of e-folding time to (2*)time step
#: Called `hdiff_efdt_ratio` inmo_diffusion_nml.f90
self.hdiff_efdt_ratio: float = hdiff_efdt_ratio
#: Options for discretizing the Smagorinsky momentum diffusion
#: Called `itype_vn_diffu` in mo_diffusion_nml.f90
type_vn_diffu: int = 1

#: Ratio of e-folding time to time step for w diffusion (NH only)
#: Called `hdiff_w_efdt_ratio` inmo_diffusion_nml.f90.
self.hdiff_w_efdt_ratio: float = hdiff_w_efdt_ratio
#: If True, compute 3D Smagorinsky diffusion coefficient
#: Called `lsmag_3d` in mo_diffusion_nml.f90
compute_3d_smag_coeff: bool = False

#: Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below
#: Called `hdiff_smag_fac` in mo_diffusion_nml.f90
self.smagorinski_scaling_factor: float = smagorinski_scaling_factor
#: Options for discretizing the Smagorinsky temperature diffusion
#: Called `itype_t_diffu` inmo_diffusion_nml.f90
type_t_diffu: int = 2

#: If True, apply truly horizontal temperature diffusion over steep slopes
#: Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90
self.apply_zdiffusion_t: bool = zdiffu_t
#: Ratio of e-folding time to (2*)time step
#: Called `hdiff_efdt_ratio` inmo_diffusion_nml.f90
hdiff_efdt_ratio: float = 36.0

#:slope threshold (temperature diffusion): is used to build up an index list for application of truly horizontal diffusion in mo_vertical_grid.f89
self.thslp_zdiffu = thslp_zdiffu
#: threshold [m] for height difference between adjacent grid points, defaults to 200m (temperature diffusion)
self.thhgtd_zdiffu = thhgtd_zdiffu
#: Ratio of e-folding time to time step for w diffusion (NH only)
#: Called `hdiff_w_efdt_ratio` inmo_diffusion_nml.f90.
hdiff_w_efdt_ratio: float = 15.0

# from other namelists:
# from parent namelist mo_nonhydrostatic_nml
#: Scaling factor for Smagorinsky diffusion at height hdiff_smag_z and below
#: Called `hdiff_smag_fac` in mo_diffusion_nml.f90
smagorinski_scaling_factor: float = 0.015

#: Number of dynamics substeps per fast-physics step
#: Called 'ndyn_substeps' in mo_nonhydrostatic_nml.f90
#: Number of dynamics substeps per fast-physics step
#: Called 'ndyn_substeps' in mo_nonhydrostatic_nml.f90
# TODO (magdalena) ndyn_substeps may dynamically increase during a model run in order to
# reduce instabilities. Need to figure out whether the parameter is the configured
# (constant!) one or the dynamical one. In the latter case it should be removed from
# DiffusionConfig and init()
ndyn_substeps: int = 5

# TODO (magdalena) ndyn_substeps may dynamically increase during a model run in order to
# reduce instabilities. Need to figure out whether the parameter is the configured
# (constant!) one or the dynamical one. In the latter case it should be removed from
# DiffusionConfig and init()
self.ndyn_substeps: int = n_substeps
#: If True, apply truly horizontal temperature diffusion over steep slopes
#: Called 'l_zdiffu_t' in mo_nonhydrostatic_nml.f90
apply_zdiffusion_t: bool = True

# namelist mo_gridref_nml.f90
#:slope threshold (temperature diffusion): is used to build up an index list for application of truly horizontal diffusion in mo_vertical_grid.f90
thslp_zdiffu: float = 0.025

#: Denominator for temperature boundary diffusion
#: Called 'denom_diffu_t' in mo_gridref_nml.f90
self.temperature_boundary_diffusion_denominator: float = (
temperature_boundary_diffusion_denom
)
#: threshold [m] for height difference between adjacent grid points, defaults to 200m (temperature diffusion)
thhgtd_zdiffu: float = 200.0

#: Denominator for velocity boundary diffusion
#: Called 'denom_diffu_v' in mo_gridref_nml.f90
self.velocity_boundary_diffusion_denominator: float = velocity_boundary_diffusion_denom
#: Denominator for velocity boundary diffusion
#: Called 'denom_diffu_v' in mo_gridref_nml.f90
velocity_boundary_diffusion_denominator: float = 200.0

# parameters from namelist: mo_interpol_nml.f90
#: Denominator for temperature boundary diffusion
#: Called 'denom_diffu_t' in mo_gridref_nml.f90
temperature_boundary_diffusion_denominator: float = 135.0

#: Parameter describing the lateral boundary nudging in limited area mode.
#:
#: Maximal value of the nudging coefficients used cell row bordering the boundary interpolation zone,
#: from there nudging coefficients decay exponentially with `nudge_efold_width` in units of cell rows.
#: Called `nudge_max_coeff` in mo_interpol_nml.f90
self.nudge_max_coeff: float = max_nudging_coeff
#: Parameter describing the lateral boundary nudging in limited area mode.
#:
#: Maximal value of the nudging coefficients used cell row bordering the boundary interpolation zone,
#: from there nudging coefficients decay exponentially with `nudge_efold_width` in units of cell rows.
#: Called `nudge_max_coeff` in mo_interpol_nml.f90
max_nudging_coefficient: float = 0.02

#: Exponential decay rate (in units of cell rows) of the lateral boundary nudging coefficients
#: Called `nudge_efold_width` in mo_interpol_nml.f90
self.nudge_efold_width: float = nudging_decay_rate
#: Exponential decay rate (in units of cell rows) of the lateral boundary nudging coefficients
#: Called `nudge_efold_width` in mo_interpol_nml.f90
nudge_efold_width: float = 2.0

#: Type of shear forcing used in turbulence
#: Called itype_shear in `mo_turbdiff_nml.f90
self.shear_type = shear_type
#: Type of shear forcing used in turbulence
#: Called itype_shear in `mo_turbdiff_nml.f90
shear_type: TurbulenceShearForcingType = TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND

def __post_init__(self):
self._validate()

def _validate(self):
"""Apply consistency checks and validation on configuration parameters."""
if self.diffusion_type != 5:
if self.diffusion_type != DiffusionType.SMAGORINSKY_4TH_ORDER:
raise NotImplementedError(
"Only diffusion type 5 = `Smagorinsky diffusion with fourth-order background "
"diffusion` is implemented"
)

if self.diffusion_type < 0:
self.apply_to_temperature = False
self.apply_to_horizontal_wind = False
self.apply_to_vertical_wind = False
if self.diffusion_type == DiffusionType.NO_DIFFUSION:
assert (
not self.apply_to_temperature
and not self.apply_to_horizontal_wind
and not self.apply_to_vertical_wind
), f"Inconsistent configuration: DiffusionType = {self.diffusion_type} - but application flags are set to: horizontal wind = '{self.apply_to_horizontal_wind}', vertical wind = '{self.apply_to_vertical_wind}', temperature = '{self.apply_to_temperature}' "

if self.shear_type not in (
TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND,
TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND,
):
raise NotImplementedError(
f"Turbulence Shear only {TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND} "
f"Turbulence Shear: only {TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND} "
f"and {TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND} "
f"implemented"
)
Expand Down Expand Up @@ -302,7 +282,7 @@ def __post_init__(self, config):
object.__setattr__(
self,
"scaled_nudge_max_coeff",
config.nudge_max_coeff * constants.DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO,
config.max_nudging_coefficient * constants.DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO,
)

def _determine_smagorinski_factor(self, config: DiffusionConfig):
Expand All @@ -314,12 +294,12 @@ def _determine_smagorinski_factor(self, config: DiffusionConfig):
It is calculated/used only in the case of diffusion_type 3 or 5
"""
match config.diffusion_type:
case 5:
case DiffusionType.SMAGORINSKY_4TH_ORDER:
(
smagorinski_factor,
smagorinski_height,
) = diffusion_type_5_smagorinski_factor(config)
case 4:
case DiffusionType.LINEAR_4TH_ORDER:
# according to mo_nh_diffusion.f90 this isn't used anywhere the factor is only
# used for diffusion_type (3,5) but the defaults are only defined for iequations=3
smagorinski_factor = (
Expand Down Expand Up @@ -704,10 +684,7 @@ def _do_diffusion_step(
offset_provider=self._grid.offset_providers,
)
log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end")
if (
self.config.shear_type
>= TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND
):
if self.config.shear_type != TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND:
log.debug(
"running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start"
)
Expand Down
19 changes: 16 additions & 3 deletions model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,10 +102,11 @@ def _get_or_initialize(experiment, backend, name):
return grid_functionality[experiment].get(name)


@pytest.mark.xfail
def test_diffusion_coefficients_with_hdiff_efdt_ratio(experiment):
config = construct_diffusion_config(experiment, ndyn_substeps=5)
config.hdiff_efdt_ratio = 1.0
config.hdiff_w_efdt_ratio = 2.0
config = diffusion.DiffusionConfig(
experiment, ndyn_substeps=5, hdiff_efdt_ratio=1.0, hdiff_w_efdt_ratio=2.0
)

params = diffusion.DiffusionParams(config)

Expand All @@ -115,6 +116,7 @@ def test_diffusion_coefficients_with_hdiff_efdt_ratio(experiment):
assert params.K4W == pytest.approx(1.0 / 72.0, abs=1e-12)


@pytest.mark.xfail
def test_diffusion_coefficients_without_hdiff_efdt_ratio(experiment):
config = construct_diffusion_config(experiment)
config.hdiff_efdt_ratio = 0.0
Expand All @@ -128,6 +130,7 @@ def test_diffusion_coefficients_without_hdiff_efdt_ratio(experiment):
assert params.K4W == 0.0


@pytest.mark.xfail
def test_smagorinski_factor_for_diffusion_type_4(experiment):
config = construct_diffusion_config(experiment, ndyn_substeps=5)
config.smagorinski_scaling_factor = 0.15
Expand All @@ -139,6 +142,7 @@ def test_smagorinski_factor_for_diffusion_type_4(experiment):
assert params.smagorinski_height is None


@pytest.mark.xfail
def test_smagorinski_heights_diffusion_type_5_are_consistent(
experiment,
):
Expand All @@ -163,6 +167,15 @@ def test_smagorinski_factor_diffusion_type_5(experiment):
assert np.all(params.smagorinski_factor >= np.zeros(len(params.smagorinski_factor)))


def vertical_grid(vertical_config: v_grid.VerticalGridConfig, grid_savepoint: sb.IconGridSavepoint):
return v_grid.VerticalGrid(
config=vertical_config,
vct_a=grid_savepoint.vct_a(),
vct_b=grid_savepoint.vct_b(),
_min_index_flat_horizontal_grad_pressure=grid_savepoint.nflat_gradp(),
)


@pytest.mark.datatest
def test_diffusion_init(
savepoint_diffusion_init,
Expand Down
Loading
Loading