from __future__ import annotations
import datetime
from os import PathLike
from pathlib import Path
from typing import Any, Literal
import numpy as np
import xarray as xr
from numpy.random import f
from ......utils import path
from ......utils.check.errclass import CifValueError
from ......utils.hdf5 import _hdf5_lock
[docs]
def get_species(self, input_type: str) -> list[str]:
"""Return the species list for a given chemistry input type"""
if input_type == "prescrconcs":
return self.chemistry.prescribed_species
elif input_type == "prodloss3d":
return self.chemistry.prodloss_species
elif input_type == "deposition":
return self.chemistry.deposition_species
elif input_type == "photorates":
return self.chemistry.jrates_varname
else:
raise CifValueError(f"Unknown chemistry input '{input_type}'")
[docs]
def make_chemfields(
self,
datastore: dict[tuple[str, str], dict[str | datetime.datetime, Any]],
input_type: Literal["prodloss3d", "prescrconcs", "deposition", "photorates"],
ddi: datetime.datetime,
runsubdir: str | PathLike,
mode: Literal["fwd", "tl", "adj"],
) -> None:
"""Write chemistry concentration fields (prescribed, deposition, etc.) for LMDZ-ico"""
for spec in get_species(self, input_type):
trid = (input_type, spec)
if trid not in datastore:
continue
tracer = datastore[trid]["tracer"]
data = datastore[trid]["data"][ddi]
if "spec" in data:
target_path = Path(runsubdir, f"{input_type}.nc")
self.chemfields.write(spec, target_path, data["spec"])
# If species is in the control vector, dump tangent-linear increments
if tracer.iscontrol and mode == "tl":
if input_type not in ("prodloss3d", "prescrconcs"):
raise CifValueError(
f"Chemistry input '{input_type}' is not supported for tangent-linear mode"
)
if "incr" not in data:
data["incr"] = data["spec"].copy(data=np.zeros(data["spec"].shape))
self.chemfields.write(f"{spec}_tl", target_path, data["incr"])
else:
dirorig = datastore[trid]["dirorig"]
fileorig = datastore[trid]["fileorig"]
origin_path = Path(dirorig, ddi.strftime(fileorig))
target_path = Path(runsubdir, f"{input_type}_{spec}.nc")
# Special legacy case for 'prodloss3d' and 'deposition' inputs
if input_type == "prodloss3d":
accepted_varnames = (spec, f"{spec}_prod")
elif input_type == "deposition":
accepted_varnames = (spec, f"Dep_{spec}")
else:
accepted_varnames = (spec,)
if not tracer.varname or tracer.varname in accepted_varnames:
# Tracer varname input argument is not specified or is identical
# to the parameter name, can link safely
path.link(origin_path, target_path)
else:
# Tracer varname input argument is different from the parameter
# name, need to rename the variable so LMDZ can finfd it
with _hdf5_lock:
with xr.open_dataset(origin_path) as ds:
if tracer.varname not in ds:
raise CifValueError(
f"Variable '{tracer.varname}' not found in file '{origin_path}'"
)
ds = ds[[tracer.varname]].rename({tracer.varname: spec})
ds.to_netcdf(target_path)
[docs]
def make_kinetic(
self,
datastore: dict[tuple[str, str], dict[str | datetime.datetime, Any]],
ddi: datetime.datetime,
runsubdir: str | PathLike,
) -> None:
"""Write the kinetic (pressure/temperature) field for one period"""
target_path = Path(runsubdir, "kinetic.nc")
for trid in datastore:
data = datastore[trid]["data"][ddi]
if "spec" in data:
_, var_name = trid
self.chemfields.write(var_name, target_path, data["spec"])
else:
input_file_list = list(set(datastore[trid]["input_files"][ddi]))
if len(input_file_list) == 0:
raise CifValueError("There is no file to link to 'kinetic.nc'")
if len(input_file_list) > 1:
raise CifValueError("There is multiple files to link to 'kinetic.nc'")
(input_file,) = input_file_list
path.link(input_file, target_path)
