nisplus 0.1.0

Neural Information Squeezer Plus core models and utilities.

NIS2026 / nisplus

nisplus is the reusable core library extracted from the NIS2026 research codebase for neural effective information and causal emergence experiments.

It provides:

• NIS+ PyTorch models for micro-to-macro representation learning.
• Training losses for stage-1 prediction/reconstruction and stage-2 EI reweighting.
• Effective-information helper utilities.
• Synthetic SIR and Kuramoto data generators used by the research experiments.

The repository still contains experiment scripts, notebooks, result artifacts, and local data assets, but the importable package lives under src/nisplus/.

Installation

From the repository root:

python -m pip install -e ".[dev]"

Verify the install:

python -c "import nisplus; print(nisplus.NISPlusNN)"
pytest -q

Core runtime dependencies are torch and numpy. The dev extra installs pytest.

Quick Start

import torch

from nisplus import NISPlusNN, stage1_loss

model = NISPlusNN(x_dim=4, h_dim=2, noise_dim=2)
x_t = torch.randn(32, 4)
x_tp1 = torch.randn(32, 4)

out = model(x_t, x_tp1)
loss, metrics = stage1_loss(out, x_t, x_tp1)
loss.backward()

print(out["h_t"].shape)
print(metrics)

For a runnable walkthrough, open:

exp/nisplus_toy_example.ipynb

Core API

Top-level imports:

from nisplus import (
    NISPlusNN,
    NISPlusMicroDynamics,
    stage1_loss,
    stage2_loss,
    multi_step_prediction_loss,
    compute_weights,
    recompute_weights,
    estimate_dei,
    estimate_dei_micro,
    compute_causal_emergence,
)

Model modules:

• NISPlusNN: encoder, decoder, forward latent dynamics, and backward latent dynamics.
• NISPlusMicroDynamics: direct micro-dynamics baseline for causal-emergence comparisons.

Useful model methods:

• model.encode(x): map micro-state x to latent macro-state h.
• model.decode(h, noise=None): decode latent state back to micro-state.
• model.latent_forward(h): predict the next latent state.
• model.latent_backward(h_next): approximate inverse latent dynamics.
• model.multi_step_predict(x_t, n_steps, n_samples): roll forward in latent space and decode predictions.

Data Utilities

SIR toy data:

from nisplus.data.sir import SIRDataset, generate_sir_dataset, get_sir_dataloaders

x_t, x_tp1, y_t, y_tp1 = generate_sir_dataset(
    n_trajectories=100,
    n_steps=3,
    sigma=0.03,
    seed=42,
)
dataset = SIRDataset(x_t, x_tp1, y_t, y_tp1)

Kuramoto toy data:

from nisplus.data.kuramoto import generate_kuramoto_dataset

x_t, x_tp1, y_t, y_tp1, meta = generate_kuramoto_dataset(
    n_trajectories=100,
    n_steps=4,
    group_sizes=(5, 5),
    seed=42,
)

YRD air-quality data and loaders remain repository-local under data/ and are not part of the first-stage nisplus package.

Training Pattern

Stage 1 trains prediction and reconstruction:

optimizer.zero_grad()
out = model(x_t, x_tp1)
loss, metrics = stage1_loss(out, x_t, x_tp1)
loss.backward()
optimizer.step()

Stage 2 adds EI-style density-ratio weighting and backward latent consistency:

from nisplus import compute_weights, stage2_loss

with torch.no_grad():
    h_all = model.encode(all_x_t)
weights = compute_weights(h_all, bandwidth=0.05, max_weight=10.0)

out = model(x_batch, x_next_batch)
loss, metrics = stage2_loss(out, x_batch, x_next_batch, weights, indices)

indices should index each batch sample into the full training-set weight vector. The bundled Dataset classes return idx for this purpose.

Effective Information Helpers

Use compute_weights for latent density-ratio weights and estimate_dei / estimate_dei_micro for dimension-averaged EI estimates:

from nisplus import compute_causal_emergence, estimate_dei, estimate_dei_micro

dei_macro, macro_info = estimate_dei(model, h_t, h_tp1, h_pred, weights)
dei_micro, micro_info = estimate_dei_micro(micro_model, x_t, x_tp1, x_pred)
delta_j = compute_causal_emergence(dei_macro, dei_micro)

EI estimation uses Jacobian calculations and can be slower than ordinary training loss evaluation.

Repository Layout

• src/nisplus/: importable NIS+ core library.
• src/nisplus/data/: bundled synthetic SIR and Kuramoto generators.
• tests/: smoke tests for the packaged API.
• exp/nisplus_toy_example.ipynb: minimal tutorial notebook.
• exp/sir/: SIR experiment scripts, notebooks, state, and runs.
• exp/kuramoto/: Kuramoto experiment scripts, notebooks, state, and runs.
• exp/yrd_air/: YRD air-quality experiment scripts and state.
• data/: local data assets and non-core data loaders.
• results/: generated result artifacts.
• docs/: paper drafts, logs, and references.

Experiment Entrypoints

Install the package in editable mode before running experiment scripts:

python -m pip install -e ".[dev]"

Main scripts:

python exp/sir/train.py
python exp/sir/ood_search.py
python exp/kuramoto/train.py
python exp/kuramoto/scale_search.py

Small SIR smoke run:

python exp/sir/train.py \
  --n_trajectories 8 \
  --n_steps 2 \
  --total_epochs 1 \
  --stage2_start 1 \
  --batch_size 4 \
  --eval_interval 1 \
  --ei_eval_interval 999 \
  --output_dir /tmp/nisplus_sir_smoke \
  --device cpu

Testing

pytest -q

The smoke tests cover package imports, model forward shapes, loss backward passes, toy data generation, and EI helper execution.