Metadata-Version: 2.4
Name: nisplus
Version: 0.1.0
Summary: Neural Information Squeezer Plus core models and utilities.
Requires-Python: >=3.11
Description-Content-Type: text/markdown
Requires-Dist: numpy
Requires-Dist: torch
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"

# 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:

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

Verify the install:

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

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

## Quick Start

```python
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:

```text
exp/nisplus_toy_example.ipynb
```

## Core API

Top-level imports:

```python
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:

```python
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:

```python
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:

```python
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:

```python
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:

```python
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:

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

Main scripts:

```bash
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:

```bash
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

```bash
pytest -q
```

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