import pandas as pd
import numpy as np
from .......utils.parallel import thread
from .......utils.check.errclass import CifError
def _to_int64(dates):
values = np.asarray(dates)
if values.dtype.kind != "M":
values = pd.to_datetime(values.ravel()).values.reshape(values.shape)
return values.astype("datetime64[ns]").view("int64")
[docs]
def adjoint(
transf, ddi, mapper, outputs, inout_datastore, onlyinit,
nthreads=1
):
if ddi not in mapper["interpol_indexes"]:
return
interpol_indexes = mapper["interpol_indexes"][ddi]
# Load input and output dates from transform
trid_ref = list(mapper["inputs"].keys())[0]
in_dates = mapper["inputs"][trid_ref]["input_dates"]
out_dates = mapper["outputs"][trid_ref]["input_dates"]
if np.size(out_dates[ddi]) == 1:
raise CifError("Output dates should not be single for adjoint interpolation")
# Work on int64 dates for fast comparisons
out_start = _to_int64(out_dates[ddi]["start_date"])
out_end = _to_int64(out_dates[ddi]["end_date"])
# Output dates from output datastore
ds_start = _to_int64(outputs[trid_ref][ddi]["metadata"]["date"])
ds_end = _to_int64(outputs[trid_ref][ddi]["metadata"]["enddate"])
# Match (start, end) pairs of the datastore with their first occurrence
# in output dates. Output dates come first in the stacked array, so the
# first occurrence of any pair present in output dates is among them
# (lexsort is much faster than np.unique with axis=0)
n_out = len(out_start)
all_start = np.concatenate([out_start, ds_start])
all_end = np.concatenate([out_end, ds_end])
order = np.lexsort((all_end, all_start))
sorted_start = all_start[order]
sorted_end = all_end[order]
new_group = np.ones(len(order), dtype=bool)
new_group[1:] = (sorted_start[1:] != sorted_start[:-1]) \
| (sorted_end[1:] != sorted_end[:-1])
# lexsort is stable, so the first row of each group is its first occurrence
group_first = order[new_group]
first_index = np.empty(len(order), dtype=int)
first_index[order] = group_first[np.cumsum(new_group) - 1]
ds_first = first_index[n_out:]
if np.any(ds_first >= n_out):
raise CifError(
"Some dates in the output datastore are not in the output dates "
"of the temporal interpolation")
# Initializing inputs
inout_datastore["inputs"][trid_ref] = {}
# Initialize metadata if not already done
transf.metadata = getattr(transf, "metadata", {})
transf.metadata[ddi] = {}
list_trids = list(mapper["inputs"].keys())
nthreads_tmp = min(nthreads, len(list_trids))
thread_intervals = np.linspace(
0, len(list_trids), nthreads_tmp + 1
).astype(int)
# Now loop over target dates
for ddtarget in interpol_indexes:
# Crop out_dates to the date interval covered by input dates
in_all = _to_int64(in_dates[ddtarget]).ravel()
in_all = in_all[in_all != np.iinfo(np.int64).min]
in_min, in_max = in_all.min(), in_all.max()
# The crop mask only depends on the (start, end) pair, so duplicates
# are either all kept or all dropped: the position of the first
# occurrence in the cropped output dates is a simple cumsum
out_mask = (out_start <= in_max) & (out_end >= in_min)
ds_mask = (ds_start <= in_max) & (ds_end >= in_min)
ds_rows = np.flatnonzero(ds_mask)
out_unique_index = (np.cumsum(out_mask) - 1)[ds_first[ds_rows]]
# Inner join of datastore rows with interpolation indexes;
# one datastore row can match several interpolation rows
inter_index = interpol_indexes[ddtarget]
inter_keys = inter_index.index.values
order = np.argsort(inter_keys, kind="stable")
sorted_keys = inter_keys[order]
lo = np.searchsorted(sorted_keys, out_unique_index, side="left")
hi = np.searchsorted(sorted_keys, out_unique_index, side="right")
counts = hi - lo
left_rows = np.repeat(np.arange(len(out_unique_index)), counts)
offsets = np.arange(counts.sum()) \
- np.repeat(np.cumsum(counts) - counts, counts)
right_rows = order[np.repeat(lo, counts) + offsets]
target_index = ds_rows[left_rows]
data_index = out_unique_index[left_rows]
weights = inter_index["weights"].values[right_rows]
tstep = inter_index["indexes"].values[right_rows].astype(int)
dtstep = np.ones(len(tstep), dtype=int)
# Update metadata for later forward run
wgt = pd.DataFrame({
("interpol_indexes", "weights"): weights,
("target_index", ""): target_index,
}, index=data_index)
transf.metadata[ddi][ddtarget] = {
"weights": wgt,
"target_mask": ds_mask
}
@thread
def thread_function(ithread):
for itrid in range(thread_intervals[ithread], thread_intervals[ithread + 1]):
trid = list_trids[itrid]
columns = {
col: values.values[target_index]
for col, values in outputs[trid][ddi].items()
}
columns[("target_index", "")] = target_index
columns[("metadata", "tstep")] = tstep
columns[("metadata", "dtstep")] = dtstep
if not onlyinit:
columns[("maindata", "adj_out")] = \
columns[("maindata", "adj_out")] * weights
inout_datastore["inputs"][trid][ddtarget] = pd.DataFrame(
columns, index=data_index)
# Apply threaded function
thread_function(range(nthreads_tmp))
