Skip to content

Commit

Permalink
Time selections can now choose to narrow the selection to a specific …
Browse files Browse the repository at this point in the history 
…tau step in each selected day.
  • Loading branch information
Tyler Coles committed Dec 17, 2024
1 parent 88c5e4a commit 7d64e22
Showing 1 changed file with 133 additions and 53 deletions.
186 changes: 133 additions & 53 deletions epymorph/time.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -353,8 +353,9 @@ class TimeStrategy:

time_frame: TimeFrame
"""The original time frame."""
selection: slice
"""A slice for selection of a subset of the time frame."""
selection: tuple[slice, int | None]
"""The selected subset of the time frame: described as a date slice
and an optional tau step index."""
grouping: TimeGrouping | None
"""A method for grouping the time series data."""
aggregation: TimeAggMethod | None
Expand All @@ -371,8 +372,9 @@ def group_format(self) -> Literal["date", "tick", "day", "other"]:
@property
def date_bounds(self) -> tuple[date, date]:
"""The bounds of the selection, given as the first and last date included."""
start = self.selection.start or 0
stop = self.selection.stop or self.time_frame.duration_days
date_slice, _ = self.selection
start = date_slice.start or 0
stop = date_slice.stop or self.time_frame.duration_days
first_date = self.time_frame.start_date + timedelta(days=start)
last_date = self.time_frame.start_date + timedelta(days=stop - 1)
return (first_date, last_date)
Expand All @@ -386,50 +388,33 @@ def to_time_frame(self) -> TimeFrame:
return TimeFrame.range(first, last)

def selection_ticks(self, taus: int) -> slice:
s = self.selection
return slice(
None if s.start is None else s.start * taus,
None if s.stop is None else s.stop * taus,
None if s.step is None else s.step * taus,
)


@dataclass(frozen=True)
class TimeSelector:
"""A utility class for selecting a subset of a time frame."""

time_frame: TimeFrame
"""The original time frame."""

def all(self) -> "TimeSelection":
"""Select the entirety of the time frame."""
return TimeSelection(self.time_frame, slice(None))

def by_slice(self, start: int, stop: int | None = None) -> "TimeSelection":
return TimeSelection(self.time_frame, slice(start, stop))

def _to_selection(self, other: TimeFrame) -> "TimeSelection":
if not self.time_frame.is_subset(other):
err = "When selecting part of a time frame you must specify a subset."
"""Converts this into a slice for which ticks are selected (by index)."""
day, tau_step = self.selection
if tau_step is not None and tau_step >= taus:
err = (
"Invalid time-axis tau step selection: this model has "
f"{taus} tau steps but you selected step index {tau_step} "
"which is out of range."
)
raise ValueError(err)
from_index = (other.start_date - self.time_frame.start_date).days
to_index = (other.end_date - self.time_frame.start_date).days + 1
return self.by_slice(from_index, to_index)

def range(self, from_date: date | str, to_date: date | str) -> "TimeSelection":
"""Subset the time frame by providing the start and end date (inclusive)."""
other = TimeFrame.range(from_date, to_date)
return self._to_selection(other)

def rangex(self, from_date: date | str, until_date: date | str) -> "TimeSelection":
"""Subset the time frame by providing the start and end date (exclusive)."""
other = TimeFrame.rangex(from_date, until_date)
return self._to_selection(other)

def year(self, year: int) -> "TimeSelection":
"""Subset the time frame to a specific year."""
other = TimeFrame.year(year)
return self._to_selection(other)
# There are two cases:
if tau_step is not None:
# If tau_step is specified, we want to select only the ticks
# in the date range which correspond to that tau step.
start = taus * (day.start or 0) + tau_step
stop = None if day.stop is None else taus * day.stop + tau_step
step = taus
return slice(start, stop, step)
else:
# If tau_step is None, then we will select every tau step
# in the date range.
# This implies it is not possible to "step" over days --
# the time series must be contiguous w.r.t. simulation days.
return slice(
None if day.start is None else day.start * taus,
None if day.stop is None else day.stop * taus,
# day.step is ignored if it is present
)


class _CanAggregate(TimeStrategy):
Expand All @@ -453,8 +438,9 @@ class TimeSelection(_CanAggregate, TimeStrategy):

time_frame: TimeFrame
"""The original time frame."""
selection: slice
"""A slice for selection of a subset of the time frame."""
selection: tuple[slice, int | None]
"""The selected subset of the time frame: described as a date slice
and an optional tau step index."""
grouping: None = field(init=False, default=None)
"""A method for grouping the time series data."""
aggregation: None = field(init=False, default=None)
Expand Down Expand Up @@ -509,8 +495,9 @@ class TimeGroup(_CanAggregate, TimeStrategy):

time_frame: TimeFrame
"""The original time frame."""
selection: slice
"""A slice for selection of a subset of the time frame."""
selection: tuple[slice, int | None]
"""The selected subset of the time frame: described as a date slice
and an optional tau step index."""
grouping: TimeGrouping
"""A method for grouping the time series data."""
aggregation: None = field(init=False, default=None)
Expand All @@ -530,8 +517,9 @@ class TimeAggregation(TimeStrategy):

time_frame: TimeFrame
"""The original time frame."""
selection: slice
"""A slice for selection of a subset of the time frame."""
selection: tuple[slice, int | None]
"""The selected subset of the time frame: described as a date slice
and an optional tau step index."""
grouping: TimeGrouping | None
"""A method for grouping the time series data."""
aggregation: TimeAggMethod
Expand All @@ -542,3 +530,95 @@ class TimeAggregation(TimeStrategy):
@override
def group_format(self):
return "tick" if self.grouping is None else self.grouping.group_format


@dataclass(frozen=True)
class TimeSelector:
"""A utility class for selecting a subset of a time frame."""

time_frame: TimeFrame
"""The original time frame."""

def all(self, step: int | None = None) -> TimeSelection:
"""Select the entirety of the time frame.
Parameters
----------
step : int, optional
if given, narrow the selection to a specific tau step (by index) within
the date range; by default include all steps
"""
return TimeSelection(self.time_frame, (slice(None), step))

def _to_selection(
self,
other: TimeFrame,
step: int | None = None,
) -> TimeSelection:
if not self.time_frame.is_subset(other):
err = "When selecting part of a time frame you must specify a subset."
raise ValueError(err)
from_index = (other.start_date - self.time_frame.start_date).days
to_index = (other.end_date - self.time_frame.start_date).days + 1
return TimeSelection(self.time_frame, (slice(from_index, to_index), step))

def range(
self,
from_date: date | str,
to_date: date | str,
step: int | None = None,
) -> TimeSelection:
"""Subset the time frame by providing the start and end date (inclusive).
Parameters
----------
from_date : date | str
the starting date of the range, as a date object or an ISO-8601 string
to_date : date | str
the last included date of the range, as a date object or an ISO-8601 string
step : int, optional
if given, narrow the selection to a specific tau step (by index) within
the date range; by default include all steps
"""
other = TimeFrame.range(from_date, to_date)
return self._to_selection(other, step)

def rangex(
self,
from_date: date | str,
until_date: date | str,
step: int | None = None,
) -> TimeSelection:
"""Subset the time frame by providing the start and end date (exclusive).
Parameters
----------
from_date : date | str
the starting date of the range, as a date object or an ISO-8601 string
until_date : date | str
the stop date date of the range (the first date excluded)
as a date object or an ISO-8601 string
step : int, optional
if given, narrow the selection to a specific tau step (by index) within
the date range; by default include all steps
"""
other = TimeFrame.rangex(from_date, until_date)
return self._to_selection(other, step)

def year(
self,
year: int,
step: int | None = None,
) -> TimeSelection:
"""Subset the time frame to a specific year.
Parameters
----------
year : int
the year to include, from January 1st through December 31st
step : int, optional
if given, narrow the selection to a specific tau step (by index) within
the date range; by default include all steps
"""
other = TimeFrame.year(year)
return self._to_selection(other, step)

0 comments on commit 7d64e22

Please sign in to comment.