squash lookup table types (#745)

* remove old LookupTable types

* merge categorical and interpolated paths

Author: rmudambi

docs/source/concepts/lookup.rst

@@ -41,55 +41,32 @@ population in a simulation.
 The lookup table system is built in layers. At the top is the
 :class:`Lookup Table <vivarium.framework.lookup.table.LookupTable>` object which
 is responsible for providing a uniform interface to the user regardless
-of the underlying implementation. From the user's perspective, it takes in
-a data set or scalar value on initialization and then lets them query against
-that data with a population index.
-
-The next layer is selected at initialization time based on the type of data
-provided. The :class:`Lookup Table <vivarium.framework.lookup.table.LookupTable>`
-picks a :class:`ScalarTable <vivarium.framework.lookup.table.ScalarTable>`
-if a single value is provided as the data, a
-:class:`CategoricalTable <vivarium.framework.lookup.table.CategoricalTable>` if a
-:class:`pandas.DataFrame` with only categorical variables is provided as the
-data, and a :class:`InterpolatedTable <vivarium.framework.lookup.table.InterpolatedTable>`
-if a :class:`pandas.DataFrame` which has at least one continuous variable is
-provided as the data.
+of the underlying data. From the user's perspective, it takes in a data set
+or scalar value on initialization and then lets them query against that data
+with a population index.
+
+At initialization time, the
+:class:`Lookup Table <vivarium.framework.lookup.table.LookupTable>` examines the
+provided data and configures itself accordingly. If the data is a scalar value
+(or list/tuple of scalars), the table simply broadcasts those values over the
+population index when called. If the data is a :class:`pandas.DataFrame`, the
+table delegates to an
+:class:`Interpolation <vivarium.framework.lookup.interpolation.Interpolation>`
+object that handles both categorical and continuous parameter lookups. The
+:class:`Interpolation <vivarium.framework.lookup.interpolation.Interpolation>`
+groups the data by any categorical (key) columns and then, for each group,
+finds the correct bin for any continuous parameters. Tables with only
+categorical parameters are simply the special case where there are no
+continuous parameters to bin on.
 
 .. note::
 
-   The :class:`InterpolatedTable <vivarium.framework.lookup.table.InterpolatedTable>`
-   is a misnomer here. It confuses the data handling strategy with the
-   underlying data representation.  A better name would be ``BinnedDataTable``
-   to indicate that it wraps data where the continuous parameters are
-   represented by bin edges in the provided data.  This would allow us
-   to easily think about and extend the lookup system to wrap data where the
-   continuous parameters are represented by points and to tables where all
-   parameters are categorical.
-
-If the underlying data is a single value or consists only of categorical variables,
-this is the last layer of abstraction. The
-:class:`ScalarTable <vivarium.framework.lookup.table.ScalarTable>` and
-:class:`CategoricalTable <vivarium.framework.lookup.table.CategoricalTable>` each
-have only one reasonable strategy which is to broadcast the value over the
-population index. If we have continuous variables and therefore an
-:class:`InterpolatedTable <vivarium.framework.lookup.table.InterpolatedTable>`,
-there are additional layers to the lookup system to allow the user to
-control the strategy for turning the population index into values based on
-the data.  The
-:class:`InterpolatedTable <vivarium.framework.lookup.table.InterpolatedTable>`
-is then responsible for turning the population index into a set of
-attributes relevant to the value production based on the structure of
-the input data and then providing those attributes to the value production
-strategy.
-
-.. note::
-
-   I'm being careful with language here.  We have objects named
-   ``Interpolation`` and ``InterpolatedTable`` though the operation they
-   perform is actually disaggregation.  If we extend the system to
-   work with point estimates for the continuous parameters, then
-   interpolation would appropriately describe what we do.  Both are
-   value production strategies based on the structure of the input data.
+   The ``Interpolation`` name is somewhat of a misnomer. For order 0
+   (the only currently supported order), the operation is really
+   disaggregation -- finding the correct bin a value belongs to rather
+   than interpolating between points. If the system is extended to work
+   with point estimates for continuous parameters, then interpolation
+   would appropriately describe the operation.
 
 More information about the value production strategies can be found in
 :ref:`here <interpolation_concept>`.
@@ -223,8 +200,8 @@ When building a lookup table from a :class:`pandas.DataFrame` using ``data_sourc
 the component automatically determines key columns, parameter columns, and value columns
 based on the data structure:
 
-- **Value columns** are assumed by the structure of the artifact to be ``["value"]``. In principle,
-  this could be configured by implementing a custom :class:`~vivarium.framework.artifact.manager.ArtifactManager`.
+- **Value columns** can be provided as an argument to :meth:`~vivarium.component.Component.build_lookup_table`
+  If value columns are not provided, it will default to ``"value"``.
 - **Parameter columns** are detected by finding columns ending in ``_start``
   that have corresponding ``_end`` columns (e.g., ``age_start``/``age_end``).
 - **Key columns** are all remaining columns that are neither value columns
@@ -296,11 +273,15 @@ integrating a lookup table into a :term:`component <Component>`, which is primar
 how they are used. Assuming you have a valid simulation object named ``sim`` and 
 the data from the above table in a :class:`pandas.DataFrame` named ``data``, you 
 can construct a lookup table in the following way, using the interface from the builder.
+You don't have to provide a name for the table, but it is recommended to do so for clarity
+and for ease of debugging. If you don't provide value column names, it will default to
+``"value"``.
+ 
 
 .. code-block:: python
 
       # value_columns implicitly set to remaining columns
-    > bmi = sim.builder.lookup.build_table(data, key_columns=['sex'], parameter_columns=['age'])
+    > bmi = sim.builder.lookup.build_table(data, name="bmi")
     > population = sim.get_population()
     > bmi(population.index).head()  # returns BMI values for the population
 
@@ -316,7 +297,7 @@ can construct a lookup table in the following way, using the interface from the
    Constructing a lookup table currently requires your data meet specific
    conditions. These are a consequence of the method the lookup table uses to
    arrive at the correct data. Specifically, your parameter columns must
-   represent bins and they must overlap.
+   represent bins and they must not overlap or have gaps.
 
 Estimating Unknown Values
 -------------------------

src/vivarium/framework/lookup/__init__.py

@@ -1,7 +1,4 @@
 from vivarium.framework.lookup.interface import LookupTableInterface
-from vivarium.framework.lookup.manager import (
-    LookupTableManager,
-    validate_build_table_parameters,
-)
+from vivarium.framework.lookup.manager import LookupTableManager
 from vivarium.framework.lookup.table import DEFAULT_VALUE_COLUMN, LookupTable
 from vivarium.types import LookupTableData, ScalarValue

src/vivarium/framework/lookup/interpolation.py

@@ -145,11 +145,6 @@ def validate_parameters(
     if data.empty:
         raise ValueError("You must supply non-empty data to create the interpolation.")
 
-    if len(continuous_parameters) < 1:
-        raise ValueError(
-            "You must supply at least one continuous parameter over which to interpolate."
-        )
-
     for p in continuous_parameters:
         if not isinstance(p, (tuple, list)) or len(p) != 3:
             raise ValueError(
@@ -160,7 +155,6 @@ def validate_parameters(
             )
 
     # break out the individual columns from binned column name lists
-    param_cols = [col for p in continuous_parameters for col in p]
     if not value_columns:
         raise ValueError(
             f"No non-parameter data. Available columns: {data.columns}, "
@@ -343,12 +337,20 @@ def __call__(self, interpolants: pd.DataFrame) -> pd.DataFrame:
         Parameters
         ----------
         interpolants
-            Data frame containing the parameters to interpolate..
+            Data frame containing the parameters to interpolate.
 
         Returns
         -------
             A table with the interpolated values for the given interpolants.
         """
+        if not self.parameter_bins:
+            # No continuous parameters — just broadcast the data values.
+            # With only categorical parameters, each sub-table has a single row.
+            return pd.DataFrame(
+                {col: self.data[col].iloc[0] for col in self.value_columns},
+                index=interpolants.index,
+            )
+
         # build a dataframe where we have the start of each parameter bin for each interpolant
         interpolant_bins = pd.DataFrame(index=interpolants.index)
 

src/vivarium/framework/lookup/manager.py

@@ -15,23 +15,14 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
-from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any
-from typing import SupportsFloat as Numeric
-from typing import overload
+from typing import TYPE_CHECKING, Any, overload
 
 import pandas as pd
 from layered_config_tree import LayeredConfigTree
 
 from vivarium.framework.event import Event
 from vivarium.framework.lifecycle import lifecycle_states
-from vivarium.framework.lookup.table import (
-    DEFAULT_VALUE_COLUMN,
-    CategoricalTable,
-    InterpolatedTable,
-    LookupTable,
-    ScalarTable,
-)
+from vivarium.framework.lookup.table import DEFAULT_VALUE_COLUMN, LookupTable
 from vivarium.manager import Manager
 from vivarium.types import LookupTableData
 
@@ -62,14 +53,14 @@ def __init__(self) -> None:
         super().__init__()
         self.tables: dict[str, LookupTable[pd.Series[Any]] | LookupTable[pd.DataFrame]] = {}
 
-    def setup(self, builder: "Builder") -> None:
+    def setup(self, builder: Builder) -> None:
         self._logger = builder.logging.get_logger(self.name)
         self._configuration = builder.configuration
-        self._pop_view_builder = builder.population.get_view
+        self._get_view = builder.population.get_view
         self.clock = builder.time.clock()
-        self._interpolation_order = builder.configuration.interpolation.order
-        self._extrapolate = builder.configuration.interpolation.extrapolate
-        self._validate = builder.configuration.interpolation.validate
+        self.interpolation_order = builder.configuration.interpolation.order
+        self.extrapolate = builder.configuration.interpolation.extrapolate
+        self.validate_interpolation = builder.configuration.interpolation.validate
         self._add_resources = builder.resources.add_resources
         self._add_constraint = builder.lifecycle.add_constraint
         self._get_current_component = builder.components.get_current_component
@@ -125,7 +116,7 @@ def build_table(
         table = self._build_table(component, data, name, value_columns)
         self._add_resources(component, table, table.required_resources)
         self._add_constraint(
-            table.call,
+            table._call,
             restrict_during=[
                 lifecycle_states.INITIALIZATION,
                 lifecycle_states.SETUP,
@@ -150,107 +141,19 @@ def _build_table(
             data = pd.DataFrame(data)
 
         value_columns_ = value_columns if value_columns else DEFAULT_VALUE_COLUMN
-        validate_build_table_parameters(data, value_columns_)
 
-        table: LookupTable[pd.Series[Any]] | LookupTable[pd.DataFrame]
-        if isinstance(data, pd.DataFrame):
-            parameter_columns, key_columns = self._get_columns(value_columns_, data)
-            if parameter_columns:
-                table = InterpolatedTable(
-                    name=name,
-                    component=component,
-                    data=data,
-                    population_view_builder=self._pop_view_builder,
-                    key_columns=key_columns,
-                    parameter_columns=parameter_columns,
-                    value_columns=value_columns_,
-                    interpolation_order=self._interpolation_order,
-                    clock=self.clock,
-                    extrapolate=self._extrapolate,
-                    validate=self._validate,
-                )
-            else:
-                table = CategoricalTable(
-                    name=name,
-                    component=component,
-                    data=data,
-                    population_view_builder=self._pop_view_builder,
-                    key_columns=key_columns,
-                    value_columns=value_columns_,
-                )
-        else:
-            table = ScalarTable(
-                name=name, component=component, data=data, value_columns=value_columns_
-            )
+        table = LookupTable(
+            name=name,
+            component=component,
+            data=data,
+            value_columns=value_columns_,
+            manager=self,
+            population_view=self._get_view(),
+        )
 
         self.tables[table.name] = table
 
         return table
 
     def __repr__(self) -> str:
         return "LookupTableManager()"
-
-    @staticmethod
-    def _get_columns(
-        value_columns: list[str] | tuple[str, ...] | str, data: pd.DataFrame
-    ) -> tuple[list[str], list[str]]:
-        if isinstance(value_columns, str):
-            value_columns = [value_columns]
-
-        all_columns = list(data.columns)
-
-        potential_parameter_columns = [
-            str(col).removesuffix("_start")
-            for col in all_columns
-            if str(col).endswith("_start")
-        ]
-        parameter_columns = []
-        bin_edge_columns = []
-        for column in potential_parameter_columns:
-            if f"{column}_end" in all_columns:
-                parameter_columns.append(column)
-                bin_edge_columns += [f"{column}_start", f"{column}_end"]
-
-        key_columns = [
-            col
-            for col in all_columns
-            if col not in value_columns and col not in bin_edge_columns
-        ]
-
-        return parameter_columns, key_columns
-
-
-def validate_build_table_parameters(
-    data: LookupTableData,
-    value_columns: list[str] | tuple[str, ...] | str,
-) -> None:
-    """Makes sure the data format agrees with the provided column layout."""
-    if (
-        data is None
-        or (isinstance(data, pd.DataFrame) and data.empty)
-        or (isinstance(data, (list, tuple)) and not data)
-    ):
-        raise ValueError("Must supply some data")
-
-    acceptable_types = (Numeric, datetime, timedelta, list, tuple, pd.DataFrame)
-    if not isinstance(data, acceptable_types):
-        raise TypeError(
-            f"The only allowable types for data are {acceptable_types}. "
-            f"You passed {type(data)}."
-        )
-
-    if isinstance(data, (list, tuple)):
-        if isinstance(value_columns, str):
-            raise ValueError(
-                "When supplying multiple values, value_columns must be a list or tuple of strings."
-            )
-        if len(value_columns) != len(data):
-            raise ValueError(
-                "The number of value columns must match the number of values."
-                f"You supplied values: {data} and value_columns: {value_columns}"
-            )
-    elif not isinstance(data, pd.DataFrame):
-        if not isinstance(value_columns, str):
-            raise ValueError(
-                "When supplying a single value, value_columns must be a string if provided."
-            )