kpnn2 0.1.0__py3-none-any.whl

kpnn2/__init__.py

@@ -0,0 +1,45 @@
+"""
+Turn a named edgelist into sparsely connected PyTorch layers
+you assemble yourself.
+
+``kpnn2`` turns a source/target edgelist into a spec so you write
+ordinary PyTorch. Pick the layout yourself: ``parse_layered``
+ranks a DAG into a ``LayeredSpec``, one ``Hop`` per layer whose
+mask carries every edge entering it, skips included, while
+``parse_adjacency`` puts every node in one state vector with
+packed edge indices (``AdjacencySpec``) and allows cycles. Use
+``PackedLinear`` on those packed indices, or
+``MaskedLinear(spec.to_mask())`` for the dense path.
+
+It is not a graph compiler: there is no ready-made model object
+and no training loop.
+"""
+
+from ._adjacency_spec import AdjacencySpec
+from ._align import align_inputs
+from ._attributions import map_node_attributions
+from ._errors import Kpnn2Error
+from ._gather import gather_hop_inputs
+from ._masked_linear import MaskedLinear
+from ._packed_linear import PackedLinear
+from ._parse import parse_layered
+from ._parse_adjacency import parse_adjacency
+from ._spec import Hop, LayeredSpec, Skip
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "parse_layered",
+    "parse_adjacency",
+    "LayeredSpec",
+    "Hop",
+    "Skip",
+    "AdjacencySpec",
+    "MaskedLinear",
+    "PackedLinear",
+    "gather_hop_inputs",
+    "align_inputs",
+    "map_node_attributions",
+    "Kpnn2Error",
+    "__version__",
+]

kpnn2/_adjacency_spec.py

@@ -0,0 +1,387 @@
+"""
+Structural blueprint for an edgelist-defined node network.
+"""
+
+from dataclasses import dataclass
+
+import pandas as pd
+import torch
+from torch import Tensor
+
+from ._layout import build_layout
+
+
+@dataclass(frozen=True)
+class AdjacencySpec:
+    """
+    Frozen blueprint from ``parse_adjacency``.
+
+    Structure only: not an ``nn.Module`` and no parameters. Every
+    node lives in one state vector and connectivity is packed as
+    source/target index tuples, so the graph may contain cycles
+    and self-loops. There is no stored square mask.
+
+    There are no depths here: no ``layer_nodes``, no ``hops``
+    tuple, and no ``skips``. This is not a one-layer
+    ``LayeredSpec``. Use ``PackedLinear`` on ``source_index`` /
+    ``target_index`` for a state update that never allocates
+    ``(n, n)``, or ``MaskedLinear(spec.to_mask())`` for the
+    dense path, and write the recurrence in ``forward()``.
+
+    Parameters
+    ----------
+    nodes : tuple[str, ...]
+        Every node name, alphabetical. This is the unit order of
+        the state vector and the row and column order of
+        ``to_mask()``.
+    input_nodes : tuple[str, ...]
+        In-degree 0 names, alphabetical. This is the column order
+        of tensors returned by ``align_inputs``.
+    output_nodes : tuple[str, ...]
+        Out-degree 0 names, alphabetical.
+    hidden_nodes : tuple[str, ...]
+        Names that are neither input nor output, alphabetical.
+    source_index : tuple[int, ...]
+        For each original edge, the column in ``nodes`` (the
+        source). Same length as ``target_index`` and as the
+        edge count. Order is canonical: lexicographic by
+        ``(source name, target name)``, identical to
+        ``to_edgelist()`` row order. Cycles and self-loops are
+        included.
+    target_index : tuple[int, ...]
+        For each original edge, the row in ``nodes`` (the
+        target). A dense square would have ``1.0`` at
+        ``[target_index[i], source_index[i]]``.
+    input_index : tuple[int, ...]
+        Position of each ``input_nodes`` name in ``nodes``.
+    output_index : tuple[int, ...]
+        Position of each ``output_nodes`` name in ``nodes``.
+
+    Notes
+    -----
+    Fields cannot be reassigned and sequences are tuples, so the
+    structure itself is fixed. There is no ``mask`` field and no
+    densifying ``mask`` property. ``to_mask()`` allocates a
+    fresh dense square on every call; mutating that tensor does
+    not change this spec. ``MaskedLinear(spec.to_mask())``
+    clones the square into a non-persistent buffer, so a layer
+    built earlier keeps its own connectivity.
+
+    ``align_inputs`` returns ``len(input_nodes)`` columns, which
+    is not the state width. Scatter that tensor into the
+    ``n``-wide state vector with ``input_index``. Input rows of
+    ``to_mask()`` are all zeros, so under the degree-aware init
+    of ``MaskedLinear`` they stay zero: writing the inputs in
+    is required, not cosmetic.
+
+    ``to_edgelist()`` returns the original edges as a
+    two-column ``source`` / ``target`` DataFrame, rows sorted
+    lexicographically, including cycle edges and self-loops.
+    ``parse_adjacency`` on that table reconstructs this spec's
+    node lists, packed indices, and input/output indices.
+
+    ``to_dict()`` returns a JSON-safe tagged dict
+    (``kpnn2_spec``, ``layout``, ``edges``). ``from_dict``
+    rebuilds this spec by calling ``parse_adjacency``.
+    ``fingerprint`` is the SHA-256 of that canonical JSON.
+    Pickle / ``torch.save`` of the dataclass is not the
+    supported interchange.
+
+    Examples
+    --------
+    An input feeding a two-node feedback core plus one output:
+
+    >>> import pandas as pd
+    >>> import kpnn2
+    >>> edgelist = pd.DataFrame(
+    ...     {
+    ...         "source": ["x", "a", "b", "a"],
+    ...         "target": ["a", "b", "a", "y"],
+    ...     }
+    ... )
+    >>> spec = kpnn2.parse_adjacency(edgelist)
+    >>> spec.nodes
+    ('a', 'b', 'x', 'y')
+    >>> spec.input_nodes, spec.output_nodes
+    (('x',), ('y',))
+    >>> spec.hidden_nodes
+    ('a', 'b')
+    >>> spec.input_index, spec.output_index
+    ((2,), (3,))
+    >>> spec.source_index
+    (0, 0, 1, 2)
+    >>> spec.target_index
+    (1, 3, 0, 0)
+    >>> tuple(spec.to_mask().shape)
+    (4, 4)
+    >>> spec.to_mask()[0].tolist()
+    [0.0, 1.0, 1.0, 0.0]
+    """
+
+    nodes: tuple[str, ...]
+    input_nodes: tuple[str, ...]
+    output_nodes: tuple[str, ...]
+    hidden_nodes: tuple[str, ...]
+    source_index: tuple[int, ...]
+    target_index: tuple[int, ...]
+    input_index: tuple[int, ...]
+    output_index: tuple[int, ...]
+
+    def __post_init__(self) -> None:
+        object.__setattr__(
+            self,
+            "nodes",
+            tuple(self.nodes),
+        )
+        object.__setattr__(
+            self,
+            "input_nodes",
+            tuple(self.input_nodes),
+        )
+        object.__setattr__(
+            self,
+            "output_nodes",
+            tuple(self.output_nodes),
+        )
+        object.__setattr__(
+            self,
+            "hidden_nodes",
+            tuple(self.hidden_nodes),
+        )
+        object.__setattr__(
+            self,
+            "source_index",
+            tuple(self.source_index),
+        )
+        object.__setattr__(
+            self,
+            "target_index",
+            tuple(self.target_index),
+        )
+        object.__setattr__(
+            self,
+            "input_index",
+            tuple(self.input_index),
+        )
+        object.__setattr__(
+            self,
+            "output_index",
+            tuple(self.output_index),
+        )
+
+    def to_mask(self) -> Tensor:
+        """
+        Allocate a dense float32 square from the packed edges.
+
+        Shape is ``(n, n)`` with ``n`` from the node layout
+        (``len(nodes)`` at width 1). The result starts at zeros;
+        each live edge sets ``1.0`` at
+        ``[target_index[i], source_index[i]]``. Every call
+        returns a fresh tensor. Mutating it does not change this
+        spec or the next ``to_mask()`` call.
+
+        Returns
+        -------
+        torch.Tensor
+            New dense connectivity square. This allocates.
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import kpnn2
+        >>> edgelist = pd.DataFrame(
+        ...     {
+        ...         "source": ["x", "a", "b", "a"],
+        ...         "target": ["a", "b", "a", "y"],
+        ...     }
+        ... )
+        >>> spec = kpnn2.parse_adjacency(edgelist)
+        >>> mask = spec.to_mask()
+        >>> tuple(mask.shape)
+        (4, 4)
+        >>> mask[0, 1].item(), mask[1, 0].item()
+        (1.0, 1.0)
+        """
+        layout = build_layout(self.nodes)
+        n_units = layout.n_units
+        mask = torch.zeros(
+            (
+                n_units,
+                n_units,
+            ),
+            dtype=torch.float32,
+        )
+        for source, target in zip(
+            self.source_index,
+            self.target_index,
+        ):
+            mask[target, source] = 1.0
+        return mask
+
+    def to_edgelist(self) -> pd.DataFrame:
+        """
+        Return this spec's edges as a two-column table.
+
+        Columns are exactly ``source`` then ``target``. Rows
+        follow the packed indices in canonical order: sorted
+        lexicographically by ``(source, target)``, one row per
+        original edge, names as strings, including cycle edges
+        and self-loops. Extra columns from the DataFrame that
+        was parsed are not reproduced.
+
+        ``parse_adjacency`` on this table reconstructs the same
+        node lists, packed indices, and input/output indices.
+
+        Returns
+        -------
+        pandas.DataFrame
+            One row per original edge.
+
+        Examples
+        --------
+        A cycle comes back as sorted pairs:
+
+        >>> import pandas as pd
+        >>> import kpnn2
+        >>> edgelist = pd.DataFrame(
+        ...     {
+        ...         "source": ["x", "a", "b", "a"],
+        ...         "target": ["a", "b", "a", "y"],
+        ...     }
+        ... )
+        >>> spec = kpnn2.parse_adjacency(edgelist)
+        >>> table = spec.to_edgelist()
+        >>> list(table.columns)
+        ['source', 'target']
+        >>> table["source"].tolist()
+        ['a', 'a', 'b', 'x']
+        >>> table["target"].tolist()
+        ['b', 'y', 'a', 'a']
+        """
+        from ._serialize import spec_to_edgelist
+
+        return spec_to_edgelist(self)
+
+    def to_dict(self) -> dict:
+        """
+        Return this spec as a JSON-safe tagged dict.
+
+        Keys are ``kpnn2_spec`` (integer ``1``), ``layout``
+        (``"adjacency"``), and ``edges`` (list of
+        ``[source, target]`` lists in the same order as
+        ``to_edgelist()`` rows, including cycle edges and
+        self-loops). The returned dict is new on every call.
+
+        Returns
+        -------
+        dict
+            Tagged edge list plus layout.
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import kpnn2
+        >>> edgelist = pd.DataFrame(
+        ...     {
+        ...         "source": ["x", "a", "b", "a"],
+        ...         "target": ["a", "b", "a", "y"],
+        ...     }
+        ... )
+        >>> spec = kpnn2.parse_adjacency(edgelist)
+        >>> payload = spec.to_dict()
+        >>> payload["kpnn2_spec"]
+        1
+        >>> payload["layout"]
+        'adjacency'
+        >>> payload["edges"]
+        [['a', 'b'], ['a', 'y'], ['b', 'a'], ['x', 'a']]
+        """
+        from ._serialize import spec_to_dict
+
+        return spec_to_dict(self)
+
+    @classmethod
+    def from_dict(cls, payload: dict) -> "AdjacencySpec":
+        """
+        Rebuild an ``AdjacencySpec`` from ``to_dict()`` output.
+
+        Calls ``parse_adjacency`` on a DataFrame built from
+        ``payload["edges"]``. Packed indices are not assembled
+        by hand. Extra unknown keys are ignored.
+
+        Parameters
+        ----------
+        payload : dict
+            A dict with ``kpnn2_spec``, ``layout``, and
+            ``edges``. ``layout`` must be ``"adjacency"``.
+
+        Returns
+        -------
+        AdjacencySpec
+            The parsed spec.
+
+        Raises
+        ------
+        Kpnn2Error
+            If ``payload`` is not a dict; ``kpnn2_spec`` is
+            missing or not ``1``; ``layout`` is missing, not a
+            known layout, or is ``"layered"``; or ``edges`` is
+            missing or not a sequence of two nonempty names.
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import kpnn2
+        >>> edgelist = pd.DataFrame(
+        ...     {
+        ...         "source": ["x", "a", "b", "a"],
+        ...         "target": ["a", "b", "a", "y"],
+        ...     }
+        ... )
+        >>> spec = kpnn2.parse_adjacency(edgelist)
+        >>> roundtrip = kpnn2.AdjacencySpec.from_dict(spec.to_dict())
+        >>> roundtrip.nodes == spec.nodes
+        True
+        """
+        from ._serialize import adjacency_spec_from_dict
+
+        return adjacency_spec_from_dict(payload)
+
+    @property
+    def fingerprint(self) -> str:
+        """
+        SHA-256 hex digest of the canonical ``to_dict()`` JSON.
+
+        The payload is ``json.dumps(self.to_dict(),
+        sort_keys=True, separators=(",", ":"),
+        ensure_ascii=False)`` encoded as UTF-8. The result is
+        64 lowercase hex characters. It is not Python
+        ``hash()``.
+
+        Returns
+        -------
+        str
+            Hex digest of the tagged spec dict.
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import kpnn2
+        >>> edgelist = pd.DataFrame(
+        ...     {
+        ...         "source": ["x", "a", "b", "a"],
+        ...         "target": ["a", "b", "a", "y"],
+        ...     }
+        ... )
+        >>> spec = kpnn2.parse_adjacency(edgelist)
+        >>> len(spec.fingerprint)
+        64
+        >>> (
+        ...     spec.fingerprint
+        ...     == kpnn2.parse_adjacency(spec.to_edgelist()).fingerprint
+        ... )
+        True
+        """
+        from ._serialize import spec_fingerprint
+
+        return spec_fingerprint(self)

kpnn2/_align.py

@@ -0,0 +1,215 @@
+"""
+Align named DataFrame columns to ``spec.input_nodes``.
+"""
+
+import pandas as pd
+import torch
+
+from ._adjacency_spec import AdjacencySpec
+from ._errors import Kpnn2Error
+from ._layout import build_layout, expand_columns
+from ._spec import LayeredSpec
+
+_TENSOR_NOT_ACCEPTED_MSG = (
+    "'data' is a tensor; a pandas DataFrame is required. "
+    "Pass a DataFrame so columns can be matched to "
+    "spec.input_nodes. Pre-ordered tensors go straight to "
+    "the model."
+)
+
+
+def align_inputs(
+    data: pd.DataFrame,
+    spec: LayeredSpec | AdjacencySpec,
+) -> torch.Tensor:
+    """
+    Return a float32 tensor whose columns follow ``spec.input_nodes``.
+
+    **DataFrame.** Required columns are ``spec.input_nodes``. Labels
+    are matched after ``str(...)``, the same conversion used for
+    edgelist node names, so an integer column ``1`` matches node
+    ``"1"``. Extra columns are ignored. Columns are reordered to
+    ``spec.input_nodes``. Missing, duplicate, or non-numeric required
+    columns raise an error.
+
+    **Tensor.** Not accepted. Raise ``Kpnn2Error``. Pre-ordered
+    tensors go straight to the model. Users who need alignment pass
+    a DataFrame.
+
+    AnnData is not supported.
+
+    Parameters
+    ----------
+    data : DataFrame
+        Feature table with named columns.
+    spec : LayeredSpec or AdjacencySpec
+        Graph structure whose ``input_nodes`` define column order.
+        Both spec types work the same way here; only
+        ``input_nodes`` is read.
+
+    Returns
+    -------
+    Tensor
+        Float32 tensor of shape
+        ``(n_samples, len(spec.input_nodes))``.
+
+    Raises
+    ------
+    Kpnn2Error
+        If ``spec`` is neither a ``LayeredSpec`` nor an
+        ``AdjacencySpec``; ``data`` is a tensor; ``data`` is not a
+        DataFrame; required DataFrame columns are missing or
+        duplicated (including after ``str`` conversion; the message
+        names the unique duplicated labels, sorted,
+        comma-separated); or required columns are non-numeric.
+
+    Notes
+    -----
+    PyTorch never sees feature names. Column ``i`` of the returned
+    tensor is ``spec.input_nodes[i]``. Use this function when the
+    table has named columns. A tensor whose columns already follow
+    ``spec.input_nodes`` goes straight to the model. Passing
+    ``DataFrame.to_numpy()`` (or any hand-stacked array) into the
+    model can silently wire the wrong features if the column order
+    differs.
+
+    The returned width is always ``len(spec.input_nodes)``. For a
+    ``LayeredSpec`` that is the width of ``hops[0].mask``, whose
+    only source layer is layer 0, so the tensor feeds the first
+    hop directly and needs no gathering. For an ``AdjacencySpec``
+    it is **not** the state width: scatter the tensor into the
+    ``len(spec.nodes)``-wide state vector with ``spec.input_index``
+    before calling ``MaskedLinear(spec.to_mask())``.
+
+    Examples
+    --------
+    Extra columns are dropped and remaining columns are reordered:
+
+    >>> import pandas as pd
+    >>> import torch
+    >>> import kpnn2
+    >>> edgelist = pd.DataFrame(
+    ...     {
+    ...         "source": ["A", "H"],
+    ...         "target": ["H", "C"],
+    ...     }
+    ... )
+    >>> spec = kpnn2.parse_layered(edgelist)
+    >>> spec.input_nodes
+    ('A',)
+    >>> df = pd.DataFrame(
+    ...     {
+    ...         "unused": [9.0, 8.0],
+    ...         "A": [0.5, 1.5],
+    ...     }
+    ... )
+    >>> x = kpnn2.align_inputs(df, spec)
+    >>> x.dtype
+    torch.float32
+    >>> tuple(x.shape)
+    (2, 1)
+    >>> x.tolist()
+    [[0.5], [1.5]]
+
+    An ``AdjacencySpec`` works the same way, but the result is
+    ``len(input_nodes)`` wide and must be scattered into the state
+    vector before it reaches ``MaskedLinear(spec.to_mask())``:
+
+    >>> cyclic = pd.DataFrame(
+    ...     {
+    ...         "source": ["x", "a", "b", "a"],
+    ...         "target": ["a", "b", "a", "y"],
+    ...     }
+    ... )
+    >>> state_spec = kpnn2.parse_adjacency(cyclic)
+    >>> inputs = pd.DataFrame({"x": [0.5, 1.5]})
+    >>> x = kpnn2.align_inputs(inputs, state_spec)
+    >>> tuple(x.shape), tuple(state_spec.to_mask().shape)
+    ((2, 1), (4, 4))
+    >>> state = torch.zeros(
+    ...     2,
+    ...     len(state_spec.nodes),
+    ... )
+    >>> state[:, state_spec.input_index] = x
+    >>> state.tolist()
+    [[0.0, 0.0, 0.5, 0.0], [0.0, 0.0, 1.5, 0.0]]
+
+    A tensor is not accepted; pass a DataFrame instead:
+
+    >>> t = torch.tensor([[0.5], [1.5]])
+    >>> kpnn2.align_inputs(t, spec)  # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    ...
+    Kpnn2Error: 'data' is a tensor; a pandas DataFrame is required. ...
+    """
+    if not isinstance(
+        spec,
+        (LayeredSpec, AdjacencySpec),
+    ):
+        raise Kpnn2Error("'spec' must be a LayeredSpec or an AdjacencySpec.")
+    if isinstance(data, torch.Tensor):
+        raise Kpnn2Error(_TENSOR_NOT_ACCEPTED_MSG)
+    if isinstance(data, pd.DataFrame):
+        return _align_dataframe(
+            data,
+            spec,
+        )
+    raise Kpnn2Error(
+        "Unsupported input data type. Expected a pandas DataFrame."
+    )
+
+
+def _align_dataframe(
+    data: pd.DataFrame,
+    spec: LayeredSpec | AdjacencySpec,
+) -> torch.Tensor:
+    """
+    Reorder numeric DataFrame columns to ``spec.input_nodes``.
+    """
+    str_columns = [str(name) for name in data.columns]
+    label_counts: dict[str, int] = {}
+    for label in str_columns:
+        label_counts[label] = label_counts.get(label, 0) + 1
+    duplicated_labels = sorted(
+        label for label, count in label_counts.items() if count > 1
+    )
+    if duplicated_labels:
+        labels_str = ", ".join(duplicated_labels)
+        raise Kpnn2Error(
+            "Input DataFrame must not contain duplicate column names "
+            "(including after converting labels to strings): "
+            f"{labels_str}."
+        )
+
+    renamed = data.copy(deep=False)
+    renamed.columns = str_columns
+
+    missing = [name for name in spec.input_nodes if name not in renamed.columns]
+    if missing:
+        missing_str = ", ".join(sorted(missing))
+        raise Kpnn2Error(
+            f"Input data is missing required feature name(s): {missing_str}."
+        )
+
+    ordered = renamed[list(spec.input_nodes)]
+    non_numeric = [
+        name
+        for name in spec.input_nodes
+        if not pd.api.types.is_numeric_dtype(ordered[name])
+    ]
+    if non_numeric:
+        non_numeric_str = ", ".join(sorted(non_numeric))
+        raise Kpnn2Error(
+            "Input data contains non-numeric feature column(s): "
+            f"{non_numeric_str}."
+        )
+
+    layout = build_layout(spec.input_nodes)
+    values = expand_columns(
+        ordered.to_numpy(copy=True),
+        layout,
+    )
+    return torch.tensor(
+        values,
+        dtype=torch.float32,
+    )