feat: Allow to update dims and coords in stan model

Author: aseyboldt

Cargo.lock

@@ -1,7 +1,10 @@
 [package]
 name = "nutpie"
-version = "0.5.1"
-authors = ["Adrian Seyboldt <[email protected]>"]
+version = "0.6.0-beta.1"
+authors = [
+    "Adrian Seyboldt <[email protected]>",
+    "PyMC Developers <[email protected]>"
+]
 edition = "2021"
 license = "MIT"
 repository = "https://github.com/pymc-devs/nutpie"

Cargo.toml

@@ -1,7 +1,7 @@
 import json
 import pathlib
 import tempfile
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any, Dict, Optional
 
 import numpy as np
@@ -58,6 +58,22 @@ def with_data(self, *, seed=None, **updates):
             model=model,
         )
 
+    def with_coords(self, **coords):
+        if self.coords is None:
+            coords_new = {}
+        else:
+            coords_new = self.coords.copy()
+        coords_new.update(coords)
+        return replace(self, coords=coords_new)
+
+    def with_dims(self, **dims):
+        if self.dims is None:
+            dims_new = {}
+        else:
+            dims_new = self.dims.copy()
+        dims_new.update(dims)
+        return replace(self, dims=dims_new)
+
     def _make_model(self, init_mean):
         if self.model is None:
             return self.with_data().model

nutpie/compile_stan.py

@@ -1,9 +1,9 @@
 [build-system]
-requires = ["maturin>=0.14,<0.15"]
+requires = ["maturin>=1.1,<2.0"]
 build-backend = "maturin"
 
 [tool.maturin]
-bindings = "pyo3"
+features = ["pyo3/extension-module"]
 
 [project]
 name = "nutpie"
@@ -12,7 +12,11 @@ authors = [{name = "PyMC Developers", email = "[email protected]"}]
 readme = "README.md"
 requires-python = ">=3.9"
 license = { text = "MIT" }
-version = "0.6.0"
+classifiers = [
+    "Programming Language :: Rust",
+    "Programming Language :: Python :: Implementation :: CPython",
+    "Programming Language :: Python :: Implementation :: PyPy",
+]
 
 dependencies = [
     "pyarrow >= 12.0.0",
@@ -30,3 +34,4 @@ all = [
     "pymc >= 5.5.0",
     "numba >= 0.57.1",
 ]
+