update results management system docstrings (#464)

* update docstrings
* pin sphinx-rtd-theme>=0.6
* various type hint updates
* change Stratification __call__() method stratify()
* move shared custom types to a new types.py module

CHANGELOG.rst

@@ -1,3 +1,7 @@
+**TBD - TBD**
+
+  - Update results-related docstrings
+
 **3.0.1- 08/20/24**
 
  - Create script to find matching dependency branches

setup.py

@@ -64,7 +64,7 @@
 
     doc_requirements = [
         "sphinx>=4.0",
-        "sphinx-rtd-theme",
+        "sphinx-rtd-theme>=0.6",
         "sphinx-click",
         "IPython",
         "matplotlib",

src/vivarium/framework/artifact/hdf.py

@@ -48,7 +48,7 @@
 import tables
 from tables.nodes import filenode
 
-PandasObj = (pd.DataFrame, pd.Series)
+_PandasObj = (pd.DataFrame, pd.Series)
 
 ####################
 # Public interface #
@@ -106,7 +106,7 @@ def write(path: Union[str, Path], entity_key: str, data: Any):
     path = _get_valid_hdf_path(path)
     entity_key = EntityKey(entity_key)
 
-    if isinstance(data, PandasObj):
+    if isinstance(data, _PandasObj):
         _write_pandas_data(path, entity_key, data)
     else:
         _write_json_blob(path, entity_key, data)
@@ -330,7 +330,7 @@ def _get_valid_hdf_path(path: Union[str, Path]) -> Path:
     return path
 
 
-def _write_pandas_data(path: Path, entity_key: EntityKey, data: Union[PandasObj]):
+def _write_pandas_data(path: Path, entity_key: EntityKey, data: Union[_PandasObj]):
     """Write data in a pandas format to an HDF file.
 
     This method currently supports :class:`pandas DataFrame` objects, with or

src/vivarium/framework/engine.py

@@ -44,8 +44,9 @@
 from vivarium.framework.randomness import RandomnessInterface
 from vivarium.framework.resource import ResourceInterface
 from vivarium.framework.results import ResultsInterface
-from vivarium.framework.time import Time, TimeInterface
+from vivarium.framework.time import TimeInterface
 from vivarium.framework.values import ValuesInterface
+from vivarium.types import Time
 
 
 class SimulationContext:

src/vivarium/framework/event.py

@@ -32,8 +32,8 @@
 import pandas as pd
 
 from vivarium.framework.lifecycle import ConstraintError
-from vivarium.framework.time import Time, Timedelta
 from vivarium.manager import Manager
+from vivarium.types import Time, Timedelta
 
 
 class Event(NamedTuple):

src/vivarium/framework/lookup/__init__.py

@@ -3,4 +3,5 @@
     LookupTableManager,
     validate_build_table_parameters,
 )
-from vivarium.framework.lookup.table import LookupTable, LookupTableData, ScalarValue
+from vivarium.framework.lookup.table import LookupTable
+from vivarium.types import LookupTableData, ScalarValue

src/vivarium/framework/lookup/interpolation.py

@@ -13,7 +13,7 @@
 import numpy as np
 import pandas as pd
 
-ParameterType = Union[List[List[str]], List[Tuple[str, str, str]]]
+_ParameterType = Union[List[List[str]], List[Tuple[str, str, str]]]
 
 
 class Interpolation:
@@ -39,7 +39,7 @@ def __init__(
         self,
         data: pd.DataFrame,
         categorical_parameters: Union[List[str], Tuple[str, ...]],
-        continuous_parameters: ParameterType,
+        continuous_parameters: _ParameterType,
         value_columns: Union[List[str], Tuple[str, ...]],
         order: int,
         extrapolate: bool,
@@ -265,7 +265,7 @@ class Order0Interp:
     def __init__(
         self,
         data,
-        continuous_parameters: ParameterType,
+        continuous_parameters: _ParameterType,
         value_columns: List[str],
         extrapolate: bool,
         validate: bool,

src/vivarium/framework/lookup/manager.py

@@ -22,10 +22,10 @@
     CategoricalTable,
     InterpolatedTable,
     LookupTable,
-    LookupTableData,
     ScalarTable,
 )
 from vivarium.manager import Manager
+from vivarium.types import LookupTableData
 
 if TYPE_CHECKING:
     from vivarium.framework.engine import Builder

src/vivarium/framework/lookup/table.py

@@ -13,17 +13,13 @@
 
 import dataclasses
 from abc import ABC, abstractmethod
-from datetime import datetime, timedelta
-from numbers import Number
 from typing import Callable, List, Tuple, Union
 
 import numpy as np
 import pandas as pd
 
 from vivarium.framework.lookup.interpolation import Interpolation
-
-ScalarValue = Union[Number, timedelta, datetime]
-LookupTableData = Union[ScalarValue, pd.DataFrame, List[ScalarValue], Tuple[ScalarValue]]
+from vivarium.types import ScalarValue
 
 
 @dataclasses.dataclass