Pufferlib-inspired API redesign: proposal + kernel core (+ Optuna parallelization)#19
Open
RaulSimpetru wants to merge 33 commits into
Open
Pufferlib-inspired API redesign: proposal + kernel core (+ Optuna parallelization)#19RaulSimpetru wants to merge 33 commits into
RaulSimpetru wants to merge 33 commits into
Conversation
Design proposal for a backend-swappable, vectorized MyoGen API: a shared flat SimState buffer stepped in lockstep by a backend-agnostic driver, with two backend seams (neural dynamics NEURON->jaxley; muscle/force analytic->FEM) and a closed-loop stage graph for afferent feedback. Buffer-first data model with neo/NWB as opt-in adapters. Maps the full current public surface onto the new layers and lists net-new work and open questions for team review.
Bite-sized TDD plan for the backend-agnostic simulation kernel (myogen/kernel): SimState shared buffer, Stage/Backend protocols, the Simulation lockstep driver with one-tick closed-loop feedback, and the buffer-first SimResult with lazy neo/NWB adapters. Seven tasks, fully CI-testable without NEURON/jaxley/FEM. NEURONBackend, real stages, jaxley and FEM are scoped as follow-on plans.
SimResult.from_state() extracts spike times in seconds, copies force and surface EMG arrays, infers grid_shape from EMG buffer dimensions, and handles missing optional buffers gracefully. Simulation.run() now returns SimResult instead of bare SimState; all kernel tests updated to read buffers through sim.state.
…rocesses Extract simulation logic and constants into _oscillating_dc_helpers.py, add _optimize_dc_worker.py subprocess worker, and rewrite the launch in 02_optimize_oscillating_dc.py to fan out N workers that all write to one shared SQLite Optuna study. Avoids NEURON h singleton / macOS spawn constraints; gallery-safe via conf.py ignore_pattern for the _ files.
Example 14 (calibrate_noise_from_real) imports h5py directly to load real recordings, but it was only available transitively. Declaring it in the nwb extra makes 'uv sync --extra nwb' sufficient to run the NWB/HDF5 examples (11, 13, 14).
Many MyoGen examples use multiprocessing/joblib (which breaks under the gallery's exec context) or form slow inter-dependent chains, so executing the gallery in CI is unreliable. Render examples as source + prose with full-code download links instead; opt into execution with MKDOCS_GALLERY_PLOT=true.
Add MuscleGeometry/JointGeometry/JointBiomechanics to the simulator API, the four missing type aliases, and a RANDOM_GENERATOR/SEED note; restore the How-to-cite, Contributing, and License sections on the home page and the MUAP simulate/access step in the quick-start (per migration-coverage audit).
Capture one representative figure per example (run standalone, where they work) and wire them as mkdocs-gallery thumbnails via per-example mkdocs_gallery_thumbnail_path directives + a logo default_thumb_file fallback. Gives the source-only gallery real pictures on the overview grid without executing examples in the docs build (which spawn-bombs on the multiprocessing examples).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This branch bundles three related-but-distinct workstreams (kept together per request). Each is independently reviewable.
1. Pufferlib-inspired API redesign — proposal + kernel core
A design proposal and the first working slice (a backend-agnostic simulation kernel) for a vectorized, backend-swappable MyoGen API.
docs/superpowers/specs/2026-06-23-pufferlib-inspired-api-design.md): 7-layer architecture around a shared flatSimStatebuffer, a closed-loop stage graph (afferent feedback), and two backend seams — neural dynamics (NEURON → jaxley) and muscle/force (analytic → FEM). Buffer-first data model with opt-in neo/NWB.myogen/kernel/):SimState,Stage/Backendprotocols, theSimulationlockstep driver with one-tick feedback, andSimResultwith lazyto_neo()/to_nwb(). 23 tests, CI-runnable with no NEURON/jaxley/FEM. Purely additive — no existing code touched.2. Optuna example parallelization (~5×)
examples/03_papers/watanabe/02now fans its 25 independent Optuna trials across worker processes sharing one SQLite study: ~2000 s → 139 s, equivalent optimum. Verified faithful (byte-identical sim/objective, deterministic per seed, independently audited).h5pyadded to thenwbextra so the NWB/HDF5 examples are fully declarable.3. Docs: Sphinx → properdocs (MkDocs)
Migrates the documentation to properdocs (mkdocs-material + mkdocstrings), matching MyoGestic / Virtual-Hand-Interface.
properdocs.yml(material theme, mkdocstrings API, section-index), flat markdown underdocs/(home, getting-started, neo-blocks), 8 mkdocstrings API pages, and a mkdocs-gallery example gallery..pydownload links, but are not executed — many examples use multiprocessing/joblib (which breaks under the gallery'sexec()context) or form slow inter-dependent chains. Opt into execution withMKDOCS_GALLERY_PLOT=true..github/workflows/docs.ymlbuilds + deploys to GitHub Pages; the Sphinx config,conf.py, rst pages, andmake/Makefileare removed.Notes for reviewers
docs.ymldeploys via GitHub Actions, so repo Settings → Pages → Source must be set to "GitHub Actions" (the old workflow pushed to agh-pagesbranch).properdocs build --strictsurfaces 8 pre-existing griffe docstring/signature warnings inmyogensource (muscle.py,afferents.py) — not introduced here; the non-strict build (used by CI) passes.