Skip to main content

JAX Field Neural Equations: a source-to-field neurophysiology engine for TFNE models.

Project description

Installation

For the future PyPI release, the intended Colab entry point is %pip install jaxfne. Until PyPI publication, use the repository or built wheel.

See docs/COLAB.md and docs/PACKAGING.md.

jaxfne

JAX Field Neural Equations (jaxfne) is a JAX-native source-to-field neurophysiology engine for Tensor-Field Neural Equations (TFNE).

import jaxfne as jtfne

cfg = (
    jtfne.configuration()
    .network(name="V1", kind="cortical_column", n=64)
    .emitter(family="izhikevich", preset="cortical_eig")
    .field(domain="laminar_column", conductivity="proxy", boundary="mean_zero_neumann", gauge="mean_zero")
    .probe(name="laminar_probe", modes=["spikes", "V_m", "CSD", "LFP"])
)

model = jtfne.construct(cfg)
sim = jtfne.simulation(duration_ms=100.0, dt_ms=0.1, seed=0)

# canonical v0.0.16+ workflow
signals = model.simulate(sim)
receipt = model.run_receipt(signals)
readout = model.compute_readout(signals, [
    jtfne.readout_spec("rate", "spike_rate_hz"),
    jtfne.readout_spec("csd",  "csd_abs_mean"),
])

Claim Boundary Note

JaxFNE is a computational scaffold, not a biological proof engine.

  • Truth Status: truth_safe_unverified.
  • Claim Level: computational_scaffold.
  • Source Calibration: uncalibrated_izhikevich_native_current.
  • Field Solver: laminar_proxy_no_pde.

Receipts and reports are validation artifacts used to document the computational state of a model. They do not constitute empirical validation, biological calibration, or mechanism proof. Calibrated physical CSD/LFP amplitudes and biological mechanism claims are explicitly forbidden until a rigorous validation gate with empirical datasets is established.

Identity

jaxfne is not primarily a neuron model, optimizer, plotting library, or data format. It is the composition layer:

Emitter -> Source -> Field -> Probe -> Objective -> Optimizer

JAX handles arrays, compilation, batching, and device execution. Jaxley can later provide detailed emitters. Optax can later provide differentiable optimizers. jaxfne handles TFNE source-to-field/readout contracts, diagnostics, invariant checks, and manifests.

Current status (v0.1.0)

v0.1.0 declares the practical OOP core freeze for the compact JAX-native TFNE scaffold. Wheel and sdist install smokes validated from final dev state. MIT LICENSE added; examples normalized to 00-06 naming convention. Truth status preserved at truth_safe_unverified. v0.1.0 is the release-candidate milestone for TestPyPI rehearsal and subsequent public PyPI release.

It still does not solve the full resistive extracellular TFNE PDE:

J_e = -sigma_e grad(phi_e)
div(J_e) = q
div(-sigma_e grad(phi_e)) = q
CSD = div(J_e)

Current field outputs are laminar proxy readouts:

source_projection_mode = proxy_no_field_solve
field_solver_status = laminar_proxy_no_pde
field_claim_level = proxy_readout_only
physical_amplitude_claim_allowed = false

v0.0.5 API surface

Task paradigm

# API scaffold for Paper 2.0 (Omission deferred)
# Standard visual omission task: 12 conditions, condition-number mapping
paradigm = jtfne.standard_visual_omission()
print(paradigm.condition_names())        # ['AAAB', 'AXAB', ..., 'RRRX']
print(paradigm.omission_conditions())    # 9 conditions with omission_position set
print(paradigm.event_codes)              # {'fx': 10, 'p1': 101, ..., 'rw': 96}
print(paradigm.analysis_windows)        # {'baseline': (-500, 0), ...}

# Paradigm is JSON-safe
import json
json.dumps(paradigm.to_dict(), allow_nan=False)

Objective / evaluation

obj = (
    jtfne.objective()
    .loss("rate_loss", target=20.0, weight=1.0, metric="spike_rate_hz_mean")
    .regularizer("vm_reg", target=-65.0, weight=0.1, metric="mean_V_m")
    .gate("rate_gate", threshold=200.0, criterion="below", metric="spike_rate_hz_mean")
)

eval_report = model.evaluate(signals, obj)
# eval_report["evaluation_status"] == "objective_evaluate_v0.0.5"
# eval_report["all_gates_pass"] in {True, False}  <- computational diagnostic only
# eval_report["truth_mode"] == "truth_safe_unverified"
# eval_report["physical_amplitude_claim_allowed"] == False

Known metrics: spike_rate_hz_mean, spike_count_total, mean_V_m, source_proxy_abs_mean, csd_proxy_abs_mean, lfp_proxy_abs_mean.

Gate criteria: below, above, equal, in_range.

Gate pass/fail is a computational diagnostic only. It does not imply empirical validation, biological calibration, or mechanism proof.

Optimizer metadata / tune scaffold

# Blackbox path (always allowed, no gradient required)
same_model, tune_report = model.tune(obj, optimizer="GSDR", steps=1)
assert same_model is model                          # model never mutated
assert tune_report["tuning_status"] == "metadata_only_v0.0.5"
assert tune_report["same_model_unchanged"] is True

# OptimizerSpec with explicit differentiability declaration
spec = jtfne.gsdr(alpha=0.7, exploration=0.05)
spec = jtfne.agsdr()
spec = jtfne.random_search()
# Optax path (requires declared_surrogate or differentiable, Optax installed)
spec = jtfne.optax_adam(learning_rate=1e-3, differentiability_status="declared_surrogate")

Model.tune() in v0.0.5 is a metadata-only scaffold. No optimization loop runs and no parameters are changed. The report documents what would happen in a future real tuning pass.

The Optax path requires explicit differentiability_status="declared_surrogate" or "differentiable". The default "not_checked" is blocked because spiking networks are not differentiable through spike resets without a surrogate gradient declaration.

Manifest with v0.0.5 metadata (compatibility alias)

manifest = model.manifest(
    signals=signals,
    readout=readout,
    paradigm=paradigm.to_dict(),
    objective={"name": obj.name, "losses": obj.losses, ...},
    evaluation=eval_report,
    tuning=tune_report,
)
# manifest["v005_claim_labels"]["empirical_validation_status"] == "not_empirically_validated"
# manifest["v005_claim_labels"]["mechanism_claim_status"] == "not_claimed"
# All v0.0.4 truth gates still present

model.manifest() is a compatibility alias retained from v0.0.4–v0.0.14. For the canonical v0.1 workflow, prefer model.run_receipt() and model.evaluate_report().

Version roadmap (Publication Event 1.0)

The roadmap centers around Paper 1.0: in-silico spectrolaminar motif, deferring omission (Paper 2.0) and global/local oddball (Paper 3.0).

Pre-v0.1 bridge lane:

  • v0.0.11 uncalibrated multi-receptor synaptic kernel
  • v0.0.12 paradigm/stimulus injection into recurrent model
  • v0.0.13 laminar population builder / source geometry metadata
  • v0.0.14 trial runner and condition-aligned outputs
  • v0.0.15 config/JaxFNEConfig standard (.jcfg.json)
  • v0.0.16 RunReceipt deterministic audit receipt
  • v0.0.17 ReadoutSpec declarative feature extraction
  • v0.0.18 ObjectiveReport structured evaluation result

Publication Event 1.0 sequence:

  • v0.0.20 semantic correctness hardening ← current release
  • v0.1.0 practical OOP core freeze (pending v0.0.20 gate)
  • v0.2.x spectrolaminar objective/readout base
  • v0.3.x generative spectrolaminar workflow
  • v0.4.0 Paper 1.0 minimum sufficient release (in-silico spectrolaminar motif)

Deferred (Post-Paper 1.0):

  • Paper 2.0: Omission mismatch workflows
  • Paper 3.0: Global/local oddball workflows
  • Dense Jaxley/NEURON compartment bridges
  • EEG/MEG physical solvers

Claims Discipline

Allowed:

  • v0.1.x provides the compact JAX-native OOP core required to build reproducible TFNE workflows.

Forbidden:

  • v0.1.x validates spectrolaminar mechanisms.
  • v0.1.x produces calibrated LFP/CSD amplitudes.
  • v0.1.x is a full simulator.

Package structure

jaxfne/
  __init__.py        public API surface
  core.py            Configuration, Model, Simulation, RuntimeConfig, Signals,
                     Probe, Objective, Paradigm, ParadigmEvent, ParadigmCondition,
                     RunReceipt, ReadoutSpec, ReadoutResult, ObjectiveReport,
                     JaxFNEConfig, ConfigValidationResult (v0.0.15–v0.0.18)
  emitters.py        Izhikevich EIG scaffold
  fields.py          laminar proxy source/probe layer and invariant diagnostics
  optim.py           OptimizerSpec, GSDR/AGSDR/random_search/Optax specs,
                     require_optax guard, legacy AGSDR class
  bridges.py         Jaxley bridge scaffold + require_jaxley guard
  io.py              strict JSON manifest, hashing, save/load, save_receipt

Development smoke

python -m compileall -q jaxfne tests examples
python -m pytest -q
python examples/minimal_eig_column.py
python examples/global_local_oddball_sketch.py
python examples/02_omission_scaffold.py
python examples/03_objective_and_tune_smoke.py

Expected: 178 passed, 0 failed.

v0.0.9 edge-list backend

v0.0.9 adds a sparse EdgeList recurrent backend using JAX pytrees, jax.lax.scan, and jax.ops.segment_sum. It is selected with runtime(recurrent_backend="edge_list"). This is a computational backend upgrade only; field output remains laminar_proxy_no_pde, source calibration remains uncalibrated, and optimizer-selected candidates do not establish empirical or mechanistic claims.

v0.0.10 synapse metadata

v0.0.10 hardens source and synapse declarations. It introduces metadata-only ReceptorSpec and SynapseSpec definitions without adding new conductance-based physical solvers or biological kernels. The backend manifest now flows EdgeList details transparently. Dense-vs-edge computations maintain statistical parity. No calibrated synapse claim and no new PDE/field/empirical/mechanism claim is made.

v0.0.11 receptor-indexed exponential synaptic kernel

v0.0.11 adds an opt-in second synaptic kernel selected via runtime(recurrent_backend="edge_list", synaptic_kernel="receptor_exponential"). The default remains synaptic_kernel="exponential", preserving the v0.0.9/v0.0.10 edge-list path. The new path keeps syn_state.shape == (n_edges,) and looks up the per-edge decay time constant from edge.receptor_index against the standard ReceptorSpec table (AMPA/GABA_A/NMDA/GABA_B). Aggregation uses jax.ops.segment_sum(weight * syn_state, post, n_neurons), so each edge contributes exactly once to its postsynaptic native recurrent input; multiple edges may legitimately converge on the same neuron. Receptor reversal potentials remain metadata-only and are not used in the current computation. Weights remain native/unphysical and the source readout remains a laminar proxy. No conductance equation, no physical-amplitude claim, no PDE upgrade, and no empirical-validation or biological-mechanism claim is introduced.

v0.0.12 native stimulus injection

v0.0.12 adds explicit event-aligned native-drive stimulus injection. It introduces the StimulusSchedule value object and stimulus_schedule() factory to build timed drive arrays from ParadigmCondition events. Injected drive is added as native (uncalibrated) current to the recurrent kernels (dense, edge_list, and receptor_exponential). This enables basic condition-aligned trial workflows for Paper 1.0. It does not implement a full Paradigm.batch() trial runner, and no cognitive omission or global-local logic is encoded. No calibrated current, physical amplitude, PDE upgrade, or empirical/mechanism claim is introduced.

v0.0.13 laminar source geometry

v0.0.13 adds explicit laminar population and source geometry metadata. It introduces the LaminarPopulation and LaminarSourceGeometry frozen dataclasses and the laminar_source_geometry() factory. An optional geometry kwarg is added to construct(cfg, *, geometry=None); when provided, the geometry's n_units_total must match the network n or a ValueError: geometry_n_units_total_mismatch is raised. Population depths use deterministic NumPy linspace — no random placement. Co-located populations (depth overlap) are allowed and not treated as errors, as this is anatomically valid for different cell types. The geometry manifest propagates into Model.manifest() under source_geometry. No new field solver, no physical-amplitude claim, no PDE upgrade, no empirical or mechanism claim is introduced.

v0.0.14 sequential trial runner

v0.0.14 adds a sequential trial runner for executing batches of condition-aligned simulations. It introduces TrialSpec, TrialBatch, TrialResult, and TrialBatchResult dataclasses, along with a trial_batch() factory for deterministic batch creation. The Model.run_trials(batch, sim) method performs a sequential loop over batch.trials, replacing sim.seed with the trial's seed and capturing any exceptions into the TrialResult. Serialization via to_dict() is strictly JSON-safe and automatically excludes large JAX arrays (V_m, spikes, sources, field) from the results, providing instead a compact Signals.summary(). This enables robust condition-aligned workflow orchestration for Paper 1.0. No parallel execution, no new field solver, no physical-amplitude claim, no PDE upgrade, and no empirical or mechanism claim is introduced.

v0.0.15 config/object standard foundation

v0.0.15 adds the first formal .jcfg.json declarative configuration standard. It introduces JaxFNEConfig, ConfigValidationResult, load_config, validate_config, config_to_simulation, config_to_geometry, config_to_configuration, config_to_trial_batch, and config_truth_boundary.

The config schema maps directly onto existing proven objects (Simulation, Configuration, LaminarSourceGeometry, TrialBatch) without introducing new kernels, field solvers, or physical-amplitude claims. Truth boundary fields are required in every config file and any escalation is a blocking validation error. Geometry depths remain normalized proxy coordinates in [0, 1] with default position_units = "relative_laminar_depth_proxy" — no physical spatial units (mm, µm) are introduced. No new PDE solver, no calibrated current, no empirical validation, and no mechanism claim is introduced.

v0.0.16 run receipt

v0.0.16 adds RunReceipt — an immutable, JSON-safe record of a single simulation run. It introduces RunReceipt, Model.run_receipt(), the module-level run_receipt() factory, and save_receipt() for writing receipts to disk with an overwrite guard.

A RunReceipt captures: config fingerprint (config_hash), simulation parameters, Signals.summary(), frozen truth gates, claim labels, and backend metadata. The receipt_id is deterministic: same configuration + seed + version always produces the same 16-char SHA256 prefix. save_receipt(receipt, path) raises ValueError("receipt_file_exists: ...") if the path already exists, unless overwrite=True is passed.

The manifest_schema_version default in io.manifest() is bumped from "0.0.4" to "0.0.16" to reflect the current package version.

No new PDE solver, no calibrated current, no empirical validation, and no mechanism claim is introduced.

v0.0.17 readout spec

v0.0.17 adds the declarative feature-extraction standard for Paper 1.0 workflows. It introduces ReadoutSpec, ReadoutResult, the readout_spec() factory, and Model.compute_readout(signals, specs).

A ReadoutSpec declares a named scalar metric to extract from Signals. Supported metrics:

Metric Description
spike_rate_hz Mean firing rate across all units (Hz)
spike_count Total spike count
mean_V_m Mean membrane voltage
csd_abs_mean Mean absolute CSD proxy
lfp_abs_mean Mean absolute LFP proxy
source_abs_mean Mean absolute source proxy

Optional time_window_ms=(start, end) and n_contacts_slice=(start, end) allow temporal and depth windowing. ReadoutResult.status is "computed", "no_field", or "unknown_metric".

_KNOWN_READOUT_METRICS is exported from both jaxfne and jaxfne.core.

No new PDE solver, no physical-amplitude claim, no empirical validation, and no mechanism claim is introduced.

v0.0.18 objective report

v0.0.18 adds ObjectiveReport — a structured, immutable result of evaluating an Objective against Signals. It introduces ObjectiveReport (frozen dataclass) and Model.evaluate_report(signals, objective, *, readout_specs=None).

evaluate_report wraps the existing Model.evaluate() output into a typed, JSON-safe object that also embeds ReadoutResult items when readout_specs are provided. ObjectiveReport.truth carries frozen conservative truth gates.

report = model.evaluate_report(signals, obj, readout_specs=[
    jtfne.readout_spec("rate", "spike_rate_hz"),
    jtfne.readout_spec("csd",  "csd_abs_mean"),
])
assert report.evaluation_status == "objective_report_v0.0.18"
assert report.truth["physical_amplitude_claim_allowed"] is False
json.dumps(report.to_dict(), allow_nan=False)

Gate pass/fail is a computational diagnostic only. No physical-amplitude claim, no empirical validation, and no mechanism claim is introduced.

Release rehearsal (no credentials required)

./scripts/release_rehearsal.sh

Runs the full pre-publish gate locally: clean build → twine check → fresh venv wheel+sdist install smokes from /tmp → pytest → all examples. No upload, no credentials, no tagging.

Colab smoke

See docs/COLAB_SMOKE_V010.md for copy-pasteable Colab cells to validate pip install jaxfne==0.1.0 from TestPyPI or real PyPI.

Truth status

truth_mode:                    truth_safe_unverified
claim_level:                   computational_scaffold
source_calibration_status:     uncalibrated_izhikevich_native_current
source_projection_mode:        proxy_no_field_solve
field_solver_status:           laminar_proxy_no_pde
field_claim_level:             proxy_readout_only
physical_amplitude_claim_allowed: false
empirical_validation_status:   not_empirically_validated
mechanism_claim_status:        not_claimed

No calibrated physical CSD/LFP/EEG/MEG amplitude claim is made. No biological mechanism is implied by gate pass/fail or tuning metadata. Optimizer success does not constitute empirical validation.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

jaxfne-0.1.1.tar.gz (125.9 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

jaxfne-0.1.1-py3-none-any.whl (56.2 kB view details)

Uploaded Python 3

File details

Details for the file jaxfne-0.1.1.tar.gz.

File metadata

  • Download URL: jaxfne-0.1.1.tar.gz
  • Upload date:
  • Size: 125.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for jaxfne-0.1.1.tar.gz
Algorithm Hash digest
SHA256 c837db51e7793ff113bddc6bbefc1312a8a2d520d49bf8867a7dcd0f93866235
MD5 d97a9a1af28b9557c1fd7ce51c2a7499
BLAKE2b-256 676a1231cc4294181ad6d51e678e250f5efe847fbd91887ba7ae1ef83b77a24f

See more details on using hashes here.

Provenance

The following attestation bundles were made for jaxfne-0.1.1.tar.gz:

Publisher: publish.yml on HNXJ/jaxfne

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file jaxfne-0.1.1-py3-none-any.whl.

File metadata

  • Download URL: jaxfne-0.1.1-py3-none-any.whl
  • Upload date:
  • Size: 56.2 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.12

File hashes

Hashes for jaxfne-0.1.1-py3-none-any.whl
Algorithm Hash digest
SHA256 16bbe47e66aa22397dd8b47b6c1a1917cefd087df7147a90cface34cb1915475
MD5 5ba8e63bc023c1eda0cc74f0da349634
BLAKE2b-256 a4ad604caa8d3d6fd2fe7134014b42c2448c6f9e964a1497833705765f16b88e

See more details on using hashes here.

Provenance

The following attestation bundles were made for jaxfne-0.1.1-py3-none-any.whl:

Publisher: publish.yml on HNXJ/jaxfne

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page