"""
=============
Interpolation
=============
Provides interpolation algorithms across tabular data for ``vivarium``
simulations.
"""
from __future__ import annotations
import math
from collections.abc import Hashable, Sequence
from typing import Any, ClassVar, NamedTuple
import numpy as np
import numpy.typing as npt
import pandas as pd
from vivarium.engine.types import has_named_row_index
_START_SUFFIX = "_start"
_END_SUFFIX = "_end"
[docs]
class ContinuousParameter(NamedTuple):
"""A continuous (binned) lookup parameter."""
name: str
"""Base name, e.g. ``"age"``; the interpolant column."""
start_column: str
"""Left-edge column name; a merge key."""
end_column: str
"""Right-edge column name."""
left_edges: npt.NDArray[Any]
"""Ascending unique left bin edges, shared across categorical groups."""
max_right: Any
"""Maximum right bin edge (the exclusive upper bound of the last bin)."""
[docs]
class Interpolation:
"""A callable that interpolates value columns over categorical/continuous parameters.
Lookup attributes are inferred from the input data: when the data has its
lookup attributes in the row index (see
``vivarium.engine.types.has_named_row_index``), the
row-index level names are the attributes; when the data is a flat
DataFrame (deprecated), the attributes are the columns not listed in
``value_columns``. Attributes whose names follow the ``<name>_start`` /
``<name>_end`` convention are paired up as continuous binned parameters;
all other attributes are treated as categorical key columns.
This naming convention is the only way ``Interpolation`` distinguishes
continuous from categorical parameters, so the class is not fully generic
— it is the interpolation primitive that
:class:`~vivarium.engine.framework.lookup.table.LookupTable` is built on.
For order 0 (currently the only supported order) a call resolves each interpolant to
the bin its continuous parameters fall in and returns that bin's values.
This requires that the key columns uniquely identify a data row and
that the bin edges are identical across every categorical group; both are
validated on construction when ``validate`` is set. With ``validate=False``
violations surface as a ``KeyError`` or ``pandas.errors.MergeError`` on
call.
"""
_FLAT_COLUMN_PREFIX: ClassVar[str] = "__lookup_col_"
"""Prefix for opaque internal value-column IDs used in interpolation."""
@property
def continuous_parameters(self) -> list[str]:
"""Lookup attributes used as binned ranges. The base name (e.g.
``"age"``) of each ``<name>_start`` / ``<name>_end`` pair."""
return [parameter.name for parameter in self._continuous_parameters]
def __init__(
self,
data: pd.DataFrame | pd.Series[Any],
value_columns: pd.Index[Any],
order: int,
extrapolate: bool,
validate: bool,
):
# TODO: allow for order 1 interpolation with binned edges
if order != 0:
raise NotImplementedError(
f"Interpolation is only supported for order 0. You specified order {order}"
)
self.value_columns: pd.Index[Any] = value_columns
"""User-facing column labels (including tuple labels from a column
``MultiIndex`` and ``None`` from a nameless Series) restored on the
output of :meth:`__call__`."""
self._internal_value_columns = [
f"{self._FLAT_COLUMN_PREFIX}{i}" for i in range(len(value_columns))
]
"""Opaque internal column IDs used in the interpolation pipeline."""
self.data: pd.DataFrame = self._reshape_data(data, value_columns)
"""Flat DataFrame the interpolation pipeline operates on. Value
columns are renamed to opaque internal IDs (see ``_FLAT_COLUMN_PREFIX``);
:attr:`value_columns` carries the original user-facing labels and is
reapplied to the output of :meth:`__call__`."""
parameter_columns = [
c for c in self.data.columns if c not in self._internal_value_columns
]
self._continuous_parameters: list[
ContinuousParameter
] = self._get_continuous_parameters(parameter_columns, self.data)
"""Continuous lookup parameters with their edge columns and shared bins."""
self.categorical_parameters: list[str] = self._get_categorical_parameters(
parameter_columns, self._continuous_parameters
)
"""Lookup attributes used to select between value rows."""
self.order: int = order
"""Order of interpolation. Only ``0`` is currently supported."""
self.extrapolate: bool = extrapolate
"""Whether to extrapolate beyond the edges of the supplied bins."""
self.validate: bool = validate
"""Whether to validate inputs on construction and on call."""
self._key_columns: list[str] = list(self.categorical_parameters) + [
parameter.start_column for parameter in self._continuous_parameters
]
# No key columns → __call__ broadcasts the single data row.
self._merge_target: pd.DataFrame | None = (
self.data[self._key_columns + self._internal_value_columns]
if self._key_columns
else None
)
if validate:
validate_parameters(
self.data,
self.categorical_parameters,
self._continuous_parameters,
self._internal_value_columns,
)
@staticmethod
def _get_continuous_parameters(
parameter_columns: list[str], data: pd.DataFrame
) -> list[ContinuousParameter]:
"""Build a ``ContinuousParameter`` for each ``<name>_start`` / ``<name>_end`` column pair."""
parameter_columns_set = set(parameter_columns)
parameters: list[ContinuousParameter] = []
for column in parameter_columns:
if str(column).endswith(_START_SUFFIX):
name = str(column).removesuffix(_START_SUFFIX)
start_column = f"{name}{_START_SUFFIX}"
end_column = f"{name}{_END_SUFFIX}"
if end_column in parameter_columns_set:
left_edges = data[start_column].drop_duplicates().sort_values()
parameters.append(
ContinuousParameter(
name=name,
start_column=start_column,
end_column=end_column,
left_edges=left_edges.to_numpy(),
max_right=data[end_column].drop_duplicates().max(),
)
)
return parameters
@staticmethod
def _get_categorical_parameters(
parameter_columns: list[str], continuous_parameters: list[ContinuousParameter]
) -> list[str]:
"""Get categorical parameter columns from the given list of parameter columns."""
bin_edge_columns = {
column
for parameter in continuous_parameters
for column in (parameter.start_column, parameter.end_column)
}
return [col for col in parameter_columns if col not in bin_edge_columns]
def _reshape_data(
self,
raw_data: pd.DataFrame | pd.Series[Any],
returned_columns: pd.Index[Any],
) -> pd.DataFrame:
"""Get the flat representation of ``data`` for interpolation."""
if has_named_row_index(raw_data):
if isinstance(raw_data, pd.Series):
flat = raw_data.to_frame(name=raw_data.name)
else:
flat = raw_data.copy(deep=False)
flat.columns = pd.Index(self._internal_value_columns)
return flat.reset_index()
# This is the deprecated path where the input data is already in flat form
assert isinstance(raw_data, pd.DataFrame) # only Series/indexed paths above
return raw_data.rename(
columns=dict(zip(list(returned_columns), self._internal_value_columns))
)
def __call__(self, interpolants: pd.DataFrame) -> pd.DataFrame:
"""Get the interpolated results for the parameters in interpolants.
Runs one :func:`numpy.digitize` per continuous parameter over the whole
population to resolve each interpolant's bin, then performs a single
left :meth:`pandas.DataFrame.merge` keyed on the categorical columns
plus each parameter's left bin edge.
Parameters
----------
interpolants
Data frame containing the parameters to interpolate.
Returns
-------
A table with the interpolated values for the given interpolants,
indexed by ``interpolants.index`` with columns
:attr:`value_columns`.
Raises
------
ValueError
If ``extrapolate`` is False and an interpolant falls outside the bins.
KeyError
If an interpolant carries a categorical value not present in the
source data.
"""
if self.validate:
validate_call_data(
interpolants, self.categorical_parameters, self._continuous_parameters
)
if interpolants.empty:
return pd.DataFrame(
index=interpolants.index, columns=self.value_columns, dtype=np.float64
)
original_index = interpolants.index
if self._merge_target is None:
# No categorical or continuous parameters: broadcast the first row.
first_row = self.data[self._internal_value_columns].iloc[0]
broadcast = pd.DataFrame(
{column: first_row[column] for column in self._internal_value_columns},
index=original_index,
)
broadcast.columns = self.value_columns
return broadcast
key_frame = pd.DataFrame(index=original_index)
for column in self.categorical_parameters:
key_frame[column] = interpolants[column].to_numpy()
for parameter in self._continuous_parameters:
values = interpolants[parameter.name]
if not self.extrapolate and (
values.min() < parameter.left_edges[0] or values.max() >= parameter.max_right
):
raise ValueError(
f"Extrapolation outside the provided bins is disabled, but "
f"parameter '{parameter.name}' has values outside its bin range "
f"[{parameter.left_edges[0]}, {parameter.max_right})."
)
# Left edge inclusive, right exclusive; out-of-range values fold to
# the nearest edge bin (below the minimum -> first, at/above the
# maximum -> last).
positions = np.digitize(values, parameter.left_edges)
positions[positions > 0] -= 1
key_frame[parameter.start_column] = parameter.left_edges[positions]
merged = key_frame.merge(
self._merge_target,
how="left",
on=self._key_columns,
validate="many_to_one",
indicator=True,
)
# A left many-to-one merge preserves left-row order, so the merged rows
# line up positionally with the interpolants.
merged.index = original_index
# Raise if interpolant's key is missing from the source
if self.categorical_parameters:
unmatched = merged["_merge"].to_numpy() == "left_only"
if unmatched.any():
unknown = merged.loc[unmatched, self.categorical_parameters].drop_duplicates()
raise KeyError(
f"Interpolants carry categorical values absent from the "
f"interpolation data (or, with validate=False, bin edges that "
f"differ across categorical groups):\n"
f"{unknown.to_string(index=False)}"
)
result = merged[self._internal_value_columns]
result.columns = self.value_columns
return result
def __repr__(self) -> str:
return "Interpolation()"
[docs]
def validate_parameters(
data: pd.DataFrame,
categorical_parameters: Sequence[str],
continuous_parameters: Sequence[ContinuousParameter],
value_columns: Sequence[Hashable],
) -> None:
"""Validate that the source data satisfies the lookup contract."""
if data.empty:
raise ValueError("You must supply non-empty data to create the interpolation.")
if not value_columns:
raise ValueError(
f"No non-parameter data. Available columns: {data.columns}, "
f"Parameter columns: "
f"{set(categorical_parameters) | {p.name for p in continuous_parameters}}"
)
required_cols = {
*categorical_parameters,
*(c for p in continuous_parameters for c in (p.start_column, p.end_column)),
*value_columns,
}
if extra_columns := list(data.columns.difference(list(required_cols))):
raise ValueError(
"Data contains extra columns not in key_columns, parameter_columns, or "
f"value_columns: {extra_columns}"
)
key_columns = list(categorical_parameters) + [
parameter.start_column for parameter in continuous_parameters
]
if key_columns:
duplicated = data.duplicated(subset=key_columns, keep=False)
if duplicated.any():
duplicate_keys = data.loc[duplicated, key_columns].drop_duplicates()
raise ValueError(
f"Interpolation data rows must be uniquely identified by the "
f"key columns {key_columns}, but multiple rows share "
f"these keys:\n{duplicate_keys.to_string(index=False)}"
)
elif len(data) > 1:
raise ValueError(
f"Interpolation data with no categorical or continuous parameters "
f"must be a single row. You provided {len(data)} rows."
)
if not continuous_parameters:
return
# Validate completeness one categorical group at a time: on the full
# multi-group table the duplicate-bin guard would trip on the repeated bins.
groups = (
[group for _, group in data.groupby(list(categorical_parameters))]
if categorical_parameters
else [data]
)
for group in groups:
check_data_complete(group, continuous_parameters)
for parameter in continuous_parameters:
reference_edges = set(data[parameter.start_column])
reference_max = data[parameter.end_column].max()
for group in groups:
if (
set(group[parameter.start_column]) != reference_edges
or group[parameter.end_column].max() != reference_max
):
raise ValueError(
f"Continuous parameter '{parameter.name}' has different bin "
f"edges across categorical groups. Interpolation requires "
f"uniform bin edges across all categorical groups."
)
[docs]
def validate_call_data(
data: pd.DataFrame,
categorical_parameters: Sequence[str],
continuous_parameters: Sequence[ContinuousParameter],
) -> None:
if not isinstance(data, pd.DataFrame):
raise TypeError(
f"Interpolations can only be called on pandas.DataFrames. You"
f"passed {type(data)}."
)
continuous_parameter_names = [p.name for p in continuous_parameters]
if not set(continuous_parameter_names) <= set(data.columns.values.tolist()):
raise ValueError(
f"The continuous continuous parameters with which you built the Interpolation must all "
f"be present in the data you call it on. The Interpolation has key "
f"columns: {continuous_parameter_names} and your data has columns: "
f"{data.columns.values.tolist()}"
)
if categorical_parameters and not set(categorical_parameters) <= set(
data.columns.values.tolist()
):
raise ValueError(
f"The key (categorical) columns with which you built the Interpolation must all"
f"be present in the data you call it on. The Interpolation has key"
f"columns: {categorical_parameters} and your data has columns: "
f"{data.columns.values.tolist()}"
)
[docs]
def check_data_complete(
data: pd.DataFrame, continuous_parameters: Sequence[ContinuousParameter]
) -> None:
"""Check that data provides complete, contiguous bins for each continuous parameter.
For each parameter, require that every combination of parameter bins is
present, that each left edge pairs with exactly one right edge, and that
the bins tile a continuous range: each bin's exclusive right edge equals
the next bin's inclusive left edge.
Raises
------
ValueError
If bins are duplicated or overlap, or if a combination of continuous
parameters is missing.
NotImplementedError
If a parameter contains non-continuous bins.
"""
if not continuous_parameters:
return
start_columns = [p.start_column for p in continuous_parameters]
# Per-parameter geometry: each parameter's distinct bins must tile a
# contiguous range. Checked before the cross-parameter grid below so a
# malformed bin reports as such rather than as a spurious missing or
# duplicate combination.
for parameter in continuous_parameters:
start_column, end_column = parameter.start_column, parameter.end_column
bins = data[[start_column, end_column]].drop_duplicates().sort_values(start_column)
starts = bins[start_column].to_numpy()
ends = bins[end_column].to_numpy()
if len(np.unique(starts)) != len(starts):
raise ValueError(
f"Parameter ('{start_column}', '{end_column}') pairs a left edge "
f"with multiple right edges; each bin start must map to one end."
)
# For contiguous bins each bin's exclusive right edge equals the next
# bin's inclusive left edge; a positive step is an overlap, negative a gap.
step = ends[:-1] - starts[1:]
if (step > 0).any():
raise ValueError(
f"Parameter ('{start_column}', '{end_column}') has overlapping "
f"bins: a right edge extends past the next bin's left edge."
)
if (step < 0).any():
raise NotImplementedError(
f"Interpolation only supports contiguous bins. Parameter "
f"('{start_column}', '{end_column}') leaves gaps between bins."
)
# Cross-parameter completeness: with each parameter a clean tiling, the rows
# must be the full cross product of the per-parameter bins, present once.
if data.duplicated(subset=start_columns).any():
raise ValueError(
f"Parameter data contains duplicate bins: multiple rows share the "
f"same left edges for "
f"{[(p.start_column, p.end_column) for p in continuous_parameters]}."
)
if len(data) != math.prod(data[col].nunique() for col in start_columns):
raise ValueError(
f"You must provide a value for every combination of "
f"{[p.name for p in continuous_parameters]}."
)