Skip to content

Commit

Permalink
Strategy API improvements (#210)
Browse files Browse the repository at this point in the history
* 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.
  • Loading branch information
JavadocMD authored Dec 17, 2024
1 parent fc8c190 commit 4cd11cd
Show file tree
Hide file tree
Showing 12 changed files with 993 additions and 179 deletions.
22 changes: 13 additions & 9 deletions USAGE.ipynb

Large diffs are not rendered by default.

510 changes: 510 additions & 0 deletions doc/devlog/2024-12-16.ipynb

Large diffs are not rendered by default.

2 changes: 2 additions & 0 deletions doc/devlog/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
15 changes: 13 additions & 2 deletions epymorph/compartment_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions epymorph/geography/custom.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
from typing_extensions import override

from epymorph.geography.scope import (
GeoGroup,
GeoGrouping,
GeoScope,
GeoSelection,
GeoSelector,
Expand Down Expand Up @@ -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]):
Expand Down
18 changes: 7 additions & 11 deletions epymorph/geography/scope.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
115 changes: 82 additions & 33 deletions epymorph/geography/us_census.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading

0 comments on commit 4cd11cd

Please sign in to comment.