code/batch_generate_v8.py

"""
CAISc Paper — v8: Production-ready with resumption + L40S + vmapped ACh sweep

KEY FEATURES:
  - L40S GPU ($1.95/hr) — best cost/performance
  - vmap across 11 ACh levels per circuit (~20s/circuit)
  - 500 circuits per container (fits in 7200s timeout)
  - RESUMPTION: checks Volume for existing circuit files, skips completed ones
  - Incremental saves + volume commits every 10 circuits
  - Designed for auto-resume observer (run_production.py)

Usage:
  modal run batch_generate_v8.py --n-circuits 5000 --ach-model v14b --batch-name v14b_5k --start-id 20000
  modal run batch_generate_v8.py --n-circuits 5000 --ach-model v14c --batch-name v14c_5k --start-id 30000
"""

import modal
import json
import math
import time

app = modal.App("caisc-prod-v8")
vol = modal.Volume.from_name("caisc-data", create_if_missing=True)
VOL_PATH = "/data"

gpu_image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "jax[cuda12]>=0.4.30,<0.5",
        "jaxley",
        "numpy",
        "scipy",
        "h5py",
    )
)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
ACH_LEVELS = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
N_ACH = len(ACH_LEVELS)
N_EXC = 160
N_INH = 40
N_TOTAL = N_EXC + N_INH
OU_MU = -0.001
OU_SIGMA = 0.001
OU_TAU = 5.0
EXC_GS = 0.000005
INH_GS = 0.00003
E_SYN_EXC = 0.0
E_SYN_INH = -80.0

def ach_mu_v14a(ach): return 0.0
def ach_syn_v14a(ach): return max(0.05, math.exp(-2.3 * ach))
def ach_mu_v14b(ach): return 0.0002 / (1.0 + math.exp(-10.0 * (ach - 0.3)))
def ach_syn_v14b(ach): return 1.0
def ach_mu_v14c(ach): return 0.0002 / (1.0 + math.exp(-10.0 * (ach - 0.3)))
def ach_syn_v14c(ach): return max(0.05, math.exp(-2.3 * ach))

ACH_MODELS = {
    "v14a": (ach_mu_v14a, ach_syn_v14a),
    "v14b": (ach_mu_v14b, ach_syn_v14b),
    "v14c": (ach_mu_v14c, ach_syn_v14c),
}

CIRCUITS_PER_CONTAINER = 500


@app.function(
    image=gpu_image,
    gpu="L40S",
    memory=32768,
    timeout=7200,  # 2hr — 500 circuits × 20s = 10000s ≈ 2.8hr, tight but with skip logic should be fine
    volumes={VOL_PATH: vol},
    retries=1,
)
def simulate_container(circuit_ids: list[int], conn_probs: list[float],
                       seeds: list[int], sim_duration_ms: float,
                       batch_name: str, ach_model: str = "v14a") -> list[dict]:
    """Process circuits with vmapped ACh sweep. Skips already-completed circuits."""
    import logging
    logging.getLogger("jaxley").setLevel(logging.ERROR)
    logging.getLogger("jax").setLevel(logging.WARNING)
    
    import builtins
    _real_print = builtins.print
    def _quiet_print(*args, **kwargs):
        msg = str(args[0]) if args else ""
        if "Added" in msg and ("external_states" in msg or "recordings" in msg):
            return
        _real_print(*args, **kwargs)
    builtins.print = _quiet_print
    
    import jax
    import jax.numpy as jnp
    import jaxley as jx
    from jaxley.channels import HH
    from jaxley.synapses import IonotropicSynapse
    from jaxley.connect import connectivity_matrix_connect
    from jaxley.integrate import build_init_and_step_fn
    from jaxley.utils.cell_utils import params_to_pstate
    from jaxley.utils.jax_utils import nested_checkpoint_scan
    import numpy as np
    import copy
    import time as t
    import json as json_mod
    import os
    
    _real_print(f"[Container] {len(circuit_ids)} circuits assigned, JAX: {jax.devices()}")
    
    # =====================================================
    # RESUMPTION: Check which circuits are already done
    # =====================================================
    outdir = f"{VOL_PATH}/sims/{batch_name}/circuits"
    os.makedirs(outdir, exist_ok=True)
    vol.reload()  # Refresh volume to see latest files
    
    existing = set()
    for fname in os.listdir(outdir):
        if fname.startswith("circuit_") and fname.endswith(".json"):
            try:
                cid = int(fname.replace("circuit_", "").replace(".json", ""))
                existing.add(cid)
            except ValueError:
                pass
    
    # Filter to only circuits not yet completed
    todo = [(cid, cp, seed) for cid, cp, seed in zip(circuit_ids, conn_probs, seeds)
            if cid not in existing]
    
    skipped = len(circuit_ids) - len(todo)
    if skipped > 0:
        _real_print(f"  SKIPPING {skipped} already-completed circuits")
    if not todo:
        _real_print(f"  All {len(circuit_ids)} circuits already done!")
        return [{"circuit_id": cid, "build_s": 0, "sim_vmap_s": 0, "total_s": 0, "n_syn": 0, "skipped": True}
                for cid in circuit_ids]
    
    _real_print(f"  TODO: {len(todo)} circuits to simulate")
    
    mu_fn, syn_fn = ACH_MODELS[ach_model]
    dt = 0.025
    n_steps = int(sim_duration_ms / dt)
    stim_len = n_steps + 1
    
    @jax.jit
    def generate_ou_all(key, mu, sigma):
        dt_s = dt / 1000.0
        tau_s = OU_TAU / 1000.0
        noise_coeff = sigma * jnp.sqrt(2.0 * dt_s / tau_s)
        decay = dt_s / tau_s
        def ou_step(x, noise_val):
            x_new = x + decay * (mu - x) + noise_coeff * noise_val
            return x_new, x_new
        keys = jax.random.split(key, N_TOTAL)
        def single_ou(subkey):
            noise = jax.random.normal(subkey, (n_steps,))
            _, trace = jax.lax.scan(ou_step, mu, noise)
            return trace
        return jax.vmap(single_ou)(keys)
    
    # Build template cells ONCE
    t0_build = t.time()
    template_cells = []
    for i in range(N_TOTAL):
        comp = jx.Compartment()
        comp.insert(HH())
        template_cells.append(comp)
    t_cell_build = t.time() - t0_build
    _real_print(f"  Cell template: {t_cell_build:.1f}s")
    
    ach_mu_effs = [OU_MU + mu_fn(ach) for ach in ACH_LEVELS]
    ach_syn_scales = [syn_fn(ach) for ach in ACH_LEVELS]
    
    all_results = []
    
    for idx, (cid, cp, seed) in enumerate(todo):
        t_start = t.time()
        rng = np.random.RandomState(seed)
        
        net = jx.Network(copy.deepcopy(template_cells))
        adjacency = rng.random((N_TOTAL, N_TOTAL)) < cp
        np.fill_diagonal(adjacency, False)
        
        connectivity_matrix_connect(
            net.cell("all"), net.cell("all"),
            IonotropicSynapse(), adjacency,
        )
        
        syn_count = int(adjacency.sum())
        if syn_count > 0:
            pre_from_adj, _ = np.where(adjacency)
            net.edges["IonotropicSynapse_gS"] = np.where(pre_from_adj < N_EXC, EXC_GS, INH_GS)
            net.edges["IonotropicSynapse_e_syn"] = np.where(pre_from_adj < N_EXC, E_SYN_EXC, E_SYN_INH)
            net.edges["IonotropicSynapse_s"] = 0.0
        
        t_build = t.time() - t_start
        mean_in_degree = max(1.0, syn_count / N_TOTAL)
        jax_key = jax.random.PRNGKey(seed)
        
        # VMAPPED ACh SWEEP
        t_sim_start = t.time()
        
        net.cell("all").record("v", verbose=False)
        dummy_stim = jnp.zeros((N_TOTAL, stim_len))
        net.cell("all").stimulate(dummy_stim, verbose=False)
        net.to_jax()
        
        rec_inds = net.recordings.rec_index.to_numpy()
        rec_states = net.recordings.state.to_numpy()
        externals_base = net.externals.copy()
        external_inds_base = net.external_inds.copy()
        for key in externals_base:
            externals_base[key] = externals_base[key].T
        
        init_fn, step_fn = build_init_and_step_fn(net)
        params = net.get_parameters()
        pstate_base = params_to_pstate(params, net.indices_set_by_trainables)
        all_params_base = net.get_all_parameters(pstate_base)
        all_states_base = net.get_all_states(pstate_base)
        all_states_base = net.append_channel_currents_to_states(all_states_base, all_params_base, dt)
        
        exc_syn_mask = (net.edges["IonotropicSynapse_e_syn"] == E_SYN_EXC).to_numpy()
        
        # Build batched params
        batched_params = {}
        for k, v in all_params_base.items():
            if k == "axial_conductances":
                batched_params[k] = {}
                for sub_k, sub_v in v.items():
                    batched_params[k][sub_k] = jnp.stack([sub_v] * N_ACH)
            else:
                batched_params[k] = jnp.stack([v] * N_ACH)
        
        if syn_count > 0:
            for ach_idx in range(N_ACH):
                syn_scale = ach_syn_scales[ach_idx]
                gs_modified = jnp.where(jnp.array(exc_syn_mask), EXC_GS * syn_scale, all_params_base["IonotropicSynapse_gS"])
                batched_params["IonotropicSynapse_gS"] = batched_params["IonotropicSynapse_gS"].at[ach_idx].set(gs_modified)
        
        batched_states = {k: jnp.stack([v] * N_ACH) for k, v in all_states_base.items()}
        
        batched_externals = {}
        for key in externals_base:
            if key == "i":
                all_stims = []
                for ach_idx in range(N_ACH):
                    ach_key = jax.random.fold_in(jax_key, ach_idx)
                    ou_traces = generate_ou_all(ach_key, ach_mu_effs[ach_idx], OU_SIGMA)
                    ou_np = np.array(ou_traces)
                    if ou_np.shape[1] < stim_len:
                        ou_np = np.concatenate([ou_np, np.zeros((N_TOTAL, stim_len - ou_np.shape[1]))], axis=1)
                    else:
                        ou_np = ou_np[:, :stim_len]
                    all_stims.append(jnp.array(ou_np.T))
                batched_externals[key] = jnp.stack(all_stims)
            else:
                batched_externals[key] = jnp.stack([externals_base[key]] * N_ACH)
        
        nsteps_to_return = stim_len
        
        def single_ach_simulate(single_params, single_states, single_externals):
            def _body_fun(state, ext_slice):
                state = step_fn(state, single_params, ext_slice, external_inds_base, dt)
                recs = jnp.asarray([state[rec_state][rec_ind] for rec_state, rec_ind in zip(rec_states, rec_inds)])
                return state, recs
            init_recs = jnp.asarray([single_states[rec_state][rec_ind] for rec_state, rec_ind in zip(rec_states, rec_inds)])
            init_recording = jnp.expand_dims(init_recs, axis=0)
            final_state, recordings = nested_checkpoint_scan(
                _body_fun, single_states, single_externals,
                length=nsteps_to_return, nested_lengths=[nsteps_to_return],
            )
            recs = jnp.concatenate([init_recording, recordings[:nsteps_to_return]], axis=0).T
            return recs
        
        vmapped_simulate = jax.vmap(single_ach_simulate)
        all_voltages = vmapped_simulate(batched_params, batched_states, batched_externals)
        
        t_sim_total = t.time() - t_sim_start
        
        # SPIKE EXTRACTION + STATS
        results = []
        warmup_steps = int(200.0 / dt)
        stable_ms = sim_duration_ms - 200.0
        
        for ach_idx, ach in enumerate(ACH_LEVELS):
            v_np = np.array(all_voltages[ach_idx])
            spike_trains = []
            for i in range(N_TOTAL):
                v_s = v_np[i, warmup_steps:]
                crossings = np.where((v_s[:-1] < 0.0) & (v_s[1:] >= 0.0))[0]
                spike_trains.append(((crossings + warmup_steps) * dt).tolist())
            
            stats = compute_statistics(spike_trains, N_EXC, N_INH, stable_ms)
            syn_scale = ach_syn_scales[ach_idx]
            mu_eff = ach_mu_effs[ach_idx]
            
            results.append({
                "circuit_id": cid, "ach_level": ach, "ach_model": ach_model,
                "n_exc": N_EXC, "n_inh": N_INH, "n_total": N_TOTAL,
                "conn_prob": cp, "n_synapses": syn_count,
                "mean_in_degree": round(mean_in_degree, 1),
                "gS_exc_effective": round(EXC_GS * syn_scale, 8),
                "ou_mu_effective": round(mu_eff, 6),
                "ou_sigma_effective": round(OU_SIGMA, 6),
                "ou_tau": OU_TAU, "sim_duration_ms": sim_duration_ms,
                "seed": seed, "sim_time_s": round(t_sim_total / N_ACH, 2),
                "statistics": stats,
            })
        
        total_time = t.time() - t_start
        
        # Save incrementally
        with open(f"{outdir}/circuit_{cid:05d}.json", "w") as f:
            json_mod.dump(results, f)
        if idx % 10 == 0 or idx == len(todo) - 1:
            vol.commit()
        
        all_results.append({
            "circuit_id": cid, "build_s": round(t_build, 1),
            "sim_vmap_s": round(t_sim_total, 1),
            "total_s": round(total_time, 1), "n_syn": syn_count,
        })
        
        if idx % 10 == 0 or idx == len(todo) - 1:
            _real_print(f"  [{idx+1}/{len(todo)}] cid={cid} build={t_build:.1f}s "
                       f"vmap={t_sim_total:.1f}s total={total_time:.1f}s")
    
    vol.commit()
    builtins.print = _real_print
    return all_results


def compute_statistics(spike_trains, n_exc, n_inh, duration_ms):
    import numpy as np
    n_total = n_exc + n_inh
    duration_s = duration_ms / 1000.0
    
    rates = [len(st) / duration_s for st in spike_trains]
    mean_rate = float(np.mean(rates)) if rates else 0.0
    mean_exc_rate = float(np.mean(rates[:n_exc])) if n_exc > 0 else 0.0
    mean_inh_rate = float(np.mean(rates[n_exc:])) if n_inh > 0 else 0.0
    
    cv_isis = []
    for st in spike_trains:
        if len(st) > 2:
            isis = np.diff(st)
            if np.mean(isis) > 0:
                cv_isis.append(float(np.std(isis) / np.mean(isis)))
    mean_cv_isi = float(np.mean(cv_isis)) if cv_isis else 0.0
    
    bin_size_ms = 50.0
    n_bins = int(duration_ms / bin_size_ms)
    mean_fano = 1.0
    if n_bins > 0:
        bin_counts = np.zeros((n_total, n_bins))
        for i, st in enumerate(spike_trains):
            for ts in st:
                bin_counts[i, min(int(ts / bin_size_ms), n_bins - 1)] += 1
        ff = [float(np.var(bin_counts[i]) / np.mean(bin_counts[i])) 
              for i in range(n_total) if np.mean(bin_counts[i]) > 0]
        mean_fano = float(np.mean(ff)) if ff else 1.0
    
    pop_bin_ms = 5.0
    n_pop_bins = int(duration_ms / pop_bin_ms)
    synchrony_index = 0.0
    pop_rate = np.zeros(max(1, n_pop_bins))
    if n_pop_bins > 0:
        for st in spike_trains:
            for ts in st:
                pop_rate[min(int(ts / pop_bin_ms), n_pop_bins - 1)] += 1
        pop_rate /= n_total
        pm = np.mean(pop_rate)
        synchrony_index = float(np.var(pop_rate) / pm) if pm > 0 else 0.0
    
    n_sample = min(100, n_exc * (n_exc - 1) // 2)
    mean_pairwise_corr = 0.0
    if n_sample > 0 and n_pop_bins > 0:
        corr_bin_ms = 10.0
        n_corr_bins = int(duration_ms / corr_bin_ms)
        binned = np.zeros((n_exc, n_corr_bins))
        for i in range(n_exc):
            for ts in spike_trains[i]:
                binned[i, min(int(ts / corr_bin_ms), n_corr_bins - 1)] += 1
        pair_rng = np.random.RandomState(0)
        corrs = []
        for _ in range(n_sample):
            i, j = pair_rng.choice(n_exc, 2, replace=False)
            if np.std(binned[i]) > 0 and np.std(binned[j]) > 0:
                r = float(np.corrcoef(binned[i], binned[j])[0, 1])
                if not np.isnan(r): corrs.append(r)
        mean_pairwise_corr = float(np.mean(corrs)) if corrs else 0.0
    
    peak_freq, total_power = 0.0, 0.0
    if n_pop_bins > 10:
        from scipy import signal
        freqs, psd = signal.welch(pop_rate, fs=1000.0/pop_bin_ms, nperseg=min(256, n_pop_bins))
        peak_freq = float(freqs[np.argmax(psd)])
        total_power = float(np.sum(psd))
    
    return {
        "mean_firing_rate": round(mean_rate, 3),
        "mean_exc_rate": round(mean_exc_rate, 3),
        "mean_inh_rate": round(mean_inh_rate, 3),
        "mean_cv_isi": round(mean_cv_isi, 3),
        "mean_fano_factor": round(mean_fano, 3),
        "synchrony_index": round(synchrony_index, 4),
        "mean_pairwise_corr": round(mean_pairwise_corr, 4),
        "peak_frequency_hz": round(peak_freq, 3),
        "total_spectral_power": round(total_power, 6),
        "n_active_neurons": sum(1 for r in rates if r > 0.5),
        "total_spikes": sum(len(st) for st in spike_trains),
    }


@app.local_entrypoint()
def main(
    n_circuits: int = 5000,
    start_id: int = 0,
    sim_duration_ms: float = 3000.0,
    batch_name: str = "prod_v8",
    ach_model: str = "v14a",
    circuits_per_gpu: int = CIRCUITS_PER_CONTAINER,
):
    import time as t
    import random
    import numpy as np
    
    random.seed(42 + start_id)
    
    all_cids, all_cps, all_seeds = [], [], []
    for i in range(n_circuits):
        cid = start_id + i
        all_cids.append(cid)
        all_cps.append(round(random.uniform(0.04, 0.10), 4))
        all_seeds.append(cid * 1000 + random.randint(0, 999))
    
    # Chunk into containers
    batches_cids, batches_cps, batches_seeds = [], [], []
    for i in range(0, n_circuits, circuits_per_gpu):
        batches_cids.append(all_cids[i:i+circuits_per_gpu])
        batches_cps.append(all_cps[i:i+circuits_per_gpu])
        batches_seeds.append(all_seeds[i:i+circuits_per_gpu])
    
    nc = len(batches_cids)
    
    print(f"\n{'='*70}")
    print(f"PRODUCTION v8 — L40S + vmapped + resumable")
    print(f"  ACh: {ach_model} | {n_circuits} circuits in {nc} containers ({circuits_per_gpu}/ea)")
    print(f"  IDs: {start_id}..{start_id+n_circuits-1} | Batch: {batch_name}")
    print(f"{'='*70}\n")
    
    t0 = t.time()
    all_results = []
    failed = 0
    
    for i, result in enumerate(simulate_container.map(
        batches_cids, batches_cps, batches_seeds,
        [sim_duration_ms] * nc, [batch_name] * nc, [ach_model] * nc,
        order_outputs=False, return_exceptions=True,
    )):
        if isinstance(result, Exception):
            failed += 1
            print(f"  FAILED container {i+1}/{nc}: {result}")
        else:
            all_results.extend(result)
        elapsed = t.time() - t0
        done_circuits = len(all_results)
        rate = done_circuits / elapsed * 3600 if elapsed > 0 else 0
        print(f"  Container {i+1}/{nc}: {done_circuits} total [{elapsed:.0f}s, ~{rate:.0f}/hr]")
    
    elapsed = t.time() - t0
    ok = len([r for r in all_results if not r.get("skipped")])
    total = len(all_results)
    
    print(f"\n{'='*70}")
    print(f"DONE: {ok} simulated, {total-ok} skipped, {failed} failed")
    print(f"  Wall time: {elapsed:.0f}s ({elapsed/60:.1f}min)")
    
    if ok > 0:
        active = [r for r in all_results if not r.get("skipped")]
        totals = [r["total_s"] for r in active]
        vmap_times = [r["sim_vmap_s"] for r in active]
        thr = ok / elapsed * 3600 if elapsed > 0 else 0
        
        print(f"  Build: {np.mean([r['build_s'] for r in active]):.1f}s avg")
        print(f"  Vmap sim: {np.mean(vmap_times):.1f}s avg")
        print(f"  Total/circuit: {np.mean(totals):.0f}s")
        print(f"  Throughput: {thr:.0f} circuits/hr")
    print(f"{'='*70}\n")