Skip to content

Add optional PyO3 Rust core acceleration#48

Merged
Adlgr87 merged 1 commit into
mainfrom
codex/identify-functions-for-rust-migration
May 30, 2026
Merged

Add optional PyO3 Rust core acceleration#48
Adlgr87 merged 1 commit into
mainfrom
codex/identify-functions-for-rust-migration

Conversation

@Adlgr87
Copy link
Copy Markdown
Owner

@Adlgr87 Adlgr87 commented May 30, 2026

Motivation

  • Move hot numeric kernels out of Python into an optional Rust core to reduce CPU overhead while preserving the existing Python API.
  • Target kernels that operate on dense NumPy arrays with deterministic clipping behavior (Langevin opinion update, active-mask computation, potential gradient) to minimize compatibility risk.
  • Provide a safe, opt-in Python Wrapper / Rust Core design that keeps RNG in Python and falls back to NumPy when the Rust extension is not available.

Description

  • Add a PyO3/Maturin Rust crate (rust_core/src/lib.rs, Cargo.toml) exposing multi_potential_gradient_rs, langevin_opinion_update_inplace, and active_mask_step_rs as deterministic numeric kernels.
  • Add a stable Python wrapper module massive_core/rust_core.py that imports the compiled massive_rust_core when present and otherwise provides NumPy fallbacks for the same functions.
  • Wire the wrapper into the runtime: IntegratedSimulator.update_agents_with_langevin now delegates opinion updates to langevin_opinion_update_inplace, and ActiveSet.step delegates to active_mask_step, preserving previous public APIs and semantics.
  • Add packaging/build files and docs (pyproject.toml, tests/test_rust_core_wrapper.py, docs/rust_core_plan_ES.md) documenting selection criteria, migration plan, and equivalence tests.

Testing

  • cargo test ran and passed the Rust unit test(s) for the crate (Rust tests: passed).
  • cargo fmt / cargo fmt --check were run and the code was formatted (format check passed after formatting).
  • maturin build --release produced a wheel and the wheel was installed successfully, and importing massive_rust_core succeeded (build + import smoke test: passed).
  • pytest tests/test_rust_core_wrapper.py and the full Python test suite pytest tests/ were executed and all tests passed (300 passed, 1 warning).

Codex Task


Summary by cubic

Add an optional Rust core (massive_rust_core) to accelerate hot numeric kernels while preserving the Python API and exact results. Falls back to numpy when the extension isn’t present.

  • New Features

    • Added Rust crate massive_rust_core (via pyo3/maturin) exposing multi_potential_gradient_rs, langevin_opinion_update_inplace, and active_mask_step_rs.
    • Added Python wrapper massive_core.rust_core that uses Rust if available, or numpy fallbacks otherwise (same signatures and clipping behavior).
    • Wired into runtime: IntegratedSimulator.update_agents_with_langevin and ActiveSet.step now delegate to the wrapper without changing public APIs.
    • Included docs and tests for equivalence; build config in pyproject.toml.
  • Migration

    • No code changes required; NumPy is used if massive_rust_core isn’t installed.
    • To enable acceleration: pip install maturin and run maturin develop --release (or maturin build --release) to install massive_rust_core.

Written for commit 4cb7c72. Summary will update on new commits.

Review in cubic

@Adlgr87 Adlgr87 merged commit 03afd31 into main May 30, 2026
3 checks passed
@Adlgr87 Adlgr87 deleted the codex/identify-functions-for-rust-migration branch May 30, 2026 23:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant