Skip to content

API Reference

Auto-generated from docstrings. For narrative explanations of why and when to use each piece, see Concepts and Usage first — this page is for looking up exact signatures.

Core

The main entry points: loading a network, computing its partition, and querying the result.

Find the linear regions of a ReLU network.

Parameters:

Name Type Description Default
source

A single network's parameters. Accepted forms: PyTorch state_dict, nn.Module, or path to a .pth / .h5 file. For per-epoch analyses, iterate over state dicts at the call site (see :func:parx.io.iter_state_dicts).

 required
data ndarray

Input points, shape (N, input_dim). Sparse methods scan the array for activation patterns; exact methods use data[0] as the DFS starting point.

 required
method str

Name of a registered region-finding method. Built-ins: "sparse_julia" (default), "exact_julia", "sparse_python", "exact_python". See :func:parx.list_methods.

 'sparse_julia'
include_output_layer bool

Include the final linear layer in the partition. Defaults to False because only hidden ReLU layers define the polyhedral partition.

 False
**method_kwargs

Forwarded verbatim to the chosen method's function.

 {}

Returns:

Type Description
Partition
Source code in src/parx/__init__.py 
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def compute_partition(
    source,
    data: np.ndarray,
    *,
    method: str = "sparse_julia",
    include_output_layer: bool = False,
    **method_kwargs,
) -> Partition:
    """Find the linear regions of a ReLU network.

    Parameters
    ----------
    source:
        A single network's parameters.  Accepted forms: PyTorch ``state_dict``,
        ``nn.Module``, or path to a ``.pth`` / ``.h5`` file.  For per-epoch
        analyses, iterate over state dicts at the call site (see
        :func:`parx.io.iter_state_dicts`).
    data:
        Input points, shape ``(N, input_dim)``.  Sparse methods scan the array
        for activation patterns; exact methods use ``data[0]`` as the DFS
        starting point.
    method:
        Name of a registered region-finding method.  Built-ins:
        ``"sparse_julia"`` (default), ``"exact_julia"``, ``"sparse_python"``,
        ``"exact_python"``.  See :func:`parx.list_methods`.
    include_output_layer:
        Include the final linear layer in the partition.  Defaults to ``False``
        because only hidden ReLU layers define the polyhedral partition.
    **method_kwargs:
        Forwarded verbatim to the chosen method's function.

    Returns
    -------
    Partition
    """
    weights, biases = load_network(source, include_output_layer=include_output_layer)
    fn = get_method(method)
    result = fn(weights, biases, data, **method_kwargs)
    return Partition.from_result(result, weights, biases)

Extract ordered (weights, biases) from a network.

Parameters:

Name Type Description Default
model

nn.Module, PyTorch state_dict dict, or path to a .pth / .h5 file.

 required
include_output_layer bool

Include the final linear layer in the returned lists. Defaults to False because only hidden ReLU layers define the polyhedral partition.

 False

Returns:

Type Description
weights, biases:

Each entry is a float64 NumPy array. Weight shape is (out_features, in_features) matching PyTorch convention.
Source code in src/parx/network.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def load_network(
    model,
    *,
    include_output_layer: bool = False,
) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """Extract ordered (weights, biases) from a network.

    Parameters
    ----------
    model:
        ``nn.Module``, PyTorch ``state_dict`` dict, or path to a ``.pth`` /
        ``.h5`` file.
    include_output_layer:
        Include the final linear layer in the returned lists.  Defaults to
        ``False`` because only hidden ReLU layers define the polyhedral
        partition.

    Returns
    -------
    weights, biases:
        Each entry is a ``float64`` NumPy array.  Weight shape is
        ``(out_features, in_features)`` matching PyTorch convention.
    """
    if isinstance(model, str | Path):
        weights, biases = _from_path(Path(model))
    elif isinstance(model, dict):
        weights, biases = _from_state_dict(model)
    else:
        weights, biases = _from_module(model)

    if not include_output_layer and len(weights) > 1:
        weights = weights[:-1]
        biases = biases[:-1]

    return weights, biases

Run a forward pass and return the input activations to a chosen Linear layer.

layer_index uses Python list indexing over all nn.Linear submodules (-1 = last, 0 = first). Returns shape (N, feature_dim) float64 array.

Source code in src/parx/network.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def extract_features(
    model,
    X: np.ndarray,
    *,
    layer_index: int = -1,
) -> np.ndarray:
    """Run a forward pass and return the input activations to a chosen Linear layer.

    ``layer_index`` uses Python list indexing over all ``nn.Linear`` submodules
    (-1 = last, 0 = first).  Returns shape ``(N, feature_dim)`` float64 array.
    """
    try:
        import torch
        import torch.nn as nn
    except ImportError as exc:
        raise ImportError("torch is required: pip install torch") from exc

    linears = [m for m in model.modules() if isinstance(m, nn.Linear)]
    if not linears:
        raise ValueError("model has no nn.Linear layers")

    target = linears[layer_index]  # natural IndexError if out of range

    captured: list = []

    def _hook(module, inputs, output):
        captured.append(inputs[0])

    handle = target.register_forward_hook(_hook)
    try:
        model.eval()
        with torch.no_grad():
            tensor = torch.tensor(X, dtype=torch.float32)
            model(tensor)
    finally:
        handle.remove()

    return captured[0].detach().cpu().numpy().astype(np.float64)

The polyhedral partition of a ReLU network's input space.

Holds a flat list of Region objects together with the network weights needed to reconstruct halfspace systems on demand. No tree structure is stored; the activation paths on each Region encode all needed topology.

Source code in src/parx/partition.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
class Partition:
    """The polyhedral partition of a ReLU network's input space.

    Holds a flat list of ``Region`` objects together with the network weights
    needed to reconstruct halfspace systems on demand.  No tree structure is
    stored; the activation paths on each Region encode all needed topology.
    """

    def __init__(
        self,
        regions: list[Region],
        weights: list[np.ndarray],
        biases: list[np.ndarray],
    ) -> None:
        self.regions = regions
        self.weights = weights
        self.biases = biases
        self.input_dim: int = weights[0].shape[1]
        self.n_layers: int = len(weights)

    def __len__(self) -> int:
        return len(self.regions)

    def __repr__(self) -> str:
        return (
            "Partition("
            f"n_regions={len(self)}, "
            f"n_layers={self.n_layers}, "
            f"input_dim={self.input_dim}"
            ")"
        )

    # ── Geometry ──────────────────────────────────────────────────────────────

    def halfspaces(
        self,
        region: Region,
        active_only: bool = False,
    ) -> tuple[np.ndarray, np.ndarray]:
        """Reconstruct the halfspace system D*x ≤ g for a region.

        Direct port of ``compute_path_geometry`` from the original Julia
        implementation.  Pure NumPy; no Julia call required.

        Parameters
        ----------
        region:
            The region whose constraints to compute.
        active_only:
            If ``True`` and ``region.active_indices`` is available (exact
            construction only), return only the non-redundant rows.

        Returns
        -------
        D : ndarray, shape (n_constraints, input_dim)
        g : ndarray, shape (n_constraints,)
            The system D*x ≤ g defines the polytope.
        """
        q_path = region.activation_path
        if not q_path:
            return np.zeros((0, self.input_dim)), np.zeros(0)

        A = np.eye(self.input_dim)  # accumulated linear map from input
        c = np.zeros(self.input_dim)  # accumulated bias from input
        D_blocks: list[np.ndarray] = []
        g_blocks: list[np.ndarray] = []

        for layer_idx, q in enumerate(q_path):
            W, b = self.weights[layer_idx], self.biases[layer_idx]
            W_hat = W @ A  # effective weight:  (out_l, input_dim)
            b_hat = W @ c + b  # effective bias:    (out_l,)

            # s[i] = -1 if neuron i is active (q=1), +1 if inactive (q=0)
            s = -2.0 * q + 1.0
            D_blocks.append(s[:, None] * W_hat)
            g_blocks.append(-(s * b_hat))

            # Propagate the affine map through active neurons only
            A = q[:, None] * W_hat  # (out_l, input_dim)
            c = q * b_hat  # (out_l,)

        D = np.vstack(D_blocks)
        g = np.concatenate(g_blocks)

        if (
            active_only
            and region.active_indices is not None
            and len(region.active_indices) > 0
        ):
            return D[region.active_indices], g[region.active_indices]

        return D, g

    # ── Local linearisation ───────────────────────────────────────────────────

    def local_affine(self, region: Region) -> tuple[np.ndarray, np.ndarray]:
        """The local affine map ``f(x) = A x + b`` for this region.

        Walks the activation path layer by layer, gating inactive neurons.
        ``A`` has shape ``(last_layer_out_dim, input_dim)`` and ``b`` has shape
        ``(last_layer_out_dim,)``.  This is the network's representation up to
        whichever final layer the partition was built with — typically the
        last hidden ReLU layer's output, post-gating.
        """
        A = np.eye(self.input_dim)
        c = np.zeros(self.input_dim)
        for layer_idx, q in enumerate(region.activation_path):
            W, b = self.weights[layer_idx], self.biases[layer_idx]
            W_hat = W @ A
            b_hat = W @ c + b
            A = q[:, None] * W_hat
            c = q * b_hat
        return A, c

    # ── Routing ───────────────────────────────────────────────────────────────

    def route(self, X: np.ndarray) -> list[Region | None]:
        """Assign each row of X to its region via a forward pass.

        Points that fall outside all known regions (only possible with
        sparse-mode partitions) are returned as ``None``.

        Time complexity: O(N·L) for the forward pass, O(N·n_regions) worst case
        for lookup (typically O(N) with the hash table).
        """
        X = np.asarray(X, dtype=float)
        N = X.shape[0]

        # Vectorised forward pass: collect activation patterns at every layer
        A = X
        q_per_layer: list[np.ndarray] = []
        for W, b in zip(self.weights, self.biases):
            Z = A @ W.T + b  # (N, out_l)
            Q = Z > 0  # (N, out_l), dtype bool
            q_per_layer.append(Q)
            A = Q * Z

        # Build hash lookup: tuple-of-bytes-per-layer → Region
        lookup: dict[tuple[bytes, ...], Region] = {}
        for r in self.regions:
            key = tuple(q.tobytes() for q in r.activation_path)
            lookup[key] = r

        results: list[Region | None] = []
        for i in range(N):
            key = tuple(
                q_per_layer[layer_idx][i].tobytes()
                for layer_idx in range(self.n_layers)
            )
            results.append(lookup.get(key))
        return results

    # ── Filtering ─────────────────────────────────────────────────────────────

    def regions_at_layer(self, layer: int) -> list[Region]:
        """Return regions whose activation path has exactly ``layer`` layers."""
        return [r for r in self.regions if r.n_layers == layer]

    # ── Construction from a method's RegionFindResult ─────────────────────────

    @classmethod
    def from_result(
        cls,
        result,
        weights: list[np.ndarray],
        biases: list[np.ndarray],
    ) -> Partition:
        """Build a Partition from a ``RegionFindResult`` (any registered method)."""
        patterns = np.asarray(result.patterns)
        offsets = np.asarray(result.offsets, dtype=np.int64)
        centroids = np.asarray(result.centroids)

        n_regions = patterns.shape[0]
        n_layers = len(offsets) - 1

        has_active = result.active_indices_flat is not None
        has_bounded = result.bounded is not None

        regions = []
        for i in range(n_regions):
            active_indices = None
            if has_active:
                a_start = int(result.active_offsets[i])
                a_stop = int(result.active_offsets[i + 1])
                active_indices = np.asarray(
                    result.active_indices_flat[a_start:a_stop], dtype=np.int32
                )
            regions.append(
                Region(
                    activation_path=[
                        patterns[i, offsets[layer_idx] : offsets[layer_idx + 1]].astype(
                            bool
                        )
                        for layer_idx in range(n_layers)
                    ],
                    centroid=centroids[i],
                    active_indices=active_indices,
                    bounded=bool(result.bounded[i]) if has_bounded else False,
                )
            )
        return cls(regions=regions, weights=weights, biases=biases)

halfspaces(region, active_only=False)

Reconstruct the halfspace system D*x ≤ g for a region.

Direct port of compute_path_geometry from the original Julia implementation. Pure NumPy; no Julia call required.

Parameters:

Name Type Description Default
region Region

The region whose constraints to compute.

 required
active_only bool

If True and region.active_indices is available (exact construction only), return only the non-redundant rows.

 False

Returns:

Name Type Description
D (ndarray, shape(n_constraints, input_dim))
g (ndarray, shape(n_constraints))

The system D*x ≤ g defines the polytope.
Source code in src/parx/partition.py 
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def halfspaces(
    self,
    region: Region,
    active_only: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
    """Reconstruct the halfspace system D*x ≤ g for a region.

    Direct port of ``compute_path_geometry`` from the original Julia
    implementation.  Pure NumPy; no Julia call required.

    Parameters
    ----------
    region:
        The region whose constraints to compute.
    active_only:
        If ``True`` and ``region.active_indices`` is available (exact
        construction only), return only the non-redundant rows.

    Returns
    -------
    D : ndarray, shape (n_constraints, input_dim)
    g : ndarray, shape (n_constraints,)
        The system D*x ≤ g defines the polytope.
    """
    q_path = region.activation_path
    if not q_path:
        return np.zeros((0, self.input_dim)), np.zeros(0)

    A = np.eye(self.input_dim)  # accumulated linear map from input
    c = np.zeros(self.input_dim)  # accumulated bias from input
    D_blocks: list[np.ndarray] = []
    g_blocks: list[np.ndarray] = []

    for layer_idx, q in enumerate(q_path):
        W, b = self.weights[layer_idx], self.biases[layer_idx]
        W_hat = W @ A  # effective weight:  (out_l, input_dim)
        b_hat = W @ c + b  # effective bias:    (out_l,)

        # s[i] = -1 if neuron i is active (q=1), +1 if inactive (q=0)
        s = -2.0 * q + 1.0
        D_blocks.append(s[:, None] * W_hat)
        g_blocks.append(-(s * b_hat))

        # Propagate the affine map through active neurons only
        A = q[:, None] * W_hat  # (out_l, input_dim)
        c = q * b_hat  # (out_l,)

    D = np.vstack(D_blocks)
    g = np.concatenate(g_blocks)

    if (
        active_only
        and region.active_indices is not None
        and len(region.active_indices) > 0
    ):
        return D[region.active_indices], g[region.active_indices]

    return D, g

local_affine(region)

The local affine map f(x) = A x + b for this region.

Walks the activation path layer by layer, gating inactive neurons. A has shape (last_layer_out_dim, input_dim) and b has shape (last_layer_out_dim,). This is the network's representation up to whichever final layer the partition was built with — typically the last hidden ReLU layer's output, post-gating.

Source code in src/parx/partition.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def local_affine(self, region: Region) -> tuple[np.ndarray, np.ndarray]:
    """The local affine map ``f(x) = A x + b`` for this region.

    Walks the activation path layer by layer, gating inactive neurons.
    ``A`` has shape ``(last_layer_out_dim, input_dim)`` and ``b`` has shape
    ``(last_layer_out_dim,)``.  This is the network's representation up to
    whichever final layer the partition was built with — typically the
    last hidden ReLU layer's output, post-gating.
    """
    A = np.eye(self.input_dim)
    c = np.zeros(self.input_dim)
    for layer_idx, q in enumerate(region.activation_path):
        W, b = self.weights[layer_idx], self.biases[layer_idx]
        W_hat = W @ A
        b_hat = W @ c + b
        A = q[:, None] * W_hat
        c = q * b_hat
    return A, c

route(X)

Assign each row of X to its region via a forward pass.

Points that fall outside all known regions (only possible with sparse-mode partitions) are returned as None.

Time complexity: O(N·L) for the forward pass, O(N·n_regions) worst case for lookup (typically O(N) with the hash table).

Source code in src/parx/partition.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def route(self, X: np.ndarray) -> list[Region | None]:
    """Assign each row of X to its region via a forward pass.

    Points that fall outside all known regions (only possible with
    sparse-mode partitions) are returned as ``None``.

    Time complexity: O(N·L) for the forward pass, O(N·n_regions) worst case
    for lookup (typically O(N) with the hash table).
    """
    X = np.asarray(X, dtype=float)
    N = X.shape[0]

    # Vectorised forward pass: collect activation patterns at every layer
    A = X
    q_per_layer: list[np.ndarray] = []
    for W, b in zip(self.weights, self.biases):
        Z = A @ W.T + b  # (N, out_l)
        Q = Z > 0  # (N, out_l), dtype bool
        q_per_layer.append(Q)
        A = Q * Z

    # Build hash lookup: tuple-of-bytes-per-layer → Region
    lookup: dict[tuple[bytes, ...], Region] = {}
    for r in self.regions:
        key = tuple(q.tobytes() for q in r.activation_path)
        lookup[key] = r

    results: list[Region | None] = []
    for i in range(N):
        key = tuple(
            q_per_layer[layer_idx][i].tobytes()
            for layer_idx in range(self.n_layers)
        )
        results.append(lookup.get(key))
    return results

regions_at_layer(layer)

Return regions whose activation path has exactly layer layers.

Source code in src/parx/partition.py 
164
165
166
def regions_at_layer(self, layer: int) -> list[Region]:
    """Return regions whose activation path has exactly ``layer`` layers."""
    return [r for r in self.regions if r.n_layers == layer]

from_result(result, weights, biases) classmethod

Build a Partition from a RegionFindResult (any registered method).

Source code in src/parx/partition.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
@classmethod
def from_result(
    cls,
    result,
    weights: list[np.ndarray],
    biases: list[np.ndarray],
) -> Partition:
    """Build a Partition from a ``RegionFindResult`` (any registered method)."""
    patterns = np.asarray(result.patterns)
    offsets = np.asarray(result.offsets, dtype=np.int64)
    centroids = np.asarray(result.centroids)

    n_regions = patterns.shape[0]
    n_layers = len(offsets) - 1

    has_active = result.active_indices_flat is not None
    has_bounded = result.bounded is not None

    regions = []
    for i in range(n_regions):
        active_indices = None
        if has_active:
            a_start = int(result.active_offsets[i])
            a_stop = int(result.active_offsets[i + 1])
            active_indices = np.asarray(
                result.active_indices_flat[a_start:a_stop], dtype=np.int32
            )
        regions.append(
            Region(
                activation_path=[
                    patterns[i, offsets[layer_idx] : offsets[layer_idx + 1]].astype(
                        bool
                    )
                    for layer_idx in range(n_layers)
                ],
                centroid=centroids[i],
                active_indices=active_indices,
                bounded=bool(result.bounded[i]) if has_bounded else False,
            )
        )
    return cls(regions=regions, weights=weights, biases=biases)

A single linear region of a ReLU network.

Defined by its activation path: the sequence of per-layer activation patterns (one bool array per hidden layer) that is shared by every input point inside this polytope.

Source code in src/parx/region.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Region:
    """A single linear region of a ReLU network.

    Defined by its activation path: the sequence of per-layer activation
    patterns (one bool array per hidden layer) that is shared by every input
    point inside this polytope.
    """

    __slots__ = ("activation_path", "centroid", "active_indices", "bounded")

    def __init__(
        self,
        activation_path: list[np.ndarray],
        centroid: np.ndarray,
        active_indices: np.ndarray | None = None,
        bounded: bool = False,
    ) -> None:
        self.activation_path = activation_path
        self.centroid = centroid
        self.active_indices = active_indices  # non-redundant rows in D*x≤g (exact only)
        self.bounded = bounded

    @property
    def n_layers(self) -> int:
        """Number of hidden layers spanned by this region's activation path."""
        return len(self.activation_path)

    def __repr__(self) -> str:
        shapes = [q.shape[0] for q in self.activation_path]
        return f"Region(layers={shapes}, bounded={self.bounded})"

n_layers property

Number of hidden layers spanned by this region's activation path.

Analysis

Neuron activity, structural complexity, and geometric size statistics — see Concepts § Outputs for what these numbers mean.

Analysis and statistics API for ReLU network partitions.

This module provides:

  • Cheap (no LP): neuron activity statistics, complexity profile.
  • Moderate (one LP per region): Chebyshev radii, size summary.
  • Expensive (Monte Carlo): volume estimation — do not call in tight loops.

All functions accept a :class:~parx.partition.Partition object and return plain NumPy arrays, scalars, or dicts. No Julia is imported here.

neuron_activity_rates(partition)

Fraction of regions where each neuron is active, per layer.

For layer l (0-based), returns an array of shape (layer_width,) where entry j is the fraction of regions with neuron j active.

Complexity: O(n_regions × total_neurons). No LP calls.

Parameters:

Name Type Description Default
partition Partition

The partition to analyse.

 required

Returns:

Type Description
dict mapping layer index (int) to ndarray of shape ``(layer_width,)``.
Source code in src/parx/analysis.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def neuron_activity_rates(partition: Partition) -> dict[int, np.ndarray]:
    """Fraction of regions where each neuron is active, per layer.

    For layer ``l`` (0-based), returns an array of shape ``(layer_width,)``
    where entry ``j`` is the fraction of regions with neuron ``j`` active.

    Complexity: O(n_regions × total_neurons).  No LP calls.

    Parameters
    ----------
    partition:
        The partition to analyse.

    Returns
    -------
    dict mapping layer index (int) to ndarray of shape ``(layer_width,)``.
    """
    if not partition.regions:
        n_layers = partition.n_layers
        return {
            layer: np.zeros(partition.weights[layer].shape[0], dtype=float)
            for layer in range(n_layers)
        }

    n_regions = len(partition.regions)
    # Determine layer widths from the first region's activation path.
    sample_path = partition.regions[0].activation_path
    n_layers = len(sample_path)

    # Accumulate sum of active indicator per layer.
    totals: list[np.ndarray] = [
        np.zeros(sample_path[layer].shape[0], dtype=float)
        for layer in range(n_layers)
    ]
    for region in partition.regions:
        for layer, q in enumerate(region.activation_path):
            totals[layer] += q.astype(float)

    return {layer: totals[layer] / n_regions for layer in range(n_layers)}

dead_neurons(partition, *, threshold=0.0)

Neurons with activity rate at or below threshold.

A neuron is "dead" when it is active in at most threshold fraction of regions. The default threshold=0.0 finds neurons that are never active across all known regions.

Parameters:

Name Type Description Default
partition Partition

The partition to analyse.

 required
threshold float

Upper bound on activity rate for a neuron to be reported. Use 0.0 for strictly dead neurons, or a small positive value (e.g. 0.05) for rarely-active neurons.

 0.0

Returns:

Type Description
List of ``(layer, neuron_index)`` pairs, sorted by layer then index.
Source code in src/parx/analysis.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def dead_neurons(
    partition: Partition,
    *,
    threshold: float = 0.0,
) -> list[tuple[int, int]]:
    """Neurons with activity rate at or below *threshold*.

    A neuron is "dead" when it is active in at most ``threshold`` fraction of
    regions.  The default ``threshold=0.0`` finds neurons that are **never**
    active across all known regions.

    Parameters
    ----------
    partition:
        The partition to analyse.
    threshold:
        Upper bound on activity rate for a neuron to be reported.  Use
        ``0.0`` for strictly dead neurons, or a small positive value (e.g.
        ``0.05``) for rarely-active neurons.

    Returns
    -------
    List of ``(layer, neuron_index)`` pairs, sorted by layer then index.
    """
    rates = neuron_activity_rates(partition)
    result: list[tuple[int, int]] = []
    for layer_idx, arr in rates.items():
        for j, rate in enumerate(arr):
            if rate <= threshold:
                result.append((layer_idx, j))
    return result

always_active_neurons(partition, *, threshold=1.0)

Neurons active in (at least) threshold fraction of regions.

The default threshold=1.0 finds neurons that are always active across all known regions.

Parameters:

Name Type Description Default
partition Partition

The partition to analyse.

 required
threshold float

Lower bound on activity rate. Use 1.0 for universally active neurons, or a value like 0.95 for nearly-always-active neurons.

 1.0

Returns:

Type Description
List of ``(layer, neuron_index)`` pairs, sorted by layer then index.
Source code in src/parx/analysis.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def always_active_neurons(
    partition: Partition,
    *,
    threshold: float = 1.0,
) -> list[tuple[int, int]]:
    """Neurons active in (at least) *threshold* fraction of regions.

    The default ``threshold=1.0`` finds neurons that are **always** active
    across all known regions.

    Parameters
    ----------
    partition:
        The partition to analyse.
    threshold:
        Lower bound on activity rate.  Use ``1.0`` for universally active
        neurons, or a value like ``0.95`` for nearly-always-active neurons.

    Returns
    -------
    List of ``(layer, neuron_index)`` pairs, sorted by layer then index.
    """
    rates = neuron_activity_rates(partition)
    result: list[tuple[int, int]] = []
    for layer_idx, arr in rates.items():
        for j, rate in enumerate(arr):
            if rate >= threshold:
                result.append((layer_idx, j))
    return result

complexity_profile(partition)

Structural complexity statistics for the partition.

Returns:

Type Description
dict with keys:
``n_regions`` : int

Total number of regions.
``n_layers`` : int

Number of hidden layers (depth of activation paths).
``input_dim`` : int

Input dimensionality of the network.
``regions_per_layer`` : list[int]

At depth l (1-based), the number of unique activation-path prefixes of length l. regions_per_layer[0] corresponds to depth 1 (after the first hidden layer).
``total_neurons`` : int

Sum of widths across all hidden layers.
``total_constraints`` : int

Sum over all regions of the number of halfspace constraints.
``mean_constraints_per_region`` : float

Average number of constraints per region.
Source code in src/parx/analysis.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def complexity_profile(partition: Partition) -> dict:
    """Structural complexity statistics for the partition.

    Returns
    -------
    dict with keys:

    ``n_regions`` : int
        Total number of regions.
    ``n_layers`` : int
        Number of hidden layers (depth of activation paths).
    ``input_dim`` : int
        Input dimensionality of the network.
    ``regions_per_layer`` : list[int]
        At depth ``l`` (1-based), the number of unique activation-path
        *prefixes* of length ``l``.  ``regions_per_layer[0]`` corresponds to
        depth 1 (after the first hidden layer).
    ``total_neurons`` : int
        Sum of widths across all hidden layers.
    ``total_constraints`` : int
        Sum over all regions of the number of halfspace constraints.
    ``mean_constraints_per_region`` : float
        Average number of constraints per region.
    """
    n_regions = len(partition.regions)
    n_layers = partition.n_layers
    input_dim = partition.input_dim
    total_neurons = sum(w.shape[0] for w in partition.weights)

    # regions_per_layer: unique prefixes at each depth.
    regions_per_layer: list[int] = []
    for depth in range(1, n_layers + 1):
        prefixes: set[tuple[bytes, ...]] = set()
        for region in partition.regions:
            prefix = tuple(q.tobytes() for q in region.activation_path[:depth])
            prefixes.add(prefix)
        regions_per_layer.append(len(prefixes))

    # total_constraints: sum of D.shape[0] for each region.
    total_constraints = 0
    for region in partition.regions:
        D, _ = partition.halfspaces(region)
        total_constraints += D.shape[0]

    mean_constraints = total_constraints / n_regions if n_regions > 0 else 0.0

    return {
        "n_regions": n_regions,
        "n_layers": n_layers,
        "input_dim": input_dim,
        "regions_per_layer": regions_per_layer,
        "total_neurons": total_neurons,
        "total_constraints": total_constraints,
        "mean_constraints_per_region": float(mean_constraints),
    }

region_chebyshev_radii(partition, *, max_radius=1000.0)

Chebyshev (largest inscribed ball) radius for every region.

One LP is solved per region via :func:parx._lp.chebyshev_center. A radius of 0.0 indicates an empty or degenerate region (a bug for exact-mode partitions). A radius at or near max_radius indicates an unbounded region.

Parameters:

Name Type Description Default
partition Partition

The partition to analyse.

 required
max_radius float

Upper bound on the radius passed to the LP solver. Unbounded regions will have their radius capped at this value.

 1000.0

Returns:

Type Description
ndarray of shape ``(n_regions,)``.
Source code in src/parx/analysis.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def region_chebyshev_radii(
    partition: Partition,
    *,
    max_radius: float = 1e3,
) -> np.ndarray:
    """Chebyshev (largest inscribed ball) radius for every region.

    One LP is solved per region via :func:`parx._lp.chebyshev_center`.
    A radius of ``0.0`` indicates an empty or degenerate region (a bug for
    exact-mode partitions).  A radius at or near ``max_radius`` indicates an
    unbounded region.

    Parameters
    ----------
    partition:
        The partition to analyse.
    max_radius:
        Upper bound on the radius passed to the LP solver.  Unbounded regions
        will have their radius capped at this value.

    Returns
    -------
    ndarray of shape ``(n_regions,)``.
    """
    radii = np.zeros(len(partition), dtype=float)
    for i, region in enumerate(partition.regions):
        D, g = partition.halfspaces(region)
        _, r = chebyshev_center(D, g, max_radius=max_radius)
        radii[i] = r
    return radii

region_size_summary(partition, *, max_radius=1000.0)

Descriptive statistics of per-region Chebyshev radii.

Calls :func:region_chebyshev_radii (one LP per region) and summarises the resulting distribution.

Parameters:

Name Type Description Default
partition Partition

The partition to analyse.

 required
max_radius float

Passed through to :func:region_chebyshev_radii.

 1000.0

Returns:

Type Description
dict with keys:
``min`` : float
``median`` : float
``mean`` : float
``max`` : float
``fraction_bounded`` : float

Fraction of regions whose radius is strictly below max_radius.
Source code in src/parx/analysis.py 
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
def region_size_summary(
    partition: Partition,
    *,
    max_radius: float = 1e3,
) -> dict:
    """Descriptive statistics of per-region Chebyshev radii.

    Calls :func:`region_chebyshev_radii` (one LP per region) and summarises
    the resulting distribution.

    Parameters
    ----------
    partition:
        The partition to analyse.
    max_radius:
        Passed through to :func:`region_chebyshev_radii`.

    Returns
    -------
    dict with keys:

    ``min`` : float
    ``median`` : float
    ``mean`` : float
    ``max`` : float
    ``fraction_bounded`` : float
        Fraction of regions whose radius is strictly below ``max_radius``.
    """
    radii = region_chebyshev_radii(partition, max_radius=max_radius)
    if len(radii) == 0:
        return {
            "min": float("nan"),
            "median": float("nan"),
            "mean": float("nan"),
            "max": float("nan"),
            "fraction_bounded": float("nan"),
        }
    fraction_bounded = float(np.mean(radii < max_radius))
    return {
        "min": float(np.min(radii)),
        "median": float(np.median(radii)),
        "mean": float(np.mean(radii)),
        "max": float(np.max(radii)),
        "fraction_bounded": fraction_bounded,
    }

region_volume_estimate(partition, region, *, n_samples=10000, seed=None)

Estimate the hypervolume of a region via rejection sampling.

.. warning:: This function is slow — it calls the halfspace system for every sample and scales poorly with dimensionality. For input_dim > 10 or large n_samples the runtime can be minutes. Do not call in tight loops; use :func:partition_volume_estimates for batch processing.

Algorithm

A bounding box is constructed from the Chebyshev ball of radius r centred at region.centroid: the axis-aligned box [x0 - r, x0 + r]^d. Points are sampled uniformly in this box and tested against D @ x <= g. The volume estimate is::

volume = acceptance_rate × (2r)^d

For unbounded regions the Chebyshev radius is capped at max_radius=1e3 internally, which can produce misleadingly large or small estimates.

Parameters:

Name Type Description Default
partition Partition

The partition containing the region.

 required
region Region

The specific region to estimate.

 required
n_samples int

Number of candidate points to draw.

 10000
seed int | None

Random seed for reproducibility.

 None

Returns:

Type Description
float — estimated hypervolume.
Source code in src/parx/analysis.py 
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
def region_volume_estimate(
    partition: Partition,
    region: Region,
    *,
    n_samples: int = 10_000,
    seed: int | None = None,
) -> float:
    """Estimate the hypervolume of a region via rejection sampling.

    .. warning::
        This function is **slow** — it calls the halfspace system for every
        sample and scales poorly with dimensionality.  For ``input_dim > 10``
        or large ``n_samples`` the runtime can be minutes.  Do not call in
        tight loops; use :func:`partition_volume_estimates` for batch
        processing.

    Algorithm
    ---------
    A bounding box is constructed from the Chebyshev ball of radius ``r``
    centred at ``region.centroid``: the axis-aligned box
    ``[x0 - r, x0 + r]^d``.  Points are sampled uniformly in this box and
    tested against ``D @ x <= g``.  The volume estimate is::

        volume = acceptance_rate × (2r)^d

    For unbounded regions the Chebyshev radius is capped at ``max_radius=1e3``
    internally, which can produce misleadingly large or small estimates.

    Parameters
    ----------
    partition:
        The partition containing the region.
    region:
        The specific region to estimate.
    n_samples:
        Number of candidate points to draw.
    seed:
        Random seed for reproducibility.

    Returns
    -------
    float — estimated hypervolume.
    """
    rng = np.random.default_rng(seed)
    D, g = partition.halfspaces(region)
    x0 = region.centroid
    d = len(x0)

    # Find the Chebyshev radius as the bounding scale.
    _, r = chebyshev_center(D, g, max_radius=1e3)
    if r == 0.0:
        return 0.0
    if not np.isfinite(r):
        # No constraints at all — polytope is all of R^d; return inf.
        return float("inf")

    # Sample uniformly in the L-infinity ball [x0 - r, x0 + r]^d.
    samples = rng.uniform(x0 - r, x0 + r, size=(n_samples, d))

    # Accept samples satisfying all halfspace constraints.
    if D.shape[0] == 0:
        acceptance_rate = 1.0
    else:
        inside = np.all(samples @ D.T <= g, axis=1)
        acceptance_rate = float(np.mean(inside))

    box_volume = (2.0 * r) ** d
    return acceptance_rate * box_volume

partition_volume_estimates(partition, *, n_samples=5000, seed=None)

Volume estimate for every region in the partition.

.. warning:: This function is very slow — it calls :func:region_volume_estimate once per region. For a partition with hundreds of regions and input_dim > 5, expect minutes of runtime.

Parameters:

Name Type Description Default
partition Partition

The partition to analyse.

 required
n_samples int

Number of Monte Carlo samples per region.

 5000
seed int | None

Base random seed. Each region uses a derived seed for reproducibility.

 None

Returns:

Type Description
ndarray of shape ``(n_regions,)`` with per-region volume estimates.
Source code in src/parx/analysis.py 
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
def partition_volume_estimates(
    partition: Partition,
    *,
    n_samples: int = 5_000,
    seed: int | None = None,
) -> np.ndarray:
    """Volume estimate for every region in the partition.

    .. warning::
        This function is **very slow** — it calls :func:`region_volume_estimate`
        once per region.  For a partition with hundreds of regions and
        ``input_dim > 5``, expect minutes of runtime.

    Parameters
    ----------
    partition:
        The partition to analyse.
    n_samples:
        Number of Monte Carlo samples per region.
    seed:
        Base random seed.  Each region uses a derived seed for reproducibility.

    Returns
    -------
    ndarray of shape ``(n_regions,)`` with per-region volume estimates.
    """
    volumes = np.zeros(len(partition), dtype=float)
    for i, region in enumerate(partition.regions):
        # Derive a per-region seed so results are reproducible regardless of
        # the order regions are processed.
        region_seed = None if seed is None else seed + i
        volumes[i] = region_volume_estimate(
            partition,
            region,
            n_samples=n_samples,
            seed=region_seed,
        )
    return volumes

complexity_over_epochs(partitions, labels=None)

Track partition complexity statistics across a sequence of epochs.

Computes :func:complexity_profile, mean Chebyshev radius (via :func:region_chebyshev_radii), and dead-neuron count for each partition in the list.

Parameters:

Name Type Description Default
partitions list[Partition]

Ordered list of partitions (one per epoch or checkpoint).

 required
labels

Optional epoch labels. If None, defaults to [0, 1, 2, ...].

 None

Returns:

Type Description
dict with keys:
``labels`` : list

Epoch identifiers.
``n_regions`` : list[int]

Number of regions per epoch.
``mean_chebyshev_radius`` : list[float]

Mean Chebyshev radius across regions per epoch.
``dead_neuron_count`` : list[int]

Number of dead neurons (activity rate == 0) per epoch.
Source code in src/parx/analysis.py 
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
def complexity_over_epochs(
    partitions: list[Partition],
    labels=None,
) -> dict[str, list]:
    """Track partition complexity statistics across a sequence of epochs.

    Computes :func:`complexity_profile`, mean Chebyshev radius (via
    :func:`region_chebyshev_radii`), and dead-neuron count for each partition
    in the list.

    Parameters
    ----------
    partitions:
        Ordered list of partitions (one per epoch or checkpoint).
    labels:
        Optional epoch labels.  If ``None``, defaults to ``[0, 1, 2, ...]``.

    Returns
    -------
    dict with keys:

    ``labels`` : list
        Epoch identifiers.
    ``n_regions`` : list[int]
        Number of regions per epoch.
    ``mean_chebyshev_radius`` : list[float]
        Mean Chebyshev radius across regions per epoch.
    ``dead_neuron_count`` : list[int]
        Number of dead neurons (activity rate == 0) per epoch.
    """
    if labels is None:
        labels = list(range(len(partitions)))
    else:
        labels = list(labels)

    n_regions_list: list[int] = []
    mean_radius_list: list[float] = []
    dead_count_list: list[int] = []

    for partition in partitions:
        profile = complexity_profile(partition)
        n_regions_list.append(profile["n_regions"])

        radii = region_chebyshev_radii(partition)
        mean_radius_list.append(float(np.mean(radii)) if len(radii) > 0 else 0.0)

        dead = dead_neurons(partition, threshold=0.0)
        dead_count_list.append(len(dead))

    return {
        "labels": labels,
        "n_regions": n_regions_list,
        "mean_chebyshev_radius": mean_radius_list,
        "dead_neuron_count": dead_count_list,
    }

Verification

Sanity checks on a computed partition: no overlaps, full coverage, routing consistency.

Partition verification utilities — both sample-based and LP-based.

The sample-based checks (check_no_overlaps, check_covers_space) probe the partition at user-supplied points. The LP-based checks (check_regions_nonempty, region_chebyshev_radii) interrogate the geometry directly through the halfspace systems. check_routing_consistency cross-checks that Partition.route agrees with the membership returned by the halfspace tests.

A correct exact-mode partition should satisfy:

  • every region has a strictly positive Chebyshev radius (check_regions_nonempty)
  • no two regions claim the same interior point (check_no_overlaps)
  • every input point is claimed by exactly one region (check_covers_space)
  • route(x) returns the region whose halfspace system contains x (check_routing_consistency)

count_region_memberships(partition, X, tol=1e-08)

For each row of X, count how many regions contain it.

A point x is considered inside region r when D @ x ≤ g + tol for every row of the region's halfspace system.

Source code in src/parx/verify.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def count_region_memberships(
    partition: Partition,
    X: np.ndarray,
    tol: float = 1e-8,
) -> np.ndarray:
    """For each row of X, count how many regions contain it.

    A point ``x`` is considered inside region ``r`` when ``D @ x ≤ g + tol``
    for every row of the region's halfspace system.
    """
    X = np.asarray(X, dtype=float)
    counts = np.zeros(len(X), dtype=int)
    for region in partition.regions:
        D, g = partition.halfspaces(region)
        contained = np.all(X @ D.T <= g + tol, axis=1)
        counts += contained
    return counts

check_no_overlaps(partition, X, tol=1e-08)

Check that no two regions share an interior point in the sample.

Two regions can share a boundary (a lower-dimensional face) which is fine, but they must not share interior points. Sample-based: only fails if a sample point satisfies two full halfspace systems simultaneously.

Source code in src/parx/verify.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
def check_no_overlaps(
    partition: Partition,
    X: np.ndarray,
    tol: float = 1e-8,
) -> tuple[bool, np.ndarray]:
    """Check that no two regions share an interior point in the sample.

    Two regions can share a boundary (a lower-dimensional face) which is fine,
    but they must not share interior points.  Sample-based: only fails if a
    sample point satisfies two full halfspace systems simultaneously.
    """
    counts = count_region_memberships(partition, X, tol=tol)
    return bool(np.all(counts <= 1)), counts

check_covers_space(partition, X, tol=1e-08)

Check that every sampled point belongs to exactly one region.

Should hold for exact-mode partitions. Sparse partitions may return counts == 0 for points in regions not covered by the data.

Source code in src/parx/verify.py 
68
69
70
71
72
73
74
75
76
77
78
79
def check_covers_space(
    partition: Partition,
    X: np.ndarray,
    tol: float = 1e-8,
) -> tuple[bool, np.ndarray]:
    """Check that every sampled point belongs to exactly one region.

    Should hold for exact-mode partitions.  Sparse partitions may return
    ``counts == 0`` for points in regions not covered by the data.
    """
    counts = count_region_memberships(partition, X, tol=tol)
    return bool(np.all(counts == 1)), counts

check_regions_nonempty(partition, *, min_radius=1e-06, max_radius=1000.0)

LP-based check that every region has a strictly positive Chebyshev radius.

Catches degenerate regions that survived construction — a class of bugs that sample-based checks miss (an empty region is never selected by a sample so it cannot create an overlap or a gap).

Returns:

Name Type Description
ok True when every region has radius ≥ ``min_radius``.
bad_indices indices of regions failing the test, sorted ascending by

radius (smallest first — most degenerate first).
radii per-region Chebyshev radii (same as ``region_chebyshev_radii``).
Source code in src/parx/verify.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def check_regions_nonempty(
    partition: Partition,
    *,
    min_radius: float = 1e-6,
    max_radius: float = 1e3,
) -> tuple[bool, np.ndarray, np.ndarray]:
    """LP-based check that every region has a strictly positive Chebyshev radius.

    Catches degenerate regions that survived construction — a class of bugs
    that sample-based checks miss (an empty region is never selected by a
    sample so it cannot create an overlap or a gap).

    Returns
    -------
    ok : True when every region has radius ≥ ``min_radius``.
    bad_indices : indices of regions failing the test, sorted ascending by
        radius (smallest first — most degenerate first).
    radii : per-region Chebyshev radii (same as ``region_chebyshev_radii``).
    """
    radii = region_chebyshev_radii(partition, max_radius=max_radius)
    mask = radii < min_radius
    bad = np.where(mask)[0]
    bad = bad[np.argsort(radii[bad])]
    return bool(len(bad) == 0), bad, radii

check_routing_consistency(partition, X, *, tol=1e-08)

Verify that route(x) matches the halfspace-membership region.

For every sample point, the region returned by :meth:Partition.route must contain that point under D x ≤ g + tol. Discrepancies indicate a mismatch between the forward-pass routing logic and the halfspace reconstruction.

Returns:

Name Type Description
ok True when every routed region contains its sample.
bad sample indices where the check failed (either ``route`` returned

None or the returned region did not contain the point).
Source code in src/parx/verify.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def check_routing_consistency(
    partition: Partition,
    X: np.ndarray,
    *,
    tol: float = 1e-8,
) -> tuple[bool, list[int]]:
    """Verify that ``route(x)`` matches the halfspace-membership region.

    For every sample point, the region returned by :meth:`Partition.route`
    must contain that point under ``D x ≤ g + tol``.  Discrepancies indicate
    a mismatch between the forward-pass routing logic and the halfspace
    reconstruction.

    Returns
    -------
    ok : True when every routed region contains its sample.
    bad : sample indices where the check failed (either ``route`` returned
        ``None`` or the returned region did not contain the point).
    """
    X = np.asarray(X, dtype=float)
    routed = partition.route(X)
    bad: list[int] = []
    for i, region in enumerate(routed):
        if region is None:
            bad.append(i)
            continue
        D, g = partition.halfspaces(region)
        if not np.all(D @ X[i] <= g + tol):
            bad.append(i)
    return len(bad) == 0, bad

sample_near_boundaries(partition, *, eps=0.001)

Generate sample points near every region's halfspace boundaries.

For each region and each halfspace D[i] · x ≤ g[i], project the region's centroid onto the boundary plane and step eps along the inward normal direction. These points sit just inside a region but near a seam shared with a neighbour — exactly where numerical overlaps or coverage gaps would manifest.

Returns an (M, input_dim) array of points. Useful when used as the X argument to :func:check_no_overlaps or :func:check_covers_space.

Source code in src/parx/verify.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
def sample_near_boundaries(
    partition: Partition,
    *,
    eps: float = 1e-3,
) -> np.ndarray:
    """Generate sample points near every region's halfspace boundaries.

    For each region and each halfspace ``D[i] · x ≤ g[i]``, project the
    region's centroid onto the boundary plane and step ``eps`` along the
    inward normal direction.  These points sit just inside a region but near
    a seam shared with a neighbour — exactly where numerical overlaps or
    coverage gaps would manifest.

    Returns an ``(M, input_dim)`` array of points.  Useful when used as the
    ``X`` argument to :func:`check_no_overlaps` or :func:`check_covers_space`.
    """
    pts: list[np.ndarray] = []
    for region in partition.regions:
        D, g = partition.halfspaces(region)
        c = region.centroid
        for i in range(D.shape[0]):
            n_i = D[i]
            nrm = np.linalg.norm(n_i)
            if nrm < 1e-10:
                continue
            # Project centroid onto the plane D[i] · x = g[i].
            offset = (n_i @ c - g[i]) / nrm
            x_plane = c - offset * (n_i / nrm)
            # Step eps inside (negative-normal direction satisfies D[i]·x ≤ g[i]).
            step = eps * (n_i / nrm)
            pts.append(x_plane - step)
    if not pts:
        return np.zeros((0, partition.input_dim))
    return np.array(pts)

Visualization

Every plotting function takes a keyword-only backend: Literal["plotly", "matplotlib"] = "plotly" argument — see Usage § Visualization for the tradeoffs between the two.

Plotly and matplotlib visualisations for Partition objects.

affine_frobenius(partition, region)

Frobenius norm ‖A‖_F of the region's local affine map.

Source code in src/parx/viz.py 
239
240
241
242
def affine_frobenius(partition: Partition, region: Region) -> float:
    """Frobenius norm ``‖A‖_F`` of the region's local affine map."""
    A, _ = partition.local_affine(region)
    return float(np.linalg.norm(A, ord="fro"))

affine_spectral(partition, region)

Spectral norm (top singular value) of A — the local Lipschitz constant.

Source code in src/parx/viz.py 
245
246
247
248
249
250
def affine_spectral(partition: Partition, region: Region) -> float:
    """Spectral norm (top singular value) of A — the local Lipschitz constant."""
    A, _ = partition.local_affine(region)
    if A.size == 0:
        return 0.0
    return float(np.linalg.svd(A, compute_uv=False)[0])

affine_det(partition, region)

Determinant of A. Raises ValueError if A is not square.

Source code in src/parx/viz.py 
253
254
255
256
257
258
def affine_det(partition: Partition, region: Region) -> float:
    """Determinant of A.  Raises ``ValueError`` if A is not square."""
    A, _ = partition.local_affine(region)
    if A.shape[0] != A.shape[1]:
        raise ValueError(f"affine_det requires square A, got shape {A.shape}")
    return float(np.linalg.det(A))

active_neuron_count(_partition, region)

Total number of active neurons across all layers of the path.

Source code in src/parx/viz.py 
261
262
263
def active_neuron_count(_partition: Partition, region: Region) -> float:
    """Total number of active neurons across all layers of the path."""
    return float(sum(int(q.sum()) for q in region.activation_path))

region_palette(partition, scheme='random')

Return one CSS colour string per region in partition.regions.

Parameters:

Name Type Description Default
partition Partition
required
scheme 'random', 'frobenius', or 'spatial'

'random' — Turbo palette indexed by region order. 'frobenius' — Viridis palette ordered by ‖A‖_F. 'spatial' — HSV colour keyed to centroid angle (hue) and radial distance from the mean centroid (saturation).

 'random'
Notes

Returned CSS strings work as the colors= argument of :func:plot_partition_2d under either backend — the matplotlib path converts them internally.

Source code in src/parx/viz.py 
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
def region_palette(partition: Partition, scheme: str = "random") -> list[str]:
    """Return one CSS colour string per region in ``partition.regions``.

    Parameters
    ----------
    partition : Partition
    scheme : 'random', 'frobenius', or 'spatial'
        ``'random'`` — Turbo palette indexed by region order.
        ``'frobenius'`` — Viridis palette ordered by ``‖A‖_F``.
        ``'spatial'`` — HSV colour keyed to centroid angle (hue) and
        radial distance from the mean centroid (saturation).

    Notes
    -----
    Returned CSS strings work as the ``colors=`` argument of
    :func:`plot_partition_2d` under either ``backend`` — the matplotlib path
    converts them internally.
    """
    regions = partition.regions
    n = len(regions)
    if scheme == "random":
        positions = [i / max(n - 1, 1) for i in range(n)]
        return list(px.colors.sample_colorscale("Turbo", positions))
    if scheme == "frobenius":
        metrics = np.array(
            [affine_frobenius(partition, r) for r in regions], dtype=float
        )
        span = metrics.max() - metrics.min()
        normed = (metrics - metrics.min()) / (span if span > 1e-12 else 1.0)
        return list(px.colors.sample_colorscale("Viridis", normed))
    if scheme == "spatial":
        return _spatial_colors(regions)
    raise ValueError(
        f"unknown scheme {scheme!r}; choose 'random', 'frobenius', or 'spatial'"
    )

plot_partition_2d(partition, x_range=None, y_range=None, pad=0.3, *, domain=None, layer=None, color_by=affine_frobenius, color_label=None, log_color=False, colorscale='Viridis', colors=None, backend='plotly')

Draw each linear region as a filled, crisp polygon.

Each region is rendered as a vector-quality go.Scatter polygon, giving sharp boundaries regardless of plot size — no pixellation artefacts. Hovering over a region shows its per-layer activation pattern.

Parameters:

Name Type Description Default
partition Partition

Must have input_dim == 2.

 required
domain ((float, float), (float, float)) or None

Explicit (x_range, y_range) bounds for the plot. When provided, these bounds are used directly and no auto-ranging is performed.

 None
x_range (float, float) or None

Axis extents. When None (default) the range is auto-computed from the arrangement of hyperplane intersection points so all regions are visible. Unbounded regions are clipped to the computed box.

 None
y_range (float, float) or None

Axis extents. When None (default) the range is auto-computed from the arrangement of hyperplane intersection points so all regions are visible. Unbounded regions are clipped to the computed box.

 None
pad float

Fractional padding added around the auto-computed bounding box.

 0.3
layer int or None

If given (1-indexed), collapse all leaf regions to their activation-path prefix of this length and plot the resulting coarser partition. layer=1 shows the regions induced by just the first ReLU layer; layer=partition.n_layers (the default) shows full leaf regions. Must satisfy 1 <= layer <= partition.n_layers.

 None
color_by callable

Callable (partition, region) -> float returning the scalar that determines each region's colour. Built-in helpers in this module: :func:affine_frobenius (default — ‖A‖_F), :func:affine_spectral (local Lipschitz constant), :func:affine_det, :func:active_neuron_count. Pass None to disable the metric mapping and revert to the discrete Turbo palette.

 :func:`affine_frobenius`
color_label str or None

Colourbar title. When None, derived from color_by.__name__.

 None
log_color bool

If True, the metric is mapped onto the colorscale via log10 (clamped at 1e-12). Useful when ‖A‖ spans orders of magnitude.

 False
colorscale str

Any Plotly colorscale name ("Viridis" default, "Turbo", "Plasma", …).

 'Viridis'
colors list[str] or None

Pre-computed CSS colour strings — one per region in partition.regions (or per plot_regions when layer is set). When provided, color_by, colorscale, and log_color are ignored and no colourbar is shown. Use :func:region_palette to generate a matching list.

 None
backend 'plotly' or 'matplotlib'

'plotly' (default) returns an interactive go.Figure with hover tooltips. 'matplotlib' returns a static matplotlib.figure.Figure (no hover text) and requires pip install 'parx[animate]'.

 'plotly'

Returns:

Type Description
Figure or Figure
Source code in src/parx/viz.py 
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
def plot_partition_2d(
    partition: Partition,
    x_range: tuple[float, float] | None = None,
    y_range: tuple[float, float] | None = None,
    pad: float = 0.3,
    *,
    domain: tuple[tuple[float, float], tuple[float, float]] | None = None,
    layer: int | None = None,
    color_by=affine_frobenius,
    color_label: str | None = None,
    log_color: bool = False,
    colorscale: str = "Viridis",
    colors: list[str] | None = None,
    backend: Literal["plotly", "matplotlib"] = "plotly",
) -> go.Figure | matplotlib.figure.Figure:
    """Draw each linear region as a filled, crisp polygon.

    Each region is rendered as a vector-quality ``go.Scatter`` polygon, giving
    sharp boundaries regardless of plot size — no pixellation artefacts.
    Hovering over a region shows its per-layer activation pattern.

    Parameters
    ----------
    partition : Partition
        Must have ``input_dim == 2``.
    domain : ((float, float), (float, float)) or None
        Explicit ``(x_range, y_range)`` bounds for the plot. When provided,
        these bounds are used directly and no auto-ranging is performed.
    x_range, y_range : (float, float) or None
        Axis extents.  When ``None`` (default) the range is auto-computed from
        the arrangement of hyperplane intersection points so all regions are
        visible.  Unbounded regions are clipped to the computed box.
    pad : float
        Fractional padding added around the auto-computed bounding box.
    layer : int or None
        If given (1-indexed), collapse all leaf regions to their activation-path
        prefix of this length and plot the resulting coarser partition.
        ``layer=1`` shows the regions induced by just the first ReLU layer;
        ``layer=partition.n_layers`` (the default) shows full leaf regions.
        Must satisfy ``1 <= layer <= partition.n_layers``.
    color_by : callable, default :func:`affine_frobenius`
        Callable ``(partition, region) -> float`` returning the scalar that
        determines each region's colour.  Built-in helpers in this module:
        :func:`affine_frobenius` (default — ``‖A‖_F``),
        :func:`affine_spectral` (local Lipschitz constant),
        :func:`affine_det`, :func:`active_neuron_count`.  Pass ``None`` to
        disable the metric mapping and revert to the discrete Turbo palette.
    color_label : str or None
        Colourbar title.  When ``None``, derived from ``color_by.__name__``.
    log_color : bool
        If ``True``, the metric is mapped onto the colorscale via ``log10``
        (clamped at ``1e-12``).  Useful when ``‖A‖`` spans orders of magnitude.
    colorscale : str
        Any Plotly colorscale name (``"Viridis"`` default, ``"Turbo"``,
        ``"Plasma"``, …).
    colors : list[str] or None
        Pre-computed CSS colour strings — one per region in
        ``partition.regions`` (or per ``plot_regions`` when ``layer`` is set).
        When provided, ``color_by``, ``colorscale``, and ``log_color`` are
        ignored and no colourbar is shown.  Use :func:`region_palette` to
        generate a matching list.
    backend : 'plotly' or 'matplotlib'
        ``'plotly'`` (default) returns an interactive ``go.Figure`` with hover
        tooltips.  ``'matplotlib'`` returns a static ``matplotlib.figure.Figure``
        (no hover text) and requires ``pip install 'parx[animate]'``.

    Returns
    -------
    go.Figure or matplotlib.figure.Figure
    """
    if backend not in ("plotly", "matplotlib"):
        raise ValueError(
            f"unknown backend {backend!r}; choose 'plotly' or 'matplotlib'"
        )

    if partition.input_dim != 2:
        raise ValueError(
            f"plot_partition_2d requires input_dim=2, got {partition.input_dim}"
        )

    if layer is not None and not (1 <= layer <= partition.n_layers):
        raise ValueError(
            f"layer must be between 1 and {partition.n_layers}, got {layer}"
        )

    # Build the list of regions to plot (possibly coarser than leaf level)
    if layer is not None:
        seen: dict[tuple[bytes, ...], Region] = {}
        for r in partition.regions:
            key = tuple(
                r.activation_path[layer_idx].tobytes() for layer_idx in range(layer)
            )
            if key not in seen:
                seen[key] = Region(
                    activation_path=r.activation_path[:layer],
                    centroid=r.centroid,
                )
        plot_regions = list(seen.values())
    else:
        plot_regions = partition.regions

    n = len(plot_regions)
    if n == 0:
        if backend == "matplotlib":
            ns = _require_matplotlib("plot_partition_2d")
            return ns.plt.figure()
        return go.Figure()

    if domain is not None:
        x_range, y_range = domain
    elif x_range is None or y_range is None:
        auto_x, auto_y = _auto_range_2d(partition, pad=pad)
        x_range = x_range or auto_x
        y_range = y_range or auto_y

    if backend == "matplotlib":
        return _plot_partition_2d_matplotlib(
            plot_regions,
            x_range,
            y_range,
            partition,
            layer,
            color_by,
            color_label,
            log_color,
            colorscale,
            colors,
        )

    if colors is not None:
        # Pre-computed palette: bypass color_by entirely, no colourbar.
        metrics = None
        scaled = None
        m_min = m_max = None
        label_str = None
    elif color_by is None:
        # Discrete Turbo palette, no colourbar, no per-region metric.
        colors = px.colors.sample_colorscale(
            "Turbo", [i / max(n - 1, 1) for i in range(n)]
        )
        metrics = None
        scaled = None
        m_min = m_max = None
        label_str = None
    else:
        metrics = np.array([color_by(partition, r) for r in plot_regions], dtype=float)
        if log_color:
            scaled = np.log10(np.maximum(metrics, 1e-12))
        else:
            scaled = metrics
        m_min, m_max = float(scaled.min()), float(scaled.max())
        if m_max - m_min < 1e-12:
            normed = np.zeros_like(scaled)
        else:
            normed = (scaled - m_min) / (m_max - m_min)
        colors = px.colors.sample_colorscale(colorscale, normed)
        label_str = color_label or getattr(color_by, "__name__", "metric")

    depth_str = f"layer {layer}" if layer is not None else "all layers"
    fig = go.Figure()
    for i, region in enumerate(plot_regions):
        D, g = partition.halfspaces(region)
        verts = _region_vertices_2d(D, g, x_range, y_range)
        if verts is None:
            continue

        xs = np.append(verts[:, 0], verts[0, 0])
        ys = np.append(verts[:, 1], verts[0, 1])

        hover_lines = []
        if metrics is not None:
            hover_lines.append(f"{label_str}: {metrics[i]:.4g}")
        hover_lines.append(_activation_label(region))
        hover_text = "<br>".join(hover_lines)

        fig.add_trace(
            go.Scatter(
                x=xs,
                y=ys,
                fill="toself",
                fillcolor=colors[i],
                line=dict(color="black", width=0.8),
                mode="lines",
                opacity=0.75,
                showlegend=False,
                name=f"region {i}",
                hovertemplate=f"{hover_text}<extra></extra>",
            )
        )

    # Add a hidden marker trace just to expose the colorbar.  The markers
    # themselves are invisible (opacity=0) and placed inside the view; Plotly
    # surfaces the colorscale via showscale=True on the marker.
    if metrics is not None:
        cb_title = label_str + (" (log10)" if log_color else "")
        fig.add_trace(
            go.Scatter(
                x=[None],
                y=[None],
                mode="markers",
                marker=dict(
                    color=[m_min, m_max],
                    colorscale=colorscale,
                    cmin=m_min,
                    cmax=m_max,
                    showscale=True,
                    opacity=0,
                    colorbar=dict(title=cb_title, thickness=14),
                ),
                hoverinfo="skip",
                showlegend=False,
            )
        )

    fig.update_layout(
        xaxis=dict(range=list(x_range), title="x₁", constrain="domain"),
        yaxis=dict(range=list(y_range), title="x₂", scaleanchor="x"),
        title=f"Linear regions — {depth_str}  ({n} regions)",
        width=620 if metrics is not None else 520,
        height=500,
    )
    return fig

plot_region_counts(partition, *, backend='plotly')

Bar chart of distinct region count at each depth.

At depth d, counts the number of unique activation-path prefixes of length d across all leaf regions. Shows how partition complexity grows layer by layer.

Parameters:

Name Type Description Default
partition Partition
required
backend 'plotly' or 'matplotlib'

'plotly' (default) returns a go.Figure; 'matplotlib' returns a matplotlib.figure.Figure and requires pip install 'parx[animate]'.

 'plotly'

Returns:

Type Description
Figure or Figure
Source code in src/parx/viz.py 
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
def plot_region_counts(
    partition: Partition,
    *,
    backend: Literal["plotly", "matplotlib"] = "plotly",
) -> go.Figure | matplotlib.figure.Figure:
    """Bar chart of distinct region count at each depth.

    At depth ``d``, counts the number of unique activation-path prefixes of
    length ``d`` across all leaf regions.  Shows how partition complexity grows
    layer by layer.

    Parameters
    ----------
    partition : Partition
    backend : 'plotly' or 'matplotlib'
        ``'plotly'`` (default) returns a ``go.Figure``; ``'matplotlib'`` returns
        a ``matplotlib.figure.Figure`` and requires ``pip install 'parx[animate]'``.

    Returns
    -------
    go.Figure or matplotlib.figure.Figure
    """
    if backend not in ("plotly", "matplotlib"):
        raise ValueError(
            f"unknown backend {backend!r}; choose 'plotly' or 'matplotlib'"
        )

    depths = list(range(1, partition.n_layers + 1))
    counts = []
    for d in depths:
        prefixes = {
            tuple(r.activation_path[layer_idx].tobytes() for layer_idx in range(d))
            for r in partition.regions
        }
        counts.append(len(prefixes))

    title = f"Region complexity by depth  (leaf total: {len(partition)})"

    if backend == "matplotlib":
        ns = _require_matplotlib("plot_region_counts")
        fig, ax = ns.plt.subplots(figsize=(5.2, 4.0))
        ax.bar(depths, counts, color="steelblue")
        ax.set_xlabel("Layer depth")
        ax.set_ylabel("Distinct regions")
        ax.set_xticks(depths)
        ax.set_title(title)
        return fig

    fig = go.Figure(go.Bar(x=depths, y=counts, marker_color="steelblue"))
    fig.update_layout(
        xaxis=dict(title="Layer depth", tickmode="linear", dtick=1),
        yaxis_title="Distinct regions",
        title=title,
    )
    return fig

plot_partition_slice(partition, free_dims, fixed_values, *, x_range=None, y_range=None, pad=0.3, color_by=affine_frobenius, color_label=None, log_color=False, colorscale='Viridis', backend='plotly')

Draw an axis-aligned 2D slice through a higher-dimensional partition.

Fixes all input dimensions except free_dims[0] and free_dims[1] to the values in fixed_values, then renders the induced 2D polyhedral partition on those two free axes. Regions whose polytopes do not intersect the slice plane are silently skipped.

Parameters:

Name Type Description Default
partition Partition

The partition to visualise. input_dim must be ≥ 2.

 required
free_dims (int, int)

The two input-dimension indices to keep free (0-indexed).

 required
fixed_values dict[int, float]

Maps every other dimension index to its fixed value. All dimensions not in free_dims must appear as keys.

 required
x_range (float, float) or None

Axis extents for the free dimensions. When None (default) the range is auto-computed from the induced 2D hyperplane intersections.

 None
y_range (float, float) or None

Axis extents for the free dimensions. When None (default) the range is auto-computed from the induced 2D hyperplane intersections.

 None
pad float

Fractional padding added around the auto-computed bounding box.

 0.3
color_by callable or None

Callable (partition, region) -> float returning the scalar used to colour each region. Pass None for a discrete Turbo palette.

 :func:`affine_frobenius`
color_label str or None

Colourbar title. When None, derived from color_by.__name__.

 None
log_color bool

If True, colour values are mapped via log10.

 False
colorscale str

Any Plotly colorscale name (default "Viridis").

 'Viridis'
backend 'plotly' or 'matplotlib'

'plotly' (default) returns a go.Figure; 'matplotlib' returns a matplotlib.figure.Figure and requires pip install 'parx[animate]'.

 'plotly'

Returns:

Type Description
Figure or Figure
Source code in src/parx/viz.py 
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
def plot_partition_slice(
    partition: Partition,
    free_dims: tuple[int, int],
    fixed_values: dict[int, float],
    *,
    x_range: tuple[float, float] | None = None,
    y_range: tuple[float, float] | None = None,
    pad: float = 0.3,
    color_by=affine_frobenius,
    color_label: str | None = None,
    log_color: bool = False,
    colorscale: str = "Viridis",
    backend: Literal["plotly", "matplotlib"] = "plotly",
) -> go.Figure | matplotlib.figure.Figure:
    """Draw an axis-aligned 2D slice through a higher-dimensional partition.

    Fixes all input dimensions except ``free_dims[0]`` and ``free_dims[1]``
    to the values in ``fixed_values``, then renders the induced 2D polyhedral
    partition on those two free axes.  Regions whose polytopes do not intersect
    the slice plane are silently skipped.

    Parameters
    ----------
    partition : Partition
        The partition to visualise.  ``input_dim`` must be ≥ 2.
    free_dims : (int, int)
        The two input-dimension indices to keep free (0-indexed).
    fixed_values : dict[int, float]
        Maps every other dimension index to its fixed value.  All dimensions
        not in ``free_dims`` must appear as keys.
    x_range, y_range : (float, float) or None
        Axis extents for the free dimensions.  When ``None`` (default) the
        range is auto-computed from the induced 2D hyperplane intersections.
    pad : float
        Fractional padding added around the auto-computed bounding box.
    color_by : callable or None, default :func:`affine_frobenius`
        Callable ``(partition, region) -> float`` returning the scalar used
        to colour each region.  Pass ``None`` for a discrete Turbo palette.
    color_label : str or None
        Colourbar title.  When ``None``, derived from ``color_by.__name__``.
    log_color : bool
        If ``True``, colour values are mapped via ``log10``.
    colorscale : str
        Any Plotly colorscale name (default ``"Viridis"``).
    backend : 'plotly' or 'matplotlib'
        ``'plotly'`` (default) returns a ``go.Figure``; ``'matplotlib'`` returns
        a ``matplotlib.figure.Figure`` and requires ``pip install 'parx[animate]'``.

    Returns
    -------
    go.Figure or matplotlib.figure.Figure
    """
    if backend not in ("plotly", "matplotlib"):
        raise ValueError(
            f"unknown backend {backend!r}; choose 'plotly' or 'matplotlib'"
        )

    free0, free1 = free_dims
    fixed_dims = [d for d in range(partition.input_dim) if d not in free_dims]
    fixed_vals = np.array([fixed_values[d] for d in fixed_dims])

    systems_2d: list[tuple[np.ndarray, np.ndarray]] = []
    plot_regions = []
    for region in partition.regions:
        D, g = partition.halfspaces(region)
        D_2d = D[:, [free0, free1]]
        g_2d = g - (D[:, fixed_dims] @ fixed_vals if fixed_dims else g * 0.0)
        systems_2d.append((D_2d, g_2d))
        plot_regions.append(region)

    if not plot_regions:
        if backend == "matplotlib":
            ns = _require_matplotlib("plot_partition_slice")
            return ns.plt.figure()
        return go.Figure()

    if x_range is None or y_range is None:
        auto_x, auto_y = _auto_range_2d_from_systems(systems_2d, pad=pad)
        x_range = x_range if x_range is not None else auto_x
        y_range = y_range if y_range is not None else auto_y

    n_visible = sum(
        1
        for D_2d, g_2d in systems_2d
        if _region_vertices_2d(D_2d, g_2d, x_range, y_range) is not None
    )
    fixed_str = ", ".join(f"x{d}={v:.3g}" for d, v in fixed_values.items())
    title = (
        f"Partition slice — free dims ({free0}, {free1}), fixed: {fixed_str}"
        f"  ({n_visible} visible regions)"
    )

    builder = (
        _build_colored_figure_matplotlib
        if backend == "matplotlib"
        else _build_colored_figure
    )
    return builder(
        plot_regions,
        systems_2d,
        x_range,
        y_range,
        partition,
        color_by,
        color_label,
        log_color,
        colorscale,
        title,
    )

plot_partition_projection(partition, projection, *, x_range=None, y_range=None, pad=0.3, color_by=affine_frobenius, color_label=None, log_color=False, colorscale='Viridis', backend='plotly')

Draw an approximate 2D projection of a higher-dimensional partition.

Projects each region's halfspace normals onto a 2D subspace defined by projection:

.. code-block:: python

D_proj = D @ projection   # (n_constraints, 2)
# g unchanged

.. warning:: This is an approximation. Projecting the dual representation (halfspace normals) does not, in general, yield the exact projected polytope — the true projection of a convex polytope requires eliminating the non-free variables (e.g. via Fourier-Motzkin), which is computationally expensive. The visualisation is useful for qualitative exploration but should not be interpreted as geometrically exact in dimensions ≥ 3.

Parameters:

Name Type Description Default
partition Partition

The partition to visualise.

 required
projection (ndarray, shape(input_dim, 2))

Linear map from input space to the 2D display plane. Columns are the two basis vectors of the subspace. Need not be orthonormal, though orthonormal columns give a shape-preserving projection.

 required
x_range (float, float) or None

Axis extents. When None (default), auto-computed from the projected hyperplane intersections.

 None
y_range (float, float) or None

Axis extents. When None (default), auto-computed from the projected hyperplane intersections.

 None
pad float

Fractional padding added around the auto-computed bounding box.

 0.3
color_by callable or None

Callable (partition, region) -> float. Pass None for a discrete Turbo palette.

 :func:`affine_frobenius`
color_label str or None

Colourbar title.

 None
log_color bool

If True, colour values are mapped via log10.

 False
colorscale str

Any Plotly colorscale name (default "Viridis").

 'Viridis'
backend 'plotly' or 'matplotlib'

'plotly' (default) returns a go.Figure; 'matplotlib' returns a matplotlib.figure.Figure and requires pip install 'parx[animate]'.

 'plotly'

Returns:

Type Description
Figure or Figure
Source code in src/parx/viz.py 
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
def plot_partition_projection(
    partition: Partition,
    projection: np.ndarray,
    *,
    x_range: tuple[float, float] | None = None,
    y_range: tuple[float, float] | None = None,
    pad: float = 0.3,
    color_by=affine_frobenius,
    color_label: str | None = None,
    log_color: bool = False,
    colorscale: str = "Viridis",
    backend: Literal["plotly", "matplotlib"] = "plotly",
) -> go.Figure | matplotlib.figure.Figure:
    """Draw an approximate 2D projection of a higher-dimensional partition.

    Projects each region's halfspace normals onto a 2D subspace defined by
    ``projection``:

    .. code-block:: python

        D_proj = D @ projection   # (n_constraints, 2)
        # g unchanged

    .. warning::
        This is an **approximation**.  Projecting the dual representation
        (halfspace normals) does not, in general, yield the exact projected
        polytope — the true projection of a convex polytope requires
        eliminating the non-free variables (e.g. via Fourier-Motzkin), which
        is computationally expensive.  The visualisation is useful for
        qualitative exploration but should not be interpreted as geometrically
        exact in dimensions ≥ 3.

    Parameters
    ----------
    partition : Partition
        The partition to visualise.
    projection : np.ndarray, shape (input_dim, 2)
        Linear map from input space to the 2D display plane.  Columns are the
        two basis vectors of the subspace.  Need not be orthonormal, though
        orthonormal columns give a shape-preserving projection.
    x_range, y_range : (float, float) or None
        Axis extents.  When ``None`` (default), auto-computed from the
        projected hyperplane intersections.
    pad : float
        Fractional padding added around the auto-computed bounding box.
    color_by : callable or None, default :func:`affine_frobenius`
        Callable ``(partition, region) -> float``.  Pass ``None`` for a
        discrete Turbo palette.
    color_label : str or None
        Colourbar title.
    log_color : bool
        If ``True``, colour values are mapped via ``log10``.
    colorscale : str
        Any Plotly colorscale name (default ``"Viridis"``).
    backend : 'plotly' or 'matplotlib'
        ``'plotly'`` (default) returns a ``go.Figure``; ``'matplotlib'`` returns
        a ``matplotlib.figure.Figure`` and requires ``pip install 'parx[animate]'``.

    Returns
    -------
    go.Figure or matplotlib.figure.Figure
    """
    if backend not in ("plotly", "matplotlib"):
        raise ValueError(
            f"unknown backend {backend!r}; choose 'plotly' or 'matplotlib'"
        )

    projection = np.asarray(projection, dtype=float)
    if projection.shape != (partition.input_dim, 2):
        raise ValueError(
            f"projection must have shape (input_dim, 2) = ({partition.input_dim}, 2), "
            f"got {projection.shape}"
        )

    systems_2d: list[tuple[np.ndarray, np.ndarray]] = []
    plot_regions = list(partition.regions)
    for region in plot_regions:
        D, g = partition.halfspaces(region)
        D_proj = D @ projection
        systems_2d.append((D_proj, g))

    if not plot_regions:
        if backend == "matplotlib":
            ns = _require_matplotlib("plot_partition_projection")
            return ns.plt.figure()
        return go.Figure()

    if x_range is None or y_range is None:
        auto_x, auto_y = _auto_range_2d_from_systems(systems_2d, pad=pad)
        x_range = x_range if x_range is not None else auto_x
        y_range = y_range if y_range is not None else auto_y

    n = len(plot_regions)
    title = f"Partition projection (approx.)  ({n} regions)"

    builder = (
        _build_colored_figure_matplotlib
        if backend == "matplotlib"
        else _build_colored_figure
    )
    return builder(
        plot_regions,
        systems_2d,
        x_range,
        y_range,
        partition,
        color_by,
        color_label,
        log_color,
        colorscale,
        title,
    )

plot_partition_pca(partition, data, *, color_by=affine_frobenius, color_label=None, log_color=False, colorscale='Viridis', backend='plotly')

Draw an approximate 2D PCA projection of a higher-dimensional partition.

Fits a 2-component PCA on data, uses the principal-component axes as the projection matrix, and delegates to :func:plot_partition_projection.

The projection is the same approximation noted in :func:plot_partition_projection — halfspace normals are projected, not the polytopes themselves.

Parameters:

Name Type Description Default
partition Partition

The partition to visualise.

 required
data (ndarray, shape(N, input_dim))

Data used to fit the PCA. Typically the training or evaluation set.

 required
color_by callable or None

Callable (partition, region) -> float. Pass None for a discrete Turbo palette.

 :func:`affine_frobenius`
color_label str or None

Colourbar title.

 None
log_color bool

If True, colour values are mapped via log10.

 False
colorscale str

Any Plotly colorscale name (default "Viridis").

 'Viridis'
backend 'plotly' or 'matplotlib'

'plotly' (default) returns a go.Figure; 'matplotlib' returns a matplotlib.figure.Figure and requires pip install 'parx[animate]'.

 'plotly'

Returns:

Type Description
Figure or Figure

Raises:

Type Description
ImportError

If scikit-learn is not installed.
Source code in src/parx/viz.py 
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
def plot_partition_pca(
    partition: Partition,
    data: np.ndarray,
    *,
    color_by=affine_frobenius,
    color_label: str | None = None,
    log_color: bool = False,
    colorscale: str = "Viridis",
    backend: Literal["plotly", "matplotlib"] = "plotly",
) -> go.Figure | matplotlib.figure.Figure:
    """Draw an approximate 2D PCA projection of a higher-dimensional partition.

    Fits a 2-component PCA on ``data``, uses the principal-component axes as
    the projection matrix, and delegates to :func:`plot_partition_projection`.

    The projection is the same approximation noted in
    :func:`plot_partition_projection` — halfspace normals are projected, not
    the polytopes themselves.

    Parameters
    ----------
    partition : Partition
        The partition to visualise.
    data : np.ndarray, shape (N, input_dim)
        Data used to fit the PCA.  Typically the training or evaluation set.
    color_by : callable or None, default :func:`affine_frobenius`
        Callable ``(partition, region) -> float``.  Pass ``None`` for a
        discrete Turbo palette.
    color_label : str or None
        Colourbar title.
    log_color : bool
        If ``True``, colour values are mapped via ``log10``.
    colorscale : str
        Any Plotly colorscale name (default ``"Viridis"``).
    backend : 'plotly' or 'matplotlib'
        ``'plotly'`` (default) returns a ``go.Figure``; ``'matplotlib'`` returns
        a ``matplotlib.figure.Figure`` and requires ``pip install 'parx[animate]'``.

    Returns
    -------
    go.Figure or matplotlib.figure.Figure

    Raises
    ------
    ImportError
        If ``scikit-learn`` is not installed.
    """
    try:
        from sklearn.decomposition import PCA
    except ImportError as e:
        raise ImportError(
            "plot_partition_pca requires scikit-learn: pip install scikit-learn"
        ) from e

    pca = PCA(n_components=2)
    pca.fit(np.asarray(data, dtype=float))
    projection = pca.components_.T  # shape (input_dim, 2)

    return plot_partition_projection(
        partition,
        projection,
        color_by=color_by,
        color_label=color_label,
        log_color=log_color,
        colorscale=colorscale,
        backend=backend,
    )

plot_halfspaces(partition, region, x_range=(-2.0, 2.0), y_range=(-2.0, 2.0), *, backend='plotly')

Draw the halfspace boundaries that define a single region.

For each row D[i,:] x = g[i] of the region's constraint system, draws the corresponding line in the 2-D plane clipped to the plot window.

Parameters:

Name Type Description Default
partition Partition

Must have input_dim == 2.

 required
region Region

The region whose halfspace boundaries to draw.

 required
x_range (float, float)

Axis extents.

 (-2.0, 2.0)
y_range (float, float)

Axis extents.

 (-2.0, 2.0)
backend 'plotly' or 'matplotlib'

'plotly' (default) returns a go.Figure; 'matplotlib' returns a matplotlib.figure.Figure and requires pip install 'parx[animate]'.

 'plotly'

Returns:

Type Description
Figure or Figure
Source code in src/parx/viz.py 
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
def plot_halfspaces(
    partition: Partition,
    region: Region,
    x_range: tuple[float, float] = (-2.0, 2.0),
    y_range: tuple[float, float] = (-2.0, 2.0),
    *,
    backend: Literal["plotly", "matplotlib"] = "plotly",
) -> go.Figure | matplotlib.figure.Figure:
    """Draw the halfspace boundaries that define a single region.

    For each row ``D[i,:] x = g[i]`` of the region's constraint system, draws
    the corresponding line in the 2-D plane clipped to the plot window.

    Parameters
    ----------
    partition : Partition
        Must have ``input_dim == 2``.
    region : Region
        The region whose halfspace boundaries to draw.
    x_range, y_range : (float, float)
        Axis extents.
    backend : 'plotly' or 'matplotlib'
        ``'plotly'`` (default) returns a ``go.Figure``; ``'matplotlib'`` returns
        a ``matplotlib.figure.Figure`` and requires ``pip install 'parx[animate]'``.

    Returns
    -------
    go.Figure or matplotlib.figure.Figure
    """
    if backend not in ("plotly", "matplotlib"):
        raise ValueError(
            f"unknown backend {backend!r}; choose 'plotly' or 'matplotlib'"
        )

    if partition.input_dim != 2:
        raise ValueError(
            f"plot_halfspaces requires input_dim=2, got {partition.input_dim}"
        )

    D, g = partition.halfspaces(region)
    xs = np.linspace(x_range[0], x_range[1], 400)

    if backend == "matplotlib":
        return _plot_halfspaces_matplotlib(D, g, region, x_range, y_range, xs)

    fig = go.Figure()

    for i in range(D.shape[0]):
        d0, d1 = D[i, 0], D[i, 1]
        gi = g[i]
        if abs(d1) > 1e-10:
            y_line = (gi - d0 * xs) / d1
            mask = (y_line >= y_range[0]) & (y_line <= y_range[1])
            if mask.any():
                fig.add_trace(
                    go.Scatter(
                        x=xs[mask],
                        y=y_line[mask],
                        mode="lines",
                        line=dict(color="rgba(80,80,80,0.45)", width=1),
                        showlegend=False,
                    )
                )
        elif abs(d0) > 1e-10:
            x_val = gi / d0
            if x_range[0] <= x_val <= x_range[1]:
                fig.add_trace(
                    go.Scatter(
                        x=[x_val, x_val],
                        y=[y_range[0], y_range[1]],
                        mode="lines",
                        line=dict(color="rgba(80,80,80,0.45)", width=1),
                        showlegend=False,
                    )
                )

    c = region.centroid
    if x_range[0] <= c[0] <= x_range[1] and y_range[0] <= c[1] <= y_range[1]:
        fig.add_trace(
            go.Scatter(
                x=[c[0]],
                y=[c[1]],
                mode="markers",
                marker=dict(size=10, color="crimson", symbol="x"),
                name="centroid",
            )
        )

    fig.update_layout(
        xaxis=dict(range=list(x_range), title="x₁"),
        yaxis=dict(range=list(y_range), title="x₂"),
        title=f"Halfspaces for region  ({D.shape[0]} constraints)",
    )
    return fig

animate_epochs(partitions, *, epoch_labels=None, x_range=None, y_range=None, pad=0.3, color_by=affine_frobenius, color_label=None, log_color=False, colorscale='Viridis', frame_duration=500, backend='plotly', figsize=(6.0, 5.0))

One frame per epoch, animated over the training/optimisation trajectory.

All frames share a fixed color scale and spatial range so that changes across epochs are visually comparable. Each partition in partitions corresponds to one epoch (frame); the list index is used as the epoch label unless epoch_labels is provided.

Parameters:

Name Type Description Default
partitions list[Partition]

One Partition per epoch, all must have input_dim == 2.

 required
epoch_labels list[str] or None

Labels shown in the slider. Defaults to ["0", "1", …].

 None
x_range (float, float) or None

Axis extents shared across all frames. Auto-computed when None.

 None
y_range (float, float) or None

Axis extents shared across all frames. Auto-computed when None.

 None
pad float

Fractional padding for auto-computed range.

 0.3
color_by callable or None

(partition, region) -> float metric. Pass None for a discrete Turbo palette (color range not shared across epochs in that case).

 affine_frobenius
color_label str or None

Colourbar title.

 None
log_color bool

Map metric via log10 before colouring.

 False
colorscale str

Any Plotly colorscale name.

 'Viridis'
frame_duration int

Milliseconds each frame is shown during playback (Plotly play/pause speed, or the matplotlib FuncAnimation frame interval).

 500
backend 'plotly' or 'matplotlib'

'plotly' (default) returns a go.Figure with an interactive play/pause button and slider. 'matplotlib' returns a matplotlib.animation.FuncAnimation instead — matplotlib has no native slider widget, so playback is via .to_jshtml() in a notebook or by saving to a file (see :func:animate_epochs_video). Requires pip install 'parx[animate]'.

 'plotly'
figsize (float, float)

Matplotlib figure size in inches. Only used when backend='matplotlib'.

 (6.0, 5.0)

Returns:

Type Description
Figure or FuncAnimation
Source code in src/parx/viz.py 
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
def animate_epochs(
    partitions: list[Partition],
    *,
    epoch_labels: list[str] | None = None,
    x_range: tuple[float, float] | None = None,
    y_range: tuple[float, float] | None = None,
    pad: float = 0.3,
    color_by=affine_frobenius,
    color_label: str | None = None,
    log_color: bool = False,
    colorscale: str = "Viridis",
    frame_duration: int = 500,
    backend: Literal["plotly", "matplotlib"] = "plotly",
    figsize: tuple[float, float] = (6.0, 5.0),
):
    """One frame per epoch, animated over the training/optimisation trajectory.

    All frames share a fixed color scale and spatial range so that changes
    across epochs are visually comparable.  Each partition in ``partitions``
    corresponds to one epoch (frame); the list index is used as the epoch label
    unless ``epoch_labels`` is provided.

    Parameters
    ----------
    partitions : list[Partition]
        One Partition per epoch, all must have ``input_dim == 2``.
    epoch_labels : list[str] or None
        Labels shown in the slider.  Defaults to ``["0", "1", …]``.
    x_range, y_range : (float, float) or None
        Axis extents shared across all frames.  Auto-computed when ``None``.
    pad : float
        Fractional padding for auto-computed range.
    color_by : callable or None
        ``(partition, region) -> float`` metric.  Pass ``None`` for a discrete
        Turbo palette (color range not shared across epochs in that case).
    color_label : str or None
        Colourbar title.
    log_color : bool
        Map metric via ``log10`` before colouring.
    colorscale : str
        Any Plotly colorscale name.
    frame_duration : int
        Milliseconds each frame is shown during playback (Plotly play/pause
        speed, or the matplotlib ``FuncAnimation`` frame interval).
    backend : 'plotly' or 'matplotlib'
        ``'plotly'`` (default) returns a ``go.Figure`` with an interactive
        play/pause button and slider.  ``'matplotlib'`` returns a
        ``matplotlib.animation.FuncAnimation`` instead — matplotlib has no
        native slider widget, so playback is via ``.to_jshtml()`` in a
        notebook or by saving to a file (see :func:`animate_epochs_video`).
        Requires ``pip install 'parx[animate]'``.
    figsize : (float, float)
        Matplotlib figure size in inches.  Only used when
        ``backend='matplotlib'``.

    Returns
    -------
    go.Figure or matplotlib.animation.FuncAnimation
    """
    if backend not in ("plotly", "matplotlib"):
        raise ValueError(
            f"unknown backend {backend!r}; choose 'plotly' or 'matplotlib'"
        )

    if not partitions:
        if backend == "matplotlib":
            raise ValueError("animate_epochs requires at least one partition")
        return go.Figure()
    if any(p.input_dim != 2 for p in partitions):
        raise ValueError(
            "animate_epochs requires all partitions to have input_dim == 2"
        )

    labels = epoch_labels or [str(i) for i in range(len(partitions))]
    if len(labels) != len(partitions):
        raise ValueError("epoch_labels length must match number of partitions")

    if x_range is None or y_range is None:
        auto_x, auto_y = _global_range(partitions, pad)
        x_range = x_range or auto_x
        y_range = y_range or auto_y

    if backend == "matplotlib":
        anim, _fig = _animate_epochs_matplotlib(
            "animate_epochs",
            partitions,
            labels,
            x_range,
            y_range,
            color_by,
            color_label,
            log_color,
            colorscale,
            frame_duration,
            figsize,
        )
        return anim

    if color_by is not None:
        m_min, m_max = _global_metric_range(partitions, color_by, log_color)
        label_str = color_label or getattr(color_by, "__name__", "metric")
        cb_title = label_str + (" (log10)" if log_color else "")
    else:
        m_min = m_max = None
        label_str = None
        cb_title = None

    # Pre-compute per-epoch polygon data
    epoch_traces: list[list[go.Scatter]] = []
    for partition in partitions:
        traces: list[go.Scatter] = []
        regions = partition.regions
        if color_by is not None:
            raw = np.array([color_by(partition, r) for r in regions], dtype=float)
            scaled = np.log10(np.maximum(raw, 1e-12)) if log_color else raw
            if m_max - m_min < 1e-12:
                normed = np.zeros_like(scaled)
            else:
                normed = (scaled - m_min) / (m_max - m_min)
            colors = px.colors.sample_colorscale(colorscale, normed)
        else:
            n = len(regions)
            colors = px.colors.sample_colorscale(
                "Turbo", [i / max(n - 1, 1) for i in range(n)]
            )

        for i, region in enumerate(regions):
            D, g = partition.halfspaces(region)
            verts = _region_vertices_2d(D, g, x_range, y_range)
            if verts is None:
                xs_poly: list = []
                ys_poly: list = []
            else:
                xs_poly = list(np.append(verts[:, 0], verts[0, 0]))
                ys_poly = list(np.append(verts[:, 1], verts[0, 1]))

            hover_lines = []
            if color_by is not None:
                hover_lines.append(f"{label_str}: {raw[i]:.4g}")
            hover_lines.append(_activation_label(region))
            hover_text = "<br>".join(hover_lines)

            traces.append(
                go.Scatter(
                    x=xs_poly,
                    y=ys_poly,
                    fill="toself",
                    fillcolor=colors[i],
                    line=dict(color="black", width=0.8),
                    mode="lines",
                    opacity=0.75,
                    showlegend=False,
                    hovertemplate=f"{hover_text}<extra></extra>",
                )
            )
        epoch_traces.append(traces)

    n_max = max(len(t) for t in epoch_traces)
    # Colorbar trace index is always n_max
    _empty = go.Scatter(
        x=[],
        y=[],
        fill="toself",
        fillcolor="rgba(0,0,0,0)",
        line=dict(color="rgba(0,0,0,0)", width=0),
        mode="lines",
        showlegend=False,
        hoverinfo="skip",
    )

    def _pad_traces(traces: list[go.Scatter]) -> list[go.Scatter]:
        padded = list(traces)
        while len(padded) < n_max:
            padded.append(_empty)
        return padded

    colorbar_trace = go.Scatter(
        x=[None],
        y=[None],
        mode="markers",
        marker=dict(
            color=[m_min, m_max] if m_min is not None else [0, 1],
            colorscale=colorscale,
            cmin=m_min,
            cmax=m_max,
            showscale=(m_min is not None),
            opacity=0,
            colorbar=dict(title=cb_title or "", thickness=14),
        ),
        hoverinfo="skip",
        showlegend=False,
    )

    # Initial figure data (first epoch)
    initial_traces = _pad_traces(epoch_traces[0]) + [colorbar_trace]
    fig = go.Figure(data=initial_traces)

    # Build frames
    frames = []
    for label, traces in zip(labels, epoch_traces):
        frame_data = _pad_traces(traces) + [colorbar_trace]
        frames.append(
            go.Frame(
                data=frame_data,
                name=label,
                traces=list(range(n_max + 1)),
            )
        )
    fig.frames = frames

    # Slider steps
    slider_steps = [
        dict(
            args=[
                [label],
                dict(
                    frame=dict(duration=frame_duration, redraw=True),
                    mode="immediate",
                    transition=dict(duration=0),
                ),
            ],
            label=label,
            method="animate",
        )
        for label in labels
    ]

    fig.update_layout(
        xaxis=dict(range=list(x_range), title="x₁", constrain="domain"),
        yaxis=dict(range=list(y_range), title="x₂", scaleanchor="x"),
        title=f"Partition evolution  ({len(partitions)} epochs)",
        width=640 if m_min is not None else 520,
        height=540,
        updatemenus=[
            dict(
                type="buttons",
                showactive=False,
                y=1.08,
                x=0.0,
                xanchor="left",
                buttons=[
                    dict(
                        label="▶ Play",
                        method="animate",
                        args=[
                            None,
                            dict(
                                frame=dict(duration=frame_duration, redraw=True),
                                fromcurrent=True,
                                transition=dict(duration=0),
                            ),
                        ],
                    ),
                    dict(
                        label="⏸ Pause",
                        method="animate",
                        args=[
                            [None],
                            dict(
                                frame=dict(duration=0, redraw=False),
                                mode="immediate",
                                transition=dict(duration=0),
                            ),
                        ],
                    ),
                ],
            )
        ],
        sliders=[
            dict(
                active=0,
                currentvalue=dict(prefix="Epoch: ", visible=True, xanchor="center"),
                pad=dict(t=60),
                steps=slider_steps,
            )
        ],
    )
    return fig

animate_epochs_video(partitions, path, *, epoch_labels=None, x_range=None, y_range=None, pad=0.3, color_by=affine_frobenius, color_label=None, log_color=False, colorscale='Viridis', fps=4, dpi=150, figsize=(6.0, 5.0))

Write an animated .gif or .mp4 of the partition across epochs.

Format is inferred from the path suffix (.gif → Pillow writer, .mp4 → ffmpeg writer).

Parameters:

Name Type Description Default
partitions list[Partition]

One Partition per epoch, all must have input_dim == 2.

 required
path str or PathLike

Output file. Suffix determines format: .gif or .mp4.

 required
epoch_labels list[str] or None

Title suffix per frame. Defaults to ["0", "1", …].

 None
x_range (float, float) or None

Axis extents shared across all frames. Auto-computed when None.

 None
y_range (float, float) or None

Axis extents shared across all frames. Auto-computed when None.

 None
pad float

Fractional padding for auto-computed range.

 0.3
color_by callable or None

(partition, region) -> float metric.

 affine_frobenius
color_label str or None

Colorbar label.

 None
log_color bool

Map metric via log10.

 False
colorscale str

Plotly colorscale name; lowercased to resolve a matplotlib colormap ("Viridis" → "viridis" etc.).

 'Viridis'
fps int

Frames per second.

 4
dpi int

Output resolution in dots per inch.

 150
figsize (float, float)

Matplotlib figure size in inches.

 (6.0, 5.0)

Raises:

Type Description
ImportError

If matplotlib is not installed.
Source code in src/parx/viz.py 
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
def animate_epochs_video(
    partitions: list[Partition],
    path,
    *,
    epoch_labels: list[str] | None = None,
    x_range: tuple[float, float] | None = None,
    y_range: tuple[float, float] | None = None,
    pad: float = 0.3,
    color_by=affine_frobenius,
    color_label: str | None = None,
    log_color: bool = False,
    colorscale: str = "Viridis",
    fps: int = 4,
    dpi: int = 150,
    figsize: tuple[float, float] = (6.0, 5.0),
) -> None:
    """Write an animated ``.gif`` or ``.mp4`` of the partition across epochs.

    Format is inferred from the ``path`` suffix (``.gif`` → Pillow writer,
    ``.mp4`` → ffmpeg writer).

    Parameters
    ----------
    partitions : list[Partition]
        One Partition per epoch, all must have ``input_dim == 2``.
    path : str or os.PathLike
        Output file.  Suffix determines format: ``.gif`` or ``.mp4``.
    epoch_labels : list[str] or None
        Title suffix per frame.  Defaults to ``["0", "1", …]``.
    x_range, y_range : (float, float) or None
        Axis extents shared across all frames.  Auto-computed when ``None``.
    pad : float
        Fractional padding for auto-computed range.
    color_by : callable or None
        ``(partition, region) -> float`` metric.
    color_label : str or None
        Colorbar label.
    log_color : bool
        Map metric via ``log10``.
    colorscale : str
        Plotly colorscale name; lowercased to resolve a matplotlib colormap
        (``"Viridis"`` → ``"viridis"`` etc.).
    fps : int
        Frames per second.
    dpi : int
        Output resolution in dots per inch.
    figsize : (float, float)
        Matplotlib figure size in inches.

    Raises
    ------
    ImportError
        If ``matplotlib`` is not installed.
    """
    ns = _require_matplotlib("animate_epochs_video")

    if not partitions:
        return
    if any(p.input_dim != 2 for p in partitions):
        raise ValueError(
            "animate_epochs_video requires all partitions to have input_dim == 2"
        )

    import os

    path = os.fspath(path)
    labels = epoch_labels or [str(i) for i in range(len(partitions))]
    if len(labels) != len(partitions):
        raise ValueError("epoch_labels length must match number of partitions")

    if x_range is None or y_range is None:
        auto_x, auto_y = _global_range(partitions, pad)
        x_range = x_range or auto_x
        y_range = y_range or auto_y

    anim, fig = _animate_epochs_matplotlib(
        "animate_epochs_video",
        partitions,
        labels,
        x_range,
        y_range,
        color_by,
        color_label,
        log_color,
        colorscale,
        1000 // fps,
        figsize,
    )

    suffix = os.path.splitext(path)[-1].lower()
    if suffix == ".gif":
        anim.save(path, writer="pillow", fps=fps, dpi=dpi)
    elif suffix == ".mp4":
        anim.save(path, writer="ffmpeg", fps=fps, dpi=dpi)
    else:
        anim.save(path, fps=fps, dpi=dpi)

    ns.plt.close(fig)

plot_feature_embedding(features, partition, X, *, method='tsne', color_by='region', title=None, backend='plotly', **embed_kwargs)

Scatter the 2D embedding of penultimate-layer features, coloured by region.

Parameters:

Name Type Description Default
features (ndarray, shape(N, D))

Pre-extracted penultimate-layer activations.

 required
partition Partition

Used to route X to linear regions via partition.route(X).

 required
X (ndarray, shape(N, input_dim))

Original input points corresponding to each row of features.

 required
method 'tsne' or 'umap'

Dimensionality reduction method forwarded to :func:_compute_embedding.

 'tsne'
color_by 'region' or np.ndarray shape (N,)

'region' colours discretely by region index (Turbo palette); an array colours continuously by scalar value (Viridis + colorbar).

 'region'
title str or None

Figure title; defaults to embedding method and region count.

 None
backend 'plotly' or 'matplotlib'

'plotly' (default) returns a go.Figure; 'matplotlib' returns a matplotlib.figure.Figure (no hover text) and requires pip install 'parx[animate]'.

 'plotly'
**embed_kwargs

Forwarded to :func:_compute_embedding.

 {}
Source code in src/parx/viz.py 
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
def plot_feature_embedding(
    features: np.ndarray,
    partition: Partition,
    X: np.ndarray,
    *,
    method: str = "tsne",
    color_by="region",
    title: str | None = None,
    backend: Literal["plotly", "matplotlib"] = "plotly",
    **embed_kwargs,
) -> go.Figure | matplotlib.figure.Figure:
    """Scatter the 2D embedding of penultimate-layer features, coloured by region.

    Parameters
    ----------
    features : np.ndarray, shape (N, D)
        Pre-extracted penultimate-layer activations.
    partition : Partition
        Used to route ``X`` to linear regions via ``partition.route(X)``.
    X : np.ndarray, shape (N, input_dim)
        Original input points corresponding to each row of ``features``.
    method : 'tsne' or 'umap'
        Dimensionality reduction method forwarded to :func:`_compute_embedding`.
    color_by : 'region' or np.ndarray shape (N,)
        ``'region'`` colours discretely by region index (Turbo palette); an
        array colours continuously by scalar value (Viridis + colorbar).
    title : str or None
        Figure title; defaults to embedding method and region count.
    backend : 'plotly' or 'matplotlib'
        ``'plotly'`` (default) returns a ``go.Figure``; ``'matplotlib'`` returns
        a ``matplotlib.figure.Figure`` (no hover text) and requires
        ``pip install 'parx[animate]'``.
    **embed_kwargs
        Forwarded to :func:`_compute_embedding`.
    """
    if backend not in ("plotly", "matplotlib"):
        raise ValueError(
            f"unknown backend {backend!r}; choose 'plotly' or 'matplotlib'"
        )

    emb = _compute_embedding(np.asarray(features, dtype=float), method, **embed_kwargs)

    routed = partition.route(X)
    region_map = {
        tuple(q.tobytes() for q in r.activation_path): i
        for i, r in enumerate(partition.regions)
    }
    region_ids = np.array(
        [
            region_map.get(tuple(q.tobytes() for q in r.activation_path), -1)
            if r is not None
            else -1
            for r in routed
        ],
        dtype=int,
    )

    if backend == "matplotlib":
        return _plot_feature_embedding_matplotlib(
            emb, region_ids, partition, method, color_by, title
        )

    fig = go.Figure()

    is_region_palette = isinstance(color_by, list) or (
        isinstance(color_by, str) and color_by == "region"
    )
    if is_region_palette:
        # Per-region discrete palette: either pre-computed list[str] or Turbo by index.
        n_regions = len(partition.regions)
        routed_mask = region_ids >= 0
        unrouted_mask = ~routed_mask
        routed_indices = np.where(routed_mask)[0]

        if isinstance(color_by, list):
            palette = color_by  # N region colors, index by region_id
            color_list = [palette[region_ids[i]] for i in routed_indices]
        else:
            color_list = px.colors.sample_colorscale(
                "Turbo",
                [region_ids[i] / max(n_regions - 1, 1) for i in routed_indices],
            )
        hover_strings = [
            f"Region {rid}<br>{_activation_label(r)}" if r is not None else "outside"
            for rid, r in zip(region_ids, routed)
        ]

        if routed_mask.any():
            fig.add_trace(
                go.Scattergl(
                    x=emb[routed_mask, 0],
                    y=emb[routed_mask, 1],
                    mode="markers",
                    marker=dict(color=color_list, size=4, opacity=0.7),
                    text=[hover_strings[i] for i in routed_indices],
                    hovertemplate="%{text}<extra></extra>",
                    showlegend=False,
                    name="routed",
                )
            )

        if unrouted_mask.any():
            fig.add_trace(
                go.Scattergl(
                    x=emb[unrouted_mask, 0],
                    y=emb[unrouted_mask, 1],
                    mode="markers",
                    marker=dict(color="lightgrey", size=5, symbol="x", opacity=0.7),
                    text=["outside known regions"] * int(unrouted_mask.sum()),
                    hovertemplate="%{text}<extra></extra>",
                    showlegend=False,
                    name="unrouted",
                )
            )
    else:
        color_vals = np.asarray(color_by, dtype=float)
        fig.add_trace(
            go.Scattergl(
                x=emb[:, 0],
                y=emb[:, 1],
                mode="markers",
                marker=dict(
                    color=color_vals,
                    colorscale="Viridis",
                    showscale=True,
                    size=4,
                    opacity=0.7,
                ),
                showlegend=False,
            )
        )

    fig.update_layout(
        xaxis_title=f"{method.upper()} 1",
        yaxis_title=f"{method.upper()} 2",
        title=title
        or f"Feature embedding ({method.upper()}, {len(partition)} regions)",
        width=600,
        height=520,
    )
    return fig

I/O & Serialization

Loading training checkpoints, and saving/loading a computed Partition without needing Julia again.

Yield (label, state_dict) pairs from epoch-indexed sources.

Accepted inputs:

  • Single state dict — a dict whose keys end in .weight / .bias (PyTorch convention). Yields (None, state_dict).
  • Mapping of epoch → state_dict — a dict whose values are themselves state dicts. Yields (epoch, state_dict) in key-sorted order.
  • Sequence/set of state dicts — yields (index, state_dict) in iteration order (sets are sorted by id — pass a list for determinism).
  • Path to a .h5 file — two layouts auto-detected:
    • flat (all parameter datasets at the file root) → single state dict
    • grouped (each top-level item is a Group containing one state dict) → one yield per group. Group names are sorted numerically when they look like integers; lexicographically otherwise.

The h5 path holds the file open for the duration of iteration; consume the generator (e.g. via list(...)) before doing anything else if you need the handles released sooner.

Source code in src/parx/io.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def iter_state_dicts(source: Any) -> Iterator[tuple[Any, dict]]:
    """Yield ``(label, state_dict)`` pairs from epoch-indexed sources.

    Accepted inputs:

    * **Single state dict** — a ``dict`` whose keys end in ``.weight`` /
      ``.bias`` (PyTorch convention).  Yields ``(None, state_dict)``.
    * **Mapping of epoch → state_dict** — a ``dict`` whose values are
      themselves state dicts.  Yields ``(epoch, state_dict)`` in key-sorted
      order.
    * **Sequence/set of state dicts** — yields ``(index, state_dict)`` in
      iteration order (sets are sorted by ``id`` — pass a list for determinism).
    * **Path to a ``.h5`` file** — two layouts auto-detected:
        - **flat** (all parameter datasets at the file root) → single state dict
        - **grouped** (each top-level item is a ``Group`` containing one state
          dict) → one yield per group.  Group names are sorted numerically when
          they look like integers; lexicographically otherwise.

    The h5 path holds the file open for the duration of iteration; consume the
    generator (e.g. via ``list(...)``) before doing anything else if you need
    the handles released sooner.
    """
    if isinstance(source, str | Path):
        path = Path(source)
        if path.suffix.lower() not in (".h5", ".hdf5"):
            raise ValueError(f"Only .h5/.hdf5 paths are supported; got {path.suffix!r}")
        yield from _iter_h5(path)
        return

    if isinstance(source, dict):
        if _looks_like_state_dict(source):
            yield None, source
            return

        # Mapping of label → state_dict
        try:
            keys = sorted(source.keys())
        except TypeError:
            keys = list(source.keys())
        for k in keys:
            sd = source[k]
            if not _looks_like_state_dict(sd):
                raise TypeError(
                    f"value at key {k!r} is not a state dict (got {type(sd).__name__})"
                )
            yield k, sd
        return

    # Anything else iterable: list, tuple, set, generator
    try:
        items = list(source)
    except TypeError as exc:
        raise TypeError(
            f"Cannot iterate state dicts from {type(source).__name__}"
        ) from exc
    for i, sd in enumerate(items):
        if not _looks_like_state_dict(sd):
            raise TypeError(f"item {i} is not a state dict (got {type(sd).__name__})")
        yield i, sd

Save a Partition to a .npz file.

Parameters:

Name Type Description Default
partition Partition

The Partition to serialise.

 required
path str | Path

Destination file path. A .npz extension is appended by np.savez if not already present.

 required
Source code in src/parx/io_partition.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def save_partition(partition: Partition, path: str | Path) -> None:
    """Save a Partition to a .npz file.

    Parameters
    ----------
    partition:
        The Partition to serialise.
    path:
        Destination file path.  A ``.npz`` extension is appended by
        ``np.savez`` if not already present.
    """
    regions = partition.regions
    n_regions = len(regions)

    # Build offsets from the activation-path widths of the first region.
    # All regions share the same layer widths, so region 0 is representative.
    if n_regions == 0:
        # Edge case: empty partition — infer layer widths from weights.
        layer_widths = [w.shape[0] for w in partition.weights]
    else:
        layer_widths = [q.shape[0] for q in regions[0].activation_path]

    n_layers = len(layer_widths)
    total_bits = sum(layer_widths)

    offsets = np.zeros(n_layers + 1, dtype=np.int64)
    for l, w in enumerate(layer_widths):  # noqa: E741
        offsets[l + 1] = offsets[l] + w

    # Reconstruct patterns matrix
    patterns = np.zeros((n_regions, total_bits), dtype=np.int8)
    for i, region in enumerate(regions):
        for l, q in enumerate(region.activation_path):  # noqa: E741
            patterns[i, offsets[l] : offsets[l + 1]] = q.astype(np.int8)

    # Reconstruct centroids matrix
    input_dim = partition.input_dim
    centroids = np.zeros((n_regions, input_dim), dtype=np.float64)
    for i, region in enumerate(regions):
        centroids[i] = region.centroid

    arrays: dict[str, np.ndarray] = {
        "version": np.int64(_FORMAT_VERSION),
        "patterns": patterns,
        "offsets": offsets,
        "centroids": centroids,
    }

    # Weights and biases
    for i, (w, b) in enumerate(zip(partition.weights, partition.biases)):
        arrays[f"weights_{i}"] = np.asarray(w, dtype=np.float64)
        arrays[f"biases_{i}"] = np.asarray(b, dtype=np.float64)

    # Optional: active_indices and bounded (exact-method partitions)
    has_active = any(r.active_indices is not None for r in regions)
    if has_active:
        # Build flat active_indices and active_offsets
        active_parts: list[np.ndarray] = []
        active_offsets = np.zeros(n_regions + 1, dtype=np.int64)
        for i, region in enumerate(regions):
            if region.active_indices is not None:
                part = np.asarray(region.active_indices, dtype=np.int32)
            else:
                part = np.empty(0, dtype=np.int32)
            active_parts.append(part)
            active_offsets[i + 1] = active_offsets[i] + len(part)

        active_indices_flat = (
            np.concatenate(active_parts)
            if active_parts
            else np.empty(0, dtype=np.int32)
        )
        arrays["active_indices_flat"] = active_indices_flat
        arrays["active_offsets"] = active_offsets

    # Always save bounded when there are regions so we can round-trip the flag.
    if n_regions > 0:
        arrays["bounded"] = np.array([r.bounded for r in regions], dtype=bool)

    np.savez(path, **arrays)

Load a Partition from a .npz file produced by :func:save_partition.

Parameters:

Name Type Description Default
path str | Path

Path to the .npz file.

 required

Returns:

Type Description
Partition

Raises:

Type Description
ValueError

If the file uses an unsupported format version.
Source code in src/parx/io_partition.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def load_partition(path: str | Path) -> Partition:
    """Load a Partition from a .npz file produced by :func:`save_partition`.

    Parameters
    ----------
    path:
        Path to the ``.npz`` file.

    Returns
    -------
    Partition

    Raises
    ------
    ValueError
        If the file uses an unsupported format version.
    """
    data = np.load(path, allow_pickle=False)

    version = int(data["version"])
    if version != _FORMAT_VERSION:
        raise ValueError(
            f"Unsupported parx partition format version: {data['version']}"
        )

    patterns = data["patterns"]
    offsets = data["offsets"]
    centroids = data["centroids"]

    # Reconstruct weights and biases
    weights: list[np.ndarray] = []
    biases: list[np.ndarray] = []
    i = 0
    while f"weights_{i}" in data:
        weights.append(data[f"weights_{i}"])
        biases.append(data[f"biases_{i}"])
        i += 1

    # Optional fields
    active_indices_flat: np.ndarray | None = (
        data["active_indices_flat"] if "active_indices_flat" in data else None
    )
    active_offsets: np.ndarray | None = (
        data["active_offsets"] if "active_offsets" in data else None
    )
    bounded: np.ndarray | None = data["bounded"] if "bounded" in data else None

    result = RegionFindResult(
        patterns=patterns,
        offsets=offsets,
        centroids=centroids,
        active_indices_flat=active_indices_flat,
        active_offsets=active_offsets,
        bounded=bounded,
    )
    return Partition.from_result(result, weights, biases)

Utilities

Optional smoke-run helper to amortise Julia's TTFX (time-to-first-X).

Julia compiles each function on first call. In parx this manifests as:

  1. ~1-2 s for the Julia runtime + LinearRegions.jl to load
  2. ~3-5 s the first time JuMP + HiGHS are touched
  3. ~0.2-2 s per method the first time it is invoked

Cost (1) is paid implicitly by ensure_julia(). Cost (2) and (3) are paid by the first call to each Julia-backed method. Total cold start from a fresh process: 5-10 seconds; subsequent calls take 0-1 s.

Call :func:precompile once at the top of a session (or notebook) if you want predictable timing from your first real compute_partition call.

precompile(*, verbose=False)

Run a tiny problem through every Julia-backed method.

Returns a dict of {method: seconds_taken} so callers can confirm the smoke run actually executed. Pure-Python methods are skipped because they do not benefit.

Parameters:

Name Type Description Default
verbose bool

If True, print per-method timing as it runs.

 False
Source code in src/parx/precompile.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def precompile(*, verbose: bool = False) -> dict[str, float]:
    """Run a tiny problem through every Julia-backed method.

    Returns a dict of ``{method: seconds_taken}`` so callers can confirm the
    smoke run actually executed.  Pure-Python methods are skipped because they
    do not benefit.

    Parameters
    ----------
    verbose : bool
        If True, print per-method timing as it runs.
    """
    from parx import compute_partition
    from parx._julia_init import ensure_julia

    ensure_julia()  # cost (1) + (2)

    weights = {
        "0.weight": np.eye(2),
        "0.bias": np.zeros(2),
        "2.weight": np.eye(2),
        "2.bias": np.zeros(2),
    }
    X = np.array([[1.0, 1.0], [-1.0, -1.0]])
    x0 = np.array([1.0, 1.0])

    timings: dict[str, float] = {}
    for method in _JULIA_METHODS:
        data: Any = X if method == "sparse_julia" else x0
        t0 = time.perf_counter()
        compute_partition(weights, data, method=method)
        dt = time.perf_counter() - t0
        timings[method] = dt
        if verbose:
            print(f"  precompiled {method:<20} {dt:.2f}s")
    return timings

All registered method names.

Source code in src/parx/methods/__init__.py 
74
75
76
def list_methods() -> list[str]:
    """All registered method names."""
    return sorted(_METHODS)

Look up a registered method by name.

Source code in src/parx/methods/__init__.py 
67
68
69
70
71
def get_method(name: str) -> MethodFn:
    """Look up a registered method by name."""
    if name not in _METHODS:
        raise ValueError(f"Unknown method {name!r}. Available: {list_methods()}")
    return _METHODS[name]

Report Julia and Python thread settings.

Returns:

Type Description
dict with keys:

julia_threads : Threads.nthreads() inside the live Julia runtime (after env vars take effect). julia_max_threads: Threads.maxthreadid() (Julia ≥ 1.10). env_JULIA_NUM_THREADS: the env var as Python sees it.
Source code in src/parx/diagnostics.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def thread_info() -> dict[str, Any]:
    """Report Julia and Python thread settings.

    Returns
    -------
    dict with keys:
        ``julia_threads``    : ``Threads.nthreads()`` inside the live Julia
                               runtime (after env vars take effect).
        ``julia_max_threads``: ``Threads.maxthreadid()`` (Julia ≥ 1.10).
        ``env_JULIA_NUM_THREADS``: the env var as Python sees it.
    """
    import os

    from parx._julia_init import ensure_julia

    jl = ensure_julia()
    return {
        "julia_threads": int(jl.seval("Threads.nthreads()")),
        "julia_max_threads": int(jl.seval("Threads.maxthreadid()")),
        "env_JULIA_NUM_THREADS": os.environ.get("JULIA_NUM_THREADS"),
    }

Time a region-finding method.

Returns {"method": ..., "n_regions": ..., "best_seconds": ..., "all_seconds": [...]}.

A warm-up run is performed and discarded to amortise JIT compilation.

Source code in src/parx/diagnostics.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def benchmark_method(
    method: str,
    state_dict_or_model,
    data: np.ndarray,
    *,
    repeats: int = 3,
    warmup: int = 1,
) -> dict[str, Any]:
    """Time a region-finding method.

    Returns
    ``{"method": ..., "n_regions": ..., "best_seconds": ..., "all_seconds": [...]}``.

    A warm-up run is performed and discarded to amortise JIT compilation.
    """
    from parx import compute_partition

    for _ in range(warmup):
        compute_partition(state_dict_or_model, data, method=method)

    timings = []
    n_regions = 0
    for _ in range(repeats):
        t0 = time.perf_counter()
        p = compute_partition(state_dict_or_model, data, method=method)
        timings.append(time.perf_counter() - t0)
        n_regions = len(p)

    return {
        "method": method,
        "n_regions": n_regions,
        "best_seconds": min(timings),
        "all_seconds": timings,
    }