nisplus 0.1.0__py3-none-any.whl

nisplus/__init__.py

@@ -0,0 +1,24 @@
+"""Core public API for NIS+."""
+
+from .ei import (
+    compute_causal_emergence,
+    compute_weights,
+    estimate_dei,
+    estimate_dei_micro,
+    recompute_weights,
+)
+from .losses import multi_step_prediction_loss, stage1_loss, stage2_loss
+from .models import NISPlusMicroDynamics, NISPlusNN
+
+__all__ = [
+    "NISPlusNN",
+    "NISPlusMicroDynamics",
+    "stage1_loss",
+    "stage2_loss",
+    "multi_step_prediction_loss",
+    "compute_weights",
+    "recompute_weights",
+    "estimate_dei",
+    "estimate_dei_micro",
+    "compute_causal_emergence",
+]

nisplus/data/__init__.py

@@ -0,0 +1,41 @@
+"""Synthetic data utilities shipped with NIS+."""
+
+from .kuramoto import (
+    KuramotoDataset,
+    angles_to_micro,
+    build_two_community_coupling,
+    community_order_parameters,
+    generate_kuramoto_dataset,
+    generate_kuramoto_dataset_multistep,
+    get_kuramoto_dataloaders,
+    global_order_parameter,
+    micro_to_community_macro,
+    simulate_kuramoto_trajectory,
+)
+from .sir import (
+    SIRDataset,
+    generate_sir_dataset,
+    get_sir_dataloaders,
+    macro_to_micro,
+    simulate_sir_trajectory,
+    sir_derivatives,
+)
+
+__all__ = [
+    "SIRDataset",
+    "sir_derivatives",
+    "simulate_sir_trajectory",
+    "macro_to_micro",
+    "generate_sir_dataset",
+    "get_sir_dataloaders",
+    "KuramotoDataset",
+    "angles_to_micro",
+    "build_two_community_coupling",
+    "community_order_parameters",
+    "global_order_parameter",
+    "micro_to_community_macro",
+    "simulate_kuramoto_trajectory",
+    "generate_kuramoto_dataset",
+    "generate_kuramoto_dataset_multistep",
+    "get_kuramoto_dataloaders",
+]

nisplus/data/kuramoto.py

@@ -0,0 +1,459 @@
+"""
+Kuramoto data generation for NIS+ experiments.
+
+This module builds a small two-community Kuramoto system:
+- 10 oscillators split into two interacting communities
+- strong intra-community coupling, weaker inter-community coupling
+- micro-state is the concatenated cosine/sine embedding of all phases
+- macro-state is the pair of community order parameters
+
+The construction makes the 2-community scale the intended macro level.
+"""
+
+import math
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+
+def wrap_angles(theta):
+    """Wrap angles to [-pi, pi)."""
+    return (theta + math.pi) % (2 * math.pi) - math.pi
+
+
+def make_group_centers(n_groups, freq_bound=0.45):
+    """Create evenly spaced natural-frequency centers for any number of communities."""
+    if n_groups == 1:
+        return (0.0,)
+    return tuple(np.linspace(-freq_bound, freq_bound, n_groups).tolist())
+
+
+def build_two_community_coupling(group_sizes=(5, 5), intra_coupling=2.2, inter_coupling=0.35):
+    """Create a block-structured Kuramoto coupling matrix."""
+    n_oscillators = sum(group_sizes)
+    coupling = np.full((n_oscillators, n_oscillators), inter_coupling / n_oscillators, dtype=float)
+
+    groups = []
+    start = 0
+    for group_size in group_sizes:
+        group = np.arange(start, start + group_size)
+        groups.append(group)
+        coupling[np.ix_(group, group)] = intra_coupling / n_oscillators
+        start += group_size
+
+    np.fill_diagonal(coupling, 0.0)
+    return coupling, groups
+
+
+def sample_natural_frequencies(groups, base_frequencies=(-0.45, 0.45), jitter=0.05, rng=None):
+    """Sample oscillator frequencies with different means for the two communities."""
+    if rng is None:
+        rng = np.random.default_rng()
+
+    if base_frequencies is None or len(base_frequencies) != len(groups):
+        base_frequencies = make_group_centers(len(groups))
+
+    omegas = np.zeros(sum(len(group) for group in groups), dtype=float)
+    for group, base in zip(groups, base_frequencies):
+        omegas[group] = base + rng.normal(0.0, jitter, size=len(group))
+    return omegas
+
+
+def kuramoto_step(theta, omegas, coupling, dt=0.05, process_noise=0.01, rng=None, groups=None):
+    """Advance one Euler-Maruyama step of the Kuramoto dynamics."""
+    if rng is None:
+        rng = np.random.default_rng()
+
+    if groups is None:
+        phase_diff = theta[None, :] - theta[:, None]
+        interaction = np.sum(coupling * np.sin(phase_diff), axis=1)
+    else:
+        interaction = np.zeros_like(theta)
+        group_order = [np.exp(1j * theta[group]).mean() for group in groups]
+        for g_idx, group in enumerate(groups):
+            theta_g = theta[group]
+            total = np.zeros(len(group), dtype=float)
+            for h_idx, other_group in enumerate(groups):
+                coeff = coupling[group[0], other_group[0]]
+                total += coeff * len(other_group) * np.imag(group_order[h_idx] * np.exp(-1j * theta_g))
+            interaction[group] = total
+
+    noise = math.sqrt(dt) * process_noise * rng.normal(size=theta.shape)
+    theta_next = theta + dt * (omegas + interaction) + noise
+    return wrap_angles(theta_next)
+
+
+def simulate_kuramoto_trajectory(theta0, omegas, coupling, dt=0.05, n_steps=4,
+                                 burn_in=30, sample_stride=4, process_noise=0.01, rng=None,
+                                 groups=None):
+    """
+    Simulate one trajectory and return sampled phase states.
+
+    The returned array has shape (n_steps + 1, n_oscillators).
+    """
+    if rng is None:
+        rng = np.random.default_rng()
+
+    theta = np.asarray(theta0, dtype=float).copy()
+
+    for _ in range(burn_in):
+        theta = kuramoto_step(
+            theta, omegas, coupling, dt=dt, process_noise=process_noise, rng=rng, groups=groups
+        )
+
+    trajectory = [theta.copy()]
+    for _ in range(n_steps):
+        for _ in range(sample_stride):
+            theta = kuramoto_step(
+                theta, omegas, coupling, dt=dt, process_noise=process_noise, rng=rng, groups=groups
+            )
+        trajectory.append(theta.copy())
+
+    return np.stack(trajectory, axis=0)
+
+
+def angles_to_micro(theta, observation_noise=0.0, rng=None):
+    """Map oscillator phases to a cosine/sine micro-state embedding."""
+    if rng is None:
+        rng = np.random.default_rng()
+
+    theta = np.asarray(theta)
+    x = np.concatenate([np.cos(theta), np.sin(theta)], axis=-1)
+    if observation_noise > 0:
+        x = x + rng.normal(0.0, observation_noise, size=x.shape)
+        x = np.clip(x, -1.2, 1.2)
+    return x
+
+
+def community_order_parameters(theta, groups):
+    """
+    Compute community-level order parameters.
+
+    Returns a vector [Re z_1, Im z_1, Re z_2, Im z_2, ...].
+    """
+    theta = np.asarray(theta)
+    features = []
+    for group in groups:
+        theta_g = theta[..., group]
+        features.append(np.cos(theta_g).mean(axis=-1))
+        features.append(np.sin(theta_g).mean(axis=-1))
+    return np.stack(features, axis=-1)
+
+
+def global_order_parameter(theta):
+    """Compute the global complex order parameter [Re z, Im z]."""
+    theta = np.asarray(theta)
+    return np.stack([np.cos(theta).mean(axis=-1), np.sin(theta).mean(axis=-1)], axis=-1)
+
+
+def micro_to_community_macro(x, groups):
+    """Recover community-level macro features from cosine/sine micro-states."""
+    x = np.asarray(x)
+    n_oscillators = x.shape[-1] // 2
+    cos_part = x[..., :n_oscillators]
+    sin_part = x[..., n_oscillators:]
+
+    features = []
+    for group in groups:
+        features.append(cos_part[..., group].mean(axis=-1))
+        features.append(sin_part[..., group].mean(axis=-1))
+    return np.stack(features, axis=-1)
+
+
+def sample_initial_phases(groups, rng, phase_gap_range=(0.9, 2.1), init_spread=0.35):
+    """Sample community-structured initial phases."""
+    base_phase = rng.uniform(-math.pi, math.pi)
+    gap = rng.uniform(*phase_gap_range) * rng.choice([-1.0, 1.0])
+    centers = [base_phase, wrap_angles(base_phase + gap)]
+
+    theta0 = np.zeros(sum(len(group) for group in groups), dtype=float)
+    for center, group in zip(centers, groups):
+        theta0[group] = wrap_angles(center + rng.normal(0.0, init_spread, size=len(group)))
+    return theta0
+
+
+def generate_kuramoto_dataset(n_trajectories=500, n_steps=4, group_sizes=(5, 5),
+                              intra_coupling=3.5, inter_coupling=0.45,
+                              base_frequencies=(-0.45, 0.45), omega_jitter=0.05,
+                              dt=0.05, burn_in=40, sample_stride=3, process_noise=0.05,
+                              init_spread=0.35, phase_gap_range=(0.9, 2.1),
+                              observation_noise=0.08, seed=42):
+    """
+    Generate a two-community Kuramoto dataset for one-step prediction.
+
+    Returns:
+        x_t, x_tp1: micro-states, shape (N, 2 * n_oscillators)
+        y_t, y_tp1: community macro-states, shape (N, 2 * n_groups)
+        meta: system metadata for visualization / evaluation
+    """
+    rng = np.random.default_rng(seed)
+
+    coupling, groups = build_two_community_coupling(
+        group_sizes=group_sizes,
+        intra_coupling=intra_coupling,
+        inter_coupling=inter_coupling,
+    )
+    omegas = sample_natural_frequencies(
+        groups,
+        base_frequencies=base_frequencies,
+        jitter=omega_jitter,
+        rng=rng,
+    )
+
+    x_t_all = []
+    x_tp1_all = []
+    y_t_all = []
+    y_tp1_all = []
+    theta_t_all = []
+    theta_tp1_all = []
+
+    for _ in range(n_trajectories):
+        theta0 = sample_initial_phases(
+            groups,
+            rng=rng,
+            phase_gap_range=phase_gap_range,
+            init_spread=init_spread,
+        )
+        trajectory = simulate_kuramoto_trajectory(
+            theta0,
+            omegas,
+            coupling,
+            dt=dt,
+            n_steps=n_steps,
+            burn_in=burn_in,
+            sample_stride=sample_stride,
+            process_noise=process_noise,
+            rng=rng,
+            groups=groups,
+        )
+
+        for t in range(n_steps):
+            theta_t = trajectory[t]
+            theta_tp1 = trajectory[t + 1]
+
+            x_t_all.append(angles_to_micro(theta_t, observation_noise=observation_noise, rng=rng))
+            x_tp1_all.append(angles_to_micro(theta_tp1, observation_noise=observation_noise, rng=rng))
+            y_t_all.append(community_order_parameters(theta_t, groups))
+            y_tp1_all.append(community_order_parameters(theta_tp1, groups))
+            theta_t_all.append(theta_t)
+            theta_tp1_all.append(theta_tp1)
+
+    meta = {
+        "groups": [group.tolist() for group in groups],
+        "omegas": omegas,
+        "coupling": coupling,
+        "global_order_t": global_order_parameter(np.stack(theta_t_all, axis=0)),
+        "global_order_tp1": global_order_parameter(np.stack(theta_tp1_all, axis=0)),
+        "theta_t": np.stack(theta_t_all, axis=0),
+        "theta_tp1": np.stack(theta_tp1_all, axis=0),
+    }
+
+    return (
+        np.stack(x_t_all, axis=0),
+        np.stack(x_tp1_all, axis=0),
+        np.stack(y_t_all, axis=0),
+        np.stack(y_tp1_all, axis=0),
+        meta,
+    )
+
+
+class KuramotoDataset(Dataset):
+    """PyTorch dataset for Kuramoto transition pairs."""
+
+    def __init__(self, x_t, x_tp1, y_t=None, y_tp1=None, x_future=None, y_future=None):
+        self.x_t = torch.FloatTensor(x_t)
+        self.x_tp1 = torch.FloatTensor(x_tp1)
+        self.y_t = torch.FloatTensor(y_t) if y_t is not None else None
+        self.y_tp1 = torch.FloatTensor(y_tp1) if y_tp1 is not None else None
+        self.x_future = torch.FloatTensor(x_future) if x_future is not None else None
+        self.y_future = torch.FloatTensor(y_future) if y_future is not None else None
+
+    def __len__(self):
+        return len(self.x_t)
+
+    def __getitem__(self, idx):
+        item = {
+            "x_t": self.x_t[idx],
+            "x_tp1": self.x_tp1[idx],
+            "idx": idx,
+        }
+        if self.y_t is not None:
+            item["y_t"] = self.y_t[idx]
+        if self.y_tp1 is not None:
+            item["y_tp1"] = self.y_tp1[idx]
+        if self.x_future is not None:
+            item["x_future"] = self.x_future[idx]
+        if self.y_future is not None:
+            item["y_future"] = self.y_future[idx]
+        return item
+
+
+def generate_kuramoto_dataset_multistep(n_trajectories=500, n_steps=6, group_sizes=(5, 5),
+                                        intra_coupling=3.5, inter_coupling=0.45,
+                                        base_frequencies=(-0.45, 0.45), omega_jitter=0.05,
+                                        dt=0.05, burn_in=40, sample_stride=3, process_noise=0.05,
+                                        init_spread=0.35, phase_gap_range=(0.9, 2.1),
+                                        observation_noise=0.08, multi_horizon=1, seed=42):
+    """
+    Generate a Kuramoto dataset with optional multi-step future windows.
+
+    Returns:
+        x_t, x_tp1, y_t, y_tp1, x_future, y_future, meta
+    """
+    rng = np.random.default_rng(seed)
+    multi_horizon = max(int(multi_horizon), 1)
+    if n_steps < multi_horizon:
+        raise ValueError(f"n_steps={n_steps} must be >= multi_horizon={multi_horizon}")
+
+    coupling, groups = build_two_community_coupling(
+        group_sizes=group_sizes,
+        intra_coupling=intra_coupling,
+        inter_coupling=inter_coupling,
+    )
+    omegas = sample_natural_frequencies(
+        groups,
+        base_frequencies=base_frequencies,
+        jitter=omega_jitter,
+        rng=rng,
+    )
+
+    x_t_all = []
+    x_tp1_all = []
+    y_t_all = []
+    y_tp1_all = []
+    x_future_all = []
+    y_future_all = []
+    theta_t_all = []
+    theta_tp1_all = []
+
+    for _ in range(n_trajectories):
+        theta0 = sample_initial_phases(
+            groups,
+            rng=rng,
+            phase_gap_range=phase_gap_range,
+            init_spread=init_spread,
+        )
+        trajectory = simulate_kuramoto_trajectory(
+            theta0,
+            omegas,
+            coupling,
+            dt=dt,
+            n_steps=n_steps,
+            burn_in=burn_in,
+            sample_stride=sample_stride,
+            process_noise=process_noise,
+            rng=rng,
+            groups=groups,
+        )
+
+        x_traj = np.stack(
+            [angles_to_micro(theta, observation_noise=observation_noise, rng=rng) for theta in trajectory],
+            axis=0,
+        )
+        y_traj = community_order_parameters(trajectory, groups)
+
+        max_start = n_steps - multi_horizon + 1
+        for t in range(max_start):
+            x_t_all.append(x_traj[t])
+            x_tp1_all.append(x_traj[t + 1])
+            y_t_all.append(y_traj[t])
+            y_tp1_all.append(y_traj[t + 1])
+            x_future_all.append(x_traj[t + 1:t + 1 + multi_horizon])
+            y_future_all.append(y_traj[t + 1:t + 1 + multi_horizon])
+            theta_t_all.append(trajectory[t])
+            theta_tp1_all.append(trajectory[t + 1])
+
+    meta = {
+        "groups": [group.tolist() for group in groups],
+        "group_sizes": list(group_sizes),
+        "n_groups": len(groups),
+        "n_oscillators": int(sum(group_sizes)),
+        "omegas": omegas,
+        "coupling": coupling,
+        "multi_horizon": multi_horizon,
+        "global_order_t": global_order_parameter(np.stack(theta_t_all, axis=0)),
+        "global_order_tp1": global_order_parameter(np.stack(theta_tp1_all, axis=0)),
+        "theta_t": np.stack(theta_t_all, axis=0),
+        "theta_tp1": np.stack(theta_tp1_all, axis=0),
+    }
+
+    return (
+        np.stack(x_t_all, axis=0),
+        np.stack(x_tp1_all, axis=0),
+        np.stack(y_t_all, axis=0),
+        np.stack(y_tp1_all, axis=0),
+        np.stack(x_future_all, axis=0),
+        np.stack(y_future_all, axis=0),
+        meta,
+    )
+
+
+def get_kuramoto_dataloaders(n_trajectories=500, n_steps=4, batch_size=128,
+                             test_ratio=0.1, seed=42, multi_horizon=1, **dataset_kwargs):
+    """Build train/test dataloaders and return dataset metadata."""
+    x_t, x_tp1, y_t, y_tp1, x_future, y_future, meta = generate_kuramoto_dataset_multistep(
+        n_trajectories=n_trajectories,
+        n_steps=n_steps,
+        multi_horizon=multi_horizon,
+        seed=seed,
+        **dataset_kwargs,
+    )
+
+    n_total = len(x_t)
+    n_test = int(n_total * test_ratio)
+    n_train = n_total - n_test
+    indices = np.random.default_rng(seed).permutation(n_total)
+    train_idx = indices[:n_train]
+    test_idx = indices[n_train:]
+
+    train_dataset = KuramotoDataset(
+        x_t[train_idx], x_tp1[train_idx], y_t[train_idx], y_tp1[train_idx],
+        x_future[train_idx], y_future[train_idx],
+    )
+    test_dataset = KuramotoDataset(
+        x_t[test_idx], x_tp1[test_idx], y_t[test_idx], y_tp1[test_idx],
+        x_future[test_idx], y_future[test_idx],
+    )
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+
+    info = {
+        "x_dim": x_t.shape[1],
+        "y_dim": y_t.shape[1],
+        "n_train": n_train,
+        "n_test": n_test,
+        "train_x_t": torch.FloatTensor(x_t[train_idx]),
+        "train_x_tp1": torch.FloatTensor(x_tp1[train_idx]),
+        "train_y_t": torch.FloatTensor(y_t[train_idx]),
+        "train_y_tp1": torch.FloatTensor(y_tp1[train_idx]),
+        "train_x_future": torch.FloatTensor(x_future[train_idx]),
+        "train_y_future": torch.FloatTensor(y_future[train_idx]),
+        "test_x_t": torch.FloatTensor(x_t[test_idx]),
+        "test_x_tp1": torch.FloatTensor(x_tp1[test_idx]),
+        "test_y_t": torch.FloatTensor(y_t[test_idx]),
+        "test_y_tp1": torch.FloatTensor(y_tp1[test_idx]),
+        "test_x_future": torch.FloatTensor(x_future[test_idx]),
+        "test_y_future": torch.FloatTensor(y_future[test_idx]),
+        "full_x_t": torch.FloatTensor(x_t),
+        "full_x_tp1": torch.FloatTensor(x_tp1),
+        "full_y_t": torch.FloatTensor(y_t),
+        "full_y_tp1": torch.FloatTensor(y_tp1),
+        "full_x_future": torch.FloatTensor(x_future),
+        "full_y_future": torch.FloatTensor(y_future),
+        "meta": meta,
+    }
+
+    return train_loader, test_loader, info
+
+
+if __name__ == "__main__":
+    train_loader, test_loader, info = get_kuramoto_dataloaders(
+        n_trajectories=32, n_steps=4, batch_size=16, multi_horizon=2
+    )
+    print(f"x_dim={info['x_dim']}, y_dim={info['y_dim']}")
+    print(f"n_train={info['n_train']}, n_test={info['n_test']}")
+    batch = next(iter(train_loader))
+    print(f"x_t batch shape: {batch['x_t'].shape}")
+    print(f"x_future batch shape: {batch['x_future'].shape}")