from __future__ import annotations from collections.abc import Iterable import numpy as np import pandas as pd ROLLING_WINDOWS_DAYS = (3, 7, 14, 30, 60, 90, 180, 365) CONTEXT_LOW_VOLTAGE = 2.45 TRAJECTORY_LAGS = tuple(range(0, 365, 14)) TRAJECTORY_COLUMNS = tuple(f"traj_lag{lag}" for lag in TRAJECTORY_LAGS) TRAJECTORY_COVERAGE_COLUMN = "traj_coverage" MIN_TRAJECTORY_COVERAGE = 0.5 TRAJECTORY_REFERENCE_OFFSET = 0 CONTEXT_COLUMNS = ( "room_voltage_mean", "room_voltage_min", "building_voltage_mean", "voltage_minus_room", "voltage_minus_building", "room_low_count", "building_low_count", "scenario_low_fraction", ) NON_FEATURE_COLUMNS = { "scenario", "scenario_time", "device_id", "battery", "building", "room", "start_time", "end_time", "effective_eol", "target_rul_days", "event_observed", # Dataset/scenario boundary fields are useful to the planner and reports, # but they are not battery-health measurements. Learning from them makes # same-scenario validation optimistic and breaks on a shifted hidden split. "censor_proxy_rul_days", "scenario_month_sin", "scenario_month_cos", # Scenario-level context, used only by the auxiliary context model. *CONTEXT_COLUMNS, # Trajectory history, used only by the degradation models. *TRAJECTORY_COLUMNS, TRAJECTORY_COVERAGE_COLUMN, } def _as_columns(timeseries: pd.DataFrame) -> pd.DataFrame: frame = timeseries if "device_id" not in frame.columns or "end_time" not in frame.columns: frame = frame.reset_index() required = {"device_id", "end_time", "voltage", "temperature"} missing = required - set(frame.columns) if missing: raise ValueError(f"Battery time series is missing columns: {sorted(missing)}") return frame[["device_id", "end_time", "voltage", "temperature"]].copy() def build_daily_features( timeseries: pd.DataFrame, windows: Iterable[int] = ROLLING_WINDOWS_DAYS, ) -> pd.DataFrame: """Build past-only daily features for all scenarios.""" frame = _as_columns(timeseries) frame["end_time"] = pd.to_datetime(frame["end_time"], errors="coerce") frame["voltage"] = pd.to_numeric(frame["voltage"], errors="coerce") frame["temperature"] = pd.to_numeric(frame["temperature"], errors="coerce") frame = frame.dropna(subset=["device_id", "end_time", "voltage"]) frame["day"] = frame["end_time"].dt.floor("D") daily = ( frame.groupby(["device_id", "day"], observed=True, sort=True) .agg( voltage_median=("voltage", "median"), voltage_mean=("voltage", "mean"), voltage_min=("voltage", "min"), voltage_max=("voltage", "max"), voltage_std=("voltage", "std"), temperature_median=("temperature", "median"), temperature_std=("temperature", "std"), observations_day=("voltage", "size"), ) .reset_index() .sort_values(["device_id", "day"]) .reset_index(drop=True) ) stable = frame[frame["temperature"].between(10.0, 30.0, inclusive="both")] stable_daily = ( stable.groupby(["device_id", "day"], observed=True, sort=True)["voltage"] .median() .rename("stable_voltage_median") .reset_index() ) daily = daily.merge(stable_daily, on=["device_id", "day"], how="left") daily["stable_voltage_median"] = daily["stable_voltage_median"].fillna(daily["voltage_median"]) grouped = daily.groupby("device_id", observed=True, sort=False) first_day = grouped["day"].transform("min") daily["history_days"] = (daily["day"] - first_day).dt.total_seconds() / 86400.0 daily["days_since_previous"] = grouped["day"].diff().dt.total_seconds() / 86400.0 for window in tuple(int(value) for value in windows): min_periods = max(2, window // 4) rolled = ( grouped["stable_voltage_median"] .rolling(window, min_periods=min_periods) .agg(["mean", "min", "max", "std"]) .reset_index(level=0, drop=True) .reindex(daily.index) ) for statistic in ("mean", "min", "max", "std"): daily[f"voltage_{statistic}_{window}d"] = rolled[statistic] daily[f"temperature_mean_{window}d"] = ( grouped["temperature_median"] .rolling(window, min_periods=min_periods) .mean() .reset_index(level=0, drop=True) .reindex(daily.index) ) lag_voltage = grouped["stable_voltage_median"].shift(window - 1) lag_day = grouped["day"].shift(window - 1) span = (daily["day"] - lag_day).dt.total_seconds() / 86400.0 daily[f"voltage_slope_{window}d"] = ( daily["stable_voltage_median"] - lag_voltage ) / span.replace(0.0, np.nan) daily[f"span_{window}d"] = span daily["voltage_recent_vs_30d"] = daily["voltage_mean_7d"] - daily["voltage_mean_30d"] daily["voltage_recent_vs_90d"] = daily["voltage_mean_14d"] - daily["voltage_mean_90d"] daily["voltage_range_day"] = daily["voltage_max"] - daily["voltage_min"] return daily def scenario_history(history: pd.DataFrame, scenario_time) -> pd.DataFrame: """Restrict raw observations to what the scenario can see, exactly as iterate_scenarios does.""" start = pd.Timestamp(scenario_time) frame = history if "end_time" not in frame.columns: frame = frame.reset_index() visible = frame[pd.to_datetime(frame["end_time"]) <= start] return visible.set_index(["device_id", "end_time"]) def scenario_history_snapshot( history: pd.DataFrame, locations: pd.DataFrame, scenario_name: str, scenario_time, windows: Iterable[int] = ROLLING_WINDOWS_DAYS, ) -> pd.DataFrame: """Features for one scenario, built only from that scenario's visible history.""" visible = scenario_history(history, scenario_time) return scenario_snapshot( build_daily_features(visible, windows=windows), locations, scenario_name, scenario_time ) def build_trajectory_matrix(daily_features: pd.DataFrame) -> dict: """Dense device-by-day smoothed voltage grid, built once per split and cached.""" cached = daily_features.attrs.get("_trajectory_matrix") if cached is not None: return cached frame = daily_features[["device_id", "day", "stable_voltage_median"]].dropna() frame = frame.sort_values(["device_id", "day"]) smooth = frame.groupby("device_id", observed=True)["stable_voltage_median"].transform( lambda values: values.rolling(7, min_periods=3).median() ) frame = frame.assign(smooth=smooth).dropna(subset=["smooth"]) if frame.empty: empty = {"index": {}, "origin": pd.Timestamp("2000-01-01"), "grid": np.zeros((0, 1), dtype=np.float32), "first": np.zeros(0, dtype=np.int64), "observed_last": np.zeros(0, dtype=np.int64)} daily_features.attrs["_trajectory_matrix"] = empty return empty devices = sorted(frame["device_id"].astype(str).unique()) index = {device: position for position, device in enumerate(devices)} origin = frame["day"].min().normalize() span = int((frame["day"].max().normalize() - origin).days) + 1 grid = np.full((len(devices), span), np.nan, dtype=np.float32) rows = frame["device_id"].astype(str).map(index).to_numpy(dtype=np.int64) columns = (frame["day"] - origin).dt.days.to_numpy(dtype=np.int64) grid[rows, columns] = frame["smooth"].to_numpy() # Carry the last reading forward to the end of the grid. Stopping at the device's # final observation would make availability depend on whether it reported again later. first = np.full(len(devices), np.iinfo(np.int64).max, dtype=np.int64) observed_last = np.full(len(devices), -1, dtype=np.int64) for position in range(len(devices)): seen = np.flatnonzero(np.isfinite(grid[position])) if seen.size == 0: continue low = int(seen[0]) segment = grid[position, low:] carry = np.maximum.accumulate( np.where(np.isfinite(segment), np.arange(segment.size), 0) ) grid[position, low:] = segment[carry] first[position] = low observed_last[position] = int(seen[-1]) built = {"index": index, "origin": origin, "grid": grid, "first": first, "observed_last": observed_last} daily_features.attrs["_trajectory_matrix"] = built return built def trajectory_bins(matrix: dict, batteries, cutoff) -> np.ndarray: """Voltage at the cutoff and every 14 days before it; never reads past the cutoff.""" grid = matrix["grid"] first = matrix["first"] rows = np.array([matrix["index"].get(str(name), -1) for name in batteries], dtype=np.int64) cut = int((pd.Timestamp(cutoff).normalize() - matrix["origin"]).days) lags = np.asarray(TRAJECTORY_LAGS, dtype=np.int64) + TRAJECTORY_REFERENCE_OFFSET columns = cut - lags[None, :] out = np.full((len(rows), lags.size), np.nan, dtype=float) known = rows >= 0 if known.any(): safe = np.clip(columns, 0, grid.shape[1] - 1) values = grid[rows[known][:, None], safe] usable = (columns >= 0) & (columns >= first[rows[known]][:, None]) out[known] = np.where(usable, values, np.nan) return out def scenario_snapshot( daily_features: pd.DataFrame, locations: pd.DataFrame, scenario_name: str, scenario_time: pd.Timestamp | str, ) -> pd.DataFrame: """Get the last feature row visible at the scenario time.""" start = pd.Timestamp(scenario_time).normalize() alive = locations.copy() if "battery" not in alive.columns: raise ValueError("locations must contain a battery column") available = daily_features[daily_features["day"] <= start] latest = available.groupby("device_id", observed=True, sort=False).tail(1) snapshot = alive.merge(latest, left_on="battery", right_on="device_id", how="left") snapshot["scenario"] = scenario_name snapshot["scenario_time"] = start snapshot["location_age_days"] = ( start - pd.to_datetime(snapshot["start_time"], errors="coerce").dt.normalize() ).dt.total_seconds() / 86400.0 snapshot["data_gap_days"] = ( start - pd.to_datetime(snapshot["day"], errors="coerce") ).dt.total_seconds() / 86400.0 snapshot["censor_proxy_rul_days"] = ( pd.to_datetime(snapshot["end_time"], errors="coerce").dt.normalize() + pd.Timedelta(days=30) - start ).dt.total_seconds() / 86400.0 snapshot["scenario_month_sin"] = np.sin(2.0 * np.pi * start.month / 12.0) snapshot["scenario_month_cos"] = np.cos(2.0 * np.pi * start.month / 12.0) snapshot["voltage_rank_global"] = snapshot["stable_voltage_median"].rank(pct=True) snapshot["voltage_rank_building"] = snapshot.groupby("building", observed=True)[ "stable_voltage_median" ].rank(pct=True) snapshot["voltage_rank_room"] = snapshot.groupby("room", observed=True)[ "stable_voltage_median" ].rank(pct=True) snapshot["building_battery_count"] = snapshot.groupby("building", observed=True)[ "battery" ].transform("size") snapshot["room_battery_count"] = snapshot.groupby("room", observed=True)["battery"].transform( "size" ) snapshot = add_context_columns(snapshot) bins = trajectory_bins( build_trajectory_matrix(daily_features), snapshot["battery"].astype(str), start ) for position, column in enumerate(TRAJECTORY_COLUMNS): snapshot[column] = bins[:, position] snapshot[TRAJECTORY_COVERAGE_COLUMN] = np.isfinite(bins).mean(axis=1) return snapshot def add_context_columns(snapshot: pd.DataFrame) -> pd.DataFrame: """Room and building degradation context for one scenario.""" voltage = pd.to_numeric(snapshot["stable_voltage_median"], errors="coerce") room = snapshot.groupby("room", observed=True)["stable_voltage_median"] building = snapshot.groupby("building", observed=True)["stable_voltage_median"] snapshot["room_voltage_mean"] = room.transform("mean") snapshot["room_voltage_min"] = room.transform("min") snapshot["building_voltage_mean"] = building.transform("mean") snapshot["voltage_minus_room"] = voltage - snapshot["room_voltage_mean"] snapshot["voltage_minus_building"] = voltage - snapshot["building_voltage_mean"] low = (voltage < CONTEXT_LOW_VOLTAGE).astype(float) snapshot["room_low_count"] = low.groupby(snapshot["room"]).transform("sum") snapshot["building_low_count"] = low.groupby(snapshot["building"]).transform("sum") snapshot["scenario_low_fraction"] = float(low.mean()) return snapshot def attach_training_targets( snapshot: pd.DataFrame, eol_times: pd.Series, unobserved_eol_days: float = 30.0, ) -> pd.DataFrame: """Attach public training targets.""" result = snapshot.copy() observed = pd.to_datetime(result["battery"].map(eol_times), errors="coerce") assumed = ( pd.to_datetime(result["end_time"], errors="coerce").dt.normalize() + pd.to_timedelta(unobserved_eol_days, unit="D") ) result["event_observed"] = observed.notna().astype("int8") result["effective_eol"] = observed.fillna(assumed) result["target_rul_days"] = ( result["effective_eol"] - pd.to_datetime(result["scenario_time"]) ).dt.total_seconds() / 86400.0 return result def numeric_feature_columns(snapshot: pd.DataFrame) -> list[str]: return [ column for column in snapshot.columns if column not in NON_FEATURE_COLUMNS and not column.lower().startswith("unnamed") and pd.api.types.is_numeric_dtype(snapshot[column]) ]