Spatio-Temporal Graph Network

"""One spatio-temporal block, transcribed the way the equations read.

Two operations, alternated:

  spatial   Z = sum_k P^k X W_k        mix each node with its k-hop neighbours
  temporal  H = A * sigmoid(B)         a causal gated convolution over time

P is the row-normalized transition matrix D^-1 A. Note carefully that P^k is
APPLIED k times and never formed - forming it makes a sparse graph dense,
which is the difference between a model that scales to a city and one that
does not. Here that is just a loop; in the fast stage it is the whole point.

Plain loops, no libraries. Adjacency as a list of neighbour lists.
"""

import math

BLANK_MASK = -1.0


def row_normalize(adjacency):
    """D^-1 A as neighbour lists of (node, weight) with weights summing to 1."""
    normalized = []
    for row in adjacency:
        total = sum(weight for _, weight in row)
        if total == 0.0:
            normalized.append([])
            continue
        normalized.append([(node, weight / total) for node, weight in row])
    return normalized


def diffusion_conv(features, transition, weights_per_hop):
    """Z[n][c_out] = sum over hops k of (P^k X)[n] . W_k[:, c_out].

    features:        N x C_in
    transition:      neighbour lists of (node, weight)
    weights_per_hop: K + 1 matrices, each C_in x C_out
    """
    num_nodes = len(features)
    num_in = len(features[0])
    num_out = len(weights_per_hop[0][0])

    output = [[0.0] * num_out for _ in range(num_nodes)]

    # hop_features starts at X (k = 0) and is advanced one hop at a time by a
    # single sparse product. K products total, no matrix power anywhere.
    hop_features = [row[:] for row in features]

    for weights in weights_per_hop:
        for node in range(num_nodes):
            for channel_out in range(num_out):
                accumulated = 0.0
                for channel_in in range(num_in):
                    accumulated += hop_features[node][channel_in] * weights[channel_in][channel_out]
                output[node][channel_out] += accumulated

        # Advance one hop: next[n] = sum over neighbours m of P[n][m] * hop[m].
        advanced = [[0.0] * num_in for _ in range(num_nodes)]
        for node in range(num_nodes):
            for neighbour, weight in transition[node]:
                for channel_in in range(num_in):
                    advanced[node][channel_in] += weight * hop_features[neighbour][channel_in]
        hop_features = advanced

    return output


def gated_temporal_conv(sequence, kernel_value, kernel_gate, bias_value, bias_gate):
    """Causal 1D convolution over time with a gated linear unit.

    sequence: T x N x C_in. Returns (T - width + 1) x N x C_out - a causal
    convolution consumes width - 1 steps of context and cannot invent them.
    """
    num_steps = len(sequence)
    num_nodes = len(sequence[0])
    num_in = len(sequence[0][0])
    width = len(kernel_value)
    num_out = len(kernel_value[0][0])
    out_steps = num_steps - width + 1

    output = [[[0.0] * num_out for _ in range(num_nodes)] for _ in range(out_steps)]

    for step in range(out_steps):
        for node in range(num_nodes):
            for channel_out in range(num_out):
                value = bias_value[channel_out]
                gate = bias_gate[channel_out]
                for offset in range(width):
                    for channel_in in range(num_in):
                        observation = sequence[step + offset][node][channel_in]
                        value += observation * kernel_value[offset][channel_in][channel_out]
                        gate += observation * kernel_gate[offset][channel_in][channel_out]
                # The gate decides which timesteps pass. sigmoid(gate) near
                # zero closes the channel for this step entirely.
                opened = 1.0 / (1.0 + math.exp(-gate))
                output[step][node][channel_out] = value * opened

    return output


def masked_mae(predictions, targets, mask):
    """Sensor networks always have gaps. An unmasked loss trains the model to
    predict whatever was used to fill them."""
    total = 0.0
    counted = 0
    for step in range(len(predictions)):
        for node in range(len(predictions[step])):
            for horizon in range(len(predictions[step][node])):
                if not mask[step][node][horizon]:
                    continue
                total += abs(predictions[step][node][horizon] - targets[step][node][horizon])
                counted += 1
    if counted == 0:
        raise ValueError("every target in the batch is masked out")
    return total / counted


def stgcn_block(sequence, adjacency, temporal_a, temporal_b, spatial_weights):
    """The STGCN sandwich: temporal, then spatial, then temporal."""
    transition = row_normalize(adjacency)

    hidden = gated_temporal_conv(sequence, *temporal_a)

    # The graph operator is per-timestep: it mixes space, not time.
    mixed = [diffusion_conv(frame, transition, spatial_weights) for frame in hidden]

    return gated_temporal_conv(mixed, *temporal_b)

"""The same block, typed, with the graph in CSR and the shape contract checked.

What changes is the representation and the boundary. The adjacency becomes a
compressed sparse row structure, which is what makes "linear in edges" a
property of the code rather than of the maths. Every shape precondition is
checked once, up front, because a silent shape mismatch in a graph model
produces a plausible forecast rather than an error.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import NamedTuple, Sequence

Matrix = list[list[float]]
Tensor3 = list[list[list[float]]]


class CsrGraph(NamedTuple):
    """Row-normalized transition matrix in compressed sparse row form.

    indptr has num_nodes + 1 entries; the neighbours of node n are
    indices[indptr[n]:indptr[n + 1]] with the matching weights.
    """

    indptr: list[int]
    indices: list[int]
    weights: list[float]
    num_nodes: int

    def out_degree(self, node: int) -> int:
        return self.indptr[node + 1] - self.indptr[node]


@dataclass(frozen=True)
class BlockConfig:
    """Frozen so a block's shape contract cannot drift after construction."""

    num_nodes: int
    in_channels: int
    hidden_channels: int
    diffusion_order: int
    temporal_width: int

    def receptive_field(self) -> int:
        """A block is two causal convolutions, so it consumes twice the context."""
        return 2 * (self.temporal_width - 1) + 1


class ShapeMismatch(ValueError):
    """Raised on a shape violation rather than broadcasting past it."""


class DegenerateGraph(ValueError):
    """Raised when the graph cannot support the requested diffusion order."""


def build_csr(edges: Sequence[tuple[int, int, float]], num_nodes: int) -> CsrGraph:
    """Build a row-normalized CSR transition matrix from a weighted edge list.

    Row normalization matters: repeated application of an unnormalized
    adjacency changes the feature scale at every hop, and the model then
    spends its capacity undoing that.
    """
    if num_nodes <= 0:
        raise ShapeMismatch("a graph needs at least one node")

    by_row: list[list[tuple[int, float]]] = [[] for _ in range(num_nodes)]
    for source, target, weight in edges:
        if not 0 <= source < num_nodes or not 0 <= target < num_nodes:
            raise ShapeMismatch(f"edge ({source}, {target}) is outside the node range")
        if weight < 0.0:
            raise ShapeMismatch("a transition weight cannot be negative")
        by_row[source].append((target, weight))

    isolated = sum(1 for row in by_row if not row)
    if isolated == num_nodes:
        raise DegenerateGraph("every node is isolated; the graph operator is the identity")

    indptr = [0]
    indices: list[int] = []
    weights: list[float] = []
    for row in by_row:
        total = math.fsum(weight for _, weight in row)
        for target, weight in row:
            indices.append(target)
            weights.append(weight / total if total > 0.0 else 0.0)
        indptr.append(len(indices))

    return CsrGraph(indptr=indptr, indices=indices, weights=weights, num_nodes=num_nodes)


def spmm(graph: CsrGraph, features: Matrix) -> Matrix:
    """One sparse product: out[n] = sum over neighbours m of P[n][m] * X[m].

    The only place the graph is touched. Applying this k times gives P^k X
    without ever forming P^k, which is what keeps the cost linear in edges.
    """
    if len(features) != graph.num_nodes:
        raise ShapeMismatch(f"expected {graph.num_nodes} rows, got {len(features)}")

    num_channels = len(features[0])
    output = [[0.0] * num_channels for _ in range(graph.num_nodes)]

    for node in range(graph.num_nodes):
        row = output[node]
        for slot in range(graph.indptr[node], graph.indptr[node + 1]):
            weight = graph.weights[slot]
            neighbour = features[graph.indices[slot]]
            for channel in range(num_channels):
                row[channel] += weight * neighbour[channel]
    return output


def diffusion_conv(
    features: Matrix,
    graph: CsrGraph,
    weights_per_hop: Sequence[Matrix],
) -> Matrix:
    """Z = sum over hops k of (P^k X) W_k, hop by hop."""
    if not weights_per_hop:
        raise ShapeMismatch("a diffusion convolution needs at least the zero-hop term")

    num_out = len(weights_per_hop[0][0])
    output = [[0.0] * num_out for _ in range(graph.num_nodes)]
    hop = features

    for hop_index, weights in enumerate(weights_per_hop):
        if len(weights) != len(hop[0]):
            raise ShapeMismatch(f"hop {hop_index} weight has the wrong input width")
        for node, row in enumerate(hop):
            target = output[node]
            for channel_in, activation in enumerate(row):
                if activation == 0.0:
                    continue
                weight_row = weights[channel_in]
                for channel_out in range(num_out):
                    target[channel_out] += activation * weight_row[channel_out]
        if hop_index + 1 < len(weights_per_hop):
            hop = spmm(graph, hop)

    return output


def sigmoid(value: float) -> float:
    """Branch on the sign so exp never overflows on a large negative input."""
    if value >= 0.0:
        return 1.0 / (1.0 + math.exp(-value))
    positive = math.exp(value)
    return positive / (1.0 + positive)


class TemporalKernel(NamedTuple):
    """Value and gate branches of one gated causal convolution."""

    value: Tensor3          # width x in_channels x out_channels
    gate: Tensor3
    value_bias: list[float]
    gate_bias: list[float]

    @property
    def width(self) -> int:
        return len(self.value)


def gated_temporal_conv(sequence: Tensor3, kernel: TemporalKernel) -> Tensor3:
    """Causal gated convolution over time. Output is width - 1 steps shorter."""
    if len(sequence) < kernel.width:
        raise ShapeMismatch(
            f"a width-{kernel.width} kernel needs at least that many steps, got {len(sequence)}"
        )

    num_nodes = len(sequence[0])
    num_out = len(kernel.value[0][0])
    out_steps = len(sequence) - kernel.width + 1
    output: Tensor3 = []

    for step in range(out_steps):
        frame = [[0.0] * num_out for _ in range(num_nodes)]
        for node in range(num_nodes):
            values = list(kernel.value_bias)
            gates = list(kernel.gate_bias)
            for offset in range(kernel.width):
                observations = sequence[step + offset][node]
                value_slab = kernel.value[offset]
                gate_slab = kernel.gate[offset]
                for channel_in, activation in enumerate(observations):
                    if activation == 0.0:
                        continue
                    value_row = value_slab[channel_in]
                    gate_row = gate_slab[channel_in]
                    for channel_out in range(num_out):
                        values[channel_out] += activation * value_row[channel_out]
                        gates[channel_out] += activation * gate_row[channel_out]
            frame[node] = [
                value * sigmoid(gate) for value, gate in zip(values, gates)
            ]
        output.append(frame)

    return output


def masked_mae(predictions: Tensor3, targets: Tensor3, observed: Sequence[Sequence[Sequence[bool]]]) -> float:
    """Mean absolute error over observed entries only.

    Guard clause first: an all-masked batch is a data-pipeline bug, and
    returning zero for it would let training proceed on nothing at all.
    """
    total = 0.0
    counted = 0
    for step, (prediction_frame, target_frame, mask_frame) in enumerate(
        zip(predictions, targets, observed)
    ):
        for node, (prediction_row, target_row, mask_row) in enumerate(
            zip(prediction_frame, target_frame, mask_frame)
        ):
            if len(prediction_row) != len(target_row):
                raise ShapeMismatch(f"step {step} node {node} horizon width disagrees")
            for prediction, target, is_observed in zip(prediction_row, target_row, mask_row):
                if not is_observed:
                    continue
                total += abs(prediction - target)
                counted += 1

    if counted == 0:
        raise ShapeMismatch("every target in the batch is masked out")
    return total / counted


def stgcn_block(
    sequence: Tensor3,
    graph: CsrGraph,
    config: BlockConfig,
    first: TemporalKernel,
    spatial: Sequence[Matrix],
    second: TemporalKernel,
) -> Tensor3:
    """Temporal, spatial, temporal. Preconditions checked before any work."""
    if len(sequence) < config.receptive_field():
        raise ShapeMismatch(
            f"block needs {config.receptive_field()} steps, got {len(sequence)}"
        )
    if len(spatial) != config.diffusion_order + 1:
        raise ShapeMismatch("one weight matrix per hop, including the zero-hop term")

    hidden = gated_temporal_conv(sequence, first)
    # The graph operator mixes space and never time: applied per frame.
    mixed = [diffusion_conv(frame, graph, spatial) for frame in hidden]
    return gated_temporal_conv(mixed, second)

"""Batched, vectorized, and fused. The hop loop stays; every other loop goes.

Three structural changes:

  1. The temporal convolution becomes one matrix product. A causal window of
     width w over T steps is a strided view of shape (T - w + 1, w, N, C),
     which reshapes to a 2-D matrix and multiplies against the flattened
     kernel. Value and gate branches are one product with 2 * C_out columns,
     so the two are fused rather than computed separately.
  2. The graph operator becomes a sparse-dense product per hop. K products,
     never a matrix power - the same rule as the literal version, but now the
     reason is visible: forming P^K would be a dense N x N allocation.
  3. The batch and time axes collapse into the leading dimension of a single
     BLAS call, so channel mixing is one GEMM instead of T separate ones.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray
from numpy.lib.stride_tricks import sliding_window_view
from scipy.sparse import csr_matrix

FLOAT = np.float32


class SpatioTemporalBlock:
    """Holds the graph and the pre-flattened kernels. Both are call-invariant."""

    def __init__(
        self,
        transition: csr_matrix,
        diffusion_order: int,
        temporal_width: int,
    ) -> None:
        if transition.shape[0] != transition.shape[1]:
            raise ValueError("the transition matrix must be square")

        # CSR with sorted indices and a single dtype: scipy dispatches to its
        # fast path only when both hold, and silently copies when they do not.
        transition = transition.astype(FLOAT)
        transition.sort_indices()
        self._transition = transition
        self._order = diffusion_order
        self._width = temporal_width
        self._num_nodes = transition.shape[0]

    def gated_temporal(
        self,
        sequence: NDArray[np.float32],
        kernel: NDArray[np.float32],
        bias: NDArray[np.float32],
    ) -> NDArray[np.float32]:
        """sequence (B, T, N, C_in) -> (B, T - w + 1, N, C_out).

        kernel is (w * C_in, 2 * C_out): value and gate side by side, so one
        GEMM produces both branches and the gate is never a second pass.
        """
        batch, steps, nodes, channels = sequence.shape
        if steps < self._width:
            raise ValueError(f"need at least {self._width} steps, got {steps}")

        # A strided view, not a copy: the (w, C) window per output step is
        # materialized only by the reshape below, and only once.
        windows = sliding_window_view(sequence, self._width, axis=1)
        # (B, T', N, C, w) -> (B, T', N, w, C) -> flat rows of w * C
        windows = np.moveaxis(windows, -1, -2)
        out_steps = steps - self._width + 1
        flat = np.ascontiguousarray(
            windows.reshape(batch * out_steps * nodes, self._width * channels)
        )

        projected = flat @ kernel
        projected += bias
        value, gate = np.split(projected, 2, axis=1)

        # In-place sigmoid on the gate, then in-place multiply: the gated
        # output reuses the value buffer and no intermediate is materialized.
        np.negative(gate, out=gate)
        np.exp(gate, out=gate)
        gate += 1.0
        np.divide(value, gate, out=value)

        return value.reshape(batch, out_steps, nodes, -1)

    def diffusion(
        self,
        frames: NDArray[np.float32],
        weights: NDArray[np.float32],
    ) -> NDArray[np.float32]:
        """frames (B, T, N, C_in), weights (K + 1, C_in, C_out) -> (B, T, N, C_out).

        The hop loop is the only Python-level loop left, and it is K long.
        """
        batch, steps, nodes, channels = frames.shape
        if nodes != self._num_nodes:
            raise ValueError(f"expected {self._num_nodes} nodes, got {nodes}")

        # (B * T, N, C) with N in the middle: the sparse product needs nodes
        # as the leading axis of a 2-D operand, so the reshape below puts
        # (N, B * T * C) contiguously and the SpMM sees one dense block.
        hop = frames.reshape(batch * steps, nodes, channels)
        out_channels = weights.shape[2]
        output = np.zeros((batch * steps, nodes, out_channels), dtype=FLOAT)

        for hop_index in range(self._order + 1):
            # einsum over (BT, N, C_in) x (C_in, C_out) is one GEMM after the
            # implicit reshape, and it accumulates straight into output.
            output += np.einsum('bnc,co->bno', hop, weights[hop_index], optimize=True)
            if hop_index < self._order:
                # One sparse product advances every batch and timestep at
                # once: transpose nodes to the front, multiply, transpose back.
                as_matrix = hop.transpose(1, 0, 2).reshape(nodes, -1)
                advanced = self._transition @ as_matrix
                hop = advanced.reshape(nodes, batch * steps, channels).transpose(1, 0, 2)
                hop = np.ascontiguousarray(hop)

        return output.reshape(batch, steps, nodes, out_channels)


def masked_mae(
    predictions: NDArray[np.float32],
    targets: NDArray[np.float32],
    observed: NDArray[np.bool_],
) -> np.float32:
    """One pass, no boolean indexing - the mask becomes a weight.

    Boolean indexing would allocate a compacted copy of every observed value;
    multiplying by the mask and dividing by its sum touches each element once.
    """
    counted = observed.sum(dtype=np.int64)
    if counted == 0:
        raise ValueError("every target in the batch is masked out")

    residual = np.subtract(predictions, targets, dtype=FLOAT)
    np.abs(residual, out=residual)
    residual *= observed
    return residual.sum(dtype=np.float64) / counted