API¶

def cartesian(self, other: Graph) -> Graph:
    """Return the directed Cartesian product over left-major node pairs.

    Left-factor edges precede right-factor edges in the returned COO order.
    The one-node edgeless graph is an exact identity. Products are associative
    as edge multisets; regrouping factors can permute COO edge order.
    """
    return Graph(
        self.nodes * other.nodes,
        [source * other.nodes + node for source in self.source for node in range(other.nodes)]
        + [node * other.nodes + source for node in range(self.nodes) for source in other.source],
        [target * other.nodes + node for target in self.target for node in range(other.nodes)]
        + [node * other.nodes + target for node in range(self.nodes) for target in other.target],
    )

def edge_values(
    self,
    values: Tensor,
    *,
    endpoint: Literal["source", "target"],
) -> Tensor:
    """Gather node values into original COO edge order."""
    self._validate_node_values(values)
    if endpoint not in ("source", "target"):
        raise ValueError("endpoint must be 'source' or 'target'")

    shape = values.shape
    order = (values.ndim - 2, *range(values.ndim - 2), values.ndim - 1)
    flat = values.permute(order).reshape(self.nodes, -1)
    output = self._csr.edge_values(flat, source=endpoint == "source")
    return output.reshape(self.edges, *shape[:-2], shape[-1]).permute(
        *range(1, values.ndim - 1),
        0,
        values.ndim - 1,
    )

def in_degree(self, *, device: str) -> Tensor:
    """Return incoming degree on one device."""
    if not isinstance(device, str):
        raise ValueError("in_degree requires one device")
    return self._csr.in_degree(device)

def mean(self, values: Tensor) -> Tensor:
    """Mean incoming values over node axis -2, with zero for empty rows."""
    self._validate_node_values(values)
    if not dtypes.is_float(values.dtype):
        raise ValueError(f"mean values must have a floating dtype, got {values.dtype}")
    assert isinstance(values.device, str)
    degree = self.in_degree(device=values.device).maximum(1).cast(values.dtype)
    degree = degree.reshape((1,) * (values.ndim - 2) + (self.nodes, 1))
    return self.sum(values) / degree

def softmax(self, edge_score: Tensor) -> Tensor:
    """Normalize scalar edge scores over each target's incoming edges."""
    if edge_score.ndim != 1 or edge_score.shape[0] != self.edges:
        raise ValueError(f"edge_score must have shape [{self.edges}], got {edge_score.shape}")
    if not dtypes.is_float(edge_score.dtype):
        raise ValueError(f"edge_score must have a floating dtype, got {edge_score.dtype}")
    if not isinstance(edge_score.device, str):
        raise ValueError("graph softmax requires one device")
    return self._csr.softmax(edge_score)

def sum(self, values: Tensor, edge_weight: Tensor | None = None) -> Tensor:
    """Sum incoming values over node axis -2 with optional shared edge weights."""
    self._validate_node_values(values)
    if edge_weight is not None:
        if edge_weight.ndim != 1 or edge_weight.shape[0] != self.edges:
            raise ValueError(f"edge_weight must have shape [{self.edges}], got {edge_weight.shape}")
        if values.dtype != edge_weight.dtype:
            raise ValueError(f"values and edge_weight must have the same dtype, got {values.dtype} and {edge_weight.dtype}")
        if edge_weight.device != values.device:
            raise ValueError("weighted graph sum requires one shared device")

    shape = values.shape
    order = (values.ndim - 2, *range(values.ndim - 2), values.ndim - 1)
    flat = values.permute(order).reshape(self.nodes, -1)
    output = self._csr.sum(flat) if edge_weight is None else self._csr.weighted_sum(flat, edge_weight)
    return output.reshape(self.nodes, *shape[:-2], shape[-1]).permute(
        *range(1, values.ndim - 1),
        0,
        values.ndim - 1,
    )

def sum_edges(self, values: Tensor) -> Tensor:
    """Sum COO-ordered edge values at their target nodes over axis -2."""
    if values.ndim < 2:
        raise ValueError(f"values must have shape [..., E, H], got {values.shape}")
    if values.shape[-2] != self.edges:
        raise ValueError(f"values must have {self.edges} edge rows, got {values.shape[-2]}")
    if not isinstance(values.device, str):
        raise ValueError("edge values require one device")

    shape = values.shape
    order = (values.ndim - 2, *range(values.ndim - 2), values.ndim - 1)
    flat = values.permute(order).reshape(self.edges, prod(shape[:-2]) * shape[-1])
    output = self._csr._segment_sum(flat)
    return output.reshape(self.nodes, *shape[:-2], shape[-1]).permute(
        *range(1, values.ndim - 1),
        0,
        values.ndim - 1,
    )

def batches(self, *, batch_size: int, history: int) -> Iterator[tuple[Tensor, Tensor]]:
    """Yield causal sequence-to-one windows as [B, L, N, F] and [B, N, Y]."""
    for name, value in (("batch_size", batch_size), ("history", history)):
        if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
            raise ValueError(f"{name} must be a positive integer")
    if history > len(self):
        raise ValueError(f"history must not exceed the {len(self)} time steps")

    windows = len(self) - history + 1
    for start in range(0, windows, batch_size):
        stop = min(start + batch_size, windows)
        yield (
            Tensor.stack(
                *(self.x[start + lag:stop + lag] for lag in range(history)),
                dim=1,
            ),
            self.y[start + history - 1:stop + history - 1],
        )

def split(self, train_ratio: float) -> tuple[StaticGraphTemporalSignal, StaticGraphTemporalSignal]:
    """Split once along time, preserving order and topology."""
    if not 0 < train_ratio < 1:
        raise ValueError("train_ratio must be between zero and one")
    train_steps = int(train_ratio * len(self))
    if train_steps == 0 or train_steps == len(self):
        raise ValueError("train and test must both be non-empty")
    return self[:train_steps], self[train_steps:]

def prefix(self, cutoff: int) -> TemporalEdges:
    """Return the events with timestamps strictly before ``cutoff``."""
    if not isinstance(cutoff, int) or isinstance(cutoff, bool):
        raise ValueError("cutoff must be an integer timestamp")
    stop = bisect_left(self.timestamp, cutoff)
    return TemporalEdges(
        self.nodes,
        self.source[:stop],
        self.target[:stop],
        self.timestamp[:stop],
    )