Add network_inspectors.py module in LANfactory

.gitignore

@@ -155,4 +155,5 @@ explorations/
 
 # marimo session/layout cache (notebook source .py is the artifact we track)
 __marimo__/
+
 uv.lock

notebooks/network_inspectors_tutorial.ipynb

@@ -0,0 +1,202 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "afbe76fa",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-25T16:45:23.814081Z",
+     "iopub.status.busy": "2026-06-25T16:45:23.814002Z",
+     "iopub.status.idle": "2026-06-25T16:45:27.700947Z",
+     "shell.execute_reply": "2026-06-25T16:45:27.700705Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "from collections.abc import Callable\n",
+    "from pathlib import Path\n",
+    "from typing import Any\n",
+    "\n",
+    "from numpy.typing import NDArray\n",
+    "\n",
+    "import ssms\n",
+    "import lanfactory\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "de69930a",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-25T16:45:27.702385Z",
+     "iopub.status.busy": "2026-06-25T16:45:27.702231Z",
+     "iopub.status.idle": "2026-06-25T16:45:28.124991Z",
+     "shell.execute_reply": "2026-06-25T16:45:28.124767Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# Specify model\n",
+    "model: str = \"angle\"\n",
+    "\n",
+    "model_dir: Path = Path(\"../data/torch_models\") / model\n",
+    "if not model_dir.exists():\n",
+    "    model_dir = Path(\"data/torch_models\") / model\n",
+    "state_dict_path: Path = next(model_dir.glob(\"*state_dict*\"))\n",
+    "network_config_path: Path = next(model_dir.glob(\"*network_config*\"))\n",
+    "\n",
+    "# get_torch_mlp needs the input dimension: params + rt + choice\n",
+    "params: list[str] = ssms.config.model_config[model][\"params\"]\n",
+    "input_dim: int = len(params) + 2\n",
+    "\n",
+    "lan_angle: Callable[[NDArray[np.float32]], Any] = (\n",
+    "    lanfactory.network_inspectors.get_torch_mlp(\n",
+    "        model_file_path=str(state_dict_path),\n",
+    "        network_config=str(network_config_path),\n",
+    "        input_dim=input_dim,\n",
+    "    )\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d183daf6",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-25T16:45:28.126199Z",
+     "iopub.status.busy": "2026-06-25T16:45:28.126064Z",
+     "iopub.status.idle": "2026-06-25T16:45:28.127874Z",
+     "shell.execute_reply": "2026-06-25T16:45:28.127639Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# pick a parameter vector for the model and repeat it across 200 rows\n",
+    "parameter_vector: NDArray[np.float32] = np.array(\n",
+    "    ssms.config.model_config[model][\"default_params\"], dtype=np.float32\n",
+    ")\n",
+    "parameter_matrix: NDArray[np.float32] = np.tile(parameter_vector, (200, 1))\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1e53e164",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-25T16:45:28.128869Z",
+     "iopub.status.busy": "2026-06-25T16:45:28.128807Z",
+     "iopub.status.idle": "2026-06-25T16:45:28.136669Z",
+     "shell.execute_reply": "2026-06-25T16:45:28.136462Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# Initialize network input\n",
+    "network_input: NDArray[np.float32] = np.zeros(\n",
+    "    (parameter_matrix.shape[0], parameter_matrix.shape[1] + 2), dtype=np.float32\n",
+    ")\n",
+    "\n",
+    "# Add reaction times\n",
+    "network_input[:, -2] = np.linspace(0, 3, parameter_matrix.shape[0])\n",
+    "\n",
+    "# Add choices\n",
+    "network_input[:, -1] = np.repeat(np.random.choice([-1, 1]), parameter_matrix.shape[0])\n",
+    "\n",
+    "# Show example output\n",
+    "print(\"Some network outputs\")\n",
+    "print(lan_angle(network_input)[:10])\n",
+    "print(\"Shape\")\n",
+    "print(lan_angle(network_input).shape)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "190c8768",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-25T16:45:28.137668Z",
+     "iopub.status.busy": "2026-06-25T16:45:28.137601Z",
+     "iopub.status.idle": "2026-06-25T16:45:35.655188Z",
+     "shell.execute_reply": "2026-06-25T16:45:35.654976Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# make 10 reproducible random parameter sets within the model's bounds\n",
+    "rng: np.random.Generator = np.random.default_rng(123)\n",
+    "lb, ub = ssms.config.model_config[model][\"param_bounds\"]\n",
+    "parameter_df: pd.DataFrame = pd.DataFrame(\n",
+    "    rng.uniform(lb, ub, size=(10, len(params))), columns=params\n",
+    ")\n",
+    "\n",
+    "lanfactory.network_inspectors.kde_vs_lan_likelihoods(\n",
+    "    parameter_df=parameter_df,\n",
+    "    model=model,\n",
+    "    torch_mlp_predict=lan_angle,\n",
+    "    n_samples=2000,\n",
+    "    n_reps=10,\n",
+    "    plot=lanfactory.network_inspectors.PlotConfig(cols=3, show=True),\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1d98202",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2026-06-25T16:45:35.656289Z",
+     "iopub.status.busy": "2026-06-25T16:45:35.656192Z",
+     "iopub.status.idle": "2026-06-25T16:45:38.024575Z",
+     "shell.execute_reply": "2026-06-25T16:45:38.024320Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# use a deterministic parameter vector for the manifold plot\n",
+    "manifold_parameter_df: pd.DataFrame = pd.DataFrame(\n",
+    "    [ssms.config.model_config[model][\"default_params\"]], columns=params\n",
+    ")\n",
+    "\n",
+    "lanfactory.network_inspectors.lan_manifold(\n",
+    "    parameter_df=manifold_parameter_df,\n",
+    "    vary_dict={\"v\": np.linspace(-2, 2, 20)},\n",
+    "    model=model,\n",
+    "    torch_mlp_predict=lan_angle,\n",
+    "    grid=lanfactory.network_inspectors.GridSpec(n_rt_steps=300, max_rt=5),\n",
+    "    plot=lanfactory.network_inspectors.PlotConfig(\n",
+    "        fig_scale=1.0, save=True, show=True\n",
+    "    ),\n",
+    ")\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "hssm (3.12.12.final.0)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -40,6 +40,7 @@ dependencies = [
     "frozendict>=2.4.6",
     "onnx>=1.17.0",
     "matplotlib>=3.10.1",
+    "seaborn>=0.13.2",
     "typer>=0.9.0",
 ]
 

src/lanfactory/__init__.py

@@ -5,4 +5,16 @@
 from . import utils
 from . import onnx
 
-__all__ = ["config", "trainers", "utils", "onnx"]
+__all__ = ["config", "trainers", "utils", "onnx", "network_inspectors"]
+
+
+def __getattr__(name):
+    # lazy import so `import lanfactory` doesn't require sklearn etc.
+    # (importlib here, since `from . import` would recurse back into __getattr__)
+    if name == "network_inspectors":
+        import importlib
+
+        module = importlib.import_module(f"{__name__}.network_inspectors")
+        globals()[name] = module
+        return module
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/lanfactory/network_inspectors/__init__.py

@@ -0,0 +1,24 @@
+"""Compare LAN vs KDE likelihoods and plot LAN manifolds (LANfactory backend).
+
+Modular layout:
+    loaders   -- build LAN predictors from the LANfactory torch backend
+    config    -- ModelSpec / PlotConfig / GridSpec config objects
+    compute   -- headless RT-grid and likelihood computation
+    plotting  -- rendering of the comparison and manifold figures
+    api       -- thin entry points wiring the above together
+"""
+
+from __future__ import annotations
+
+from .api import kde_vs_lan_likelihoods, lan_manifold
+from .config import GridSpec, ModelSpec, PlotConfig
+from .loaders import get_torch_mlp
+
+__all__ = [
+    "get_torch_mlp",
+    "kde_vs_lan_likelihoods",
+    "lan_manifold",
+    "ModelSpec",
+    "PlotConfig",
+    "GridSpec",
+]

src/lanfactory/network_inspectors/api.py

@@ -0,0 +1,115 @@
+"""Thin entry points wiring config, computation, and plotting together."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+import logging
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from numpy.typing import NDArray
+
+from .compute import (
+    build_manifold,
+    evaluate_kde,
+    evaluate_network,
+    make_manifold_grid,
+    make_rt_choice_grid,
+    simulate_ground_truth,
+)
+from .config import GridSpec, ModelSpec, PlotConfig
+from .plotting import LikelihoodResult, plot_kde_vs_lan, plot_manifold
+
+logger = logging.getLogger(__name__)
+
+
+def kde_vs_lan_likelihoods(
+    parameter_df: pd.DataFrame,
+    model: str,
+    torch_mlp_predict: Callable[[NDArray[np.float32]], Any],
+    n_samples: int = 10,
+    n_reps: int = 10,
+    grid: GridSpec | None = None,
+    plot: PlotConfig | None = None,
+) -> None:
+    """Compare kernel density estimates from simulation data with LAN output.
+
+    parameter_df: one model-compatible parameter vector per row.
+    model: model name. torch_mlp_predict: predict_on_batch from get_torch_mlp.
+    n_samples/n_reps: samples per KDE / KDEs per subplot.
+    grid: optional GridSpec. plot: optional PlotConfig.
+    """
+    if parameter_df is None or model is None or torch_mlp_predict is None:
+        raise ValueError(
+            "parameter_df, model, and torch_mlp_predict are required; build the"
+            " predictor with get_torch_mlp()."
+        )
+    if not isinstance(parameter_df, pd.DataFrame):
+        raise TypeError("parameter_df must be a pandas.DataFrame.")
+
+    spec = ModelSpec.from_model(model, predictor=torch_mlp_predict)
+    cfg = plot or PlotConfig()
+    grid_arr = make_rt_choice_grid(spec, grid)
+
+    results: list[LikelihoodResult] = []
+    for i in range(parameter_df.shape[0]):
+        params = parameter_df.iloc[i, :].values
+        lan_like = np.exp(evaluate_network(spec, params, grid_arr))
+        kdes = [
+            np.exp(
+                evaluate_kde(simulate_ground_truth(spec, params, n_samples), grid_arr)
+            )
+            for _ in range(n_reps)
+        ]
+        results.append({"lan": lan_like, "kdes": kdes})
+
+    return plot_kde_vs_lan(grid_arr, results, spec, cfg)
+
+
+def lan_manifold(
+    parameter_df: pd.DataFrame | np.ndarray | None = None,
+    vary_dict: dict[str, Any] | None = None,
+    model: str = "ddm",
+    torch_mlp_predict: Callable[[NDArray[np.float32]], Any] | None = None,
+    grid: GridSpec | None = None,
+    plot: PlotConfig | None = None,
+) -> None:
+    """Plot LAN likelihoods as a 3D manifold while sweeping one parameter.
+
+    parameter_df: parameter vector (first row used). vary_dict: {param: values}.
+    model: model name. torch_mlp_predict: predict_on_batch from get_torch_mlp.
+    grid: optional GridSpec. plot: optional PlotConfig.
+    """
+    if parameter_df is None or torch_mlp_predict is None:
+        raise ValueError(
+            "parameter_df and torch_mlp_predict are required; build the predictor"
+            " with get_torch_mlp()."
+        )
+    if vary_dict is None:
+        vary_dict = {"v": [-1.0, -0.75, -0.5, -0.25, 0, 0.25, 0.5, 0.75, 1.0]}
+
+    spec = ModelSpec.from_model(model, predictor=torch_mlp_predict)
+
+    if spec.n_choices != 2:
+        raise ValueError(
+            "lan_manifold currently supports only 2-choice models; "
+            f"got {spec.n_choices} choices."
+        )
+
+    if isinstance(parameter_df, pd.DataFrame):
+        if parameter_df.shape[0] > 0:
+            logger.info("Using only the first row of the supplied parameter array.")
+        parameters = np.squeeze(
+            parameter_df.iloc[0, :][spec.params].values.astype(np.float32)
+        )
+    else:
+        parameters = np.asarray(parameter_df, dtype=np.float32)
+
+    vary_name = list(vary_dict.keys())[0]
+    vary_values = np.asarray(vary_dict[vary_name])
+
+    grid_arr = make_manifold_grid(grid or GridSpec())
+    manifold = build_manifold(spec, parameters, vary_name, vary_values, grid_arr)
+
+    return plot_manifold(manifold, spec, vary_name, plot or PlotConfig())