""" 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")