import pandas import numpy as np import datetime from pywr import _core # Constants SECONDS_IN_DAY = 3600 * 24 class Timestepper(object): def __init__(self, start="2015-01-01", end="2015-12-31", delta=1): self.start = start self.end = end self.delta = delta self._last_length = None self._periods = None self._deltas = None self._current = None self._next = None self.setup() self.reset() self._dirty = True def __iter__(self, ): return self def __len__(self, ): return len(self._periods) @property def dirty(self): return self._dirty def setup(self): periods = self.datetime_index # Compute length of each period deltas = periods.to_timestamp(how='e') - periods.to_timestamp(how='s') # Round to nearest second deltas = np.round(deltas.total_seconds()) # Convert to days deltas = deltas / SECONDS_IN_DAY self._periods = periods self._deltas = deltas self.reset() self._dirty = False def reset(self, start=None): """ Reset the timestepper If start is None it resets to the original self.start, otherwise start is used as the new starting point. """ self._current = None current_length = len(self) if start is None: current_index = 0 else: # Determine which period the start time is within for index, period in enumerate(self._periods): if period.start_time <= start < period.end_time: current_index = index break else: raise ValueError('New starting position is outside the range of the model timesteps.') self._next = _core.Timestep(self._periods[current_index], current_index, self._deltas[current_index]) length_changed = self._last_length != current_length self._last_length = current_length return length_changed def __next__(self, ): return self.next() def next(self, ): self._current = current = self._next if current.index >= len(self._periods): raise StopIteration() # Increment to next timestep next_index = current.index + 1 if next_index >= len(self._periods): # The final time-step is one offset beyond the end of the model. # Here we compute its delta and create the object. final_period = current.period + self.offset delta = final_period.end_time - final_period.start_time delta = np.round(delta.total_seconds()) delta = delta / SECONDS_IN_DAY self._next = _core.Timestep(final_period, next_index, delta) else: self._next = _core.Timestep(self._periods[next_index], next_index, self._deltas[next_index]) # Return this timestep return current @property def start(self): return self._start @start.setter def start(self, value): if isinstance(value, pandas.Timestamp): self._start = value else: self._start = pandas.to_datetime(value) self._dirty = True @property def start_period(self): return pandas.Period(self.start, freq=self.freq) @property def end(self): return self._end @end.setter def end(self, value): if isinstance(value, pandas.Timestamp): self._end = value else: self._end = pandas.to_datetime(value) self._dirty = True @property def end_period(self): return pandas.Period(self.end, freq=self.freq) @property def delta(self): return self._delta @delta.setter def delta(self, value): self._delta = value self._dirty = True @property def freq(self): d = self._delta if isinstance(d, int): freq = '{}D'.format(d) elif isinstance(d, datetime.timedelta): freq = '{}D'.format(d.days) else: freq = d return freq @property def offset(self): return pandas.tseries.frequencies.to_offset(self.freq) @property def current(self): """The current timestep If iteration has not begun this will return None. """ return self._current @property def datetime_index(self): """ Return a `pandas.DatetimeIndex` using the start, end and delta of this object This is useful for creating `pandas.DataFrame` objects from Model results """ return pandas.period_range(self.start, self.end, freq=self.freq) def __repr__(self): start = self.start.strftime("%Y-%m-%d") end = self.end.strftime("%Y-%m-%d") return "<Timestepper start=\"{}\" end=\"{}\" freq=\"{}\">".format( start, end, self.freq )