From 3b4140d8a3d0e64d8564029bd0b676887a9af4f0 Mon Sep 17 00:00:00 2001 From: Virgile Andreani Date: Mon, 13 Jan 2025 23:59:28 +0100 Subject: [PATCH] Make a proper Python wrapper over the library bindings --- pixi.lock | 2 +- pyproject.toml | 2 +- python/rebop/__init__.py | 52 +------------------------ python/rebop/gillespie.py | 80 +++++++++++++++++++++++++++++++++++++++ python/rebop/rebop.pyi | 64 ------------------------------- src/lib.rs | 8 +--- 6 files changed, 86 insertions(+), 122 deletions(-) create mode 100644 python/rebop/gillespie.py delete mode 100644 python/rebop/rebop.pyi diff --git a/pixi.lock b/pixi.lock index febdc47..261b998 100644 --- a/pixi.lock +++ b/pixi.lock @@ -561,7 +561,7 @@ packages: - pypi: . name: rebop version: 0.8.3 - sha256: 2d8d2e804394db8f6a9ee034476d5251962f5b4b421000f19c3b43b948eb5654 + sha256: 7cb212f5f191e381f42ccb7547893bf37b925c5b8007146a58883f6d58c6b2bc requires_dist: - xarray>=2023.1 - pytest>=8.3.4,<9 ; extra == 'dev' diff --git a/pyproject.toml b/pyproject.toml index ce7c4e8..b7a2349 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ dev = [ [tool.maturin] python-source = "python" -module-name = "rebop.rebop" +module-name = "rebop._lib" features = ["pyo3/extension-module"] [tool.ruff] diff --git a/python/rebop/__init__.py b/python/rebop/__init__.py index aa69bf0..8a1ac11 100644 --- a/python/rebop/__init__.py +++ b/python/rebop/__init__.py @@ -1,52 +1,4 @@ -from __future__ import annotations - -from collections.abc import Sequence -from typing import TypeAlias - -import numpy as np -import xarray as xr - -from .rebop import Gillespie, __version__ # type: ignore[attr-defined] - -SeedLike: TypeAlias = int | np.integer | Sequence[int] | np.random.SeedSequence -RNGLike: TypeAlias = np.random.Generator | np.random.BitGenerator - +from rebop._lib import __version__ +from rebop.gillespie import Gillespie __all__ = ("Gillespie", "__version__") - - -def run_xarray( # noqa: PLR0913 too many parameters in function definition - self: Gillespie, - init: dict[str, int], - tmax: float, - nb_steps: int, - *, - rng: RNGLike | SeedLike | None = None, - sparse: bool = False, - var_names: Sequence[str] | None = None, -) -> xr.Dataset: - """Run the system until `tmax` with `nb_steps` steps. - - The initial configuration is specified in the dictionary `init`. - Returns an xarray Dataset. - """ - rng_ = np.random.default_rng(rng) - seed = rng_.integers(np.iinfo(np.uint64).max, dtype=np.uint64) - times, result = self._run( - init, - tmax, - nb_steps, - seed=seed, - sparse=sparse, - var_names=var_names, - ) - ds = xr.Dataset( - data_vars={ - name: xr.DataArray(values, dims="time", coords={"time": times}) - for name, values in result.items() - }, - ) - return ds - - -Gillespie.run = run_xarray # type: ignore[method-assign] diff --git a/python/rebop/gillespie.py b/python/rebop/gillespie.py new file mode 100644 index 0000000..0fe60af --- /dev/null +++ b/python/rebop/gillespie.py @@ -0,0 +1,80 @@ +"""rebop is a fast stochastic simulator for well-mixed chemical reaction networks.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TypeAlias + +import numpy as np +import xarray as xr + +from rebop import _lib + +__all__ = ("Gillespie", "__version__") + + +__version__: str = _lib.__version__ + +SeedLike: TypeAlias = int | np.integer | Sequence[int] | np.random.SeedSequence +RNGLike: TypeAlias = np.random.Generator | np.random.BitGenerator + + +class Gillespie: + """Reaction system composed of species and reactions.""" + + def __init__(self) -> None: + """Initialize a solver.""" + self.gillespie = _lib.Gillespie() + + def add_reaction( + self, + rate: float, + reactants: list[str], + products: list[str], + reverse_rate: float | None = None, + ) -> None: + """Add a Law of Mass Action reaction to the system. + + The forward reaction rate is `rate`, while `reactants` and `products` + are lists of respectively reactant names and product names. + Add the reverse reaction with the rate `reverse_rate` if it is not + `None`. + """ + self.gillespie.add_reaction(rate, reactants, products, reverse_rate) + + def __str__(self) -> str: + """Return a textual representation of the reaction system.""" + return str(self.gillespie) + + def run( # noqa: PLR0913 too many parameters in function definition + self, + init: dict[str, int], + tmax: float, + nb_steps: int, + *, + rng: RNGLike | SeedLike | None = None, + sparse: bool = False, + var_names: Sequence[str] | None = None, + ) -> xr.Dataset: + """Run the system until `tmax` with `nb_steps` steps. + + The initial configuration is specified in the dictionary `init`. + Returns an xarray Dataset. + """ + rng_ = np.random.default_rng(rng) + seed = rng_.integers(np.iinfo(np.uint64).max, dtype=np.uint64) + times, result = self.gillespie.run( + init, + tmax, + nb_steps, + seed=seed, + sparse=sparse, + var_names=var_names, + ) + ds = xr.Dataset( + data_vars={ + name: xr.DataArray(values, dims="time", coords={"time": times}) + for name, values in result.items() + }, + ) + return ds diff --git a/python/rebop/rebop.pyi b/python/rebop/rebop.pyi deleted file mode 100644 index 38bcb4b..0000000 --- a/python/rebop/rebop.pyi +++ /dev/null @@ -1,64 +0,0 @@ -from collections.abc import Sequence -from typing import TypeAlias - -import numpy as np -import xarray - -SeedLike: TypeAlias = int | np.integer | Sequence[int] | np.random.SeedSequence -RNGLike: TypeAlias = np.random.Generator | np.random.BitGenerator - -class Gillespie: - """Reaction system composed of species and reactions.""" - - def add_reaction( - self, - /, - rate: float, - reactants: list[str], - products: list[str], - reverse_rate: float | None = None, - ) -> None: - """Add a Law of Mass Action reaction to the system. - - The forward reaction rate is `rate`, while `reactants` and `products` are - lists of respectively reactant names and product names. - Add the reverse reaction with the rate `reverse_rate` if it is not `None`. - """ - - def nb_reactions(self, /) -> int: - """Number of reactions currently in the system.""" - - def nb_species(self, /) -> int: - """Number of species currently in the system.""" - - def _run( - self, - init: dict[str, int], - tmax: float, - nb_steps: int, - *, - seed: int | None = None, - sparse: bool = False, - var_names: Sequence[str] | None = None, - ) -> tuple[np.ndarray, dict[str, np.ndarray]]: - """Run the system until `tmax` with `nb_steps` steps. - - The initial configuration is specified in the dictionary `init`. - Returns numpy arrays. - """ - - def run( - self, - init: dict[str, int], - tmax: float, - nb_steps: int, - *, - rng: RNGLike | SeedLike | None = None, - sparse: bool = False, - var_names: Sequence[str] | None = None, - ) -> xarray.Dataset: - """Run the system until `tmax` with `nb_steps` steps. - - The initial configuration is specified in the dictionary `init`. - Returns an xarray Dataset. - """ diff --git a/src/lib.rs b/src/lib.rs index fab2883..294c02c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -260,10 +260,6 @@ impl Gillespie { Ok(self.species.len()) } /// Add a Law of Mass Action reaction to the system. - /// - /// The forward reaction rate is `rate`, while `reactants` and `products` are lists of - /// respectively reactant names and product names. Add the reverse reaction with the rate - /// `reverse_rate` if it is not `None`. #[pyo3(signature = (rate, reactants, products, reverse_rate=None))] fn add_reaction( &mut self, @@ -304,7 +300,7 @@ impl Gillespie { /// If `nb_steps` is `0`, then returns all reactions, ending with the first that happens at /// or after `tmax`. #[pyo3(signature = (init, tmax, nb_steps, seed=None, sparse=false, var_names=None))] - fn _run( + fn run( &self, init: HashMap, tmax: f64, @@ -405,7 +401,7 @@ impl Gillespie { } #[pymodule] -fn rebop(m: &Bound<'_, PyModule>) -> PyResult<()> { +fn _lib(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add("__version__", env!("CARGO_PKG_VERSION"))?; m.add_class::()?; Ok(())