def compute_change(
before: Checkpoint,
after: Checkpoint,
tendency_variables: Set[str],
storage_variables: Set[str],
name: str,
timestep: float,
):
diags = {}
delp_before = before[DELP]
delp_after = after[DELP]
# Compute statistics
for variable in tendency_variables:
diag_name = f"tendency_of_{variable}_due_to_{name}"
diags[diag_name] = (after[variable] - before[variable]) / timestep
if "units" in before[variable].attrs:
diags[diag_name].attrs["units"] = before[variable].units + "/s"
for variable in storage_variables:
path_before = vcm.mass_integrate(before[variable], delp_before, "z")
path_after = vcm.mass_integrate(after[variable], delp_after, "z")
diag_name = f"storage_of_{variable}_path_due_to_{name}"
diags[diag_name] = (path_after - path_before) / timestep
if "units" in before[variable].attrs:
diags[diag_name].attrs[
"units"] = before[variable].units + " kg/m**2/s"
mass_change = (delp_after - delp_before).sum("z") / timestep
mass_change.attrs["units"] = "Pa/s"
diags[f"storage_of_mass_due_to_{name}"] = mass_change
return diags
def compute_baseline_diagnostics(state: State) -> Diagnostics:
return dict(
water_vapor_path=vcm.mass_integrate(state[SPHUM], state[DELP], dim="z")
.assign_attrs(units="mm")
.assign_attrs(description="column integrated water vapor"),
physics_precip=(state[PHYSICS_PRECIP_RATE])
.assign_attrs(units="kg/m^2/s")
.assign_attrs(
description="surface precipitation rate due to parameterized physics"
),
)
def water_vapor_path(self):
try:
return self._mapper["water_vapor_path"]
except KeyError:
da = vcm.mass_integrate(
self._mapper["specific_humidity"],
self._mapper["pressure_thickness_of_atmospheric_layer"],
dim="z",
)
return da.assign_attrs(
{"long_name": "column integrated water vapor", "units": "mm"}
)
def compute_ml_momentum_diagnostics(state: State, tendency: State) -> Diagnostics:
delp = state[DELP]
dQu = tendency.get("dQu", xr.zeros_like(delp))
dQv = tendency.get("dQv", xr.zeros_like(delp))
column_integrated_dQu = vcm.mass_integrate(dQu, delp, "z")
column_integrated_dQv = vcm.mass_integrate(dQv, delp, "z")
return dict(
dQu=dQu.assign_attrs(units="m s^-2").assign_attrs(
description="zonal wind tendency due to ML"
),
dQv=dQv.assign_attrs(units="m s^-2").assign_attrs(
description="meridional wind tendency due to ML"
),
column_integrated_dQu_stress=column_integrated_dQu.assign_attrs(
units="Pa", description="column integrated zonal wind tendency due to ML",
),
column_integrated_dQv_stress=column_integrated_dQv.assign_attrs(
units="Pa",
description="column integrated meridional wind tendency due to ML",
),
)
def _apply_physics(self) -> Diagnostics:
self._log_debug(f"Physics Step (apply)")
self._fv3gfs.apply_physics()
micro = self._fv3gfs.get_diagnostic_by_name(
"tendency_of_specific_humidity_due_to_microphysics").data_array
delp = self._state[DELP]
return {
"storage_of_specific_humidity_path_due_to_microphysics":
vcm.mass_integrate(micro, delp, "z"),
"evaporation":
self._state["evaporation"],
"cnvprcp_after_physics":
self._fv3gfs.get_diagnostic_by_name("cnvprcp").data_array,
"total_precip_after_physics":
self._state[TOTAL_PRECIP],
}
def insert_column_integrated_vars(
ds: xr.Dataset, column_integrated_vars: Sequence[str]
) -> xr.Dataset:
"""Insert column integrated (<*>) terms,
really a wrapper around vcm.calc.thermo funcs"""
for var in column_integrated_vars:
column_integrated_name = f"column_integrated_{var}"
if "Q1" in var:
da = vcm.column_integrated_heating_from_isochoric_transition(
ds[var], ds[DELP]
)
elif "Q2" in var:
da = -vcm.minus_column_integrated_moistening(ds[var], ds[DELP])
da = da.assign_attrs(
{"long_name": "column integrated moistening", "units": "mm/day"}
)
else:
da = vcm.mass_integrate(ds[var], ds[DELP], dim="z")
ds = ds.assign({column_integrated_name: da})
return ds
def __call__(self, time, state):
diagnostics: Diagnostics = {}
delp = state[DELP]
prediction: State = predict(self.model, state)
tendency, state_updates = {}, {}
for key, value in prediction.items():
if is_state_update_variable(key, state):
state_updates[key] = value
else:
tendency[key] = value
for name in state_updates.keys():
diagnostics[name] = state_updates[name]
dQ1_initial = tendency.get("dQ1", xr.zeros_like(state[SPHUM]))
dQ2_initial = tendency.get("dQ2", xr.zeros_like(state[SPHUM]))
dQ1_updated, dQ2_updated = non_negative_sphum(
state[SPHUM],
dQ1_initial,
dQ2_initial,
dt=self.timestep,
)
if "dQ1" in tendency:
if self.hydrostatic:
heating = vcm.column_integrated_heating_from_isobaric_transition(
dQ1_updated - tendency["dQ1"], delp, "z")
else:
heating = vcm.column_integrated_heating_from_isochoric_transition(
dQ1_updated - tendency["dQ1"], delp, "z")
heating = heating.assign_attrs(
long_name=
"Change in ML column heating due to non-negative specific "
"humidity limiter")
diagnostics.update({
"column_integrated_dQ1_change_non_neg_sphum_constraint":
heating
})
tendency.update({"dQ1": dQ1_updated})
if "dQ2" in tendency:
moistening = vcm.mass_integrate(dQ2_updated - tendency["dQ2"],
delp,
dim="z")
moistening = moistening.assign_attrs(
units="kg/m^2/s",
long_name=
"Change in ML column moistening due to non-negative specific "
"humidity limiter",
)
diagnostics.update({
"column_integrated_dQ2_change_non_neg_sphum_constraint":
moistening
})
tendency.update({"dQ2": dQ2_updated})
diagnostics["specific_humidity_limiter_active"] = xr.where(
dQ2_initial != dQ2_updated, 1, 0)
return (
tendency,
diagnostics,
state_updates,
)
def compute_diagnostics(
state: State, tendency: State, label: str, hydrostatic: bool
) -> Diagnostics:
delp = state[DELP]
temperature_tendency_name = "dQ1"
humidity_tendency_name = "dQ2"
temperature_tendency = tendency.get(temperature_tendency_name, xr.zeros_like(delp))
humidity_tendency = tendency.get(humidity_tendency_name, xr.zeros_like(delp))
# compute column-integrated diagnostics
if hydrostatic:
net_heating = vcm.column_integrated_heating_from_isobaric_transition(
temperature_tendency, delp, "z"
)
else:
net_heating = vcm.column_integrated_heating_from_isochoric_transition(
temperature_tendency, delp, "z"
)
diags: Diagnostics = {
f"net_moistening_due_to_{label}": vcm.mass_integrate(
humidity_tendency, delp, dim="z"
).assign_attrs(
units="kg/m^2/s",
description=f"column integrated moisture tendency due to {label}",
),
f"column_heating_due_to_{label}": net_heating.assign_attrs(
units="W/m^2"
).assign_attrs(description=f"column integrated heating due to {label}"),
}
delp_tendency = STATE_NAME_TO_TENDENCY[DELP]
if delp_tendency in tendency:
net_mass_tendency = vcm.mass_integrate(
xr.ones_like(tendency[delp_tendency]), tendency[delp_tendency], dim="z"
).assign_attrs(
units="kg/m^2/s",
description=f"column-integrated mass tendency due to {label}",
)
diags[f"net_mass_tendency_due_to_{label}"] = net_mass_tendency
# add 3D tendencies to diagnostics
if label == "nudging":
diags_3d: Mapping[Hashable, xr.DataArray] = {
f"{TENDENCY_TO_STATE_NAME[k]}_tendency_due_to_nudging": v
for k, v in tendency.items()
}
elif label == "machine_learning":
diags_3d = {
"dQ1": temperature_tendency.assign_attrs(units="K/s").assign_attrs(
description=f"air temperature tendency due to {label}"
),
"dQ2": humidity_tendency.assign_attrs(units="kg/kg/s").assign_attrs(
description=f"specific humidity tendency due to {label}"
),
}
diags.update(diags_3d)
# add 3D state to diagnostics for backwards compatibility
diags.update({TEMP: state[TEMP], SPHUM: state[SPHUM], DELP: state[DELP]})
return diags