Skip to main content

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

Project description

jaxfne

JAX-native TFNE (Tensor-Field Neural Emitter) workflows for reproducible computational neurophysiology.

PyPI · Docs · GitHub · Issues


What is jaxfne?

jaxfne is a compact JAX-native framework for composing neural simulations from modular operators:

Emitter (neuron state) → Source (membrane current) → Field (proxy/solved) → Probe (readout) → Objective

Primary use: Build reproducible laminar-field proxy simulations with deterministic PRNG, JSON-safe outputs, and clear claim boundaries.

Not a biological simulator. jaxfne is a computational-scaffold framework for teaching, prototyping, and experimenting with neural-field source models. All outputs are proxies unless explicitly validated against empirical data.


Quick Start

Install

pip install jaxfne

Optional JAX acceleration:

pip install -e '.[jax]'

Optional development/visualization:

pip install -e '.[dev,viz]'

Minimal Example

import jaxfne as jtfne

cfg = (
    jtfne.configuration()
    .network(name="V1", kind="cortical_column", n=100)
    .emitter(family="izhikevich", preset="cortical_eig")
    .field(domain="laminar_column", conductivity="proxy", boundary="mean_zero_neumann")
    .probe(name="laminar_16ch", modes=["spikes", "V_m", "source", "CSD"], n_contacts=16)
)

model = jtfne.construct(cfg)
signals = model.simulate(jtfne.simulation(duration_ms=100.0, dt_ms=0.1, seed=0))
readouts = model.compute_readout(signals, [
    jtfne.readout_spec("rate", "spike_rate_hz"),
    jtfne.readout_spec("csd", "csd_abs_mean"),
])

manifest = model.manifest(signals, readouts)
print(f"Simulation complete: {signals.V_m.shape[0]} timesteps, {signals.V_m.shape[1]} neurons")
print(f"Source status: {manifest['source_calibration_status']}")
print(f"Field status: {manifest['field_solver_status']}")

The Pipeline

1. Emitter: Neural Dynamics

Declare neuron model (Izhikevich or custom) and recurrent connectivity:

.emitter(family="izhikevich", preset="cortical_eig")

Output: State vector $z(t)$ and native membrane current $I(t)$ [time, neurons]

Status: Izhikevich presets are provided; no biological calibration claimed (computational scaffold)

2. Source: Spatial Projection

Project neural current into space (laminar probe contacts or voxels):

.field(domain="laminar_column", conductivity="proxy", ...)

Output: Source density $q(x,t)$ [time, contacts]

Status: Proxy projection using anatomical position and Izhikevich native current; no empirical validation

3. Field: Field Approximation

Current default: proxy CSD (no PDE solve).

field_solver_status = "laminar_proxy_no_pde"

CSD and LFP are computed from source without solving the Poisson equation. Conductivity is metadata-only.

Future (v0.2.27+): Optional field diagnostics; physical conductivity (v0.3.x)

4. Probe: Multimodal Readouts

Extract metrics from emitter state and field:

Operator Output Meaning
Spikes (SPK) Binary spike raster [T, N] Action potentials (thresholded state)
Voltage (V_m) Membrane voltage [T, N] Membrane potential state
Source Transmembrane current [T, X] Spatial source density
LFP-proxy Local field potential [T, X] Proxy; not physical units
CSD-proxy Current-source density [T, X] Proxy; spatial divergence of source
EEG-proxy Electroencephalogram [T, N_channels] Proxy; not physical units
MEG-proxy Magnetoencephalogram [T, N_channels] Proxy; not physical units
EMM-proxy Metabolic-like cost [T] Relative activity intensity (NOT biological metabolism)

All readouts are proxies unless explicitly solved and validated.

5. Objective & Optimization

Declare optimization targets and run GSDR/AGSDR (custom optimizers; Optax optional):

objectives = [
    jtfne.objective(name="spike_rate", target=10.0, metric="spike_rate_hz"),
    jtfne.objective(name="mean_voltage", target=-50.0, metric="mean_V_m"),
]

Validation

Fast validation (every commit, ~1 minute)

python -m compileall -q jaxfne tests examples scripts
PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. python -m pytest -q --tb=line

Results: 806 tests passed, 5 skipped. Examples 00–06 run. Build validation passes.

Extended validation (release, ~5–10 minutes)

python scripts/run_all_tutorials.py --smoke --write-figures
python scripts/validate_tutorial_outputs.py outputs/

Runs large tutorials (examples 02–07) with deterministic figures and asset hashing. See CI policy.


Documentation Map

Topic Document Purpose
Equations & Math Mathematical Glossary Flow TFNE equations (emitter, source, field, probe) with term glossaries and claim boundaries
Source Detail Source/Field Equations Source modes, forbidden double-counting pattern, field metadata, code examples
Architecture Computation Basis TFNE as collapsible tensor-field scaffold; extensibility doctrine
Probe Operators Probe Operators Eight multimodal operators, claim boundaries per operator
I/O & Manifests Output Bundles Signals, Manifest, ReadoutResult schema and JSON-safe contracts
Bridges & Interop Jaxley Interop Convert Jaxley voltage traces to jaxfne Signals
Scope & Limits Scope and Limitations What jaxfne claims and does not claim
Full Docs jaxfne.readthedocs.io API reference, tutorials, changelog

Roadmap

Version Phase Content Status
v0.2.24 Foundation Audit Audited contracts, verified solver status, updated language ✓ Released
v0.2.25 Docs-First Mathematical glossary, source/field doctrine, computation basis 🔄 In Progress
v0.2.26 Extensibility Documented future bases (layer projection, spectral), multi-area scaffolds 📋 Planned
v0.2.27 Diagnostics Optional Poisson-admissibility solver with conservation validation (no physical conductivity yet) 📋 Planned
v0.3.0 Physical Field Calibrated conductivity, solved field with empirical geometry, EEG/MEG/field calibration 📋 Planned

Current phase: v0.2.25 docs-first. No new solvers or diagnostics in v0.2.x; conservation material documented as future.


Claim Status

truth_mode: truth_safe_unverified
claim_level: computational_scaffold
physical_amplitude_claim_allowed: False

jaxfne is not a biological simulator. All outputs are computational proxies:

  • Izhikevich native current is a mathematical dynamics model, not empirically calibrated membrane current
  • Source projection uses declared anatomy but is not validated against measured sources
  • Field (proxy) is NOT a solved Poisson equation; CSD/LFP are kernel-based approximations
  • Readout proxies (LFP, CSD, EEG, MEG, EMM) are relative metrics, not physical units
  • Optimization is mathematical fitness; success ≠ biological plausibility

When to use jaxfne:

  • Teaching neural-field concepts
  • Prototyping source-field models
  • Benchmarking optimization strategies
  • Validating model consistency (future: conservation diagnostics)

When NOT to use jaxfne:

  • Making biological claims without separate empirical validation
  • Publishing simulation results as if they are real neural data
  • Claiming physical conductivity without calibration
  • Interpreting metabolic cost (EMM-proxy) as biological metabolism

License

MIT License.


Contributing

Issues, feature requests, and pull requests welcome. See CONTRIBUTING.

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.2.30.tar.gz (900.3 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.2.30-py3-none-any.whl (74.6 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: jaxfne-0.2.30.tar.gz
  • Upload date:
  • Size: 900.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.14.4

File hashes

Hashes for jaxfne-0.2.30.tar.gz
Algorithm Hash digest
SHA256 064ffaf4f3b18eb6662ced441ff29945bb27dd52a2885b7a10b2451c48bba95e
MD5 7258c0b28fe181297adff40803fe799d
BLAKE2b-256 ecb36e65c33e9a5c4cead59b5305cb306a0e5e5360eff11d39368a4e39d4375d

See more details on using hashes here.

File details

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

File metadata

  • Download URL: jaxfne-0.2.30-py3-none-any.whl
  • Upload date:
  • Size: 74.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.14.4

File hashes

Hashes for jaxfne-0.2.30-py3-none-any.whl
Algorithm Hash digest
SHA256 a886213a2ff2a8d5d77f513e5361dd0ed1e8889714794e52d58ed34b81320114
MD5 cc804838cbd7dbf85f29fa40b1a5a73e
BLAKE2b-256 15ad4e0213e2c4b30dcbac8780ed9b4fb6cec7281b1ddf85ed47d46e2e3b5880

See more details on using hashes here.

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