import copy
import datetime
import glob
import os
import tracemalloc
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_ad(
self, mode, dataobs,
fp_header_nest, fp_header_glob, spec, ddi
):
"""Compute the adjoint flux sensitivity from FLEXPART/STILT footprints.
For each observation in *dataobs*, reads the footprint grid and
accumulates the adjoint flux sensitivity (``adj_out`` departure) into
the in-memory flux store ``self.dataflx``, indexed by period and
grid cell.
Args:
self: Lagrangian model plugin instance.
mode (str): ``'adj'``.
dataobs: CIF observation data-store for the period (contains
``'adj_out'`` departure values).
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.
"""
nlat, nlon = self.domain.zlat.shape
ndates = len(self.input_dates[ddi])
# Loading fluxes
flux = self.dataflx[("flux", spec)][ddi]["spec"]
flx_dates = pd.DatetimeIndex(flux.time.values).to_pydatetime()
# Execute parallel threads
nthreads = self.nthreads
thread_intervals = np.linspace(0, len(dataobs), nthreads + 1).astype(int)
@thread
def thread_function(ithread):
flx_sensit_tmp = np.zeros((ndates + 1, 1, nlat, nlon))
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, flx_sensit_tmp, flx_dates, obs_i
)
return flx_sensit_tmp
flx_sensit = sum(thread_function(range(nthreads)))
# Flush fluxes
self.dataflx[("flux", spec)][ddi]["spec"] = None
return flx_sensit
[docs]
def process_obs_row(
self, dataobs, ithread,
fp_header_nest, fp_header_glob,
flux, flx_sensit_tmp, flx_dates, obs_i
):
"""Compute the adjoint flux sensitivity for a single observation row.
Reads the footprint grids for the observation at *obs_i*, weights them
by ``adj_out``, and accumulates the result into *flx_sensit_tmp*.
Args:
self: Lagrangian model plugin instance.
dataobs: CIF observation data-store for the period.
ithread (int): thread index.
fp_header_nest: nested-domain FLEXPART header.
fp_header_glob: outer-domain FLEXPART header.
flux (dict): forward flux arrays (for background subtraction).
flx_sensit_tmp: thread-local accumulator for adjoint flux sensitivity.
flx_dates: dates of the flux grid.
obs_i (int): row index into the observation data-store.
"""
row = dataobs["metadata"].iloc[obs_i]
mainrow = dataobs["maindata"].iloc[obs_i]
station = row.station
network = row.network
spec = row.parameter
molarmass = getattr(self.chemistry.acspecies, spec.upper()).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])
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")
# Conversion of footprints
grid_nest *= self.coeff * self.mmair / molarmass
# 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
nest_sensit = grid_nest.T[:ngrid].reshape(ngrid, -1)
# Find time steps to compare
inds_flx = (np.argmin(
(np.array(gtime)[:, np.newaxis]
- flx_dates[np.newaxis, :]) >= datetime.timedelta(0),
axis=1) - 1) % len(flx_dates)
zeros = np.zeros((inds_flx.size, self.domain.zlon_in.size),
dtype=int)
np.add.at(
flx_sensit_tmp,
(inds_flx.reshape(-1, 1),
zeros,
zeros,
np.arange(self.domain.zlon_in.size)[np.newaxis, :]),
nest_sensit * mainrow.adj_out,
)
# 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
glob_sensit = grid_glob.T[:ngrid_glob].reshape(ngrid_glob, -1)
# Find time steps to compare
inds_flx = (np.argmin(
(np.array(gtime_glob)[:, np.newaxis]
- flx_dates[np.newaxis, :]) >= datetime.timedelta(0),
axis=1) - 1) % len(flx_dates)
zeros = np.zeros(
(inds_flx.size, self.domain.zlon.size - self.domain.zlon_in.size),
dtype=int)
np.add.at(
flx_sensit_tmp,
(inds_flx.reshape(-1, 1),
zeros,
zeros,
np.arange(self.domain.zlon_in.size,
self.domain.zlon.size)[np.newaxis, :]),
glob_sensit * mainrow.adj_out,
)
