Merge branch 'main' into jitter_scale

aphc14 · aphc14 · commit 5f3f5ec5cd45 · 2025-01-19T19:12:52.000+11:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -48,7 +48,7 @@ repos:
 #        - --exclude=binder/
 #        - --exclude=versioneer.py
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.8.4
+  rev: v0.9.1
   hooks:
     - id: ruff
       args: [--fix, --show-fixes]
diff --git a/docs/source/learn/core_notebooks/pymc_pytensor.ipynb b/docs/source/learn/core_notebooks/pymc_pytensor.ipynb
@@ -1849,7 +1849,7 @@
     "print(\n",
     "    f\"\"\"\n",
     "mu_value -> {scipy.stats.norm.logpdf(x=0, loc=0, scale=2)}\n",
-    "sigma_log_value -> {- 10 + scipy.stats.halfnorm.logpdf(x=np.exp(-10), loc=0, scale=3)}\n",
+    "sigma_log_value -> {-10 + scipy.stats.halfnorm.logpdf(x=np.exp(-10), loc=0, scale=3)}\n",
     "x_value -> {scipy.stats.norm.logpdf(x=0, loc=0, scale=np.exp(-10))}\n",
     "\"\"\"\n",
     ")"
diff --git a/pymc/backends/zarr.py b/pymc/backends/zarr.py
@@ -15,14 +15,11 @@
 from typing import Any
 
 import arviz as az
-import numcodecs
 import numpy as np
 import xarray as xr
-import zarr
 
 from arviz.data.base import make_attrs
 from arviz.data.inference_data import WARMUP_TAG
-from numcodecs.abc import Codec
 from pytensor.tensor.variable import TensorVariable
 
 import pymc
@@ -44,11 +41,23 @@
 from pymc.util import UNSET, _UnsetType, get_default_varnames, is_transformed_name
 
 try:
+    import numcodecs
+    import zarr
+
+    from numcodecs.abc import Codec
+    from zarr import Group
     from zarr.storage import BaseStore, default_compressor
     from zarr.sync import Synchronizer
 
     _zarr_available = True
 except ImportError:
+    from typing import TYPE_CHECKING, TypeVar
+
+    if not TYPE_CHECKING:
+        Codec = TypeVar("Codec")
+        Group = TypeVar("Group")
+        BaseStore = TypeVar("BaseStore")
+        Synchronizer = TypeVar("Synchronizer")
     _zarr_available = False
 
 
@@ -243,7 +252,7 @@ def flush(self):
 
 def get_initial_fill_value_and_codec(
     dtype: Any,
-) -> tuple[FILL_VALUE_TYPE, np.typing.DTypeLike, numcodecs.abc.Codec | None]:
+) -> tuple[FILL_VALUE_TYPE, np.typing.DTypeLike, Codec | None]:
     _dtype = np.dtype(dtype)
     fill_value: FILL_VALUE_TYPE = None
     codec = None
@@ -366,27 +375,27 @@ def groups(self) -> list[str]:
         return [str(group_name) for group_name, _ in self.root.groups()]
 
     @property
-    def posterior(self) -> zarr.Group:
+    def posterior(self) -> Group:
         return self.root.posterior
 
     @property
-    def unconstrained_posterior(self) -> zarr.Group:
+    def unconstrained_posterior(self) -> Group:
         return self.root.unconstrained_posterior
 
     @property
-    def sample_stats(self) -> zarr.Group:
+    def sample_stats(self) -> Group:
         return self.root.sample_stats
 
     @property
-    def constant_data(self) -> zarr.Group:
+    def constant_data(self) -> Group:
         return self.root.constant_data
 
     @property
-    def observed_data(self) -> zarr.Group:
+    def observed_data(self) -> Group:
         return self.root.observed_data
 
     @property
-    def _sampling_state(self) -> zarr.Group:
+    def _sampling_state(self) -> Group:
         return self.root._sampling_state
 
     def init_trace(
@@ -646,12 +655,12 @@ def init_sampling_state_group(self, tune: int, chains: int):
 
     def init_group_with_empty(
         self,
-        group: zarr.Group,
+        group: Group,
         var_dtype_and_shape: dict[str, tuple[StatDtype, StatShape]],
         chains: int,
         draws: int,
         extra_var_attrs: dict | None = None,
-    ) -> zarr.Group:
+    ) -> Group:
         group_coords: dict[str, Any] = {"chain": range(chains), "draw": range(draws)}
         for name, (_dtype, shape) in var_dtype_and_shape.items():
             fill_value, dtype, object_codec = get_initial_fill_value_and_codec(_dtype)
@@ -689,8 +698,8 @@ def init_group_with_empty(
             array.attrs.update({"_ARRAY_DIMENSIONS": [dim]})
         return group
 
-    def create_group(self, name: str, data_dict: dict[str, np.ndarray]) -> zarr.Group | None:
-        group: zarr.Group | None = None
+    def create_group(self, name: str, data_dict: dict[str, np.ndarray]) -> Group | None:
+        group: Group | None = None
         if data_dict:
             group_coords = {}
             group = self.root.create_group(name=name, overwrite=True)
diff --git a/pymc/data.py b/pymc/data.py
@@ -257,7 +257,7 @@ def determine_coords(
     if isinstance(value, np.ndarray) and dims is not None:
         if len(dims) != value.ndim:
             raise pm.exceptions.ShapeError(
-                "Invalid data shape. The rank of the dataset must match the " "length of `dims`.",
+                "Invalid data shape. The rank of the dataset must match the length of `dims`.",
                 actual=value.shape,
                 expected=value.ndim,
             )
diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py
@@ -992,8 +992,7 @@ def get_mu_lam_phi(mu, lam, phi):
                     return mu, lam, lam / mu
 
         raise ValueError(
-            "Wald distribution must specify either mu only, "
-            "mu and lam, mu and phi, or lam and phi."
+            "Wald distribution must specify either mu only, mu and lam, mu and phi, or lam and phi."
         )
 
     def logp(value, mu, lam, alpha):
@@ -1603,8 +1602,7 @@ def dist(cls, kappa=None, mu=None, b=None, q=None, *args, **kwargs):
     def get_kappa(cls, kappa=None, q=None):
         if kappa is not None and q is not None:
             raise ValueError(
-                "Incompatible parameterization. Either use "
-                "kappa or q to specify the distribution."
+                "Incompatible parameterization. Either use kappa or q to specify the distribution."
             )
         elif q is not None:
             if isinstance(q, Variable):
@@ -3483,7 +3481,7 @@ def get_nu_b(cls, nu, b, sigma):
         elif nu is not None and b is None:
             b = nu / sigma
             return nu, b, sigma
-        raise ValueError("Rice distribution must specify either nu" " or b.")
+        raise ValueError("Rice distribution must specify either nu or b.")
 
     def support_point(rv, size, nu, sigma):
         nu_sigma_ratio = -(nu**2) / (2 * sigma**2)
diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py
@@ -247,7 +247,9 @@ class MvNormal(Continuous):
         data = np.random.multivariate_normal(mu, true_cov, 10)
 
         sd_dist = pm.Exponential.dist(1.0, shape=3)
-        chol, corr, stds = pm.LKJCholeskyCov("chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True)
+        chol, corr, stds = pm.LKJCholeskyCov(
+            "chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True
+        )
         vals = pm.MvNormal("vals", mu=mu, chol=chol, observed=data)
 
     For unobserved values it can be better to use a non-centered
@@ -2793,9 +2795,9 @@ def dist(cls, sigma=1.0, n_zerosum_axes=None, support_shape=None, **kwargs):
 
         support_shape = pt.as_tensor(support_shape, dtype="int64", ndim=1)
 
-        assert n_zerosum_axes == pt.get_vector_length(
-            support_shape
-        ), "support_shape has to be as long as n_zerosum_axes"
+        assert n_zerosum_axes == pt.get_vector_length(support_shape), (
+            "support_shape has to be as long as n_zerosum_axes"
+        )
 
         return super().dist([sigma, support_shape], **kwargs)
 
diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py
@@ -328,8 +328,7 @@ def power_spectral_density(self, omega: TensorLike) -> TensorVariable:
         check = Counter([isinstance(factor, Covariance) for factor in self._factor_list])
         if check.get(True, 0) >= 2:
             raise NotImplementedError(
-                "The power spectral density of products of covariance "
-                "functions is not implemented."
+                "The power spectral density of products of covariance functions is not implemented."
             )
         return reduce(mul, self._merge_factors_psd(omega))
 
diff --git a/pymc/gp/util.py b/pymc/gp/util.py
@@ -211,8 +211,7 @@ def plot_gp_dist(
         samples_kwargs = {}
     if np.any(np.isnan(samples)):
         warnings.warn(
-            "There are `nan` entries in the [samples] arguments. "
-            "The plot will not contain a band!",
+            "There are `nan` entries in the [samples] arguments. The plot will not contain a band!",
             UserWarning,
         )
 
diff --git a/pymc/sampling/jax.py b/pymc/sampling/jax.py
@@ -108,8 +108,7 @@ def _replace_shared_variables(graph: list[TensorVariable]) -> list[TensorVariabl
 
     if any(var.default_update is not None for var in shared_variables):
         raise ValueError(
-            "Graph contains shared variables with default_update which cannot "
-            "be safely replaced."
+            "Graph contains shared variables with default_update which cannot be safely replaced."
         )
 
     replacements = {var: pt.constant(var.get_value(borrow=True)) for var in shared_variables}
@@ -360,7 +359,7 @@ def _sample_blackjax_nuts(
         map_fn = jax.vmap
     else:
         raise ValueError(
-            "Only supporting the following methods to draw chains:" ' "parallel" or "vectorized"'
+            "Only supporting the following methods to draw chains: 'parallel' or 'vectorized'"
         )
 
     if chains == 1:
diff --git a/pymc/sampling/mcmc.py b/pymc/sampling/mcmc.py
@@ -41,7 +41,6 @@
 from rich.theme import Theme
 from threadpoolctl import threadpool_limits
 from typing_extensions import Protocol
-from zarr.storage import MemoryStore
 
 import pymc as pm
 
@@ -80,6 +79,11 @@
 )
 from pymc.vartypes import discrete_types
 
+try:
+    from zarr.storage import MemoryStore
+except ImportError:
+    MemoryStore = type("MemoryStore", (), {})
+
 sys.setrecursionlimit(10000)
 
 __all__ = [
@@ -996,7 +1000,7 @@ def _sample_return(
         total_draws = draws_per_chain.sum()
 
         _log.info(
-            f'Sampling {n_chains} chain{"s" if n_chains > 1 else ""} for {desired_tune:_d} desired tune and {desired_draw:_d} desired draw iterations '
+            f"Sampling {n_chains} chain{'s' if n_chains > 1 else ''} for {desired_tune:_d} desired tune and {desired_draw:_d} desired draw iterations "
             f"(Actually sampled {total_n_tune:_d} tune and {total_draws:_d} draws total) "
             f"took {t_sampling:.0f} seconds."
         )
@@ -1058,8 +1062,8 @@ def _sample_return(
 
     n_chains = len(mtrace.chains)
     _log.info(
-        f'Sampling {n_chains} chain{"s" if n_chains > 1 else ""} for {n_tune:_d} tune and {n_draws:_d} draw iterations '
-        f"({n_tune*n_chains:_d} + {n_draws*n_chains:_d} draws total) "
+        f"Sampling {n_chains} chain{'s' if n_chains > 1 else ''} for {n_tune:_d} tune and {n_draws:_d} draw iterations "
+        f"({n_tune * n_chains:_d} + {n_draws * n_chains:_d} draws total) "
         f"took {t_sampling:.0f} seconds."
     )
 
diff --git a/pymc/sampling/population.py b/pymc/sampling/population.py
@@ -386,9 +386,9 @@ def _prepare_iter_population(
 
     # 2. Set up the steppers
     steppers: list[Step] = []
-    assert (
-        len(rngs) == nchains
-    ), f"There must be one random Generator per chain. Got {len(rngs)} instead of {nchains}"
+    assert len(rngs) == nchains, (
+        f"There must be one random Generator per chain. Got {len(rngs)} instead of {nchains}"
+    )
     for c, rng in enumerate(rngs):
         # need independent samplers for each chain
         # it is important to copy the actual steppers (but not the delta_logp)
diff --git a/pymc/step_methods/compound.py b/pymc/step_methods/compound.py
@@ -282,9 +282,9 @@ def sampling_state(self) -> DataClassState:
 
     @sampling_state.setter
     def sampling_state(self, state: DataClassState):
-        assert isinstance(
-            state, self._state_class
-        ), f"Invalid sampling state class {type(state)}. Expected {self._state_class}"
+        assert isinstance(state, self._state_class), (
+            f"Invalid sampling state class {type(state)}. Expected {self._state_class}"
+        )
         for method, state_method in zip(self.methods, state.methods):
             method.sampling_state = state_method
 
diff --git a/pymc/step_methods/state.py b/pymc/step_methods/state.py
@@ -90,9 +90,9 @@ def sampling_state(self) -> DataClassState:
     @sampling_state.setter
     def sampling_state(self, state: DataClassState):
         state_class = self._state_class
-        assert isinstance(
-            state, state_class
-        ), f"Encountered invalid state class '{state.__class__}'. State must be '{state_class}'"
+        assert isinstance(state, state_class), (
+            f"Encountered invalid state class '{state.__class__}'. State must be '{state_class}'"
+        )
         for field in fields(state_class):
             is_tensor_name = field.metadata.get("tensor_name", False)
             state_val = deepcopy(getattr(state, field.name))
diff --git a/pymc/testing.py b/pymc/testing.py
@@ -964,9 +964,9 @@ def check_rv_size(self):
                 assert actual == expected_symbolic == expected
 
     def validate_tests_list(self):
-        assert len(self.checks_to_run) == len(
-            set(self.checks_to_run)
-        ), "There are duplicates in the list of checks_to_run"
+        assert len(self.checks_to_run) == len(set(self.checks_to_run)), (
+            "There are duplicates in the list of checks_to_run"
+        )
 
 
 def seeded_scipy_distribution_builder(dist_name: str) -> Callable:
diff --git a/pymc/variational/opvi.py b/pymc/variational/opvi.py
@@ -710,9 +710,9 @@ class Group(WithMemoization):
 
     @classmethod
     def register(cls, sbcls):
-        assert (
-            frozenset(sbcls.__param_spec__) not in cls.__param_registry
-        ), "Duplicate __param_spec__"
+        assert frozenset(sbcls.__param_spec__) not in cls.__param_registry, (
+            "Duplicate __param_spec__"
+        )
         cls.__param_registry[frozenset(sbcls.__param_spec__)] = sbcls
         assert sbcls.short_name not in cls.__name_registry, "Duplicate short_name"
         cls.__name_registry[sbcls.short_name] = sbcls
@@ -1234,7 +1234,7 @@ def __init__(self, groups, model=None):
         for g in groups:
             if g.group is None:
                 if rest is not None:
-                    raise GroupError("More than one group is specified for " "the rest variables")
+                    raise GroupError("More than one group is specified for the rest variables")
                 else:
                     rest = g
             else:
diff --git a/pymc/variational/updates.py b/pymc/variational/updates.py
@@ -1006,7 +1006,7 @@ def norm_constraint(tensor_var, max_norm, norm_axes=None, epsilon=1e-7):
     elif ndim in [3, 4, 5]:  # Conv{1,2,3}DLayer
         sum_over = tuple(range(1, ndim))
     else:
-        raise ValueError(f"Unsupported tensor dimensionality {ndim}." "Must specify `norm_axes`")
+        raise ValueError(f"Unsupported tensor dimensionality {ndim}. Must specify `norm_axes`")
 
     dtype = np.dtype(pytensor.config.floatX).type
     norms = pt.sqrt(pt.sum(pt.sqr(tensor_var), axis=sum_over, keepdims=True))
diff --git a/tests/distributions/test_multivariate.py b/tests/distributions/test_multivariate.py
@@ -1531,14 +1531,14 @@ class TestZeroSumNormal:
     def assert_zerosum_axes(self, random_samples, axes_to_check, check_zerosum_axes=True):
         if check_zerosum_axes:
             for ax in axes_to_check:
-                assert np.isclose(
-                    random_samples.mean(axis=ax), 0
-                ).all(), f"{ax} is a zerosum_axis but is not summing to 0 across all samples."
+                assert np.isclose(random_samples.mean(axis=ax), 0).all(), (
+                    f"{ax} is a zerosum_axis but is not summing to 0 across all samples."
+                )
         else:
             for ax in axes_to_check:
-                assert not np.isclose(
-                    random_samples.mean(axis=ax), 0
-                ).all(), f"{ax} is not a zerosum_axis, but is nonetheless summing to 0 across all samples."
+                assert not np.isclose(random_samples.mean(axis=ax), 0).all(), (
+                    f"{ax} is not a zerosum_axis, but is nonetheless summing to 0 across all samples."
+                )
 
     @pytest.mark.parametrize(
         "dims, n_zerosum_axes",
diff --git a/tests/gp/test_hsgp_approx.py b/tests/gp/test_hsgp_approx.py
@@ -135,9 +135,9 @@ def test_mean_invariance(self):
         with model:
             pm.set_data({"X": x_new})
 
-        assert np.allclose(
-            gp._X_center, original_center
-        ), "gp._X_center should not change after updating data for out-of-sample predictions."
+        assert np.allclose(gp._X_center, original_center), (
+            "gp._X_center should not change after updating data for out-of-sample predictions."
+        )
 
     def test_parametrization(self):
         err_msg = (
@@ -188,9 +188,9 @@ def test_parametrization_drop_first(self, model, cov_func, X1, drop_first):
 
             n_coeffs = model.f1_hsgp_coeffs.type.shape[0]
             if drop_first:
-                assert (
-                    n_coeffs == n_basis - 1
-                ), f"one basis vector should have been dropped, {n_coeffs}"
+                assert n_coeffs == n_basis - 1, (
+                    f"one basis vector should have been dropped, {n_coeffs}"
+                )
             else:
                 assert n_coeffs == n_basis, "one was dropped when it shouldn't have been"
 
diff --git a/tests/test_data.py b/tests/test_data.py

Original file line number	Diff line number	Diff line change
`@@ -257,7 +257,7 @@ def determine_coords(`
`257`	`257`	`if isinstance(value, np.ndarray) and dims is not None:`
`258`	`258`	`if len(dims) != value.ndim:`
`259`	`259`	`raise pm.exceptions.ShapeError(`
`260`		- "Invalid data shape. The rank of the dataset must match the " "length of `dims`.",
	`260`	+ "Invalid data shape. The rank of the dataset must match the length of `dims`.",
`261`	`261`	`actual=value.shape,`
`262`	`262`	`expected=value.ndim,`
`263`	`263`	`)`
Original file line number	Diff line number	Diff line change
`@@ -328,8 +328,7 @@ def power_spectral_density(self, omega: TensorLike) -> TensorVariable:`
`328`	`328`	`check = Counter([isinstance(factor, Covariance) for factor in self._factor_list])`
`329`	`329`	`if check.get(True, 0) >= 2:`
`330`	`330`	`raise NotImplementedError(`
`331`		`- "The power spectral density of products of covariance "`
`332`		`- "functions is not implemented."`
	`331`	`+ "The power spectral density of products of covariance functions is not implemented."`
`333`	`332`	`)`
`334`	`333`	`return reduce(mul, self._merge_factors_psd(omega))`
`335`	`334`
Original file line number	Diff line number	Diff line change
`@@ -211,8 +211,7 @@ def plot_gp_dist(`
`211`	`211`	`samples_kwargs = {}`
`212`	`212`	`if np.any(np.isnan(samples)):`
`213`	`213`	`warnings.warn(`
`214`		- "There are `nan` entries in the [samples] arguments. "
`215`		`- "The plot will not contain a band!",`
	`214`	+ "There are `nan` entries in the [samples] arguments. The plot will not contain a band!",
`216`	`215`	`UserWarning,`
`217`	`216`	`)`
`218`	`217`