Strategy API improvements (#210)

* Time selections can now choose to narrow the selection to a specific tau step in each selected day.

* Output line plots no longer implicitly group by day.
Now by default, event plots will often display "sawtooth" characteristics especially noticeable when the simulation is using tau steps of varying sizes. The old smooth graph can be obtained by choosing to group events by day and aggregating.

* Output tables functions support time-axis group-and-aggregate. Previously, only time selection was allowed.

* Axis strategy `.group_by()` methods have been re-standardized as `.group()`, and tend to accept both grouping objects as well as a short list of strings mapping to commonly-used grouping methods that make sense for the situation.

* Make an "agg" method available for all axis selections.
Standardizes things a bit and provides convenience; "agg" is a commonly available method in Python data science libraries.

* Minor bug and graphing fixes.

doc/devlog/README.md

@@ -74,6 +74,8 @@ This folder is a handy place to put Jupyter notebooks or other documents which h
 | 2024-10-09.ipynb | Tyler | | Comparing methods for group-and-aggregate on time-series data. |
 | 2024-10-11.ipynb | Tyler | | Comparing methods for group-and-aggregate on geo-series data. |
 | 2024-12-03-v0.7.ipynb | Tyler | | Comparing v0.6 to v0.7 via example. (See others in this series.) |
+| 2024-12-06.ipynb | Tyler | | Visualizing movement models. |
+| 2024-12-16.ipynb | Tyler | | Demonstrates table output tools with time-axis group-and-aggregate. |
 
 ## Contributing
 

epymorph/compartment_model.py

@@ -794,10 +794,17 @@ class QuantityGroupResult(NamedTuple):
 
 
 class QuantityGrouping(NamedTuple):
+    """Describes how to group simulation output quantities (events and compartments).
+    The default combines any quantity whose names match exactly. This is common in
+    multistrata models where events from several strata impact one transition.
+    You can also choose to group across strata and subscript differences.
+    Setting `strata` or `subscript` to True means those elements of quantity names
+    (if they exist) are ignored for the purposes of matching."""
+
     strata: bool
     """True to combine quantities across strata."""
     subscript: bool
-    """The to combine quantities across subscript."""
+    """True to combine quantities across subscript."""
 
     def _strip(self, name: _N) -> _N:
         if self.strata:
@@ -1036,9 +1043,13 @@ def labels(self) -> Sequence[str]:
         groups, _ = self.grouping.map(self.selected)
         return [g.name.full for g in groups]
 
+    def agg(self, agg: Literal["sum"]) -> "QuantityAggregation":
+        """Combine grouped quantities using the named aggregation."""
+        return QuantityAggregation(self.ipm, self.selection, self.grouping, agg)
+
     def sum(self) -> "QuantityAggregation":
         """Combine grouped quantities by adding their values."""
-        return QuantityAggregation(self.ipm, self.selection, self.grouping, "sum")
+        return self.agg("sum")
 
 
 @dataclass(frozen=True)

epymorph/geography/custom.py

@@ -5,6 +5,8 @@
 from typing_extensions import override
 
 from epymorph.geography.scope import (
+    GeoGroup,
+    GeoGrouping,
     GeoScope,
     GeoSelection,
     GeoSelector,
@@ -41,6 +43,9 @@ def select(self) -> "CustomSelector":
 class CustomSelection(GeoSelection[CustomScope]):
     """A GeoSelection on a CustomScope."""
 
+    def group(self, grouping: GeoGrouping) -> GeoGroup[CustomScope]:
+        return GeoGroup(self.scope, self.selection, grouping)
+
 
 @dataclass(frozen=True)
 class CustomSelector(GeoSelector[CustomScope, CustomSelection]):

epymorph/geography/scope.py

@@ -97,25 +97,21 @@ def map(self, node_ids: NDArray[np.str_]) -> NDArray[np.str_]:
 
 
 class _CanGeoAggregate(GeoStrategy[GeoScopeT_co]):
-    def _agg(self, agg: GeoAggMethod) -> "GeoAggregation[GeoScopeT_co]":
-        return GeoAggregation(
-            self.scope,
-            self.selection,
-            self.grouping,
-            agg,
-        )
+    def agg(self, agg: Literal["sum", "min", "max"]) -> "GeoAggregation[GeoScopeT_co]":
+        """Apply the named aggregation for each geo node group."""
+        return GeoAggregation(self.scope, self.selection, self.grouping, agg)
 
     def sum(self) -> "GeoAggregation[GeoScopeT_co]":
-        """Perform a sum for each geo node groups."""
-        return self._agg("sum")
+        """Perform a sum for each geo node group."""
+        return self.agg("sum")
 
     def min(self) -> "GeoAggregation[GeoScopeT_co]":
         """Take the min value for each geo node group."""
-        return self._agg("min")
+        return self.agg("min")
 
     def max(self) -> "GeoAggregation[GeoScopeT_co]":
         """Take the max value for each geo node group."""
-        return self._agg("max")
+        return self.agg("max")
 
 
 @dataclass(frozen=True)

epymorph/geography/us_census.py

@@ -430,49 +430,98 @@ class StateSelection(GeoSelection[CensusScope]):
 
 
 @dataclass(frozen=True)
-class CountySelection(GeoSelection[CensusScope]):
-    def group_by(self, grouping: Literal["state"]) -> GeoGroup[CensusScope]:
-        return GeoGroup(
-            self.scope,
-            self.selection,
-            CensusGrouping(grouping),
-        )
+class CensusGrouping(GeoGrouping):
+    granularity: CensusGranularityName
+
+    @override
+    def map(self, node_ids: NDArray[np.str_]) -> NDArray[np.str_]:
+        gran = CensusGranularity.of(self.granularity)
+        return np.array([gran.truncate(g) for g in node_ids], dtype=np.str_)
 
 
 @dataclass(frozen=True)
-class TractSelection(GeoSelection[CensusScope]):
-    def group_by(self, grouping: Literal["state", "county"]) -> GeoGroup[CensusScope]:
-        return GeoGroup(
-            self.scope,
-            self.selection,
-            CensusGrouping(grouping),
-        )
+class CountySelection(GeoSelection[CensusScope]):
+    def group(self, grouping: Literal["state"] | GeoGrouping) -> GeoGroup[CensusScope]:
+        """
+        Groups the geo series using the specified grouping.
+
+        Parameters
+        ----------
+        grouping : Literal["state"] | GeoGrouping
+            The grouping to use. You can specify a supported string value --
+            all of which act as shortcuts for common GeoGrouping instances --
+            or you can provide a GeoGrouping instance to perform custom grouping.
+
+            The shortcut values are:
+            - "state": equivalent to epymorph.geography.us_census.CensusGrouping("state")
+        """  # noqa: E501
+        match grouping:
+            case "state":
+                grouping = CensusGrouping(grouping)
+            case _:
+                pass
+        return GeoGroup(self.scope, self.selection, grouping)
 
 
 @dataclass(frozen=True)
-class BlockGroupSelection(GeoSelection[CensusScope]):
-    def group_by(
-        self, grouping: Literal["state", "county", "tract"]
+class TractSelection(GeoSelection[CensusScope]):
+    def group(
+        self, grouping: Literal["state", "county"] | GeoGrouping
     ) -> GeoGroup[CensusScope]:
-        return GeoGroup(
-            self.scope,
-            self.selection,
-            CensusGrouping(grouping),
-        )
+        """
+        Groups the geo series using the specified grouping.
+
+        Parameters
+        ----------
+        grouping : Literal["state", "county"] | GeoGrouping
+            The grouping to use. You can specify a supported string value --
+            all of which act as shortcuts for common GeoGrouping instances --
+            or you can provide a GeoGrouping instance to perform custom grouping.
+
+            The shortcut values are:
+            - "state": equivalent to epymorph.geography.us_census.CensusGrouping("state")
+            - "county": equivalent to epymorph.geography.us_census.CensusGrouping("county")
+        """  # noqa: E501
+        match grouping:
+            case "state":
+                grouping = CensusGrouping(grouping)
+            case "county":
+                grouping = CensusGrouping(grouping)
+            case _:
+                pass
+        return GeoGroup(self.scope, self.selection, grouping)
 
 
 @dataclass(frozen=True)
-class CensusGrouping(GeoGrouping):
-    granularity: CensusGranularityName
-
-    @override
-    def map(
-        self,
-        node_ids: NDArray[np.str_],
-    ) -> NDArray[np.str_]:
-        gran = CensusGranularity.of(self.granularity)
-        # return gran.truncate_array(node_ids)
-        return np.array([gran.truncate(g) for g in node_ids], dtype=np.str_)
+class BlockGroupSelection(GeoSelection[CensusScope]):
+    def group(
+        self, grouping: Literal["state", "county", "tract"] | GeoGrouping
+    ) -> GeoGroup[CensusScope]:
+        """
+        Groups the geo series using the specified grouping.
+
+        Parameters
+        ----------
+        grouping : Literal["state", "county", "tract"] | GeoGrouping
+            The grouping to use. You can specify a supported string value --
+            all of which act as shortcuts for common GeoGrouping instances --
+            or you can provide a GeoGrouping instance to perform custom grouping.
+
+            The shortcut values are:
+            - "state": equivalent to epymorph.geography.us_census.CensusGrouping("state")
+            - "county": equivalent to epymorph.geography.us_census.CensusGrouping("county")
+            - "tract": equivalent to epymorph.geography.us_census.CensusGrouping("tract")
+        """  # noqa: E501
+        match grouping:
+            case "state":
+                grouping = CensusGrouping(grouping)
+            case "county":
+                grouping = CensusGrouping(grouping)
+            case "tract":
+                grouping = CensusGrouping(grouping)
+            case _:
+                pass
+        return GeoGroup(self.scope, self.selection, grouping)
 
 
 # NOTE: these Selector classes form a hierarchy of shared functionality.