source/generate_data.py

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
from datasets import Dataset

PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"


def build_device(device_id: int, steps: int, seed: int) -> pd.DataFrame:
    rng = np.random.default_rng(seed + device_id)
    time = np.arange(steps)
    phase = device_id * 0.37
    load = 0.55 + 0.27 * np.sin(time / 93 + phase) + 0.08 * np.sin(time / 17)
    actuator = np.clip(48 + load * 42 + rng.normal(0, 1.5, steps), 0, 100)
    temperature = 34 + load * 26 + rng.normal(0, 0.7, steps) + device_id * 0.12
    pressure = 72 + load * 15 + rng.normal(0, 0.8, steps)
    vibration = 0.8 + load**2 * 2.4 + rng.normal(0, 0.08, steps)
    current = 5.2 + load * 13 + rng.normal(0, 0.25, steps)
    flow = 18 + actuator * 0.42 + rng.normal(0, 0.6, steps)
    packet_rate = 110 + load * 38 + rng.normal(0, 4, steps)
    command_rate = 4 + np.abs(np.gradient(actuator)) * 0.38 + rng.normal(0, 0.2, steps)
    labels = np.zeros(steps, dtype=np.int8)
    anomaly_type = np.full(steps, "normal", dtype=object)

    anomaly_names = [
        "sensor_drift",
        "actuator_mismatch",
        "vibration_fault",
        "pressure_spike",
        "network_flood",
    ]
    starts = rng.choice(np.arange(250, steps - 250), size=18, replace=False)
    for index, start in enumerate(starts):
        length = int(rng.integers(18, 55))
        stop = min(start + length, steps)
        kind = anomaly_names[index % len(anomaly_names)]
        labels[start:stop] = 1
        anomaly_type[start:stop] = kind
        ramp = np.linspace(0, 1, stop - start)
        if kind == "sensor_drift":
            temperature[start:stop] += 8 * ramp
        elif kind == "actuator_mismatch":
            flow[start:stop] -= 12 + actuator[start:stop] * 0.12
        elif kind == "vibration_fault":
            vibration[start:stop] += 3.5 + rng.normal(0, 0.5, stop - start)
        elif kind == "pressure_spike":
            pressure[start:stop] += 18 * np.sin(np.linspace(0, np.pi, stop - start))
        else:
            packet_rate[start:stop] += 260 + rng.normal(0, 18, stop - start)
            command_rate[start:stop] += 14

    return pd.DataFrame(
        {
            "device_id": device_id,
            "time": time,
            "temperature": temperature,
            "pressure": pressure,
            "vibration": vibration,
            "current": current,
            "flow": flow,
            "packet_rate": packet_rate,
            "command_rate": command_rate,
            "actuator_position": actuator,
            "label": labels,
            "anomaly_type": anomaly_type,
        }
    )


def add_features(frame: pd.DataFrame) -> pd.DataFrame:
    sensors = [
        "temperature",
        "pressure",
        "vibration",
        "current",
        "flow",
        "packet_rate",
        "command_rate",
        "actuator_position",
    ]
    groups = frame.groupby("device_id", sort=False)
    for column in sensors:
        frame[f"{column}_delta"] = groups[column].diff().fillna(0)
        rolling = groups[column].rolling(24, min_periods=4)
        mean = rolling.mean().reset_index(level=0, drop=True)
        std = rolling.std().reset_index(level=0, drop=True).fillna(1).clip(lower=1e-3)
        frame[f"{column}_z24"] = ((frame[column] - mean) / std).fillna(0)
    frame["flow_actuator_residual"] = frame["flow"] - (
        18 + frame["actuator_position"] * 0.42
    )
    frame["power_proxy"] = frame["current"] * frame["actuator_position"] / 100
    return frame


def main() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    frame = pd.concat(
        [build_device(device, steps=5000, seed=2026) for device in range(12)],
        ignore_index=True,
    )
    frame = add_features(frame)
    split = np.where(
        frame["device_id"] <= 7,
        "train",
        np.where(frame["device_id"] <= 9, "validation", "test"),
    )
    frame["split"] = split
    manifest = {}
    for split_name in ["train", "validation", "test"]:
        split_frame = frame[frame["split"] == split_name].reset_index(drop=True)
        path = DATA_DIR / f"{split_name}.parquet"
        Dataset.from_pandas(split_frame, preserve_index=False).to_parquet(path)
        manifest[split_name] = {
            "rows": len(split_frame),
            "devices": sorted(split_frame["device_id"].unique().tolist()),
            "anomaly_rate": float(split_frame["label"].mean()),
            "path": path.name,
        }
    (DATA_DIR / "manifest.json").write_text(
        json.dumps(manifest, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()