koopman-graph 0.1.0__py3-none-any.whl

koopman_graph/__init__.py

@@ -0,0 +1,48 @@
+"""KoopmanGraph: topology-aware Koopman dynamics on graphs.
+
+Public API
+----------
+``GraphKoopmanModel``
+    End-to-end encode → Koopman advance → decode model.
+``GNNEncoder``, ``GATEncoder``
+    GNN encoders for latent lifting.
+``GNNDecoder``
+    GNN decoder for physical reconstruction.
+``KoopmanOperator``
+    Learnable finite-dimensional Koopman matrix.
+``GraphSnapshotSequence``
+    Container for time-ordered PyG graph snapshots.
+``ForwardConsistencyLoss``
+    Latent-space linear evolution consistency loss.
+``BackwardConsistencyLoss``
+    Latent-space inverse linear evolution consistency loss.
+``FitHistory``
+    Training history returned by :meth:`~koopman_graph.model.GraphKoopmanModel.fit`.
+``LossWeights``
+    Weights for reconstruction and consistency loss terms.
+``__version__``
+    Package version string.
+"""
+
+from koopman_graph.data import GraphSnapshotSequence
+from koopman_graph.decoder import GNNDecoder
+from koopman_graph.encoder import GATEncoder, GNNEncoder
+from koopman_graph.losses import BackwardConsistencyLoss, ForwardConsistencyLoss
+from koopman_graph.model import GraphKoopmanModel
+from koopman_graph.operator import KoopmanOperator
+from koopman_graph.training import FitHistory, LossWeights
+
+__all__ = [
+    "BackwardConsistencyLoss",
+    "FitHistory",
+    "ForwardConsistencyLoss",
+    "GATEncoder",
+    "GNNDecoder",
+    "GNNEncoder",
+    "GraphKoopmanModel",
+    "GraphSnapshotSequence",
+    "KoopmanOperator",
+    "LossWeights",
+    "__version__",
+]
+__version__ = "0.1.0"

koopman_graph/data.py

@@ -0,0 +1,253 @@
+"""Utilities for spatiotemporal graph snapshot sequences."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch_geometric.data import Data
+
+ArrayLike = Tensor | np.ndarray
+
+
+def _as_tensor(value: ArrayLike, *, dtype: torch.dtype | None = None) -> Tensor:
+    """Convert an array-like value to a :class:`torch.Tensor`.
+
+    Parameters
+    ----------
+    value : Tensor or ndarray
+        Input array or tensor.
+    dtype : torch.dtype, optional
+        Target dtype. When ``value`` is already a tensor, conversion is applied
+        only if the dtypes differ.
+
+    Returns
+    -------
+    Tensor
+        Tensor representation of ``value``.
+    """
+    if isinstance(value, Tensor):
+        if dtype is not None and value.dtype != dtype:
+            return value.to(dtype=dtype)
+        return value
+    return torch.as_tensor(value, dtype=dtype)
+
+
+def _validate_shared_topology(snapshots: Sequence[Data]) -> None:
+    """Verify that all snapshots share node count, features, and topology.
+
+    Parameters
+    ----------
+    snapshots : sequence of Data
+        Graph snapshots to validate.
+
+    Raises
+    ------
+    ValueError
+        If the sequence is empty or any snapshot differs in ``edge_index``,
+        node count, or feature dimension from the first snapshot.
+    """
+    if not snapshots:
+        msg = "GraphSnapshotSequence requires at least one snapshot"
+        raise ValueError(msg)
+
+    reference = snapshots[0]
+    ref_edge_index = reference.edge_index
+    ref_num_nodes = reference.num_nodes
+    ref_in_channels = reference.x.shape[1]
+
+    for idx, snapshot in enumerate(snapshots[1:], start=1):
+        if snapshot.num_nodes != ref_num_nodes:
+            msg = (
+                f"Snapshot {idx} has {snapshot.num_nodes} nodes, "
+                f"expected {ref_num_nodes}"
+            )
+            raise ValueError(msg)
+        if snapshot.x.shape[1] != ref_in_channels:
+            msg = (
+                f"Snapshot {idx} has feature dimension {snapshot.x.shape[1]}, "
+                f"expected {ref_in_channels}"
+            )
+            raise ValueError(msg)
+        if not torch.equal(snapshot.edge_index, ref_edge_index):
+            msg = f"Snapshot {idx} has a different edge_index than snapshot 0"
+            raise ValueError(msg)
+
+
+class GraphSnapshotSequence:
+    """Container for a time-ordered sequence of PyG ``Data`` graph snapshots.
+
+    All snapshots must share the same ``edge_index``, node count, and feature
+    dimension. Downstream training APIs should require at least two snapshots;
+    construction here allows a single snapshot for inspection or prediction-only
+    workflows.
+
+    Attributes
+    ----------
+    snapshots : list of Data
+        Time-ordered PyG graph snapshots sharing the same topology.
+    edge_index : Tensor
+        Shared edge index, shape ``(2, num_edges)``.
+    num_nodes : int
+        Number of nodes in the graph topology.
+    num_timesteps : int
+        Length of the temporal sequence.
+    in_channels : int
+        Node feature dimension shared across snapshots.
+    """
+
+    def __init__(self, snapshots: Sequence[Data]) -> None:
+        """Initialize from a sequence of graph snapshots.
+
+        Parameters
+        ----------
+        snapshots : sequence of Data
+            Time-ordered graph snapshots. Validated for shared topology on
+            construction.
+        """
+        snapshot_list = list(snapshots)
+        _validate_shared_topology(snapshot_list)
+        self._snapshots = snapshot_list
+
+    @classmethod
+    def from_arrays(
+        cls,
+        node_features: ArrayLike,
+        edge_index: ArrayLike,
+        *,
+        dtype: torch.dtype = torch.float32,
+    ) -> GraphSnapshotSequence:
+        """Build a sequence from node feature arrays and a shared topology.
+
+        Parameters
+        ----------
+        node_features : array-like
+            Array with shape ``(num_timesteps, num_nodes, in_channels)``.
+        edge_index : array-like
+            Shared edge index with shape ``(2, num_edges)``.
+        dtype : torch.dtype, optional
+            Floating dtype used when converting numpy inputs to torch tensors.
+            Default is ``torch.float32``.
+
+        Returns
+        -------
+        :class:`~koopman_graph.data.GraphSnapshotSequence`
+            Validated snapshot sequence.
+
+        Raises
+        ------
+        ValueError
+            If ``node_features`` or ``edge_index`` have invalid shape.
+        """
+        features = _as_tensor(node_features, dtype=dtype)
+        edges = _as_tensor(edge_index, dtype=torch.long)
+
+        if features.ndim != 3:
+            msg = (
+                f"node_features must have shape "
+                f"(num_timesteps, num_nodes, in_channels), got {tuple(features.shape)}"
+            )
+            raise ValueError(msg)
+        if edges.ndim != 2 or edges.shape[0] != 2:
+            msg = f"edge_index must have shape (2, num_edges), got {tuple(edges.shape)}"
+            raise ValueError(msg)
+        if features.shape[0] < 1:
+            msg = "node_features must contain at least one timestep"
+            raise ValueError(msg)
+
+        snapshots = [
+            Data(x=features[t], edge_index=edges) for t in range(features.shape[0])
+        ]
+        return cls(snapshots)
+
+    @property
+    def snapshots(self) -> list[Data]:
+        """Return the underlying list of graph snapshots.
+
+        Returns
+        -------
+        list of Data
+            Time-ordered PyG graph snapshots.
+        """
+        return self._snapshots
+
+    @property
+    def edge_index(self) -> Tensor:
+        """Return the shared edge index for all snapshots.
+
+        Returns
+        -------
+        Tensor
+            Edge index with shape ``(2, num_edges)``.
+        """
+        return self._snapshots[0].edge_index
+
+    @property
+    def num_nodes(self) -> int:
+        """Return the number of nodes in the graph topology.
+
+        Returns
+        -------
+        int
+            Node count shared across all snapshots.
+        """
+        return int(self._snapshots[0].num_nodes)
+
+    @property
+    def num_timesteps(self) -> int:
+        """Return the number of timesteps in the sequence.
+
+        Returns
+        -------
+        int
+            Length of the temporal sequence.
+        """
+        return len(self._snapshots)
+
+    @property
+    def in_channels(self) -> int:
+        """Return the node feature dimension.
+
+        Returns
+        -------
+        int
+            Feature dimension shared across all snapshots.
+        """
+        return int(self._snapshots[0].x.shape[1])
+
+    def __len__(self) -> int:
+        """Return the number of timesteps in the sequence.
+
+        Returns
+        -------
+        int
+            Same value as :attr:`num_timesteps`.
+        """
+        return len(self._snapshots)
+
+    def __getitem__(self, index: int) -> Data:
+        """Return the graph snapshot at ``index``.
+
+        Parameters
+        ----------
+        index : int
+            Timestep index.
+
+        Returns
+        -------
+        Data
+            Graph snapshot at the requested timestep.
+        """
+        return self._snapshots[index]
+
+    def __iter__(self) -> Iterator[Data]:
+        """Iterate over graph snapshots in temporal order.
+
+        Yields
+        ------
+        Data
+            Graph snapshot at each timestep.
+        """
+        return iter(self._snapshots)

koopman_graph/datasets/__init__.py

@@ -0,0 +1,31 @@
+"""Dataset utilities for KoopmanGraph benchmarks and examples.
+
+Public benchmarks
+-----------------
+SyntheticDynamicGraphBenchmark
+    Reproducible Laplacian-diffusion dynamics on path/ring topologies.
+GridDynamicGraphBenchmark
+    Laplacian diffusion on a 4-connected 2D lattice graph.
+AnisotropicAdvectionGridBenchmark
+    Directional advection on a grid with asymmetric neighbor weights.
+IEEE118DynamicBenchmark
+    IEEE 118-bus MATPOWER topology with simulated voltage/load dynamics.
+MetrLaTrafficBenchmark
+    METR-LA traffic-speed sensor graph with cached speed snapshots.
+"""
+
+from koopman_graph.datasets.grid import (
+    AnisotropicAdvectionGridBenchmark,
+    GridDynamicGraphBenchmark,
+)
+from koopman_graph.datasets.ieee118 import IEEE118DynamicBenchmark
+from koopman_graph.datasets.metr_la import MetrLaTrafficBenchmark
+from koopman_graph.datasets.synthetic import SyntheticDynamicGraphBenchmark
+
+__all__ = [
+    "AnisotropicAdvectionGridBenchmark",
+    "GridDynamicGraphBenchmark",
+    "IEEE118DynamicBenchmark",
+    "MetrLaTrafficBenchmark",
+    "SyntheticDynamicGraphBenchmark",
+]

koopman_graph/datasets/dynamics.py

@@ -0,0 +1,309 @@
+"""Shared Laplacian diffusion dynamics for benchmark datasets."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+from koopman_graph.data import GraphSnapshotSequence
+
+InitialStateName = Literal["random", "ones"]
+
+
+def normalized_step_operator(
+    edge_index: Tensor,
+    num_nodes: int,
+    diffusion_rate: float,
+    *,
+    dtype: torch.dtype,
+) -> Tensor:
+    """Build one-step Laplacian diffusion operator ``I - alpha * L_sym``.
+
+    The symmetric normalized Laplacian is ``L_sym = I - D^{-1/2} A D^{-1/2}``.
+    The returned operator is ``(1 - alpha) * I + alpha * D^{-1/2} A D^{-1/2}``.
+
+    Parameters
+    ----------
+    edge_index : Tensor
+        Edge index with shape ``(2, num_edges)``.
+    num_nodes : int
+        Number of graph nodes.
+    diffusion_rate : float
+        Diffusion strength ``alpha`` in ``[0, 1]``.
+    dtype : torch.dtype
+        Floating dtype for the dense operator.
+
+    Returns
+    -------
+    Tensor
+        Step operator with shape ``(num_nodes, num_nodes)``.
+    """
+    row, col = edge_index
+    deg = torch.zeros(num_nodes, dtype=dtype)
+    deg.index_add_(0, row, torch.ones(row.size(0), dtype=dtype))
+    deg_inv_sqrt = deg.pow(-0.5)
+    deg_inv_sqrt[torch.isinf(deg_inv_sqrt)] = 0.0
+
+    norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
+    adj = torch.zeros((num_nodes, num_nodes), dtype=dtype)
+    adj[row, col] = norm
+
+    eye = torch.eye(num_nodes, dtype=dtype)
+    return (1.0 - diffusion_rate) * eye + diffusion_rate * adj
+
+
+def make_generator(seed: int | None) -> torch.Generator | None:
+    """Return a seeded torch generator when ``seed`` is provided.
+
+    Parameters
+    ----------
+    seed : int or None
+        Random seed. When ``None``, no generator is created.
+
+    Returns
+    -------
+    torch.Generator or None
+        Seeded generator, or ``None`` when ``seed`` is ``None``.
+    """
+    if seed is None:
+        return None
+    generator = torch.Generator()
+    generator.manual_seed(seed)
+    return generator
+
+
+def validate_diffusion_generation_params(
+    *,
+    decay_rate: float,
+    noise_std: float,
+    diffusion_rate: float | None = None,
+    initial_state: InitialStateName | None = None,
+) -> None:
+    """Validate shared diffusion benchmark generation parameters.
+
+    Parameters
+    ----------
+    decay_rate : float
+        Global amplitude decay applied each diffusion step.
+    noise_std : float
+        Standard deviation of additive Gaussian noise.
+    diffusion_rate : float or None, optional
+        Laplacian diffusion strength in ``[0, 1]``. Validated when provided.
+    initial_state : {"random", "ones"} or None, optional
+        Initial node feature pattern. Validated when provided.
+
+    Raises
+    ------
+    ValueError
+        If any parameter is outside its allowed range.
+    """
+    if diffusion_rate is not None and not 0.0 <= diffusion_rate <= 1.0:
+        msg = f"diffusion_rate must be in [0, 1], got {diffusion_rate}"
+        raise ValueError(msg)
+    if decay_rate <= 0.0:
+        msg = f"decay_rate must be > 0, got {decay_rate}"
+        raise ValueError(msg)
+    if noise_std < 0.0:
+        msg = f"noise_std must be >= 0, got {noise_std}"
+        raise ValueError(msg)
+    if initial_state is not None and initial_state not in {"random", "ones"}:
+        msg = f"initial_state must be 'random' or 'ones', got {initial_state!r}"
+        raise ValueError(msg)
+
+
+def initial_node_features(
+    num_nodes: int,
+    in_channels: int,
+    initial_state: InitialStateName,
+    *,
+    generator: torch.Generator | None,
+    dtype: torch.dtype,
+) -> Tensor:
+    """Create the initial node feature matrix for diffusion benchmarks.
+
+    Parameters
+    ----------
+    num_nodes : int
+        Number of graph nodes.
+    in_channels : int
+        Node feature dimension.
+    initial_state : {"random", "ones"}
+        Feature initialization pattern.
+    generator : torch.Generator or None
+        Optional RNG used when ``initial_state="random"``.
+    dtype : torch.dtype
+        Floating dtype for the returned tensor.
+
+    Returns
+    -------
+    Tensor
+        Initial node features with shape ``(num_nodes, in_channels)``.
+    """
+    if initial_state == "ones":
+        return torch.ones((num_nodes, in_channels), dtype=dtype)
+    return torch.randn(
+        (num_nodes, in_channels),
+        generator=generator,
+        dtype=dtype,
+    )
+
+
+def apply_laplacian_diffusion_step(
+    state: Tensor,
+    step_operator: Tensor,
+    decay_rate: float,
+) -> Tensor:
+    """Advance node features by one Laplacian diffusion step.
+
+    Parameters
+    ----------
+    state : Tensor
+        Current node features with shape ``(num_nodes, in_channels)``.
+    step_operator : Tensor
+        One-step diffusion operator with shape ``(num_nodes, num_nodes)``.
+    decay_rate : float
+        Global amplitude decay applied to the diffused state.
+
+    Returns
+    -------
+    Tensor
+        Updated node features with the same shape as ``state``.
+    """
+    return decay_rate * (step_operator @ state)
+
+
+def add_gaussian_noise(
+    state: Tensor,
+    noise_std: float,
+    *,
+    generator: torch.Generator | None,
+    dtype: torch.dtype,
+) -> Tensor:
+    """Add optional isotropic Gaussian noise to node features.
+
+    Parameters
+    ----------
+    state : Tensor
+        Node features to perturb.
+    noise_std : float
+        Standard deviation of additive Gaussian noise. No noise is added when
+        ``noise_std <= 0``.
+    generator : torch.Generator or None
+        Optional RNG for drawing noise samples.
+    dtype : torch.dtype
+        Floating dtype for generated noise.
+
+    Returns
+    -------
+    Tensor
+        Perturbed node features with the same shape as ``state``.
+    """
+    if noise_std <= 0.0:
+        return state
+    noise = torch.randn(state.shape, generator=generator, dtype=dtype)
+    return state + noise_std * noise
+
+
+def laplacian_diffusion_rollout(
+    *,
+    edge_index: Tensor,
+    num_nodes: int,
+    num_timesteps: int,
+    in_channels: int,
+    diffusion_rate: float,
+    decay_rate: float,
+    noise_std: float,
+    initial_state: InitialStateName,
+    dtype: torch.dtype,
+    generator: torch.Generator | None,
+    initial_features: Tensor | None = None,
+) -> Tensor:
+    """Simulate Laplacian diffusion dynamics and return stacked features.
+
+    Parameters
+    ----------
+    edge_index : Tensor
+        Shared graph topology.
+    num_nodes : int
+        Number of nodes in the graph.
+    num_timesteps : int
+        Number of temporal snapshots to generate.
+    in_channels : int
+        Node feature dimension.
+    diffusion_rate : float
+        Laplacian diffusion strength in ``[0, 1]``.
+    decay_rate : float
+        Global amplitude decay applied each step.
+    noise_std : float
+        Standard deviation of additive Gaussian noise.
+    initial_state : {"random", "ones"}
+        Initial node feature pattern when ``initial_features`` is ``None``.
+    dtype : torch.dtype
+        Floating dtype for generated features.
+    generator : torch.Generator or None
+        Optional RNG for initial state and noise.
+    initial_features : Tensor or None, optional
+        Explicit initial node features with shape ``(num_nodes, in_channels)``.
+
+    Returns
+    -------
+    Tensor
+        Node features with shape ``(num_timesteps, num_nodes, in_channels)``.
+    """
+    step_operator = normalized_step_operator(
+        edge_index,
+        num_nodes,
+        diffusion_rate,
+        dtype=dtype,
+    )
+    state = (
+        initial_features.clone()
+        if initial_features is not None
+        else initial_node_features(
+            num_nodes,
+            in_channels,
+            initial_state,
+            generator=generator,
+            dtype=dtype,
+        )
+    )
+
+    snapshots = [state.clone()]
+    for _ in range(num_timesteps - 1):
+        state = apply_laplacian_diffusion_step(state, step_operator, decay_rate)
+        state = add_gaussian_noise(
+            state,
+            noise_std,
+            generator=generator,
+            dtype=dtype,
+        )
+        snapshots.append(state.clone())
+
+    return torch.stack(snapshots, dim=0)
+
+
+def diffusion_sequence_from_features(
+    features: Tensor,
+    edge_index: Tensor,
+    *,
+    dtype: torch.dtype = torch.float32,
+) -> GraphSnapshotSequence:
+    """Wrap rollout features and topology in a validated snapshot sequence.
+
+    Parameters
+    ----------
+    features : Tensor
+        Node features with shape ``(num_timesteps, num_nodes, in_channels)``.
+    edge_index : Tensor
+        Shared edge index with shape ``(2, num_edges)``.
+    dtype : torch.dtype, optional
+        Floating dtype used when building snapshots. Default is ``torch.float32``.
+
+    Returns
+    -------
+    :class:`~koopman_graph.data.GraphSnapshotSequence`
+        Validated time-ordered snapshot sequence.
+    """
+    return GraphSnapshotSequence.from_arrays(features, edge_index, dtype=dtype)