import filecmp
import copy
import os
import shutil
import pathlib
import numpy as np
import pandas as pd
import xarray as xr
import datetime
from logging import debug
from netCDF4 import Dataset
from ......utils import path
from ......utils.hdf5 import _hdf5_lock
from .utils import replace_dates
[docs]
def make_fluxes(self, datastore, runsubdir, datei, mode):
"""Make AEMISSIONS.nc and BEMISSIONS.nc files for CHIMERE.
Use chemical scheme to check which species is needed and either take it
from the datastore (i.e. when defined in the control vector), or take it
from prescribed emissions
Args:
self (pycif.utils.classes.fluxes.Flux): Flux plugin with all
attributes
datastore (dict): information on flux species
runsubdir (str): directory to the current run
nho (int): number of hours in the run
mode (str): running mode: 'fwd', 'tl' or 'adj'
"""
# Replace name when batch ensemble computing
datastore = {
(trid[0], trid[1].replace("__sample#", "")): datastore[trid]
for trid in datastore
if trid[0] in ["flux", "bioflux"]
}
# List of dates for which emissions are needed
list_dates = pd.date_range(datei, periods=self.nhours + 1, freq="h")
# Getting the right emissions
# Loop on all anthropogenic and biogenic species
# If in datastore, take data, otherwise, link to original A/B EMISSIONS
list_trid = [("flux", spec)
for spec in self.chemistry.emis_species.attributes] \
+ [("bioflux", spec)
for spec in self.chemistry.bio_species.attributes]
data2dump = {}
data2dump_tl = {}
for trid in list_trid:
spec = trid[1]
emis_type = trid[0]
# Update ref_trid if in ensemble mode
ref_trid = (emis_type, spec)
perturbed_from_reference = False
if hasattr(self, "perturbed_species"):
if spec in self.perturbed_species:
ref_trcr = self.perturbed_species[spec]
ref_trid = (emis_type, ref_trcr)
# If spec not explicitly defined in datastore,
# fetch general component information if available
if trid in datastore:
pass
elif trid in [(k[0], k[1].replace("__sample#", "")) for k in datastore]:
trid_ind = [
(k[0], k[1].replace("__sample#", "")) for k in datastore
].index(trid)
trid = list(datastore.keys())[trid_ind]
# If trid comes from an unperturbed input, but in ensemble mode
# needs to fetch info from the reference species
elif ref_trid in datastore:
trid = ref_trid
perturbed_from_reference = True
elif trid not in datastore and (emis_type, "") in datastore:
trid = (emis_type, "")
else:
continue
# Bio or anthro file
if emis_type == "flux":
file_emisout = f"{runsubdir}/AEMISSIONS.nc"
file_emisincrout = f"{runsubdir}/AEMISSIONS.increment.nc"
else:
file_emisout = f"{runsubdir}/BEMISSIONS.nc"
file_emisincrout = f"{runsubdir}/BEMISSIONS.increment.nc"
tracer = datastore[trid]
tracer_data = tracer["data"][datei]
# If no data is provided, just copy from original file
if "spec" not in tracer_data:
dirorig = tracer["dirorig"]
fileorig = tracer["fileorig"]
fileemis = datei.strftime(f"{dirorig}/{fileorig}")
# If does not exist, just link
linked = False
if not os.path.isfile(file_emisout):
path.link(fileemis, file_emisout)
linked = True
# Otherwise, check for difference
if not linked or perturbed_from_reference:
# First check that files are binary identical
if not filecmp.cmp(fileemis, file_emisout):
# Now check that spec emissions are correct
with _hdf5_lock:
with Dataset(file_emisout, "r") as fout:
with Dataset(fileemis, "r") as fin:
emisin = fin.variables[spec][:]
overwrite_variable = True
if spec in fout.variables:
emisout = fin.variables[spec][:]
overwrite_variable = ~np.all(emisin == emisout)
emisin = xr.DataArray(
emisin,
coords={"time": list_dates},
dims=("time", "lev", "lat", "lon"),
)
# If emission file is still a link, should be copied
# to be able to modify it locally
if overwrite_variable:
if os.path.islink(file_emisout):
file_orig = pathlib.Path(file_emisout).resolve()
os.unlink(file_emisout)
shutil.copy(file_orig, file_emisout)
# Now saves to buffer variable before saving
data2dump[spec] = copy.deepcopy(emisin)
# # Now writes the new value for the corresponding species
# flx_plg.write(spec, file_emisout, emisin)
# Repeat operations for tangent linear
if mode == "tl":
# If does not exist, copy
if not os.path.isfile(file_emisincrout):
shutil.copy(fileemis, file_emisincrout)
# Set variable to 0 if exists
with _hdf5_lock:
with Dataset(file_emisincrout, "a") as fout:
if spec in fout.variables:
fout.variables[spec][:] = 0.0
continue
# Otherwise, add new species at zero
flx_incr = xr.DataArray(
np.zeros(
(
len(list_dates),
self.nlevemis if emis_type == "flux"
else self.nlevemis_bio,
self.domain.nlat,
self.domain.nlon,
)
),
coords={"time": list_dates},
dims=("time", "lev", "lat", "lon"),
)
# Now saves to buffer variable before saving
data2dump_tl[spec] = copy.deepcopy(flx_incr)
# flx_plg.write(spec, file_emisincrout, flx_incr)
else:
# Replace existing link by copy of original file to modify it
path.copyfromlink(file_emisout)
# Put in dataset and write to input
flx_fwd = tracer_data["spec"]
# flx_plg.write(spec, file_emisout, flx_fwd)
data2dump[spec] = copy.deepcopy(flx_fwd)
if mode == "tl":
path.copyfromlink(file_emisincrout)
flx_tl = tracer_data.get("incr", 0.0 * flx_fwd)
# flx_plg.write(spec, file_emisincrout, flx_tl)
data2dump_tl[spec] = copy.deepcopy(flx_tl)
emis_type_ref = emis_type
flx_plg = self.flux if emis_type == "flux" else self.bioflux
# Now dump buffer data if any
if len(data2dump) > 0:
flx_plg.write(
list(data2dump.keys()), file_emisout, data2dump, comp_type=emis_type_ref
)
# Now dump buffer data if any
if len(data2dump_tl) > 0:
flx_plg.write(
list(data2dump_tl.keys()), file_emisincrout, data2dump_tl, comp_type=emis_type_ref
)
# Check that the dates are consistent with what CHIMERE expects
replace_dates(file_emisout, list_dates, self.ignore_input_dates)
if mode == "tl":
replace_dates(file_emisincrout, list_dates, self.ignore_input_dates)
