Add transforming adaptation for stan

python/nutpie/compile_stan.py

@@ -1,16 +1,18 @@
+from functools import partial
 import json
 import tempfile
 from dataclasses import dataclass, replace
 from importlib.util import find_spec
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Optional, Callable
 
 import numpy as np
 import pandas as pd
 from numpy.typing import NDArray
 
 from nutpie import _lib
 from nutpie.sample import CompiledModel
+from nutpie.transform_adapter import make_transform_adapter
 
 
 class _NumpyArrayEncoder(json.JSONEncoder):
@@ -28,6 +30,7 @@ class CompiledStanModel(CompiledModel):
     library: Any
     model: Any
     model_name: Optional[str] = None
+    _transform_adapt_args: dict | None = None
 
     def with_data(self, *, seed=None, **updates):
         if self.data is None:
@@ -42,7 +45,15 @@ def with_data(self, *, seed=None, **updates):
         else:
             data_json = None
 
-        model = _lib.StanModel(self.library, seed, data_json)
+        kwargs = self._transform_adapt_args
+        if kwargs is None:
+            kwargs = {}
+        make_adapter = partial(
+            make_transform_adapter(**kwargs),
+            logp_fn=None,
+        )
+
+        model = _lib.StanModel(self.library, seed, data_json, make_adapter)
         coords = self._coords
         if coords is None:
             coords = {}
@@ -75,6 +86,9 @@ def with_dims(self, **dims):
         dims_new.update(dims)
         return replace(self, dims=dims_new)
 
+    def with_transform_adapt(self, **kwargs):
+        return replace(self, _transform_adapt_args=kwargs).with_data()
+
     def _make_model(self, init_mean):
         if self.model is None:
             return self.with_data().model