""" Regularization of a jittery, gappy input series onto a strict time grid. The live grid advances monotonically: each frontier point is finalized once, in order, as measured, imputed, or forecast. Late measured arrivals may re-emit corrections that upgrade prior imputed/forecast points still held in history (without rewinding the frontier): - measured data within the jitter tolerance is snapped onto the grid, - interior gaps (bounded by later measured data) are filled by the imputer, - grid points with no data by the deadline (grid_time + lag_time) are filled by the forecaster. """ import logging import threading import time from datetime import datetime, timedelta, timezone from math import floor from typing import Callable, Dict, List, Optional, Tuple from .config import ChannelConfig from .history import HistoryProvider, RegularizedHistoryStore from .io import RedisStreamSink, RedisStreamSource from .sample import Sample from .tools import Forecaster, Imputer class TimeGridRegularizer: """ Consumes irregular measured samples and produces a strictly regular sequence of samples on the grid `epoch + k * interval`. """ # Epoch for snapping to the grid. _EPOCH = datetime(1970, 1, 1, tzinfo=timezone.utc) def __init__(self, config: ChannelConfig, imputer: Imputer, forecaster: Forecaster, history_provider: Optional[HistoryProvider], logger: logging.Logger, start_time: Optional[datetime] = None ) -> None: """ Initialize the time grid regularizer. Args: config: channel configuration object imputer: imputer to use object forecaster: forecaster to use object history_provider: optional TimescaleDB history provider object logger: logger to use object start_time: optional start time to use (default=None) """ self._config = config self._imputer = imputer self._forecaster = forecaster self._history_provider = history_provider self._logger = logger # Snapped samples waiting to be emitted, keyed by grid timestamp. Also keeps # the jitter distance so that the closest sample wins for duplicate grid points. self._pending: Dict[datetime, Tuple[Sample, timedelta]] = {} # Late measured upgrades for already-finalized imputed/forecast points, # drained by poll() ahead of newly finalized frontier samples. self._corrections: List[Sample] = [] # First grid point of interest: data older than (start - lag_time) was already # published before this service started, so the grid begins right after it. start_time = start_time or datetime.now(timezone.utc) self._next_grid_ts = self._snap_to_grid(start_time, config.polling_interval, floor) - config.lag_time self._logger.debug(f'snapping start_time = {start_time.isoformat()} to _next_grid_ts = {self._next_grid_ts.isoformat()}') # Initialize the regularized history store. self._history = RegularizedHistoryStore(config.window) def bootstrap_history( self, sink: RedisStreamSink, source: RedisStreamSource, stop_event: threading.Event, ) -> bool: """ Optionally waits for the Redis source, then fetches raw DB history, replays regularization, seeds in-memory history, and publishes bootstrap emissions to Redis. Returns False if startup should abort (source wait failed / stopped); True otherwise. """ hp_config = self._config.history_provider # Skip bootstrap if history provider is not configured. if hp_config is None or self._history_provider is None: self._logger.info('Skipping history bootstrap') return True # Wait for source data to be available if configured. if hp_config.init_when_source_available: if source.wait_until_available(stop_event): self._logger.info('Source data is available, waiting a moment before starting the history bootstrap ...') time.sleep(hp_config.bootstrap_delay.total_seconds()) else: self._logger.info('Channel stopped before source data was available') return False # Compute the bootstrap window. bootstrap_start = self._next_grid_ts - self._config.window bootstrap_end = datetime.now(timezone.utc) # Reset the grid to the start of the bootstrap window. self._next_grid_ts = bootstrap_start self._pending.clear() self._corrections.clear() # Get the raw history from the history provider. history_raw = self._history_provider.get_history( hp_config.dp_name, dp_location_code=hp_config.dp_location_code, dp_device_id=hp_config.dp_device_id, dp_data_provider=hp_config.dp_data_provider, start=bootstrap_start, end=bootstrap_end ) self._logger.info(f'Retrieved history: start={bootstrap_start.isoformat()}, end={bootstrap_end.isoformat()}, samples={len(history_raw)}') # Add the raw history to the regularizer. for sample in sorted(history_raw, key=lambda s: s.timestamp): self.add(sample) # Poll the regularizer. This will finalize the grid points that can be emitted. if not self.poll(now=self._next_grid_ts, interval=self._config.window): self._logger.warning(f'No samples emitted during bootstrap') return True # Emit the history samples to the sink. sink.emit(self._history.samples) self._logger.info(f'Bootstrap complete: {len(self._history.samples)} samples emitted') return True def add(self, sample: Sample) -> None: """ Registers an incoming measured sample, snapping it onto the grid. """ # Snap the sample to the grid. grid_ts = self._snap_to_grid(sample.timestamp, self._config.update_interval) # Calculate the distance from the sample to the grid point. distance = abs(sample.timestamp - grid_ts) # Log the sample addition. self._logger.debug( f'adding sample: timestamp = {sample.timestamp.isoformat()}, value = {sample.value}, quality = {sample.quality} ' f'-> grid timestamp = {grid_ts.isoformat()}' ) # Drop the sample if it is outside the jitter tolerance. if distance > self._config.jitter_tolerance: self._logger.warning( f'Dropping sample at {sample.timestamp.isoformat()}: ' f'{distance} away from nearest grid point {grid_ts.isoformat()}') return # Add the sample to the history if it is within the jitter tolerance. if grid_ts < self._next_grid_ts: # Get the prior sample at the grid point. prior = self._history.get(grid_ts) # If the prior sample is an imputed or forecast, queue a correction. if prior is not None and prior.quality in ('imputed', 'forecast'): # Create a correction sample. correction = Sample(timestamp=grid_ts, value=sample.value, quality='measured') # Add the correction to the history. self._history.extend([correction]) # Add the correction to the corrections list. self._corrections.append(correction) self._logger.info( f'Queuing late measured correction at {grid_ts.isoformat()} ' f'(upgrading {prior.quality}, current grid point: {self._next_grid_ts.isoformat()})' ) else: self._logger.warning( f'Dropping late sample at {grid_ts.isoformat()} ' f'(current grid point: {self._next_grid_ts.isoformat()})' ) return # Add the sample to the pending list. existing = self._pending.get(grid_ts) # If the sample is not in the pending list, or if it is closer to the # grid point than the existing sample, add it to the pending list. if existing is None or distance < existing[1]: self._pending[grid_ts] = (sample, distance) def poll(self, now: Optional[datetime] = None, interval: Optional[timedelta] = None) -> List[Sample]: """ Returns all grid points that can be finalized, in order. A grid point is finalized when measured data for it (or beyond it) has arrived, or when `now` has passed its deadline `grid_time + lag_time`. """ self._logger.debug('polling now') # Get the current time and interval. now = now or datetime.now(timezone.utc) interval = interval or self._config.polling_interval # Get the horizon. horizon = self._next_grid_ts + interval # Get the corrections. corrections = self._corrections self._corrections = [] # Initialize the list of finalized samples. finalized: List[Sample] = [] while self._next_grid_ts <= horizon: # Get the next grid point. grid_ts = self._next_grid_ts # Get the pending entry for the grid point. pending_entry = self._pending.get(grid_ts) # If the pending entry exists, process it. if pending_entry is not None: # Process the pending entry as a measured sample. self._logger.debug('processing pending entry (measured)') sample = pending_entry[0] finalized.append(Sample(timestamp=grid_ts, value=sample.value, quality='measured')) elif any(ts > grid_ts for ts in self._pending): # Process the interior gap as an imputed sample. self._logger.debug('processing interior gap (imputing)') next_sample = self._pending[min(ts for ts in self._pending if ts > grid_ts)][0] value = self._imputer.impute(grid_ts, self._history_for(grid_ts), next_sample, self._config) self._logger.debug(f'imputed value {value} for interior gap at {grid_ts.isoformat()}') finalized.append(Sample(timestamp=grid_ts, value=value, quality='imputed')) elif now >= grid_ts + self._config.lag_time: # Process the deadline gap as a forecasted sample. self._logger.debug('processing deadline gap (forecasting)') value = self._forecaster.forecast(grid_ts, self._history_for(grid_ts), self._config) self._logger.debug(f'forecast value {value} for grid point {grid_ts.isoformat()} ' f'(no data by deadline)') finalized.append(Sample(timestamp=grid_ts, value=value, quality='forecast')) else: break # Remove the grid point from the pending list. self._pending.pop(grid_ts, None) # Update the next grid point. self._next_grid_ts += self._config.update_interval # Extend the history with the finalized samples. self._history.extend(finalized) # Trim the history to the current time. self._history.trim(now) # Get the emitted samples. emitted = corrections + finalized self._logger.debug(f'after polling: _next_grid_ts = {self._next_grid_ts.isoformat()}') try: self._logger.debug(f'history: start = {self._history.samples[0].timestamp.isoformat()}, ' f'end = {self._history.samples[-1].timestamp.isoformat()}') except IndexError: self._logger.debug('no history samples available') # Get the pending keys. try: pending_keys = sorted(list(self._pending.keys())) self._logger.debug(f'pending: start = {pending_keys[0].isoformat()}, ' f'end = {pending_keys[-1].isoformat()}') except IndexError: self._logger.debug('no pending samples available') self._logger.info(f'number of emitted samples: {len(emitted)} ' f'(corrections={len(corrections)}, finalized={len(finalized)})') return emitted def _history_for(self, grid_ts: datetime) -> List[Sample]: """ Returns the in-memory regularized history window ending at grid_ts. """ return self._history.window(grid_ts - self._config.window, grid_ts) def _snap_to_grid(self, ts: datetime, interval: timedelta, snap_func: Callable[[float], float] = round ) -> datetime: """ Snap the timestamp to the grid. """ # Calculate the number of intervals since the epoch. k = snap_func((ts - self._EPOCH) / interval) # Return the snapped timestamp. return self._EPOCH + k * interval