supressing ArviZ warning and allowing step selection

pymc_experimental/model_builder.py

@@ -15,6 +15,7 @@
 
 import hashlib
 import json
+import warnings
 from abc import abstractmethod
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
@@ -290,12 +291,18 @@ def sample_model(self, **kwargs):
  raise RuntimeError(
  "The model hasn't been built yet, call .build_model() first or call .fit() instead."
  )
-
  with self.model:
  sampler_args = {**self.sampler_config, **kwargs}
- idata = pm.sample(**sampler_args)
- idata.extend(pm.sample_prior_predictive())
- idata.extend(pm.sample_posterior_predictive(idata))
+ if "step" in sampler_args:
+ step_function_name = sampler_args["step"]
+ step_function = getattr(pm, step_function_name)
+ sampler_args["step"] = step_function()
+ idata = pm.sample(**sampler_args)
+ idata.extend(pm.sample_prior_predictive())
+ idata.extend(pm.sample_posterior_predictive(idata))
+ sampler_args["step"] = step_function_name
+ else:
+ idata = pm.sample(**sampler_args)
 
  idata = self.set_idata_attrs(idata)
  return idata
@@ -496,7 +503,13 @@ def fit(
  X_df = pd.DataFrame(X, columns=X.columns)
  combined_data = pd.concat([X_df, y], axis=1)
  assert all(combined_data.columns), "All columns must have non-empty names"
- self.idata.add_groups(fit_data=combined_data.to_xarray()) # type: ignore
+ with warnings.catch_warnings():
+ warnings.filterwarnings(
+ "ignore",
+ category=UserWarning,
+ message="The group fit_data is not defined in the InferenceData scheme",
+ )
+ self.idata.add_groups(fit_data=combined_data.to_xarray()) # type: ignore
 
  return self.idata # type: ignore
 

pymc_experimental/tests/test_model_builder.py

@@ -237,3 +237,12 @@ def test_id():
  ).hexdigest()[:16]
 
  assert model_builder.id == expected_id
+
+
+def test_step_selection_in_sample_config(toy_X, toy_y):
+ sampler_config = {
+ "step": "Slice",
+ }
+ model = test_ModelBuilder(sampler_config=sampler_config)
+ model.fit(toy_X, toy_y)
+ assert model.idata is not None