import copy
import datetime
import glob
import os
from logging import debug
import numpy as np
import pandas as pd
from ......utils.parallel import thread
from ...utils.read import read_footprint_grid
from ......utils.check.errclass import CifError
[docs]
def flux_contribution(
self, mode, dataobs,
fp_header_nest, fp_header_glob, spec, ddi,
batch_computation=False
):
"""Compute the flux–concentration contribution for one species and period.
Convolves pre-loaded flux fields with FLEXPART/STILT footprints to
produce simulated concentration increments for each observation in
*dataobs*. For TL mode, also convolves the flux increment.
Args:
self: Lagrangian model plugin instance.
mode (str): ``'fwd'`` or ``'tl'``.
dataobs: CIF observation data-store for the period.
fp_header_nest: FLEXPART header for the nested-domain footprints.
fp_header_glob: FLEXPART header for the outer-domain footprints.
spec (str): species name.
ddi (datetime): sub-simulation period start.
batch_computation (bool): if ``True``, process all ensemble members
in a single pass.
Returns:
dataobs updated with simulated ``'spec'`` (and optionally ``'incr'``)
columns.
"""
# Loading fluxes
flux = {spec: {"spec": self.dataflx[("flux", spec)][ddi]["spec"]}}
if mode == "tl" and "incr" in self.dataflx[("flux", spec)][ddi]:
flux[spec]["incr"] = self.dataflx[("flux", spec)][ddi]["incr"]
# Initialize debug columns if needed
if self.dump_debug:
dataobs[spec][("flexpart", "nest")] = 0.
dataobs[spec][("flexpart", "nest_tl")] = 0.
dataobs[spec][("flexpart", "glob")] = 0.
dataobs[spec][("flexpart", "glob_tl")] = 0.
# Specific treatment for batch computing
sample_species = [spec]
if batch_computation:
ref_species = self.perturbed_species[spec]
sample_species = [
s for s in self.chemistry.acspecies.attributes
if self.perturbed_species[s] == ref_species
]
flux = {
s: {"spec": self.dataflx[("flux", s)][ddi]["spec"]}
for s in sample_species
}
if mode == "tl" and "incr" in self.dataflx[("flux", spec)][ddi]:
for s in sample_species:
flux[s]["incr"] = self.dataflx[("flux", s)][ddi]["incr"]
if self.dump_debug:
for s in sample_species:
dataobs[s][("flexpart", "nest")] = 0.
dataobs[s][("flexpart", "nest_tl")] = 0.
dataobs[s][("flexpart", "glob")] = 0.
dataobs[s][("flexpart", "glob_tl")] = 0.
# Execute parallel threads
nthreads = self.nthreads
nobs = len(dataobs[sample_species[0]])
thread_intervals = np.linspace(0, nobs, nthreads + 1).astype(int)
@thread
def thread_function(ithread):
for obs_i in range(thread_intervals[ithread], thread_intervals[ithread + 1]):
process_obs_row(
self, dataobs, ithread,
fp_header_nest, fp_header_glob,
flux, obs_i, sample_species, batch_computation
)
thread_function(range(nthreads))
# Flush fluxes
for s in sample_species:
self.dataflx[("flux", s)][ddi]["spec"] = None
self.dataflx[("flux", s)][ddi]["incr"] = None
[docs]
def process_obs_row(self, dataobs, ithread,
fp_header_nest, fp_header_glob,
flux, obs_i, sample_species, batch_computation):
"""Compute the flux–concentration contribution for a single observation row.
Reads the nested-domain and (optionally) outer-domain footprint grids
for the observation at row *obs_i*, convolves them with *flux*, and
accumulates the result into *dataobs*.
Args:
self: Lagrangian model plugin instance.
dataobs (dict): species-keyed observation data-stores.
ithread (int): thread index (used for debug logging).
fp_header_nest: FLEXPART header for the nested-domain footprints.
fp_header_glob: FLEXPART header for the outer-domain footprints.
flux (dict): forward flux arrays keyed by species.
obs_i (int): row index into the observation data-store.
sample_species (list): species names to process (may include
ensemble-sample variants).
batch_computation (bool): ensemble batch mode flag.
"""
ref_spec = sample_species[0]
row = dataobs[ref_spec]["metadata"].iloc[obs_i]
station = row.station
network = row.network
molarmass = getattr(self.chemistry.acspecies, ref_spec).molarmass
# Translate station name if needed
if hasattr(self, "station_name_dict"):
station = self.dict_station_name[station.upper()].upper()
# Infer folder structure
subdir = row.date.strftime(self.footprint_dir_format)
release_date = row.date - pd.to_timedelta(self.release_shift)
file_date = release_date.strftime(self.footprint_date_format)
# Read nested grids
runsubdir_nest = os.path.join(
self.run_dir_nest,
self.footprint_stat_subdir_format.format(
stat=station.upper(), network=network.upper()),
subdir)
file_name = self.file_nest_format.format(
date=file_date, stat=station.upper(),
network=network.upper())
list_valid = glob.glob(os.path.join(runsubdir_nest, file_name))
if list_valid == []:
debug(f"WARNING: file not found: {os.path.join(runsubdir_nest, file_name)}")
return
elif len(list_valid) > 1:
raise CifError(
f"Multiple files fit the specified format {self.file_ini_format}. "
f"This can be related to the use of a wildcard... "
f"Please check your yml"
)
file_name = os.path.basename(list_valid[0])
if self.preloaded_footprints.get("loaded_file", "") == list_valid[0] \
and self.reload_footprints:
debug("Using preloaded concentrations from "
f"{file_name} for station {station}")
for s in sample_species:
sample_flux = flux[s]
for data_id in sample_flux:
dataobs[s].iloc[
obs_i, dataobs[s].columns.get_loc(("maindata", data_id))
] = self.preloaded_footprints[data_id][s]
else:
debug(f"Thread #{ithread}: Reading {file_name} for station {station}")
grid_nest, gtime, ngrid, valid_file = \
read_footprint_grid(self,
runsubdir_nest, file_name, release_date, fp_header_nest,
numscale=self.numscale, stilt=self.footprint_type == "STILT")
if self.reload_footprints:
self.preloaded_footprints = {
"loaded_file": list_valid[0],
"data": (
copy.deepcopy(grid_nest),
copy.deepcopy(gtime),
copy.deepcopy(ngrid),
copy.deepcopy(valid_file)
),
"spec": {},
"incr": {}
}
# Conversion of footprints
grid_nest *= self.coeff * self.mmair / molarmass
# Multiply footprints with fluxes for fwd and tl
for s in sample_species:
sample_flux = flux[s]
for data_id in sample_flux:
dataflx = sample_flux[data_id]
flx_dates = pd.DatetimeIndex(
dataflx.time.values).to_pydatetime()
# Check that dates are compatible
if np.any(np.array(gtime) < flx_dates.min()):
raise CifError(
f"Footprints span beyond the first available flux date: \nMin footprint date: {np.min(gtime)}\nMin flux date: {flx_dates.min()}")
nest_sim = 0
if valid_file:
inds_flx = (np.argmin(
(np.array(gtime)[:, np.newaxis]
- flx_dates[np.newaxis, :]) >= datetime.timedelta(0),
axis=1) - 1) % len(flx_dates)
# Apply decay if any
decay_coef = np.ones((ngrid, 1))
if hasattr(self, "exp_decay"):
exp_decay = self.exp_decay
halflife = pd.to_timedelta(
exp_decay.halflife) / np.log(2)
decay_coef = np.exp(
-np.maximum(0, (row.date -
np.array(gtime)) / halflife)
)[:, np.newaxis]
if exp_decay.inverse_decay:
decay_coef = 1 - decay_coef
# Compute sum of contributions
nest_sim = np.nansum((
grid_nest.T[:ngrid].reshape(ngrid, -1) * decay_coef
* dataflx[inds_flx, 0, 0, :self.domain.zlon_in.size])
.values)
# Filling simulation
dataobs[s].iloc[
obs_i, dataobs[s].columns.get_loc(("maindata", data_id))
] = nest_sim
print(s, obs_i, nest_sim, dataobs[s].iloc[
obs_i, dataobs[s].columns.get_loc(("maindata", data_id))
])
# Dump debug
if self.dump_debug:
if data_id == "spec":
sim_col = "nest"
else:
sim_col = "nest_tl"
dataobs[s].iloc[
obs_i, dataobs[s].columns.get_loc(
("flexpart", sim_col))
] = nest_sim
if self.reload_footprints:
self.preloaded_footprints[data_id][s] = nest_sim
# Read global footprints
# TODO: read correction factor dry air!
if not self.domain.nested:
return
runsubdir_glob = os.path.join(
self.run_dir_glob, station.upper(), subdir)
file_name = self.file_glob_format.format(
date=file_date, stat=station.upper(), network=network.upper())
list_valid = glob.glob(os.path.join(runsubdir_nest, file_name))
if list_valid == []:
debug(f"WARNING: file not found: {os.path.join(runsubdir_nest, file_name)}")
return
elif len(list_valid) > 1:
raise CifError(
f"Multiple files fit the specified format {self.file_ini_format}. "
f"This can be related to the use of a wildcard... "
f"Please check your yml"
)
file_name = os.path.basename(list_valid[0])
debug(f"Thread #{ithread}: Reading {file_name} for station {station}")
grid_glob, gtime_glob, ngrid_glob, valid_file = \
read_footprint_grid(self,
runsubdir_glob, file_name, release_date, fp_header_glob,
numscale=self.numscale)
# Conversion of footprints
grid_glob *= self.coeff * self.mmair / molarmass
# Keep only valid grids
grid_glob = grid_glob.T[:ngrid_glob].reshape(ngrid_glob, -1)
# Multiply footprints with fluxes for fwd and tl
for s in sample_species:
sample_flux = flux[s]
for data_id in sample_flux:
dataflx = sample_flux[data_id]
flx_dates = pd.DatetimeIndex(
dataflx.time.values).to_pydatetime()
inds_flx_glob = (np.argmin(
(np.array(gtime_glob)[:, np.newaxis]
- flx_dates[np.newaxis, :]) >= datetime.timedelta(0),
axis=1) - 1) % len(flx_dates)
# Multiply by fluxes
glob_sim = (
grid_glob
* dataflx[inds_flx_glob, 0, 0,
self.domain.zlon_in.size:].values
)
# Remove nest domain
glob_sim[:, self.domain.raveled_indexes_glob] = 0.
glob_sim = glob_sim.sum()
# Filling simulation
dataobs[s].iloc[
obs_i, dataobs[s].columns.get_loc(("maindata", data_id))
] += glob_sim
# Dump debug
if self.dump_debug:
if data_id == "spec":
sim_col = "glob"
else:
sim_col = "glob_tl"
dataobs[sample_species].iloc[
obs_i, dataobs[s].columns.get_loc(("flexpart", sim_col))
] = glob_sim
