Cyclic graph example¶

This page is an end-to-end workflow for a cyclic knowledge graph represented as an edgelist: a prior with a feedback loop. You parse it, train a small model, and interpret named nodes.

Read Feedforward example first. That tutorial is a feedforward parse_layered() neural network on a directed acyclic graph (DAG). This page is what changes when the graph has a cycle.

What you get here:

• a tiny graph where two hidden nodes exchange a feedback edge
• parse_adjacency(), which allows cycles and self-loops
• why parse_layered() cannot parse this graph (it needs a DAG)
• one shared MaskedLinear over all nodes, applied a fixed number of steps so a signal can travel around the cycle
• training on simulated data, then Captum scores on named nodes
• controls that switch which input is informative, including the path that only reaches the output through feedback

A cycle has no depth, so there are no layers to apply in order. The loop is ordinary PyTorch instead: one shared MaskedLinear over the full state vector, a single vector over all nodes, applied for a fixed number of steps, with the input values written in before each step. Weights are shared across steps. kpnn2 does not unroll time or pick the step count.

Every step sees the same x, so this is a cyclic neural network on a static feature table, not a sequence model. It is not a recurrent neural network (RNN) cell such as a long short-term memory (LSTM) or a gated recurrent unit (GRU), and not a converge-until loop. For a changing x_t, see Time-series example.

In [ ]:

import random

import numpy as np
import pandas as pd
import seaborn as sns
import torch
from graphviz import Digraph
from IPython.display import display
from matplotlib import pyplot as plt
from torch import nn

import kpnn2

import random import numpy as np import pandas as pd import seaborn as sns import torch from graphviz import Digraph from IPython.display import display from matplotlib import pyplot as plt from torch import nn import kpnn2

In [ ]:

from matplotlib_inline.backend_inline import (
    set_matplotlib_formats,
)

random.seed(42)
np.random.seed(42)
torch.manual_seed(42)

set_matplotlib_formats(
    "svg",
    bbox_inches=None,
)
_FONT_SIZE = 14
_FIG_W = 7.0
_plot_rc = {
    "font.size": _FONT_SIZE,
    "axes.titlesize": _FONT_SIZE,
    "axes.labelsize": _FONT_SIZE,
    "xtick.labelsize": _FONT_SIZE,
    "ytick.labelsize": _FONT_SIZE,
    "legend.fontsize": _FONT_SIZE,
    "legend.title_fontsize": _FONT_SIZE,
    "lines.linewidth": 1.0,
    "font.family": "Liberation Sans",
    "text.color": "black",
    "axes.labelcolor": "black",
    "xtick.color": "black",
    "ytick.color": "black",
    "axes.edgecolor": "black",
    "axes.grid": False,
    "axes.facecolor": "none",
    "figure.facecolor": "none",
    "savefig.facecolor": "none",
    "savefig.edgecolor": "none",
    "svg.fonttype": "none",
}
plt.rcParams.update(_plot_rc)
sns.set_theme(
    style="ticks",
    rc=_plot_rc,
)


def _new_fig(height):
    return plt.figure(figsize=(_FIG_W, height))


def _style_ax(ax):
    ax.tick_params(labelsize=_FONT_SIZE)
    ax.title.set_fontsize(_FONT_SIZE)
    ax.xaxis.label.set_size(_FONT_SIZE)
    ax.yaxis.label.set_size(_FONT_SIZE)
    ax.set_facecolor("none")
    ax.grid(False)

from matplotlib_inline.backend_inline import ( set_matplotlib_formats, ) random.seed(42) np.random.seed(42) torch.manual_seed(42) set_matplotlib_formats( "svg", bbox_inches=None, ) _FONT_SIZE = 14 _FIG_W = 7.0 _plot_rc = { "font.size": _FONT_SIZE, "axes.titlesize": _FONT_SIZE, "axes.labelsize": _FONT_SIZE, "xtick.labelsize": _FONT_SIZE, "ytick.labelsize": _FONT_SIZE, "legend.fontsize": _FONT_SIZE, "legend.title_fontsize": _FONT_SIZE, "lines.linewidth": 1.0, "font.family": "Liberation Sans", "text.color": "black", "axes.labelcolor": "black", "xtick.color": "black", "ytick.color": "black", "axes.edgecolor": "black", "axes.grid": False, "axes.facecolor": "none", "figure.facecolor": "none", "savefig.facecolor": "none", "savefig.edgecolor": "none", "svg.fonttype": "none", } plt.rcParams.update(_plot_rc) sns.set_theme( style="ticks", rc=_plot_rc, ) def _new_fig(height): return plt.figure(figsize=(_FIG_W, height)) def _style_ax(ax): ax.tick_params(labelsize=_FONT_SIZE) ax.title.set_fontsize(_FONT_SIZE) ax.xaxis.label.set_size(_FONT_SIZE) ax.yaxis.label.set_size(_FONT_SIZE) ax.set_facecolor("none") ax.grid(False)

Define the graph¶

This is the smallest graph in which feedback matters. Two hidden nodes feed each other, and each has its own input feature:

• input_signal is the input to node_a
• input_noise is the input to node_b
• node_a and node_b form a feedback loop
• node_a connects to the output node output

The edgelist format is the one parse_layered() also takes: a source column and a target column, nothing else, each row an edge in the direction of computation. parse_adjacency() additionally allows self-loops (a row where source == target), which this graph does not use.

The first walkthrough treats input_signal as the only informative feature. A later section repeats the same architecture under positive and negative controls that vary which inputs carry signal.

In [ ]:

edgelist = pd.DataFrame(
    {
        "source": [
            "input_signal",
            "input_noise",
            "node_a",
            "node_b",
            "node_a",
        ],
        "target": [
            "node_a",
            "node_b",
            "node_b",
            "node_a",
            "output",
        ],
    }
)

edgelist

edgelist = pd.DataFrame( { "source": [ "input_signal", "input_noise", "node_a", "node_b", "node_a", ], "target": [ "node_a", "node_b", "node_b", "node_a", "output", ], } ) edgelist

Out[ ]:

	source	target
0	input_signal	node_a
1	input_noise	node_b
2	node_a	node_b
3	node_b	node_a
4	node_a	output

In [ ]:

spec = kpnn2.parse_adjacency(edgelist)

spec.nodes, spec.input_nodes, spec.hidden_nodes, spec.output_nodes

spec = kpnn2.parse_adjacency(edgelist) spec.nodes, spec.input_nodes, spec.hidden_nodes, spec.output_nodes

Out[ ]:

(('input_noise', 'input_signal', 'node_a', 'node_b', 'output'),
 ('input_noise', 'input_signal'),
 ('node_a', 'node_b'),
 ('output',))

In [ ]:

dot = Digraph()
dot.attr(rankdir="LR")

for node_name in spec.nodes:
    dot.node(node_name, node_name)

for row in edgelist.itertuples(index=False):
    dot.edge(str(row.source), str(row.target))

dot

dot = Digraph() dot.attr(rankdir="LR") for node_name in spec.nodes: dot.node(node_name, node_name) for row in edgelist.itertuples(index=False): dot.edge(str(row.source), str(row.target)) dot

Out[ ]:

Shared masked update, unrolled for a fixed number of steps¶

The model is one layer applied several times. Each application of that layer propagates values one edge along the graph, so two applications carry a value from node_a around the loop and back.

parse_adjacency() already did the structural work:

1. named node sets, with inputs at in-degree 0 and outputs at out-degree 0: spec.input_nodes, spec.hidden_nodes, spec.output_nodes
2. packed indices over spec.nodes — source_index / target_index, one entry per edge; spec.to_mask() is the dense square for this small graph
3. where the inputs and outputs sit inside that state vector: spec.input_index, spec.output_index

The rest is ordinary PyTorch:

4. wrap spec.to_mask() in MaskedLinear
5. apply it for a fixed number of steps, writing the input columns in with index_copy before each step
6. read the output nodes and apply a linear head

This graph is small, so MaskedLinear is appropriate. See PackedLinear when n_nodes is large enough that the square would hurt RAM.

Writing the inputs in is required, not cosmetic. Input nodes have in-degree 0, so their rows of to_mask() are all zeros and their fan_in is 0. Under the degree-aware initialization of MaskedLinear those rows stay at zero forever, so nothing would ever write an input value into the state vector unless the loop does it.

You interpret the model at those unrolled steps. Each step's new state passes through its own nn.Identity, a tap that Captum can hook. The taps hold no weights: every step uses the same self.core, a normal child of the module. The first walkthrough uses two steps. Later controls use three, so the feedback path can reach the output.

The step has no activation. output is the only unit the head reads, and a ReLU on it passes no gradient once it sits below zero for every sample, so training stalls. When you add an activation, put the module between self.core and the tap, as in Layered vs. Adjacency.

In [ ]:

class CyclicNet(nn.Module):
    def __init__(
        self,
        spec: kpnn2.AdjacencySpec,
        n_steps: int,
    ):
        super().__init__()
        self.spec = spec
        self.core = kpnn2.MaskedLinear(spec.to_mask())
        self.register_buffer(
            "input_index",
            torch.as_tensor(spec.input_index),
        )
        self.taps = nn.ModuleList([nn.Identity() for _ in range(n_steps)])
        self.head = nn.Linear(1, 1)

    def forward(self, x):
        state = x.new_zeros(
            x.shape[0],
            self.core.in_features,
        )
        for tap in self.taps:
            state = state.index_copy(
                -1,
                self.input_index,
                x,
            )
            state = tap(self.core(state))
        output = state[:, self.spec.output_index]
        return self.head(output)

class CyclicNet(nn.Module): def __init__( self, spec: kpnn2.AdjacencySpec, n_steps: int, ): super().__init__() self.spec = spec self.core = kpnn2.MaskedLinear(spec.to_mask()) self.register_buffer( "input_index", torch.as_tensor(spec.input_index), ) self.taps = nn.ModuleList([nn.Identity() for _ in range(n_steps)]) self.head = nn.Linear(1, 1) def forward(self, x): state = x.new_zeros( x.shape[0], self.core.in_features, ) for tap in self.taps: state = state.index_copy( -1, self.input_index, x, ) state = tap(self.core(state)) output = state[:, self.spec.output_index] return self.head(output)

In [ ]:

model = CyclicNet(
    spec,
    n_steps=2,
)
spec.nodes, spec.input_index, len(model.taps)

model = CyclicNet( spec, n_steps=2, ) spec.nodes, spec.input_index, len(model.taps)

Out[ ]:

(('input_noise', 'input_signal', 'node_a', 'node_b', 'output'), (0, 1), 2)

Simulate data, align features, and train¶

The data are simulated so the right answer is known: the two classes differ by a shift on input_signal, while input_noise has the same distribution in both. align_inputs() then returns a column index into the feature names. Train and test share that index here because both tables use the same column order.

In [ ]:

def simulate_binary_data(
    n_per_class,
    feature_names,
    informative_features,
    rng,
):
    n_features = len(feature_names)
    n_samples = 2 * n_per_class
    x = rng.normal(
        0.0,
        1.0,
        size=(n_samples, n_features),
    )
    y = np.array([0] * n_per_class + [1] * n_per_class)
    informative_idx = [
        feature_names.index(name) for name in informative_features
    ]
    for class_value, shift in enumerate([-1.5, 1.5]):
        class_mask = y == class_value
        for feature_idx in informative_idx:
            x[class_mask, feature_idx] += shift
    x_df = pd.DataFrame(
        x,
        columns=feature_names,
    )
    y_series = pd.Series(y, name="label")
    return x_df, y_series

def simulate_binary_data( n_per_class, feature_names, informative_features, rng, ): n_features = len(feature_names) n_samples = 2 * n_per_class x = rng.normal( 0.0, 1.0, size=(n_samples, n_features), ) y = np.array([0] * n_per_class + [1] * n_per_class) informative_idx = [ feature_names.index(name) for name in informative_features ] for class_value, shift in enumerate([-1.5, 1.5]): class_mask = y == class_value for feature_idx in informative_idx: x[class_mask, feature_idx] += shift x_df = pd.DataFrame( x, columns=feature_names, ) y_series = pd.Series(y, name="label") return x_df, y_series

In [ ]:

rng = np.random.default_rng(42)

x_train_df, y_train = simulate_binary_data(
    n_per_class=80,
    feature_names=list(spec.input_nodes),
    informative_features=["input_signal"],
    rng=rng,
)

x_test_df, y_test = simulate_binary_data(
    n_per_class=40,
    feature_names=list(spec.input_nodes),
    informative_features=["input_signal"],
    rng=rng,
)

x_train_df.head()

rng = np.random.default_rng(42) x_train_df, y_train = simulate_binary_data( n_per_class=80, feature_names=list(spec.input_nodes), informative_features=["input_signal"], rng=rng, ) x_test_df, y_test = simulate_binary_data( n_per_class=40, feature_names=list(spec.input_nodes), informative_features=["input_signal"], rng=rng, ) x_train_df.head()

Out[ ]:

	input_noise	input_signal
0	0.304717	-2.539984
1	0.750451	-0.559435
2	-1.951035	-2.802180
3	0.127840	-1.816243
4	-0.016801	-2.353044

In [ ]:

plot_df = x_train_df.copy()
plot_df["label"] = y_train.values
plot_df["class"] = plot_df["label"].map({0: "class 0", 1: "class 1"})

long_df = plot_df.melt(
    id_vars=["label", "class"],
    value_vars=list(spec.input_nodes),
    var_name="feature",
    value_name="value",
)

feature_type_map = {
    "input_signal": "informative",
    "input_noise": "uninformative",
}
long_df["feature_type"] = long_df["feature"].map(feature_type_map)

n_cols = 2
g = sns.catplot(
    data=long_df,
    x="class",
    y="value",
    col="feature",
    col_wrap=n_cols,
    hue="class",
    kind="box",
    width=0.35,
    sharey=False,
    height=_FIG_W / n_cols,
    aspect=1.0,
    legend=False,
)
g.fig.set_size_inches(_FIG_W, g.fig.get_figheight())
g.fig.suptitle(
    "Simulated training data: informative vs. uninformative features",
    fontsize=_FONT_SIZE,
    y=0.98,
)
g.fig.subplots_adjust(
    top=0.82,
    bottom=0.16,
    left=0.1,
    right=0.98,
    wspace=0.3,
)
for ax, feature in zip(g.axes.flat, spec.input_nodes):
    feature_type = feature_type_map[feature]
    ax.set_title(f"{feature}\n({feature_type})")
    ax.set_xlabel("")
    ax.set_ylabel("value")
    _style_ax(ax)
g.fig.patch.set_facecolor("none")
plt.show()

plot_df = x_train_df.copy() plot_df["label"] = y_train.values plot_df["class"] = plot_df["label"].map({0: "class 0", 1: "class 1"}) long_df = plot_df.melt( id_vars=["label", "class"], value_vars=list(spec.input_nodes), var_name="feature", value_name="value", ) feature_type_map = { "input_signal": "informative", "input_noise": "uninformative", } long_df["feature_type"] = long_df["feature"].map(feature_type_map) n_cols = 2 g = sns.catplot( data=long_df, x="class", y="value", col="feature", col_wrap=n_cols, hue="class", kind="box", width=0.35, sharey=False, height=_FIG_W / n_cols, aspect=1.0, legend=False, ) g.fig.set_size_inches(_FIG_W, g.fig.get_figheight()) g.fig.suptitle( "Simulated training data: informative vs. uninformative features", fontsize=_FONT_SIZE, y=0.98, ) g.fig.subplots_adjust( top=0.82, bottom=0.16, left=0.1, right=0.98, wspace=0.3, ) for ax, feature in zip(g.axes.flat, spec.input_nodes): feature_type = feature_type_map[feature] ax.set_title(f"{feature}\n({feature_type})") ax.set_xlabel("") ax.set_ylabel("value") _style_ax(ax) g.fig.patch.set_facecolor("none") plt.show()

In [ ]:

# Train and test share column order, so one col is valid.
col = kpnn2.align_inputs(x_train_df.columns, spec)
x_train = torch.as_tensor(
    x_train_df.to_numpy()[:, col],
    dtype=torch.float32,
)
x_test = torch.as_tensor(
    x_test_df.to_numpy()[:, col],
    dtype=torch.float32,
)

y_train_tensor = torch.tensor(
    y_train.values.reshape(-1, 1),
    dtype=torch.float32,
)

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=0.05,
)
loss_fn = nn.BCEWithLogitsLoss()

n_epochs = 100
loss_history = []
model.train()
for _ in range(n_epochs):
    optimizer.zero_grad()
    logits = model(x_train)
    loss = loss_fn(logits, y_train_tensor)
    loss.backward()
    optimizer.step()
    loss_history.append(loss.item())

print(
    "Loss before first update:",
    round(loss_history[0], 4),
)
print(
    "Loss after final update: ",
    round(loss_history[-1], 4),
)

_new_fig(2.8)
plt.plot(
    range(1, n_epochs + 1),
    loss_history,
    linewidth=1,
)
plt.xlabel("Epoch")
plt.ylabel("Binary cross-entropy loss")
plt.title("Training loss")
plt.tight_layout()
plt.show()

# Train and test share column order, so one col is valid. col = kpnn2.align_inputs(x_train_df.columns, spec) x_train = torch.as_tensor( x_train_df.to_numpy()[:, col], dtype=torch.float32, ) x_test = torch.as_tensor( x_test_df.to_numpy()[:, col], dtype=torch.float32, ) y_train_tensor = torch.tensor( y_train.values.reshape(-1, 1), dtype=torch.float32, ) optimizer = torch.optim.Adam( model.parameters(), lr=0.05, ) loss_fn = nn.BCEWithLogitsLoss() n_epochs = 100 loss_history = [] model.train() for _ in range(n_epochs): optimizer.zero_grad() logits = model(x_train) loss = loss_fn(logits, y_train_tensor) loss.backward() optimizer.step() loss_history.append(loss.item()) print( "Loss before first update:", round(loss_history[0], 4), ) print( "Loss after final update: ", round(loss_history[-1], 4), ) _new_fig(2.8) plt.plot( range(1, n_epochs + 1), loss_history, linewidth=1, ) plt.xlabel("Epoch") plt.ylabel("Binary cross-entropy loss") plt.title("Training loss") plt.tight_layout() plt.show()

Loss before first update: 0.8023
Loss after final update:  0.1257

Interpret nodes¶

Only input_signal is informative in this first walkthrough, so node_a is the expected answer.

Captum LayerConductance scores each unrolled step. Then map_node_attributions() labels the (batch, n_nodes) tensor as xarray with the full state-vector names. Convert with .to_pandas() for the tables below. kpnn2 does not import Captum. For the mapper itself, see Mapping attributions.

There is no layer= argument here. An AdjacencySpec has no depths, so the node axis is the whole state vector (spec.nodes) and the result carries no scalar layer coordinate.

Those scores are signed, and a node can matter at one step and not at another. The summary below therefore keeps, per node, the signed value from the step of largest magnitude, and takes .abs() only for ranking. It covers the hidden nodes (node_a, node_b); use the per-step tables to inspect the inputs and the output as well.

In [ ]:

from captum.attr import LayerConductance

model.eval()
node_attr_by_site = {}
for index, tap in enumerate(
    model.taps,
    start=1,
):
    conductor = LayerConductance(
        model,
        tap,
    )
    attributions = conductor.attribute(
        x_test,
        target=0,
    )
    node_attr_by_site[f"step_{index}"] = kpnn2.map_node_attributions(
        attributions=attributions,
        spec=spec,
    ).to_pandas()

list(node_attr_by_site.keys())

from captum.attr import LayerConductance model.eval() node_attr_by_site = {} for index, tap in enumerate( model.taps, start=1, ): conductor = LayerConductance( model, tap, ) attributions = conductor.attribute( x_test, target=0, ) node_attr_by_site[f"step_{index}"] = kpnn2.map_node_attributions( attributions=attributions, spec=spec, ).to_pandas() list(node_attr_by_site.keys())

Out[ ]:

['step_1', 'step_2']

In [ ]:

for site_id, site_df in node_attr_by_site.items():
    print(site_id)
    display(site_df.head())

for site_id, site_df in node_attr_by_site.items(): print(site_id) display(site_df.head())

step_1

node	input_noise	input_signal	node_a	node_b	output
observation
0	0.0	0.0	-5.821504	0.0	0.0
1	0.0	0.0	-2.579521	0.0	0.0
2	0.0	0.0	-2.648712	0.0	0.0
3	0.0	0.0	-5.395569	0.0	0.0
4	0.0	0.0	-6.949645	0.0	0.0

step_2

node	input_noise	input_signal	node_a	node_b	output
observation
0	0.0	0.0	0.0	0.0	-5.821504
1	0.0	0.0	0.0	0.0	-2.579521
2	0.0	0.0	0.0	0.0	-2.648712
3	0.0	0.0	0.0	0.0	-5.395569
4	0.0	0.0	0.0	0.0	-6.949644

In [ ]:

def peak_hidden_summary(site_tables, hidden_nodes):
    stacked = np.stack(
        [table[hidden_nodes].to_numpy() for table in site_tables],
        axis=0,
    )
    pick = np.abs(stacked).argmax(axis=0)
    n_batch, n_hidden = pick.shape
    values = np.empty((n_batch, n_hidden))
    rows = np.arange(n_batch)
    for j in range(n_hidden):
        values[:, j] = stacked[
            pick[:, j],
            rows,
            j,
        ]
    return pd.DataFrame(
        values,
        columns=hidden_nodes,
    )


node_importance = peak_hidden_summary(
    list(node_attr_by_site.values()),
    list(spec.hidden_nodes),
)
node_importance.head()

def peak_hidden_summary(site_tables, hidden_nodes): stacked = np.stack( [table[hidden_nodes].to_numpy() for table in site_tables], axis=0, ) pick = np.abs(stacked).argmax(axis=0) n_batch, n_hidden = pick.shape values = np.empty((n_batch, n_hidden)) rows = np.arange(n_batch) for j in range(n_hidden): values[:, j] = stacked[ pick[:, j], rows, j, ] return pd.DataFrame( values, columns=hidden_nodes, ) node_importance = peak_hidden_summary( list(node_attr_by_site.values()), list(spec.hidden_nodes), ) node_importance.head()

Out[ ]:

	node_a	node_b
0	-5.821504	0.0
1	-2.579521	0.0
2	-2.648712	0.0
3	-5.395569	0.0
4	-6.949645	0.0

In [ ]:

summary_a_only = (
    node_importance.abs()
    .mean(axis=0)
    .sort_values(ascending=True)
    .rename("mean_abs_attribution")
    .reset_index()
    .rename(columns={"index": "node"})
)

_new_fig(2.4)
sns.barplot(
    data=summary_a_only,
    x="mean_abs_attribution",
    y="node",
    color="steelblue",
    width=0.3,
)
plt.xlabel("Mean absolute node attribution")
plt.ylabel("Node")
plt.title("Summary node importance on the test set")
plt.tight_layout()
plt.show()

summary_a_only = ( node_importance.abs() .mean(axis=0) .sort_values(ascending=True) .rename("mean_abs_attribution") .reset_index() .rename(columns={"index": "node"}) ) _new_fig(2.4) sns.barplot( data=summary_a_only, x="mean_abs_attribution", y="node", color="steelblue", width=0.3, ) plt.xlabel("Mean absolute node attribution") plt.ylabel("Node") plt.title("Summary node importance on the test set") plt.tight_layout() plt.show()

Positive and negative controls¶

The walkthrough above recovers importance on node_a when only input_signal is informative, while node_b stays near zero. That contrast is incomplete: it does not show whether importance tracks the signal or the architecture, nor what happens when the other input carries signal, when both do, or when neither does.

Hidden nodes and the output stay fixed; each of the four cases below rebuilds the edgelist and names its inputs by role:

Case	Input → node_a	Input → node_b	Expected
A only	input_signal	input_noise	a high, b ~0
Both	input_signal_a	input_signal_b	both high
B only	input_noise	input_signal	both high
Neither	input_noise_a	input_noise_b	both ~0

A only reuses the walkthrough results above. The remaining cases build CyclicNet with three steps (step_1 to step_3 in the tables), because the feedback path input → node_b → node_a → output is three edges long.

In [ ]:

def plot_graph(edgelist, input_a, input_b):
    dot = Digraph()
    dot.attr(rankdir="LR")
    for node_name in [input_a, input_b]:
        dot.node(node_name, node_name)
    for node_name in ["node_a", "node_b"]:
        dot.node(node_name, node_name)
    dot.node("output", "output")
    for row in edgelist.itertuples(index=False):
        dot.edge(
            str(row.source),
            str(row.target),
        )
    return dot


def make_edgelist(input_a, input_b):
    return pd.DataFrame(
        {
            "source": [
                input_a,
                input_b,
                "node_a",
                "node_b",
                "node_a",
            ],
            "target": [
                "node_a",
                "node_b",
                "node_b",
                "node_a",
                "output",
            ],
        }
    )


def run_cyclic_case(
    edgelist,
    informative_features,
    n_epochs=100,
):
    torch.manual_seed(42)
    rng = np.random.default_rng(42)
    case_spec = kpnn2.parse_adjacency(edgelist)
    case_model = CyclicNet(
        case_spec,
        n_steps=3,
    )
    x_train_df, y_train = simulate_binary_data(
        n_per_class=80,
        feature_names=list(case_spec.input_nodes),
        informative_features=informative_features,
        rng=rng,
    )
    x_test_df, y_test = simulate_binary_data(
        n_per_class=40,
        feature_names=list(case_spec.input_nodes),
        informative_features=informative_features,
        rng=rng,
    )
    # Same feature_names on train and test, so one col is valid.
    col = kpnn2.align_inputs(
        x_train_df.columns,
        case_spec,
    )
    x_train = torch.as_tensor(
        x_train_df.to_numpy()[:, col],
        dtype=torch.float32,
    )
    x_test = torch.as_tensor(
        x_test_df.to_numpy()[:, col],
        dtype=torch.float32,
    )
    y_train_tensor = torch.tensor(
        y_train.values.reshape(-1, 1),
        dtype=torch.float32,
    )
    optimizer = torch.optim.Adam(
        case_model.parameters(),
        lr=0.05,
    )
    loss_fn = nn.BCEWithLogitsLoss()
    loss_history = []
    case_model.train()
    for _ in range(n_epochs):
        optimizer.zero_grad()
        logits = case_model(x_train)
        loss = loss_fn(logits, y_train_tensor)
        loss.backward()
        optimizer.step()
        loss_history.append(loss.item())
    case_model.eval()
    site_tables = []
    for tap in case_model.taps:
        conductor = LayerConductance(
            case_model,
            tap,
        )
        attributions = conductor.attribute(
            x_test,
            target=0,
        )
        site_tables.append(
            kpnn2.map_node_attributions(
                attributions=attributions,
                spec=case_spec,
            ).to_pandas()
        )
    peak = peak_hidden_summary(
        site_tables,
        list(case_spec.hidden_nodes),
    )
    summary = (
        peak.abs()
        .mean(axis=0)
        .sort_values(ascending=True)
        .rename("mean_abs_attribution")
        .reset_index()
        .rename(columns={"index": "node"})
    )
    return {
        "loss_history": loss_history,
        "summary": summary,
    }


def plot_case_results(
    case_name,
    result,
    attr_xlim=None,
):
    fig, axes = plt.subplots(
        1,
        2,
        figsize=(_FIG_W, 3.2),
        gridspec_kw={"width_ratios": [1.1, 1.0]},
    )
    axes[0].plot(
        range(1, len(result["loss_history"]) + 1),
        result["loss_history"],
        linewidth=1,
    )
    axes[0].set_xlabel("Epoch")
    axes[0].set_ylabel("Binary cross-entropy loss")
    axes[0].set_title(f"{case_name}: training loss")
    sns.barplot(
        data=result["summary"],
        x="mean_abs_attribution",
        y="node",
        color="steelblue",
        width=0.3,
        ax=axes[1],
    )
    axes[1].set_xlabel("Mean absolute node attribution")
    axes[1].set_ylabel("Node")
    axes[1].set_title(f"{case_name}: node importance")
    if attr_xlim is not None:
        axes[1].set_xlim(attr_xlim)
    _style_ax(axes[0])
    _style_ax(axes[1])
    fig.tight_layout()
    plt.show()


graph_a_only = plot_graph(
    edgelist,
    "input_signal",
    "input_noise",
)
result_a_only = {
    "loss_history": loss_history,
    "summary": summary_a_only,
}

edgelist_both = make_edgelist(
    "input_signal_a",
    "input_signal_b",
)
edgelist_b_only = make_edgelist(
    "input_noise",
    "input_signal",
)
edgelist_neither = make_edgelist(
    "input_noise_a",
    "input_noise_b",
)

graph_both = plot_graph(
    edgelist_both,
    "input_signal_a",
    "input_signal_b",
)
graph_b_only = plot_graph(
    edgelist_b_only,
    "input_noise",
    "input_signal",
)
graph_neither = plot_graph(
    edgelist_neither,
    "input_noise_a",
    "input_noise_b",
)

result_both = run_cyclic_case(
    edgelist=edgelist_both,
    informative_features=[
        "input_signal_a",
        "input_signal_b",
    ],
)
result_b_only = run_cyclic_case(
    edgelist=edgelist_b_only,
    informative_features=["input_signal"],
)
result_neither = run_cyclic_case(
    edgelist=edgelist_neither,
    informative_features=[],
)

xmax = max(
    result_a_only["summary"]["mean_abs_attribution"].max(),
    result_both["summary"]["mean_abs_attribution"].max(),
    result_b_only["summary"]["mean_abs_attribution"].max(),
    result_neither["summary"]["mean_abs_attribution"].max(),
)
attr_xlim = (0, float(xmax) * 1.05)

def plot_graph(edgelist, input_a, input_b): dot = Digraph() dot.attr(rankdir="LR") for node_name in [input_a, input_b]: dot.node(node_name, node_name) for node_name in ["node_a", "node_b"]: dot.node(node_name, node_name) dot.node("output", "output") for row in edgelist.itertuples(index=False): dot.edge( str(row.source), str(row.target), ) return dot def make_edgelist(input_a, input_b): return pd.DataFrame( { "source": [ input_a, input_b, "node_a", "node_b", "node_a", ], "target": [ "node_a", "node_b", "node_b", "node_a", "output", ], } ) def run_cyclic_case( edgelist, informative_features, n_epochs=100, ): torch.manual_seed(42) rng = np.random.default_rng(42) case_spec = kpnn2.parse_adjacency(edgelist) case_model = CyclicNet( case_spec, n_steps=3, ) x_train_df, y_train = simulate_binary_data( n_per_class=80, feature_names=list(case_spec.input_nodes), informative_features=informative_features, rng=rng, ) x_test_df, y_test = simulate_binary_data( n_per_class=40, feature_names=list(case_spec.input_nodes), informative_features=informative_features, rng=rng, ) # Same feature_names on train and test, so one col is valid. col = kpnn2.align_inputs( x_train_df.columns, case_spec, ) x_train = torch.as_tensor( x_train_df.to_numpy()[:, col], dtype=torch.float32, ) x_test = torch.as_tensor( x_test_df.to_numpy()[:, col], dtype=torch.float32, ) y_train_tensor = torch.tensor( y_train.values.reshape(-1, 1), dtype=torch.float32, ) optimizer = torch.optim.Adam( case_model.parameters(), lr=0.05, ) loss_fn = nn.BCEWithLogitsLoss() loss_history = [] case_model.train() for _ in range(n_epochs): optimizer.zero_grad() logits = case_model(x_train) loss = loss_fn(logits, y_train_tensor) loss.backward() optimizer.step() loss_history.append(loss.item()) case_model.eval() site_tables = [] for tap in case_model.taps: conductor = LayerConductance( case_model, tap, ) attributions = conductor.attribute( x_test, target=0, ) site_tables.append( kpnn2.map_node_attributions( attributions=attributions, spec=case_spec, ).to_pandas() ) peak = peak_hidden_summary( site_tables, list(case_spec.hidden_nodes), ) summary = ( peak.abs() .mean(axis=0) .sort_values(ascending=True) .rename("mean_abs_attribution") .reset_index() .rename(columns={"index": "node"}) ) return { "loss_history": loss_history, "summary": summary, } def plot_case_results( case_name, result, attr_xlim=None, ): fig, axes = plt.subplots( 1, 2, figsize=(_FIG_W, 3.2), gridspec_kw={"width_ratios": [1.1, 1.0]}, ) axes[0].plot( range(1, len(result["loss_history"]) + 1), result["loss_history"], linewidth=1, ) axes[0].set_xlabel("Epoch") axes[0].set_ylabel("Binary cross-entropy loss") axes[0].set_title(f"{case_name}: training loss") sns.barplot( data=result["summary"], x="mean_abs_attribution", y="node", color="steelblue", width=0.3, ax=axes[1], ) axes[1].set_xlabel("Mean absolute node attribution") axes[1].set_ylabel("Node") axes[1].set_title(f"{case_name}: node importance") if attr_xlim is not None: axes[1].set_xlim(attr_xlim) _style_ax(axes[0]) _style_ax(axes[1]) fig.tight_layout() plt.show() graph_a_only = plot_graph( edgelist, "input_signal", "input_noise", ) result_a_only = { "loss_history": loss_history, "summary": summary_a_only, } edgelist_both = make_edgelist( "input_signal_a", "input_signal_b", ) edgelist_b_only = make_edgelist( "input_noise", "input_signal", ) edgelist_neither = make_edgelist( "input_noise_a", "input_noise_b", ) graph_both = plot_graph( edgelist_both, "input_signal_a", "input_signal_b", ) graph_b_only = plot_graph( edgelist_b_only, "input_noise", "input_signal", ) graph_neither = plot_graph( edgelist_neither, "input_noise_a", "input_noise_b", ) result_both = run_cyclic_case( edgelist=edgelist_both, informative_features=[ "input_signal_a", "input_signal_b", ], ) result_b_only = run_cyclic_case( edgelist=edgelist_b_only, informative_features=["input_signal"], ) result_neither = run_cyclic_case( edgelist=edgelist_neither, informative_features=[], ) xmax = max( result_a_only["summary"]["mean_abs_attribution"].max(), result_both["summary"]["mean_abs_attribution"].max(), result_b_only["summary"]["mean_abs_attribution"].max(), result_neither["summary"]["mean_abs_attribution"].max(), ) attr_xlim = (0, float(xmax) * 1.05)

Case: A only¶

Only the input to node_a (input_signal) is informative. The input to node_b is input_noise. Importance should concentrate on node_a. Results are taken from the walkthrough above.

In [ ]:

display(graph_a_only)
plot_case_results("A only", result_a_only, attr_xlim)

display(graph_a_only) plot_case_results("A only", result_a_only, attr_xlim)

Case: both inputs informative¶

Both inputs carry class signal (input_signal_a → node_a, input_signal_b → node_b). Both hidden nodes should receive importance.

In [ ]:

display(graph_both)
plot_case_results(
    "Both informative",
    result_both,
    attr_xlim,
)

display(graph_both) plot_case_results( "Both informative", result_both, attr_xlim, )

Case: B only¶

Only the input to node_b (input_signal) is informative. The input to node_a is input_noise. Because only node_a connects to output, the signal can reach the output only along the path input_signal → node_b → node_a → output. Both hidden nodes should therefore receive importance.

In [ ]:

display(graph_b_only)
plot_case_results("B only", result_b_only, attr_xlim)

display(graph_b_only) plot_case_results("B only", result_b_only, attr_xlim)

Case: neither input informative¶

Both inputs are uninformative (input_noise_a, input_noise_b). The model should not recover strong, consistent importance on either hidden node.

In [ ]:

display(graph_neither)
plot_case_results(
    "Neither informative",
    result_neither,
    attr_xlim,
)

display(graph_neither) plot_case_results( "Neither informative", result_neither, attr_xlim, )

Control summary¶

Across the four cases:

• A only (input_signal / input_noise): importance concentrates on node_a; node_b stays near zero
• Both (input_signal_a / input_signal_b): both hidden nodes receive importance
• B only (input_noise → node_a, input_signal → node_b): loss decreases and both hidden nodes receive importance via the feedback path to the output
• Neither (input_noise_a / input_noise_b): loss remains near chance; neither hidden node shows strong recovered importance

Low importance on node_b in the A-only setting therefore reflects its uninformative input, not an inability to score that node. Recovering importance through the feedback path also requires enough unrolled steps for input → node_b → node_a → output to reach the output.

Where to go next¶

• Time-series example writes a new x_t in at each step, so a self-loop carries state across a sequence.
• Transformer example uses the same packed indices as an attention allow-set.
• Layered vs. Adjacency compares this layout with parse_layered().
• PackedLinear replaces MaskedLinear when the (n, n) square would strain RAM.