From 80b3ebb9fc9da5cc6cef08a953df7f340ea24c47 Mon Sep 17 00:00:00 2001 From: Jesse Grabowski Date: Fri, 26 Jul 2024 20:13:46 +0800 Subject: [PATCH 1/4] Update pre-commit to use ruff --- .pre-commit-config.yaml | 37 +++++++-------------------- pyproject.toml | 56 +++++++++++++++++++++++++++++++++-------- 2 files changed, 55 insertions(+), 38 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 70fe11db9..78c4a366b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.5.0 + rev: v4.6.0 hooks: - id: check-merge-conflict - id: check-toml @@ -10,37 +10,18 @@ repos: - id: no-commit-to-branch args: [--branch, main] - id: trailing-whitespace -- repo: https://github.com/PyCQA/isort - rev: 5.13.2 + +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.5.5 hooks: - - id: isort - name: isort -- repo: https://github.com/asottile/pyupgrade - rev: v3.15.0 - hooks: - - id: pyupgrade - args: [--py37-plus] -- repo: https://github.com/psf/black - rev: 24.1.1 - hooks: - - id: black - - id: black-jupyter -- repo: https://github.com/PyCQA/pylint - rev: v3.0.3 - hooks: - - id: pylint - args: [--rcfile=.pylintrc] - files: ^pymc_experimental/ + - id: ruff + args: [ --fix, --unsafe-fixes, --exit-non-zero-on-fix ] + - id: ruff-format + types_or: [ python, pyi, jupyter ] + - repo: https://github.com/MarcoGorelli/madforhooks rev: 0.4.1 hooks: - id: no-print-statements exclude: _version.py files: ^pymc_experimental/ -- repo: local - hooks: - - id: no-relative-imports - name: No relative imports - entry: from \.[\.\w]* import - types: [python] - language: pygrep diff --git a/pyproject.toml b/pyproject.toml index 2a5ee3247..1fb70104c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,9 +22,6 @@ filterwarnings =[ 'ignore:jax\.tree_map is deprecated:DeprecationWarning', ] -[tool.black] -line-length = 100 - [tool.coverage.report] exclude_lines = [ "pragma: nocover", @@ -32,11 +29,50 @@ exclude_lines = [ "if TYPE_CHECKING:", ] -[tool.isort] -profile = "black" -# lines_between_types = 1 +[tool.ruff] +line-length = 100 +target-version = "py310" + +[tool.ruff.lint] +select = ["D", "E", "F", "I", "UP", "W", "RUF"] +ignore = [ + "E501", + "RUF001", # String contains ambiguous character (such as Greek letters) + "RUF002", # Docstring contains ambiguous character (such as Greek letters) + "RUF012", # Mutable class attributes should be annotated with `typing.ClassVar` + "D100", + "D101", + "D102", + "D103", + "D104", + "D105", + "D107", + "D200", + "D202", + "D203", + "D204", + "D205", + "D209", + "D212", + "D213", + "D301", + "D400", + "D401", + "D403", + "D413", + "D415", + "D417", +] + +[tool.ruff.lint.isort] +lines-between-types = 1 -[tool.nbqa.mutate] -isort = 1 -black = 1 -pyupgrade = 1 +[tool.ruff.lint.per-file-ignores] +'tests/*.py' = [ + 'F841', # Unused variable warning for test files -- common in pymc model declarations + 'D106' # Missing docstring for public method -- unittest test subclasses don't need docstrings +] +'tests/statespace/*.py' = [ + 'F401', # Unused import warning for test files -- this check removes imports of fixtures + 'F811' # Redefine while unused -- this check fails on imported fixtures +] From 9aee194b0f41b4c7b14637c49593c90d5f79a841 Mon Sep 17 00:00:00 2001 From: Jesse Grabowski Date: Sat, 27 Jul 2024 17:34:10 +0800 Subject: [PATCH 2/4] Run ruff with new configuration --- notebooks/SARMA Example.ipynb | 4 +- .../Structural Timeseries Modeling.ipynb | 6 +- pymc_experimental/distributions/continuous.py | 9 +- pymc_experimental/distributions/discrete.py | 1 + .../distributions/histogram_utils.py | 13 +- .../distributions/multivariate/r2d2m2cp.py | 42 +++--- pymc_experimental/distributions/timeseries.py | 6 +- pymc_experimental/gp/latent_approx.py | 14 +- pymc_experimental/inference/fit.py | 1 - pymc_experimental/inference/laplace.py | 8 +- pymc_experimental/inference/pathfinder.py | 4 +- pymc_experimental/inference/smc/sampling.py | 12 +- pymc_experimental/linearmodel.py | 14 +- pymc_experimental/model/marginal_model.py | 20 +-- .../model/transforms/autoreparam.py | 33 ++--- pymc_experimental/model_builder.py | 56 ++++---- .../preprocessing/standard_scaler.py | 1 + .../statespace/core/representation.py | 27 ++-- .../statespace/core/statespace.py | 125 +++++++++--------- .../statespace/filters/distributions.py | 8 +- .../statespace/filters/kalman_filter.py | 14 +- .../statespace/filters/kalman_smoother.py | 7 +- .../statespace/filters/utilities.py | 1 + .../statespace/models/SARIMAX.py | 13 +- pymc_experimental/statespace/models/VARMAX.py | 19 +-- .../statespace/models/structural.py | 34 ++--- .../statespace/models/utilities.py | 4 +- .../statespace/utils/data_tools.py | 9 +- pymc_experimental/utils/linear_cg.py | 15 ++- pymc_experimental/utils/pivoted_cholesky.py | 5 +- pymc_experimental/utils/prior.py | 24 ++-- pymc_experimental/utils/spline.py | 5 +- setup.py | 3 +- setupegg.py | 2 - tests/distributions/test_continuous.py | 1 + tests/distributions/test_discrete.py | 1 + .../test_discrete_markov_chain.py | 1 + tests/model/test_marginal_model.py | 17 +-- tests/statespace/test_SARIMAX.py | 3 +- tests/statespace/test_VARMAX.py | 3 +- tests/statespace/test_distributions.py | 5 +- tests/statespace/test_kalman_filter.py | 3 +- tests/statespace/test_representation.py | 1 + tests/statespace/test_statespace.py | 1 + tests/statespace/test_statespace_JAX.py | 1 + tests/statespace/test_structural.py | 32 ++--- .../utilities/statsmodel_local_level.py | 2 +- tests/statespace/utilities/test_helpers.py | 1 + tests/test_blackjax_smc.py | 1 + tests/test_laplace.py | 3 - tests/test_model_builder.py | 7 +- tests/test_prior_from_trace.py | 1 + tests/test_splines.py | 1 + tests/utils.py | 2 +- 54 files changed, 345 insertions(+), 301 deletions(-) diff --git a/notebooks/SARMA Example.ipynb b/notebooks/SARMA Example.ipynb index 66981fb28..ab85c9dc2 100644 --- a/notebooks/SARMA Example.ipynb +++ b/notebooks/SARMA Example.ipynb @@ -1554,7 +1554,7 @@ " hdi_forecast.coords[\"time\"].values,\n", " *hdi_forecast.isel(observed_state=0).values.T,\n", " alpha=0.25,\n", - " color=\"tab:blue\"\n", + " color=\"tab:blue\",\n", " )\n", "ax.set_title(\"Porcupine Graph of 10-Period Forecasts (parameters estimated on all data)\")\n", "plt.show()" @@ -2692,7 +2692,7 @@ " *forecast_hdi.values.T,\n", " label=\"Forecast 94% HDI\",\n", " color=\"tab:orange\",\n", - " alpha=0.25\n", + " alpha=0.25,\n", ")\n", "ax.legend()\n", "plt.show()" diff --git a/notebooks/Structural Timeseries Modeling.ipynb b/notebooks/Structural Timeseries Modeling.ipynb index 2c62ecaef..01dc8939b 100644 --- a/notebooks/Structural Timeseries Modeling.ipynb +++ b/notebooks/Structural Timeseries Modeling.ipynb @@ -1657,7 +1657,7 @@ " nile.index,\n", " *component_hdi.smoothed_posterior.sel(state=state).values.T,\n", " color=\"tab:blue\",\n", - " alpha=0.15\n", + " alpha=0.15,\n", " )\n", " axis.set_title(state.title())" ] @@ -1706,7 +1706,7 @@ " *hdi.smoothed_posterior.sum(dim=\"state\").values.T,\n", " color=\"tab:blue\",\n", " alpha=0.15,\n", - " label=\"HDI 94%\"\n", + " label=\"HDI 94%\",\n", ")\n", "ax.legend()\n", "plt.show()" @@ -2750,7 +2750,7 @@ "ax.fill_between(\n", " blossom_data.index,\n", " *hdi_post.predicted_posterior_observed.isel(observed_state=0).values.T,\n", - " alpha=0.25\n", + " alpha=0.25,\n", ")\n", "blossom_data.plot(ax=ax)" ] diff --git a/pymc_experimental/distributions/continuous.py b/pymc_experimental/distributions/continuous.py index 6c2a57002..3ed1f6f5b 100644 --- a/pymc_experimental/distributions/continuous.py +++ b/pymc_experimental/distributions/continuous.py @@ -19,10 +19,9 @@ The imports from pymc are not fully replicated here: add imports as necessary. """ -from typing import Tuple, Union - import numpy as np import pytensor.tensor as pt + from pymc import ChiSquared, CustomDist from pymc.distributions import transforms from pymc.distributions.dist_math import check_parameters @@ -39,7 +38,7 @@ class GenExtremeRV(RandomVariable): name: str = "Generalized Extreme Value" signature = "(),(),()->()" dtype: str = "floatX" - _print_name: Tuple[str, str] = ("Generalized Extreme Value", "\\operatorname{GEV}") + _print_name: tuple[str, str] = ("Generalized Extreme Value", "\\operatorname{GEV}") def __call__(self, mu=0.0, sigma=1.0, xi=0.0, size=None, **kwargs) -> TensorVariable: return super().__call__(mu, sigma, xi, size=size, **kwargs) @@ -47,11 +46,11 @@ def __call__(self, mu=0.0, sigma=1.0, xi=0.0, size=None, **kwargs) -> TensorVari @classmethod def rng_fn( cls, - rng: Union[np.random.RandomState, np.random.Generator], + rng: np.random.RandomState | np.random.Generator, mu: np.ndarray, sigma: np.ndarray, xi: np.ndarray, - size: Tuple[int, ...], + size: tuple[int, ...], ) -> np.ndarray: # Notice negative here, since remainder of GenExtreme is based on Coles parametrization return stats.genextreme.rvs(c=-xi, loc=mu, scale=sigma, random_state=rng, size=size) diff --git a/pymc_experimental/distributions/discrete.py b/pymc_experimental/distributions/discrete.py index 3934baa88..0bfe8e7c4 100644 --- a/pymc_experimental/distributions/discrete.py +++ b/pymc_experimental/distributions/discrete.py @@ -14,6 +14,7 @@ import numpy as np import pymc as pm + from pymc.distributions.dist_math import betaln, check_parameters, factln, logpow from pymc.distributions.shape_utils import rv_size_is_none from pytensor import tensor as pt diff --git a/pymc_experimental/distributions/histogram_utils.py b/pymc_experimental/distributions/histogram_utils.py index 7824d08ca..44b9a8242 100644 --- a/pymc_experimental/distributions/histogram_utils.py +++ b/pymc_experimental/distributions/histogram_utils.py @@ -13,10 +13,9 @@ # limitations under the License. -from typing import Dict - import numpy as np import pymc as pm + from numpy.typing import ArrayLike __all__ = ["quantile_histogram", "discrete_histogram", "histogram_approximation"] @@ -24,7 +23,7 @@ def quantile_histogram( data: ArrayLike, n_quantiles=1000, zero_inflation=False -) -> Dict[str, ArrayLike]: +) -> dict[str, ArrayLike]: try: import xhistogram.core except ImportError as e: @@ -34,7 +33,7 @@ def quantile_histogram( import dask.dataframe except ImportError: dask = None - if dask and isinstance(data, (dask.dataframe.Series, dask.dataframe.DataFrame)): + if dask and isinstance(data, dask.dataframe.Series | dask.dataframe.DataFrame): data = data.to_dask_array(lengths=True) if zero_inflation: zeros = (data == 0).sum(0) @@ -67,7 +66,7 @@ def quantile_histogram( return result -def discrete_histogram(data: ArrayLike, min_count=None) -> Dict[str, ArrayLike]: +def discrete_histogram(data: ArrayLike, min_count=None) -> dict[str, ArrayLike]: try: import xhistogram.core except ImportError as e: @@ -78,7 +77,7 @@ def discrete_histogram(data: ArrayLike, min_count=None) -> Dict[str, ArrayLike]: except ImportError: dask = None - if dask and isinstance(data, (dask.dataframe.Series, dask.dataframe.DataFrame)): + if dask and isinstance(data, dask.dataframe.Series | dask.dataframe.DataFrame): data = data.to_dask_array(lengths=True) mid, count_uniq = np.unique(data, return_counts=True) if min_count is not None: @@ -153,7 +152,7 @@ def histogram_approximation(name, dist, *, observed, **h_kwargs): import dask.dataframe except ImportError: dask = None - if dask and isinstance(observed, (dask.dataframe.Series, dask.dataframe.DataFrame)): + if dask and isinstance(observed, dask.dataframe.Series | dask.dataframe.DataFrame): observed = observed.to_dask_array(lengths=True) if np.issubdtype(observed.dtype, np.integer): histogram = discrete_histogram(observed, **h_kwargs) diff --git a/pymc_experimental/distributions/multivariate/r2d2m2cp.py b/pymc_experimental/distributions/multivariate/r2d2m2cp.py index f6214a7bd..534992299 100644 --- a/pymc_experimental/distributions/multivariate/r2d2m2cp.py +++ b/pymc_experimental/distributions/multivariate/r2d2m2cp.py @@ -14,7 +14,7 @@ from collections import namedtuple -from typing import Sequence, Tuple, Union +from collections.abc import Sequence import numpy as np import pymc as pm @@ -26,8 +26,8 @@ def _psivar2musigma( psi: pt.TensorVariable, explained_var: pt.TensorVariable, - psi_mask: Union[pt.TensorLike, None], -) -> Tuple[pt.TensorVariable, pt.TensorVariable]: + psi_mask: pt.TensorLike | None, +) -> tuple[pt.TensorVariable, pt.TensorVariable]: sign = pt.sign(psi - 0.5) if psi_mask is not None: # any computation might be ignored for ~psi_mask @@ -55,7 +55,7 @@ def _R2D2M2CP_beta( psi: pt.TensorVariable, *, psi_mask, - dims: Union[str, Sequence[str]], + dims: str | Sequence[str], centered=False, ) -> pt.TensorVariable: """R2D2M2CP beta prior. @@ -120,7 +120,7 @@ def _R2D2M2CP_beta( def _broadcast_as_dims( *values: np.ndarray, dims: Sequence[str], -) -> Union[Tuple[np.ndarray, ...], np.ndarray]: +) -> tuple[np.ndarray, ...] | np.ndarray: model = pm.modelcontext(None) shape = [len(model.coords[d]) for d in dims] ret = tuple(np.broadcast_to(v, shape) for v in values) @@ -135,7 +135,7 @@ def _psi_masked( positive_probs_std: pt.TensorLike, *, dims: Sequence[str], -) -> Tuple[Union[pt.TensorLike, None], pt.TensorVariable]: +) -> tuple[pt.TensorLike | None, pt.TensorVariable]: if not ( isinstance(positive_probs, pt.Constant) and isinstance(positive_probs_std, pt.Constant) ): @@ -172,10 +172,10 @@ def _psi_masked( def _psi( positive_probs: pt.TensorLike, - positive_probs_std: Union[pt.TensorLike, None], + positive_probs_std: pt.TensorLike | None, *, dims: Sequence[str], -) -> Tuple[Union[pt.TensorLike, None], pt.TensorVariable]: +) -> tuple[pt.TensorLike | None, pt.TensorVariable]: if positive_probs_std is not None: mask, psi = _psi_masked( positive_probs=pt.as_tensor(positive_probs), @@ -194,9 +194,9 @@ def _psi( def _phi( - variables_importance: Union[pt.TensorLike, None], - variance_explained: Union[pt.TensorLike, None], - importance_concentration: Union[pt.TensorLike, None], + variables_importance: pt.TensorLike | None, + variance_explained: pt.TensorLike | None, + importance_concentration: pt.TensorLike | None, *, dims: Sequence[str], ) -> pt.TensorVariable: @@ -210,7 +210,7 @@ def _phi( variables_importance = pt.as_tensor(variables_importance) if importance_concentration is not None: variables_importance *= importance_concentration - return pm.Dirichlet("phi", variables_importance, dims=broadcast_dims + [dim]) + return pm.Dirichlet("phi", variables_importance, dims=[*broadcast_dims, dim]) elif variance_explained is not None: if len(model.coords[dim]) <= 1: raise TypeError("Can't use variance explained with less than two variables") @@ -218,7 +218,7 @@ def _phi( else: phi = _broadcast_as_dims(1.0, dims=dims) if importance_concentration is not None: - return pm.Dirichlet("phi", importance_concentration * phi, dims=broadcast_dims + [dim]) + return pm.Dirichlet("phi", importance_concentration * phi, dims=[*broadcast_dims, dim]) else: return phi @@ -233,12 +233,12 @@ def R2D2M2CP( *, dims: Sequence[str], r2: pt.TensorLike, - variables_importance: Union[pt.TensorLike, None] = None, - variance_explained: Union[pt.TensorLike, None] = None, - importance_concentration: Union[pt.TensorLike, None] = None, - r2_std: Union[pt.TensorLike, None] = None, - positive_probs: Union[pt.TensorLike, None] = 0.5, - positive_probs_std: Union[pt.TensorLike, None] = None, + variables_importance: pt.TensorLike | None = None, + variance_explained: pt.TensorLike | None = None, + importance_concentration: pt.TensorLike | None = None, + r2_std: pt.TensorLike | None = None, + positive_probs: pt.TensorLike | None = 0.5, + positive_probs_std: pt.TensorLike | None = None, centered: bool = False, ) -> R2D2M2CPOut: """R2D2M2CP Prior. @@ -413,7 +413,7 @@ def R2D2M2CP( year = {2023} } """ - if not isinstance(dims, (list, tuple)): + if not isinstance(dims, list | tuple): dims = (dims,) *broadcast_dims, dim = dims input_sigma = pt.as_tensor(input_sigma) @@ -438,7 +438,7 @@ def R2D2M2CP( r2, phi, psi, - dims=broadcast_dims + [dim], + dims=[*broadcast_dims, dim], centered=centered, psi_mask=mask, ) diff --git a/pymc_experimental/distributions/timeseries.py b/pymc_experimental/distributions/timeseries.py index 91da141ac..0e8659915 100644 --- a/pymc_experimental/distributions/timeseries.py +++ b/pymc_experimental/distributions/timeseries.py @@ -1,10 +1,10 @@ import warnings -from typing import List, Union import numpy as np import pymc as pm import pytensor import pytensor.tensor as pt + from pymc.distributions.dist_math import check_parameters from pymc.distributions.distribution import ( Distribution, @@ -26,7 +26,7 @@ from pytensor.tensor.random.op import RandomVariable -def _make_outputs_info(n_lags: int, init_dist: Distribution) -> List[Union[Distribution, dict]]: +def _make_outputs_info(n_lags: int, init_dist: Distribution) -> list[Distribution | dict]: """ Two cases are needed for outputs_info in the scans used by DiscreteMarkovRv. If n_lags = 1, we need to throw away the first dimension of init_dist_ or else markov_chain will have shape (steps, 1, *batch_size) instead of @@ -142,7 +142,7 @@ def dist(cls, P=None, logit_P=None, steps=None, init_dist=None, n_lags=1, **kwar if init_dist is not None: if not isinstance(init_dist, TensorVariable) or not isinstance( - init_dist.owner.op, (RandomVariable, SymbolicRandomVariable) + init_dist.owner.op, RandomVariable | SymbolicRandomVariable ): raise ValueError( f"Init dist must be a distribution created via the `.dist()` API, " diff --git a/pymc_experimental/gp/latent_approx.py b/pymc_experimental/gp/latent_approx.py index ddcbb845d..346e5895f 100644 --- a/pymc_experimental/gp/latent_approx.py +++ b/pymc_experimental/gp/latent_approx.py @@ -12,11 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. from functools import partial -from typing import Optional import numpy as np import pymc as pm import pytensor.tensor as pt + from pymc.gp.util import JITTER_DEFAULT, stabilize from pytensor.tensor.linalg import cholesky, solve_triangular @@ -33,7 +33,7 @@ class ProjectedProcess(pm.gp.Latent): ## AKA: DTC def __init__( self, - n_inducing: Optional[int] = None, + n_inducing: int | None = None, *, mean_func=pm.gp.mean.Zero(), cov_func=pm.gp.cov.Constant(0.0), @@ -59,20 +59,22 @@ def prior( self, name: str, X: np.ndarray, - X_inducing: Optional[np.ndarray] = None, + X_inducing: np.ndarray | None = None, jitter: float = JITTER_DEFAULT, **kwargs, ) -> np.ndarray: """ Builds the GP prior with optional inducing points locations. - Parameters: + Parameters + ---------- - name: Name for the GP variable. - X: Input data. - X_inducing: Optional. Inducing points for the GP. - jitter: Jitter to ensure numerical stability. - Returns: + Returns + ------- - GP function """ # Check if X is a numpy array @@ -137,7 +139,7 @@ def __init__( super().__init__(mean_func=mean_func, cov_func=cov_func) def _build_prior(self, name, X, jitter=1e-6, **kwargs): - mu = self.mean_func(X) + self.mean_func(X) Kxx = pm.gp.util.stabilize(self.cov_func(X), jitter) vals, vecs = pt.linalg.eigh(Kxx) ## NOTE: REMOVED PRECISION CUTOFF diff --git a/pymc_experimental/inference/fit.py b/pymc_experimental/inference/fit.py index 71dfb9f8b..9b1f891ed 100644 --- a/pymc_experimental/inference/fit.py +++ b/pymc_experimental/inference/fit.py @@ -40,7 +40,6 @@ def fit(method, **kwargs): return fit_pathfinder(**kwargs) if method == "laplace": - from pymc_experimental.inference.laplace import laplace return laplace(**kwargs) diff --git a/pymc_experimental/inference/laplace.py b/pymc_experimental/inference/laplace.py index 1508b6e88..7d7beb59b 100644 --- a/pymc_experimental/inference/laplace.py +++ b/pymc_experimental/inference/laplace.py @@ -13,13 +13,14 @@ # limitations under the License. import warnings + from collections.abc import Sequence -from typing import Optional import arviz as az import numpy as np import pymc as pm import xarray as xr + from arviz import dict_to_dataset from pymc.backends.arviz import ( coords_and_dims_for_inferencedata, @@ -33,9 +34,9 @@ def laplace( vars: Sequence[Variable], - draws: Optional[int] = 1000, + draws: int | None = 1000, model=None, - random_seed: Optional[RandomSeed] = None, + random_seed: RandomSeed | None = None, progressbar=True, ): """ @@ -72,7 +73,6 @@ def laplace( Examples -------- - >>> import numpy as np >>> import pymc as pm >>> import arviz as az diff --git a/pymc_experimental/inference/pathfinder.py b/pymc_experimental/inference/pathfinder.py index 5e5533dd6..89e621c88 100644 --- a/pymc_experimental/inference/pathfinder.py +++ b/pymc_experimental/inference/pathfinder.py @@ -14,13 +14,13 @@ import collections import sys -from typing import Optional import arviz as az import blackjax import jax import numpy as np import pymc as pm + from packaging import version from pymc.backends.arviz import coords_and_dims_for_inferencedata from pymc.blocking import DictToArrayBijection, RaveledVars @@ -63,7 +63,7 @@ def convert_flat_trace_to_idata( def fit_pathfinder( samples=1000, - random_seed: Optional[RandomSeed] = None, + random_seed: RandomSeed | None = None, postprocessing_backend="cpu", model=None, **pathfinder_kwargs, diff --git a/pymc_experimental/inference/smc/sampling.py b/pymc_experimental/inference/smc/sampling.py index 93f8a8a32..99488c85c 100644 --- a/pymc_experimental/inference/smc/sampling.py +++ b/pymc_experimental/inference/smc/sampling.py @@ -15,13 +15,16 @@ import logging import time import warnings -from typing import Callable, Dict, NamedTuple, Optional, cast + +from collections.abc import Callable +from typing import NamedTuple, cast import arviz as az import blackjax import jax import jax.numpy as jnp import numpy as np + from blackjax.smc.resampling import systematic from pymc import draw, modelcontext, to_inference_data from pymc.backends import NDArray @@ -39,7 +42,7 @@ def sample_smc_blackjax( kernel: str = "HMC", target_essn: float = 0.5, num_mcmc_steps: int = 10, - inner_kernel_params: Optional[dict] = None, + inner_kernel_params: dict | None = None, model=None, iterations_to_diagnose: int = 100, ): @@ -319,6 +322,7 @@ def add_to_inference_data( ): """ Adds several SMC parameters into the az.InferenceData result + Parameters ---------- inference_data: arviz object to add attributes to. @@ -389,7 +393,7 @@ def logp_fn_wrap(particles): return logp_fn_wrap -def initialize_population(model, draws, random_seed) -> Dict[str, np.ndarray]: +def initialize_population(model, draws, random_seed) -> dict[str, np.ndarray]: with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning, message="The effect of Potentials") @@ -405,7 +409,7 @@ def initialize_population(model, draws, random_seed) -> Dict[str, np.ndarray]: names = [model.rvs_to_values[rv].name for rv in model.free_RVs] dict_prior = {k: np.stack(v) for k, v in zip(names, prior_values)} - return cast(Dict[str, np.ndarray], dict_prior) + return cast(dict[str, np.ndarray], dict_prior) def var_map_from_model(model, initial_point) -> dict: diff --git a/pymc_experimental/linearmodel.py b/pymc_experimental/linearmodel.py index 0c4237dab..08c3e0331 100644 --- a/pymc_experimental/linearmodel.py +++ b/pymc_experimental/linearmodel.py @@ -1,5 +1,3 @@ -from typing import Dict, Optional, Union - import numpy as np import pandas as pd import pymc as pm @@ -8,7 +6,9 @@ class LinearModel(ModelBuilder): - def __init__(self, model_config: Dict = None, sampler_config: Dict = None, nsamples=100): + def __init__( + self, model_config: dict | None = None, sampler_config: dict | None = None, nsamples=100 + ): self.nsamples = nsamples super().__init__(model_config, sampler_config) @@ -38,7 +38,7 @@ def get_default_sampler_config(): } @property - def _serializable_model_config(self) -> Dict: + def _serializable_model_config(self) -> dict: return self.model_config @property @@ -83,7 +83,7 @@ def build_model(self, X: pd.DataFrame, y: pd.Series): y_model = pm.Deterministic("y_model", intercept + slope * x, dims="observation") # observed data - y_hat = pm.Normal( + pm.Normal( "y_hat", y_model, sigma=obs_error, @@ -94,14 +94,14 @@ def build_model(self, X: pd.DataFrame, y: pd.Series): self._data_setter(X, y) - def _data_setter(self, X: pd.DataFrame, y: Optional[Union[pd.DataFrame, pd.Series]] = None): + def _data_setter(self, X: pd.DataFrame, y: pd.DataFrame | pd.Series | None = None): with self.model: pm.set_data({"x": X.squeeze()}) if y is not None: pm.set_data({"y_data": y.squeeze()}) def _generate_and_preprocess_model_data( - self, X: Union[pd.DataFrame, pd.Series], y: pd.Series + self, X: pd.DataFrame | pd.Series, y: pd.Series ) -> None: """ Generate model data for linear regression. diff --git a/pymc_experimental/model/marginal_model.py b/pymc_experimental/model/marginal_model.py index ead9a362b..220e542f6 100644 --- a/pymc_experimental/model/marginal_model.py +++ b/pymc_experimental/model/marginal_model.py @@ -1,9 +1,12 @@ import warnings -from typing import Sequence, Union + +from collections.abc import Sequence +from typing import Union import numpy as np import pymc import pytensor.tensor as pt + from arviz import InferenceData, dict_to_dataset from pymc import SymbolicRandomVariable from pymc.backends.arviz import coords_and_dims_for_inferencedata, dataset_to_point_list @@ -60,7 +63,6 @@ class MarginalModel(Model): Examples -------- - Marginalize over a single variable .. code-block:: python @@ -276,7 +278,7 @@ def marginalize( raise NotImplementedError( "Marginalization for DiscreteMarkovChain with non-matrix transition probability is not supported" ) - elif not isinstance(rv_op, (Bernoulli, Categorical, DiscreteUniform)): + elif not isinstance(rv_op, Bernoulli | Categorical | DiscreteUniform): raise NotImplementedError( f"Marginalization of RV with distribution {rv_to_marginalize.owner.op} is not supported" ) @@ -292,7 +294,7 @@ def marginalize( self.clone()._marginalize(user_warnings=True) def _to_transformed(self): - "Create a function from the untransformed space to the transformed space" + """Create a function from the untransformed space to the transformed space""" transformed_rvs = [] transformed_names = [] @@ -417,7 +419,7 @@ def transform_input(inputs): marginalized_rv = m.vars_to_clone[marginalized_rv] m.unmarginalize([marginalized_rv]) dependent_vars = find_conditional_dependent_rvs(marginalized_rv, m.basic_RVs) - joint_logps = m.logp(vars=[marginalized_rv] + dependent_vars, sum=False) + joint_logps = m.logp(vars=[marginalized_rv, *dependent_vars], sum=False) marginalized_value = m.rvs_to_values[marginalized_rv] other_values = [v for v in m.value_vars if v is not marginalized_value] @@ -570,8 +572,7 @@ def find_conditional_input_rvs(output_rvs, all_rvs): return [ var for var in ancestors(output_rvs, blockers=blockers) - if var in blockers - or (var.owner is None and not isinstance(var, (Constant, SharedVariable))) + if var in blockers or (var.owner is None and not isinstance(var, Constant | SharedVariable)) ] @@ -611,7 +612,7 @@ def is_elemwise_subgraph(rv_to_marginalize, other_input_rvs, output_rvs): ) for o in node.outputs ] - blocker_candidates = [rv_to_marginalize] + other_input_rvs + non_elemwise_blockers + blocker_candidates = [rv_to_marginalize, *other_input_rvs, *non_elemwise_blockers] blockers = [var for var in blocker_candidates if var not in output_rvs] truncated_inputs = [ @@ -619,7 +620,7 @@ def is_elemwise_subgraph(rv_to_marginalize, other_input_rvs, output_rvs): for var in ancestors(output_rvs, blockers=blockers) if ( var in blockers - or (var.owner is None and not isinstance(var, (Constant, SharedVariable))) + or (var.owner is None and not isinstance(var, Constant | SharedVariable)) ) ] @@ -804,7 +805,6 @@ def logp_fn(marginalized_rv_const, *non_sequences): @_logprob.register(DiscreteMarginalMarkovChainRV) def marginal_hmm_logp(op, values, *inputs, **kwargs): - marginalized_rvs_node = op.make_node(*inputs) inner_rvs = clone_replace( op.inner_outputs, diff --git a/pymc_experimental/model/transforms/autoreparam.py b/pymc_experimental/model/transforms/autoreparam.py index bb3996459..a48cc6540 100644 --- a/pymc_experimental/model/transforms/autoreparam.py +++ b/pymc_experimental/model/transforms/autoreparam.py @@ -1,12 +1,13 @@ +from collections.abc import Sequence from dataclasses import dataclass from functools import singledispatch -from typing import Dict, List, Optional, Sequence, Tuple, Union import numpy as np import pymc as pm import pytensor import pytensor.tensor as pt import scipy.special + from pymc.distributions import SymbolicRandomVariable from pymc.exceptions import NotConstantValueError from pymc.logprob.transforms import Transform @@ -39,10 +40,10 @@ class VIP: \end{align*} """ - _logit_lambda: Dict[str, pytensor.tensor.sharedvar.TensorSharedVariable] + _logit_lambda: dict[str, pytensor.tensor.sharedvar.TensorSharedVariable] @property - def variational_parameters(self) -> List[pytensor.tensor.sharedvar.TensorSharedVariable]: + def variational_parameters(self) -> list[pytensor.tensor.sharedvar.TensorSharedVariable]: r"""Return raw :math:`\operatorname{logit}(\lambda_k)` for custom optimization. Examples @@ -109,7 +110,7 @@ def truncate_all_lambda(self, value: float): ) self.truncate_lambda(**truncate) - def get_lambda(self) -> Dict[str, np.ndarray]: + def get_lambda(self) -> dict[str, np.ndarray]: r"""Get :math:`\lambda_k` that are currently used by the model. Returns @@ -122,7 +123,7 @@ def get_lambda(self) -> Dict[str, np.ndarray]: for name, shared in self._logit_lambda.items() } - def set_lambda(self, **kwargs: Dict[str, Union[np.ndarray, float]]): + def set_lambda(self, **kwargs: dict[str, np.ndarray | float]): r"""Set :math:`\lambda_k` per variable.""" for key, value in kwargs.items(): logit_lam = scipy.special.logit(value) @@ -133,7 +134,7 @@ def set_lambda(self, **kwargs: Dict[str, Union[np.ndarray, float]]): ) shared.set_value(fill) - def set_all_lambda(self, value: Union[np.ndarray, float]): + def set_all_lambda(self, value: np.ndarray | float): r"""Set :math:`\lambda_k` globally.""" config = dict.fromkeys( self._logit_lambda.keys(), @@ -169,9 +170,9 @@ def vip_reparam_node( op: RandomVariable, node: Apply, name: str, - dims: List[Variable], - transform: Optional[Transform], -) -> Tuple[ModelDeterministic, ModelNamed]: + dims: list[Variable], + transform: Transform | None, +) -> tuple[ModelDeterministic, ModelNamed]: if not isinstance(node.op, RandomVariable | SymbolicRandomVariable): raise TypeError("Op should be RandomVariable type") rv = node.default_output() @@ -204,8 +205,8 @@ def _vip_reparam_node( op: RandomVariable, node: Apply, name: str, - dims: List[Variable], - transform: Optional[Transform], + dims: list[Variable], + transform: Transform | None, lam: pt.TensorVariable, ) -> ModelDeterministic: raise NotImplementedError @@ -216,8 +217,8 @@ def _( op: pm.Normal, node: Apply, name: str, - dims: List[Variable], - transform: Optional[Transform], + dims: list[Variable], + transform: Transform | None, lam: pt.TensorVariable, ) -> ModelDeterministic: rng, size, loc, scale = node.inputs @@ -251,8 +252,8 @@ def _( op: pm.Exponential, node: Apply, name: str, - dims: List[Variable], - transform: Optional[Transform], + dims: list[Variable], + transform: Transform | None, lam: pt.TensorVariable, ) -> ModelDeterministic: rng, size, scale = node.inputs @@ -287,7 +288,7 @@ def _( def vip_reparametrize( model: pm.Model, var_names: Sequence[str], -) -> Tuple[pm.Model, VIP]: +) -> tuple[pm.Model, VIP]: r"""Repametrize Model using Variationally Informed Parametrization (VIP). .. math:: diff --git a/pymc_experimental/model_builder.py b/pymc_experimental/model_builder.py index 6e2a256e6..6e712e5d7 100644 --- a/pymc_experimental/model_builder.py +++ b/pymc_experimental/model_builder.py @@ -16,15 +16,17 @@ import hashlib import json import warnings + from abc import abstractmethod from pathlib import Path -from typing import Any, Dict, List, Optional, Union +from typing import Any import arviz as az import numpy as np import pandas as pd import pymc as pm import xarray as xr + from pymc.util import RandomState # If scikit-learn is available, use its data validator @@ -51,8 +53,8 @@ class ModelBuilder: def __init__( self, - model_config: Dict = None, - sampler_config: Dict = None, + model_config: dict | None = None, + sampler_config: dict | None = None, ): """ Initializes model configuration and sampler configuration for the model @@ -65,6 +67,7 @@ def __init__( dictionary of parameters that initialise model configuration. Class-default defined by the user default_model_config method. sampler_config : Dictionary, optional dictionary of parameters that initialise sampler configuration. Class-default defined by the user default_sampler_config method. + Examples -------- >>> class MyModel(ModelBuilder): @@ -79,7 +82,7 @@ def __init__( self.model_config = model_config # parameters for priors etc. self.model = None # Set by build_model - self.idata: Optional[az.InferenceData] = None # idata is generated during fitting + self.idata: az.InferenceData | None = None # idata is generated during fitting self.is_fitted_ = False def _validate_data(self, X, y=None): @@ -91,8 +94,8 @@ def _validate_data(self, X, y=None): @abstractmethod def _data_setter( self, - X: Union[np.ndarray, pd.DataFrame], - y: Union[np.ndarray, pd.DataFrame, List] = None, + X: np.ndarray | pd.DataFrame, + y: np.ndarray | pd.DataFrame | list = None, ) -> None: """ Sets new data in the model. @@ -104,8 +107,8 @@ def _data_setter( y : array, shape (n_obs,) The target values (real numbers). - Returns: - ---------- + Returns + ------- None Examples @@ -137,10 +140,11 @@ def output_var(self): @staticmethod @abstractmethod - def get_default_model_config() -> Dict: + def get_default_model_config() -> dict: """ Returns a class default config dict for model builder if no model_config is provided on class initialization Useful for understanding structure of required model_config to allow its customization by users + Examples -------- >>> @staticmethod @@ -166,10 +170,11 @@ def get_default_model_config() -> Dict: @staticmethod @abstractmethod - def get_default_sampler_config(self) -> Dict: + def get_default_sampler_config(self) -> dict: """ Returns a class default sampler dict for model builder if no sampler_config is provided on class initialization Useful for understanding structure of required sampler_config to allow its customization by users + Examples -------- >>> @staticmethod @@ -190,7 +195,7 @@ def get_default_sampler_config(self) -> Dict: @abstractmethod def _generate_and_preprocess_model_data( - self, X: Union[pd.DataFrame, pd.Series], y: pd.Series + self, X: pd.DataFrame | pd.Series, y: pd.Series ) -> None: """ Applies preprocessing to the data before fitting the model. @@ -200,7 +205,8 @@ def _generate_and_preprocess_model_data( In case of optional parameters being passed into the model, this method should implement the conditional logic responsible for correct handling of the optional parameters, and including them into the dataset. - Parameters: + Parameters + ---------- X : array, shape (n_obs, n_features) y : array, shape (n_obs,) @@ -385,7 +391,7 @@ def save(self, fname: str) -> None: raise RuntimeError("The model hasn't been fit yet, call .fit() first") @classmethod - def _model_config_formatting(cls, model_config: Dict) -> Dict: + def _model_config_formatting(cls, model_config: dict) -> dict: """ Because of json serialization, model_config values that were originally tuples or numpy are being encoded as lists. This function converts them back to tuples and numpy arrays to ensure correct id encoding. @@ -421,6 +427,7 @@ def load(cls, fname: str): ------ ValueError If the inference data that is loaded doesn't match with the model. + Examples -------- >>> class MyModel(ModelBuilder): @@ -453,9 +460,9 @@ def load(cls, fname: str): def fit( self, X: pd.DataFrame, - y: Optional[pd.Series] = None, + y: pd.Series | None = None, progressbar: bool = True, - predictor_names: List[str] = None, + predictor_names: list[str] | None = None, random_seed: RandomState = None, **kwargs: Any, ) -> az.InferenceData: @@ -484,6 +491,7 @@ def fit( ------- self : az.InferenceData returns inference data of the fitted model. + Examples -------- >>> model = MyModel() @@ -520,7 +528,7 @@ def fit( def predict( self, - X_pred: Union[np.ndarray, pd.DataFrame, pd.Series], + X_pred: np.ndarray | pd.DataFrame | pd.Series, extend_idata: bool = True, **kwargs, ) -> np.ndarray: @@ -529,7 +537,7 @@ def predict( for each input row is the expected output value, computed as the mean of MCMC samples. Parameters - --------- + ---------- X_pred : array-like if sklearn is available, otherwise array, shape (n_pred, n_features) The input data used for prediction. extend_idata : Boolean determining whether the predictions should be added to inference data object. @@ -568,7 +576,7 @@ def sample_prior_predictive( self, X_pred, y_pred=None, - samples: Optional[int] = None, + samples: int | None = None, extend_idata: bool = False, combined: bool = True, **kwargs, @@ -577,7 +585,7 @@ def sample_prior_predictive( Sample from the model's prior predictive distribution. Parameters - --------- + ---------- X_pred : array, shape (n_pred, n_features) The input data used for prediction using prior distribution. samples : int @@ -621,7 +629,7 @@ def sample_posterior_predictive(self, X_pred, extend_idata, combined, **kwargs): Sample from the model's posterior predictive distribution. Parameters - --------- + ---------- X_pred : array, shape (n_pred, n_features) The input data used for prediction using prior distribution.. extend_idata : Boolean determining whether the predictions should be added to inference data object. @@ -666,7 +674,7 @@ def set_params(self, **params): @property @abstractmethod - def _serializable_model_config(self) -> Dict[str, Union[int, float, Dict]]: + def _serializable_model_config(self) -> dict[str, int | float | dict]: """ Converts non-serializable values from model_config to their serializable reversable equivalent. Data types like pandas DataFrame, Series or datetime aren't JSON serializable, @@ -679,7 +687,7 @@ def _serializable_model_config(self) -> Dict[str, Union[int, float, Dict]]: def predict_proba( self, - X_pred: Union[np.ndarray, pd.DataFrame, pd.Series], + X_pred: np.ndarray | pd.DataFrame | pd.Series, extend_idata: bool = True, combined: bool = False, **kwargs, @@ -689,7 +697,7 @@ def predict_proba( def predict_posterior( self, - X_pred: Union[np.ndarray, pd.DataFrame, pd.Series], + X_pred: np.ndarray | pd.DataFrame | pd.Series, extend_idata: bool = True, combined: bool = True, **kwargs, @@ -698,7 +706,7 @@ def predict_posterior( Generate posterior predictive samples on unseen data. Parameters - --------- + ---------- X_pred : array-like if sklearn is available, otherwise array, shape (n_pred, n_features) The input data used for prediction. extend_idata : Boolean determining whether the predictions should be added to inference data object. diff --git a/pymc_experimental/preprocessing/standard_scaler.py b/pymc_experimental/preprocessing/standard_scaler.py index 26026447a..b16a7ca71 100644 --- a/pymc_experimental/preprocessing/standard_scaler.py +++ b/pymc_experimental/preprocessing/standard_scaler.py @@ -1,4 +1,5 @@ import pandas as pd + from sklearn.base import BaseEstimator, TransformerMixin from sklearn.preprocessing import StandardScaler diff --git a/pymc_experimental/statespace/core/representation.py b/pymc_experimental/statespace/core/representation.py index 41c316134..c257378a3 100644 --- a/pymc_experimental/statespace/core/representation.py +++ b/pymc_experimental/statespace/core/representation.py @@ -1,5 +1,6 @@ import copy -from typing import Optional, Type, Union + +from typing import Union import numpy as np import pytensor @@ -56,7 +57,7 @@ class PytensorRepresentation: It should potentially be removed in favor of the closed-form diffuse initialization. Notes - ---------- + ----- A linear statespace system is defined by two equations: .. math:: @@ -119,7 +120,7 @@ class PytensorRepresentation: sliced away unless specifically requested by the user. See the examples for details. Examples - ---------- + -------- .. code:: python from pymc_experimental.statespace.core.representation import PytensorRepresentation @@ -174,8 +175,8 @@ def __init__( k_endog: int, k_states: int, k_posdef: int, - design: Optional[np.ndarray] = None, - obs_intercept: Optional[np.ndarray] = None, + design: np.ndarray | None = None, + obs_intercept: np.ndarray | None = None, obs_cov=None, transition=None, state_intercept=None, @@ -223,8 +224,8 @@ def _validate_key(self, key: KeyLike) -> None: if key not in self.shapes: raise IndexError(f"{key} is an invalid state space matrix name") - def _update_shape(self, key: KeyLike, value: Union[np.ndarray, pt.Variable]) -> None: - if isinstance(value, (pt.TensorConstant, pt.TensorVariable)): + def _update_shape(self, key: KeyLike, value: np.ndarray | pt.Variable) -> None: + if isinstance(value, pt.TensorConstant | pt.TensorVariable): shape = value.type.shape else: shape = value.shape @@ -239,14 +240,14 @@ def _update_shape(self, key: KeyLike, value: Union[np.ndarray, pt.Variable]) -> # Add time dimension dummy if none present if key not in NEVER_TIME_VARYING: if len(shape) == 2 and key not in VECTOR_VALUED: - shape = (1,) + shape + shape = (1, *shape) elif len(shape) == 1: - shape = (1,) + shape + shape = (1, *shape) self.shapes[key] = shape def _add_time_dim_to_slice( - self, name: str, slice_: Union[list[int], tuple[int]], n_dim: int + self, name: str, slice_: list[int] | tuple[int], n_dim: int ) -> tuple[int | slice, ...]: # Case 1: There is never a time dim. No changes needed. if name in NEVER_TIME_VARYING: @@ -263,13 +264,13 @@ def _add_time_dim_to_slice( return (0,) + tuple(slice_) + empty_slice * n_omitted @staticmethod - def _validate_key_and_get_type(key: KeyLike) -> Type[str]: + def _validate_key_and_get_type(key: KeyLike) -> type[str]: if isinstance(key, tuple) and not isinstance(key[0], str): raise IndexError("First index must the name of a valid state space matrix.") return type(key) - def _validate_matrix_shape(self, name: str, X: Union[np.ndarray, pt.TensorVariable]) -> None: + def _validate_matrix_shape(self, name: str, X: np.ndarray | pt.TensorVariable) -> None: time_dim, *expected_shape = self.shapes[name] expected_shape = tuple(expected_shape) shape = X.shape if isinstance(X, np.ndarray) else X.type.shape @@ -407,7 +408,7 @@ def __getitem__(self, key: KeyLike) -> pt.TensorVariable: else: raise IndexError("First index must the name of a valid state space matrix.") - def __setitem__(self, key: KeyLike, value: Union[float, int, np.ndarray, pt.Variable]) -> None: + def __setitem__(self, key: KeyLike, value: float | int | np.ndarray | pt.Variable) -> None: _type = type(key) # Case 1: key is a string: we are setting an entire matrix. diff --git a/pymc_experimental/statespace/core/statespace.py b/pymc_experimental/statespace/core/statespace.py index 9e7b9fa47..2fb8de148 100644 --- a/pymc_experimental/statespace/core/statespace.py +++ b/pymc_experimental/statespace/core/statespace.py @@ -1,11 +1,14 @@ import logging -from typing import Any, Callable, Optional, Sequence, Union + +from collections.abc import Callable, Sequence +from typing import Any import numpy as np import pandas as pd import pymc as pm import pytensor import pytensor.tensor as pt + from arviz import InferenceData from pymc.model import modelcontext from pymc.model.transform.optimization import freeze_dims_and_data @@ -220,10 +223,10 @@ def __init__( verbose: bool = True, measurement_error: bool = False, ): - self._fit_mode: Optional[str] = None - self._fit_coords: Optional[dict[str, Sequence[str]]] = None - self._fit_dims: Optional[dict[str, Sequence[str]]] = None - self._fit_data: Optional[pt.TensorVariable] = None + self._fit_mode: str | None = None + self._fit_coords: dict[str, Sequence[str]] | None = None + self._fit_dims: dict[str, Sequence[str]] | None = None + self._fit_data: pt.TensorVariable | None = None self._needs_exog_data = None self._exog_names = [] @@ -240,7 +243,7 @@ def __init__( self.ssm = PytensorRepresentation(k_endog, k_states, k_posdef) # This will be populated with PyMC random matrices after calling _insert_random_variables - self.subbed_ssm: Optional[list[pt.TensorVariable]] = None + self.subbed_ssm: list[pt.TensorVariable] | None = None if filter_type.lower() not in FILTER_FACTORY.keys(): raise NotImplementedError( @@ -491,7 +494,7 @@ def make_and_register_variable( return placeholder def make_and_register_data( - self, name: str, shape: Union[int, tuple[int]], dtype: str = floatX + self, name: str, shape: int | tuple[int], dtype: str = floatX ) -> Variable: r""" Helper function to create a pytensor symbolic variable and register it in the _name_to_data dictionary @@ -539,7 +542,7 @@ def make_symbolic_graph(self) -> None: Every statespace model needs to implement this function. Examples - ---------- + -------- As an example, consider an ARMA(2,2) model, which has five parameters (excluding the initial state distribution): 2 AR parameters (:math:`\rho_1` and :math:`\rho_2`), 2 MA parameters (:math:`\theta_1` and :math:`theta_2`), and a single innovation covariance (:math:`\\sigma`). A common way of writing this statespace is: @@ -586,9 +589,7 @@ def make_symbolic_graph(self) -> None: """ raise NotImplementedError("The make_symbolic_statespace method has not been implemented!") - def _get_matrix_shape_and_dims( - self, name: str - ) -> tuple[Optional[tuple[int]], Optional[tuple[str]]]: + def _get_matrix_shape_and_dims(self, name: str) -> tuple[tuple[int] | None, tuple[str] | None]: """ Get the shape and dimensions of a matrix associated with the specified name. @@ -616,8 +617,8 @@ def _get_matrix_shape_and_dims( data_len = len(self._fit_data) if name in self.kalman_filter.seq_names: - shape = (data_len,) + self.ssm[SHORT_NAME_TO_LONG[name]].type.shape - dims = (TIME_DIM,) + dims + shape = (data_len, *self.ssm[SHORT_NAME_TO_LONG[name]].type.shape) + dims = (TIME_DIM, *dims) else: shape = self.ssm[SHORT_NAME_TO_LONG[name]].type.shape @@ -749,7 +750,7 @@ def _register_matrices_with_pymc_model(self) -> list[pt.TensorVariable]: has_dims = dims is not None if matrix.ndim == time_varying_ndim and has_dims: - dims = (TIME_DIM,) + dims + dims = (TIME_DIM, *dims) x = pm.Deterministic(name, matrix, dims=dims) registered_matrices.append(x) @@ -786,11 +787,11 @@ def _register_kalman_filter_outputs_with_pymc_model(outputs: tuple[pt.TensorVari def build_statespace_graph( self, - data: Union[np.ndarray, pd.DataFrame, pt.TensorVariable], + data: np.ndarray | pd.DataFrame | pt.TensorVariable, register_data: bool = True, - mode: Optional[str] = None, - missing_fill_value: Optional[float] = None, - cov_jitter: Optional[float] = JITTER_DEFAULT, + mode: str | None = None, + missing_fill_value: float | None = None, + cov_jitter: float | None = JITTER_DEFAULT, save_kalman_filter_outputs_in_idata: bool = False, ) -> None: """ @@ -872,7 +873,7 @@ def build_statespace_graph( smooth_states, smooth_covariances = self._build_smoother_graph( filtered_states, filtered_covariances, self.unpack_statespace(), mode=mode ) - all_kf_outputs = states + [smooth_states] + covs + [smooth_covariances] + all_kf_outputs = [*states, smooth_states, *covs, smooth_covariances] self._register_kalman_filter_outputs_with_pymc_model(all_kf_outputs) obs_dims = FILTER_OUTPUT_DIMS["predicted_observed_state"] @@ -896,7 +897,7 @@ def _build_smoother_graph( filtered_states: pt.TensorVariable, filtered_covariances: pt.TensorVariable, matrices, - mode: Optional[str] = None, + mode: str | None = None, cov_jitter=JITTER_DEFAULT, ): """ @@ -1084,17 +1085,20 @@ def _sample_conditional( group_idata = getattr(idata, group) with pm.Model(coords=self._fit_coords) as forward_model: - [ - x0, - P0, - c, - d, - T, - Z, - R, - H, - Q, - ], grouped_outputs = self._kalman_filter_outputs_from_dummy_graph(data=data) + ( + [ + x0, + P0, + c, + d, + T, + Z, + R, + H, + Q, + ], + grouped_outputs, + ) = self._kalman_filter_outputs_from_dummy_graph(data=data) for name, (mu, cov) in zip(FILTER_OUTPUT_TYPES, grouped_outputs): dummy_ll = pt.zeros_like(mu) @@ -1154,9 +1158,9 @@ def _sample_unconditional( self, idata: InferenceData, group: str, - steps: Optional[int] = None, + steps: int | None = None, use_data_time_dim: bool = False, - random_seed: Optional[RandomState] = None, + random_seed: RandomState | None = None, **kwargs, ): """ @@ -1266,7 +1270,7 @@ def _sample_unconditional( return idata_unconditional.posterior_predictive def sample_conditional_prior( - self, idata: InferenceData, random_seed: Optional[RandomState] = None, **kwargs + self, idata: InferenceData, random_seed: RandomState | None = None, **kwargs ) -> InferenceData: """ Sample from the conditional prior; that is, given parameter draws from the prior distribution, @@ -1296,7 +1300,7 @@ def sample_conditional_prior( return self._sample_conditional(idata, "prior", random_seed, **kwargs) def sample_conditional_posterior( - self, idata: InferenceData, random_seed: Optional[RandomState] = None, **kwargs + self, idata: InferenceData, random_seed: RandomState | None = None, **kwargs ): """ Sample from the conditional posterior; that is, given parameter draws from the posterior distribution, @@ -1327,9 +1331,9 @@ def sample_conditional_posterior( def sample_unconditional_prior( self, idata: InferenceData, - steps: Optional[int] = None, + steps: int | None = None, use_data_time_dim: bool = False, - random_seed: Optional[RandomState] = None, + random_seed: RandomState | None = None, **kwargs, ) -> InferenceData: """ @@ -1380,9 +1384,9 @@ def sample_unconditional_prior( def sample_unconditional_posterior( self, idata: InferenceData, - steps: Optional[int] = None, + steps: int | None = None, use_data_time_dim: bool = False, - random_seed: Optional[RandomState] = None, + random_seed: RandomState | None = None, **kwargs, ) -> InferenceData: """ @@ -1491,11 +1495,11 @@ def sample_statespace_matrices( def forecast( self, idata: InferenceData, - start: Union[int, pd.Timestamp], - periods: int = None, - end: Union[int, pd.Timestamp] = None, + start: int | pd.Timestamp, + periods: int | None = None, + end: int | pd.Timestamp = None, filter_output="smoothed", - random_seed: Optional[RandomState] = None, + random_seed: RandomState | None = None, **kwargs, ) -> InferenceData: """ @@ -1604,19 +1608,20 @@ def forecast( cov_dims = ["data_time", ALL_STATE_DIM, ALL_STATE_AUX_DIM] with pm.Model(coords=temp_coords) as forecast_model: - [ - x0, - P0, - c, - d, - T, - Z, - R, - H, - Q, - ], grouped_outputs = self._kalman_filter_outputs_from_dummy_graph( - data_dims=["data_time", OBS_STATE_DIM] - ) + ( + [ + x0, + P0, + c, + d, + T, + Z, + R, + H, + Q, + ], + grouped_outputs, + ) = self._kalman_filter_outputs_from_dummy_graph(data_dims=["data_time", OBS_STATE_DIM]) group_idx = FILTER_OUTPUT_TYPES.index(filter_output) mu, cov = grouped_outputs[group_idx] @@ -1667,11 +1672,11 @@ def impulse_response_function( idata, n_steps: int = 40, use_posterior_cov: bool = True, - shock_size: Optional[Union[float, np.ndarray]] = None, - shock_cov: Optional[np.ndarray] = None, - shock_trajectory: Optional[np.ndarray] = None, + shock_size: float | np.ndarray | None = None, + shock_cov: np.ndarray | None = None, + shock_trajectory: np.ndarray | None = None, orthogonalize_shocks: bool = False, - random_seed: Optional[RandomState] = None, + random_seed: RandomState | None = None, **kwargs, ): """ diff --git a/pymc_experimental/statespace/filters/distributions.py b/pymc_experimental/statespace/filters/distributions.py index edcc00c6e..0a2aa221f 100644 --- a/pymc_experimental/statespace/filters/distributions.py +++ b/pymc_experimental/statespace/filters/distributions.py @@ -2,6 +2,7 @@ import pymc as pm import pytensor import pytensor.tensor as pt + from pymc import intX from pymc.distributions.dist_math import check_parameters from pymc.distributions.distribution import Continuous, SymbolicRandomVariable @@ -41,7 +42,7 @@ def make_signature(sequence_names): for matrix in sequence_names: base_shape = matrix_to_shape[matrix] - matrix_to_shape[matrix] = (time,) + base_shape + matrix_to_shape[matrix] = (time, *base_shape) signature = ",".join(["(" + ",".join(shapes) + ")" for shapes in matrix_to_shape.values()]) @@ -66,6 +67,7 @@ class MvNormalSVD(MvNormal): try: import jax.random + from pytensor.link.jax.dispatch.random import jax_sample_fn @jax_sample_fn.register(MvNormalSVDRV) @@ -233,7 +235,7 @@ def step_fn(*args): step_fn, outputs_info=[init_dist_], sequences=None if len(sequences) == 0 else sequences, - non_sequences=non_sequences + [rng], + non_sequences=[*non_sequences, rng], n_steps=steps, mode=mode, strict=True, @@ -290,8 +292,6 @@ def __new__( latent_dims = [time_dim, state_dim] obs_dims = [time_dim, obs_dim] - matrices = () - latent_obs_combined = _LinearGaussianStateSpace( f"{name}_combined", a0, diff --git a/pymc_experimental/statespace/filters/kalman_filter.py b/pymc_experimental/statespace/filters/kalman_filter.py index 87fdc746a..887eef837 100644 --- a/pymc_experimental/statespace/filters/kalman_filter.py +++ b/pymc_experimental/statespace/filters/kalman_filter.py @@ -1,9 +1,9 @@ from abc import ABC -from typing import Optional import numpy as np import pytensor import pytensor.tensor as pt + from pytensor.compile.mode import get_mode from pytensor.graph.basic import Variable from pytensor.raise_op import Assert @@ -74,10 +74,10 @@ def __init__(self, mode=None): self.n_posdef = None self.n_endog = None - self.eye_states: Optional[TensorVariable] = None - self.eye_posdef: Optional[TensorVariable] = None - self.eye_endog: Optional[TensorVariable] = None - self.missing_fill_value: Optional[float] = None + self.eye_states: TensorVariable | None = None + self.eye_posdef: TensorVariable | None = None + self.eye_endog: TensorVariable | None = None + self.missing_fill_value: float | None = None self.cov_jitter = None def initialize_eyes(self, R: TensorVariable, Z: TensorVariable) -> None: @@ -263,7 +263,7 @@ def build_graph( results, updates = pytensor.scan( self.kalman_step, - sequences=[data] + sequences, + sequences=[data, *sequences], outputs_info=[None, a0, None, None, P0, None, None], non_sequences=non_sequences, name="forward_kalman_pass", @@ -531,7 +531,7 @@ def kalman_step(self, *args) -> tuple: y, a, P, c, d, T, Z, R, H, Q. See the docstring for the kalman filter class for details. Returns - ---------- + ------- a_filtered : TensorVariable Best linear estimate of hidden states given all information up to and including the present observation, a[t | t]. diff --git a/pymc_experimental/statespace/filters/kalman_smoother.py b/pymc_experimental/statespace/filters/kalman_smoother.py index 32581fa4f..68e69b27b 100644 --- a/pymc_experimental/statespace/filters/kalman_smoother.py +++ b/pymc_experimental/statespace/filters/kalman_smoother.py @@ -1,7 +1,6 @@ -from typing import Optional - import pytensor import pytensor.tensor as pt + from pytensor.compile import get_mode from pytensor.tensor.nlinalg import matrix_dot @@ -19,7 +18,7 @@ class KalmanSmoother: """ - def __init__(self, mode: Optional[str] = None): + def __init__(self, mode: str | None = None): self.mode = mode self.cov_jitter = JITTER_DEFAULT self.seq_names = [] @@ -84,7 +83,7 @@ def build_graph( smoother_result, updates = pytensor.scan( self.smoother_step, - sequences=[filtered_states[:-1], filtered_covariances[:-1]] + sequences, + sequences=[filtered_states[:-1], filtered_covariances[:-1], *sequences], outputs_info=[a_last, P_last], non_sequences=non_sequences, go_backwards=True, diff --git a/pymc_experimental/statespace/filters/utilities.py b/pymc_experimental/statespace/filters/utilities.py index cd1af6a32..8686d42a4 100644 --- a/pymc_experimental/statespace/filters/utilities.py +++ b/pymc_experimental/statespace/filters/utilities.py @@ -1,4 +1,5 @@ import pytensor.tensor as pt + from pytensor.tensor.nlinalg import matrix_dot from pymc_experimental.statespace.core.representation import ( diff --git a/pymc_experimental/statespace/models/SARIMAX.py b/pymc_experimental/statespace/models/SARIMAX.py index 840eb0acb..29efaf3f1 100644 --- a/pymc_experimental/statespace/models/SARIMAX.py +++ b/pymc_experimental/statespace/models/SARIMAX.py @@ -1,7 +1,9 @@ -from typing import Any, Optional, Sequence, Tuple +from collections.abc import Sequence +from typing import Any import numpy as np import pytensor.tensor as pt + from pytensor.tensor.slinalg import solve_discrete_lyapunov from pymc_experimental.statespace.core.statespace import PyMCStateSpace, floatX @@ -91,7 +93,6 @@ class BayesianSARIMA(PyMCStateSpace): Notes ----- - The ARIMAX model is a univariate time series model that posits the future evolution of a stationary time series will be a function of its past values, together with exogenous "innovations" and their past history. The model is described by its "order", a 3-tuple (p, d, q), that are: @@ -172,8 +173,8 @@ class BayesianSARIMA(PyMCStateSpace): def __init__( self, - order: Tuple[int, int, int], - seasonal_order: Optional[Tuple[int, int, int, int]] = None, + order: tuple[int, int, int], + seasonal_order: tuple[int, int, int, int] | None = None, stationary_initialization: bool = True, filter_type: str = "standard", state_structure: str = "fast", @@ -514,14 +515,14 @@ def make_symbolic_graph(self) -> None: ) # Set up the state covariance matrix - state_cov_idx = ("state_cov",) + np.diag_indices(self.k_posdef) + state_cov_idx = ("state_cov", *np.diag_indices(self.k_posdef)) state_cov = self.make_and_register_variable( "sigma_state", shape=() if self.k_posdef == 1 else (self.k_posdef,), dtype=floatX ) self.ssm[state_cov_idx] = state_cov**2 if self.measurement_error: - obs_cov_idx = ("obs_cov",) + np.diag_indices(self.k_endog) + obs_cov_idx = ("obs_cov", *np.diag_indices(self.k_endog)) obs_cov = self.make_and_register_variable( "sigma_obs", shape=() if self.k_endog == 1 else (self.k_endog,), dtype=floatX ) diff --git a/pymc_experimental/statespace/models/VARMAX.py b/pymc_experimental/statespace/models/VARMAX.py index 0942d6db3..3dab978a0 100644 --- a/pymc_experimental/statespace/models/VARMAX.py +++ b/pymc_experimental/statespace/models/VARMAX.py @@ -1,8 +1,10 @@ -from typing import Any, Sequence, Tuple +from collections.abc import Sequence +from typing import Any import numpy as np import pytensor import pytensor.tensor as pt + from pytensor.tensor.slinalg import solve_discrete_lyapunov from pymc_experimental.statespace.core.statespace import PyMCStateSpace @@ -72,7 +74,6 @@ class BayesianVARMAX(PyMCStateSpace): Notes ----- - The VARMA model is a multivariate extension of the SARIMAX model. Given a set of timeseries :math:`\{x_t\}_{t=0}^T`, with :math:`x_t = \begin{bmatrix} x_{1,t} & x_{2,t} & \cdots & x_{k,t} \end{bmatrix}^T`, a VARMA models each series as a function of the histories of all series. Specifically, denoting the AR-MA order as (p, q), a VARMA can be @@ -140,9 +141,9 @@ class BayesianVARMAX(PyMCStateSpace): def __init__( self, - order: Tuple[int, int], - endog_names: list[str] = None, - k_endog: int = None, + order: tuple[int, int], + endog_names: list[str] | None = None, + k_endog: int | None = None, stationary_initialization: bool = False, filter_type: str = "standard", measurement_error: bool = False, @@ -307,7 +308,7 @@ def make_symbolic_graph(self) -> None: self.ssm["initial_state_cov", :, :] = P0 # Design matrix is a truncated identity (first k_obs states observed) - self.ssm[("design",) + np.diag_indices(self.k_endog)] = 1 + self.ssm[("design", *np.diag_indices(self.k_endog))] = 1 # Transition matrix has 4 blocks: # Upper left: AR coefs (k_obs, k_obs * min(p, 1)) @@ -320,14 +321,14 @@ def make_symbolic_graph(self) -> None: slice(self.k_endog, self.k_endog * self.p), slice(0, self.k_endog * (self.p - 1)), ) - self.ssm[("transition",) + idx] = np.eye(self.k_endog * (self.p - 1)) + self.ssm[("transition", *idx)] = np.eye(self.k_endog * (self.p - 1)) if self.q > 1: idx = ( slice(-self.k_endog * (self.q - 1), None), slice(-self.k_endog * self.q, -self.k_endog), ) - self.ssm[("transition",) + idx] = np.eye(self.k_endog * (self.q - 1)) + self.ssm[("transition", *idx)] = np.eye(self.k_endog * (self.q - 1)) if self.p > 0: ar_param_idx = ("transition", slice(0, self.k_endog), slice(0, self.k_endog * self.p)) @@ -364,7 +365,7 @@ def make_symbolic_graph(self) -> None: self.ssm["selection", slice(self.k_endog * -self.q, end), :] = np.eye(self.k_endog) if self.measurement_error: - obs_cov_idx = ("obs_cov",) + np.diag_indices(self.k_endog) + obs_cov_idx = ("obs_cov", *np.diag_indices(self.k_endog)) sigma_obs = self.make_and_register_variable( "sigma_obs", shape=(self.k_endog,), dtype=floatX ) diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index d5985adcb..1e81f92ad 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/pymc_experimental/statespace/models/structural.py @@ -1,12 +1,15 @@ import functools as ft import logging + from abc import ABC -from typing import Any, Optional, Sequence, Union +from collections.abc import Sequence +from typing import Any import numpy as np import pytensor import pytensor.tensor as pt import xarray as xr + from pytensor import Variable from pymc_experimental.statespace.core import PytensorRepresentation @@ -349,7 +352,7 @@ def __init__( shock_names=None, param_names=None, exog_names=None, - representation: Optional[PytensorRepresentation] = None, + representation: PytensorRepresentation | None = None, measurement_error=False, combine_hidden_states=True, component_from_sum=False, @@ -689,6 +692,7 @@ class LevelTrendComponent(Component): Level and trend component of a structural time series model Parameters + ---------- __________ order : int @@ -788,8 +792,8 @@ class LevelTrendComponent(Component): def __init__( self, - order: Union[int, list[int]] = 2, - innovations_order: Optional[Union[int, list[int]]] = None, + order: int | list[int] = 2, + innovations_order: int | list[int] | None = None, name: str = "LevelTrend", ): if innovations_order is None: @@ -875,7 +879,6 @@ class MeasurementError(Component): Parameters ---------- - name: str, optional Name of the observed data. Default is "obs". @@ -1066,6 +1069,7 @@ class TimeSeasonality(Component): seasonal pattern (``season_length = 7``). If None, states will be numbered ``[State_0, ..., State_s]`` + Notes ----- A seasonal effect is any pattern that repeats every fixed interval. Although there are many possible ways to @@ -1156,8 +1160,8 @@ def __init__( self, season_length: int, innovations: bool = True, - name: Optional[str] = None, - state_names: Optional[list] = None, + name: str | None = None, + state_names: list | None = None, pop_state: bool = True, ): if name is None: @@ -1444,8 +1448,8 @@ class CycleComponent(Component): def __init__( self, - name: str = None, - cycle_length: int = None, + name: str | None = None, + cycle_length: int | None = None, estimate_cycle_length: bool = False, dampen: bool = False, innovations: bool = True, @@ -1549,9 +1553,9 @@ def populate_component_properties(self): class RegressionComponent(Component): def __init__( self, - k_exog: Optional[int] = None, - name: Optional[str] = "Exogenous", - state_names: Optional[list[str]] = None, + k_exog: int | None = None, + name: str | None = "Exogenous", + state_names: list[str] | None = None, innovations=False, ): self.innovations = innovations @@ -1574,7 +1578,7 @@ def __init__( ) @staticmethod - def _get_state_names(k_exog: Optional[int], state_names: Optional[list[str]], name: str): + def _get_state_names(k_exog: int | None, state_names: list[str] | None, name: str): if k_exog is None and state_names is None: raise ValueError("Must specify at least one of k_exog or state_names") if state_names is not None and k_exog is not None: @@ -1587,7 +1591,7 @@ def _get_state_names(k_exog: Optional[int], state_names: Optional[list[str]], na return k_exog, state_names - def _handle_input_data(self, k_exog: int, state_names: Optional[list[str]], name) -> int: + def _handle_input_data(self, k_exog: int, state_names: list[str] | None, name) -> int: k_exog, state_names = self._get_state_names(k_exog, state_names, name) self.state_names = state_names @@ -1634,7 +1638,7 @@ def populate_component_properties(self) -> None: "dims": (TIME_DIM, "exog_state"), }, } - self.coords = {f"exog_state": self.state_names} + self.coords = {"exog_state": self.state_names} if self.innovations: self.param_names += [f"sigma_beta_{self.name}"] diff --git a/pymc_experimental/statespace/models/utilities.py b/pymc_experimental/statespace/models/utilities.py index 40f38093e..7fe935d0e 100644 --- a/pymc_experimental/statespace/models/utilities.py +++ b/pymc_experimental/statespace/models/utilities.py @@ -359,14 +359,14 @@ def conform_time_varying_and_time_invariant_matrices(A, B): if T_A == 1: A_out = pt.repeat(A, B.shape[0], axis=0) - A_out = pt.specify_shape(A_out, (T_B,) + tuple(A_dims)) + A_out = pt.specify_shape(A_out, (T_B, *tuple(A_dims))) A_out.name = A.name return A_out, B if T_B == 1: B_out = pt.repeat(B, A.shape[0], axis=0) - B_out = pt.specify_shape(B_out, (T_A,) + tuple(B_dims)) + B_out = pt.specify_shape(B_out, (T_A, *tuple(B_dims))) B_out.name = B.name return A, B_out diff --git a/pymc_experimental/statespace/utils/data_tools.py b/pymc_experimental/statespace/utils/data_tools.py index 67ab3d89a..855faee0a 100644 --- a/pymc_experimental/statespace/utils/data_tools.py +++ b/pymc_experimental/statespace/utils/data_tools.py @@ -5,6 +5,7 @@ import pymc as pm import pytensor import pytensor.tensor as pt + from pymc import ImputationWarning, modelcontext from pytensor.tensor.sharedvar import TensorSharedVariable @@ -22,7 +23,7 @@ def get_data_dims(data): - if not isinstance(data, (pt.TensorVariable, TensorSharedVariable)): + if not isinstance(data, pt.TensorVariable | TensorSharedVariable): return data_name = getattr(data, "name", None) @@ -150,7 +151,7 @@ def mask_missing_values_in_data(values, missing_fill_value=None): ) impute_message = ( - f"Provided data contains missing values and" + "Provided data contains missing values and" " will be automatically imputed as hidden states" " during Kalman filtering." ) @@ -163,11 +164,11 @@ def mask_missing_values_in_data(values, missing_fill_value=None): def register_data_with_pymc( data, n_obs, obs_coords, register_data=True, missing_fill_value=None, data_dims=None ): - if isinstance(data, (pt.TensorVariable, TensorSharedVariable)): + if isinstance(data, pt.TensorVariable | TensorSharedVariable): values, index = preprocess_tensor_data(data, n_obs, obs_coords) elif isinstance(data, np.ndarray): values, index = preprocess_numpy_data(data, n_obs, obs_coords) - elif isinstance(data, (pd.DataFrame, pd.Series)): + elif isinstance(data, pd.DataFrame | pd.Series): values, index = preprocess_pandas_data(data, n_obs, obs_coords) else: raise ValueError("Data should be one of pytensor tensor, numpy array, or pandas dataframe") diff --git a/pymc_experimental/utils/linear_cg.py b/pymc_experimental/utils/linear_cg.py index 49457cad5..93b64e577 100644 --- a/pymc_experimental/utils/linear_cg.py +++ b/pymc_experimental/utils/linear_cg.py @@ -71,7 +71,10 @@ def linear_cg( initial_guess = np.zeros_like(rhs) if preconditioner is None: - preconditioner = lambda x: x + + def preconditioner(x): + return x + precond = False else: precond = True @@ -128,7 +131,7 @@ def linear_cg( residual_inner_prod = residual.T @ precond_residual # define storage matrices - mul_storage = np.zeros_like(residual) + np.zeros_like(residual) alpha = np.zeros((*batch_shape, 1, rhs.shape[-1])) beta = np.zeros_like(alpha) is_zero = np.zeros((*batch_shape, 1, rhs.shape[-1])) @@ -262,13 +265,11 @@ def linear_cg( result = result * rhs_norm if not tolerance_reached and n_iter > 0: raise RuntimeError( - "CG terminated in {} iterations with average residual norm {}" - " which is larger than the tolerance of {} specified by" + f"CG terminated in {k + 1} iterations with average residual norm {residual_norm.mean()}" + f" which is larger than the tolerance of {tolerance} specified by" " gpytorch.settings.cg_tolerance." " If performance is affected, consider raising the maximum number of CG iterations by running code in" - " a gpytorch.settings.max_cg_iterations(value) context.".format( - k + 1, residual_norm.mean(), tolerance - ) + " a gpytorch.settings.max_cg_iterations(value) context." ) if n_tridiag: diff --git a/pymc_experimental/utils/pivoted_cholesky.py b/pymc_experimental/utils/pivoted_cholesky.py index 69ea9cd7b..756d328db 100644 --- a/pymc_experimental/utils/pivoted_cholesky.py +++ b/pymc_experimental/utils/pivoted_cholesky.py @@ -1,12 +1,15 @@ try: import torch + from gpytorch.utils.permutation import apply_permutation except ImportError as e: raise ImportError("PyTorch and GPyTorch not found.") from e import numpy as np -pp = lambda x: np.array2string(x, precision=4, floatmode="fixed") + +def pp(x): + return np.array2string(x, precision=4, floatmode="fixed") def pivoted_cholesky(mat: np.matrix, error_tol=1e-6, max_iter=np.inf): diff --git a/pymc_experimental/utils/prior.py b/pymc_experimental/utils/prior.py index 30d4e9507..15833c9e7 100644 --- a/pymc_experimental/utils/prior.py +++ b/pymc_experimental/utils/prior.py @@ -13,24 +13,26 @@ # limitations under the License. -from typing import Dict, List, Optional, Sequence, Tuple, TypedDict, Union +from collections.abc import Sequence +from typing import TypedDict import arviz import numpy as np import pymc as pm import pytensor.tensor as pt + from pymc.logprob.transforms import Transform class ParamCfg(TypedDict): name: str - transform: Optional[Transform] - dims: Optional[Union[str, Tuple[str]]] + transform: Transform | None + dims: str | tuple[str] | None class ShapeInfo(TypedDict): # shape might not match slice due to a transform - shape: Tuple[int] # transformed shape + shape: tuple[int] # transformed shape slice: slice @@ -41,13 +43,13 @@ class VarInfo(TypedDict): class FlatInfo(TypedDict): data: np.ndarray - info: List[VarInfo] + info: list[VarInfo] -def _arg_to_param_cfg(key, value: Optional[Union[ParamCfg, Transform, str, Tuple]] = None): +def _arg_to_param_cfg(key, value: ParamCfg | Transform | str | tuple | None = None): if value is None: cfg = ParamCfg(name=key, transform=None, dims=None) - elif isinstance(value, Tuple): + elif isinstance(value, tuple): cfg = ParamCfg(name=key, transform=None, dims=value) elif isinstance(value, str): cfg = ParamCfg(name=value, transform=None, dims=None) @@ -62,8 +64,8 @@ def _arg_to_param_cfg(key, value: Optional[Union[ParamCfg, Transform, str, Tuple def _parse_args( - var_names: Sequence[str], **kwargs: Union[ParamCfg, Transform, str, Tuple] -) -> Dict[str, ParamCfg]: + var_names: Sequence[str], **kwargs: ParamCfg | Transform | str | tuple +) -> dict[str, ParamCfg]: results = dict() for var in var_names: results[var] = _arg_to_param_cfg(var) @@ -133,8 +135,8 @@ def prior_from_idata( name="trace_prior_", *, var_names: Sequence[str] = (), - **kwargs: Union[ParamCfg, Transform, str, Tuple] -) -> Dict[str, pt.TensorVariable]: + **kwargs: ParamCfg | Transform | str | tuple, +) -> dict[str, pt.TensorVariable]: """ Create a prior from posterior using MvNormal approximation. diff --git a/pymc_experimental/utils/spline.py b/pymc_experimental/utils/spline.py index 2e4db5e75..921c8ec6f 100644 --- a/pymc_experimental/utils/spline.py +++ b/pymc_experimental/utils/spline.py @@ -18,6 +18,7 @@ import pytensor.sparse as ps import pytensor.tensor as pt import scipy.interpolate + from pytensor.graph.op import Apply, Op @@ -43,9 +44,9 @@ def make_node(self, *inputs) -> Apply: eval_points, k, d = map(pt.as_tensor, inputs) if not (eval_points.ndim == 1 and np.issubdtype(eval_points.dtype, np.floating)): raise TypeError("eval_points should be a vector of floats") - if not k.type in pt.int_types: + if k.type not in pt.int_types: raise TypeError("k should be integer") - if not d.type in pt.int_types: + if d.type not in pt.int_types: raise TypeError("degree should be integer") if self.sparse: out_type = ps.SparseTensorType("csr", eval_points.dtype)() diff --git a/setup.py b/setup.py index 92c0ea397..d37bff646 100644 --- a/setup.py +++ b/setup.py @@ -13,6 +13,7 @@ # limitations under the License. import itertools + from codecs import open from os.path import dirname, join, realpath @@ -65,8 +66,6 @@ import os -from setuptools import find_packages, setup - def read_version(): here = os.path.abspath(os.path.dirname(__file__)) diff --git a/setupegg.py b/setupegg.py index 888a65c9b..168c05b99 100755 --- a/setupegg.py +++ b/setupegg.py @@ -17,7 +17,5 @@ A setup.py script to use setuptools, which gives egg goodness, etc. """ -from setuptools import setup - with open("setup.py") as s: exec(s.read()) diff --git a/tests/distributions/test_continuous.py b/tests/distributions/test_continuous.py index ced0745b3..07d34d5a9 100644 --- a/tests/distributions/test_continuous.py +++ b/tests/distributions/test_continuous.py @@ -18,6 +18,7 @@ import pytest import scipy.stats.distributions as sp + # test support imports from pymc from pymc.testing import ( BaseTestDistributionRandom, diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py index 60885908f..41e3adfd8 100644 --- a/tests/distributions/test_discrete.py +++ b/tests/distributions/test_discrete.py @@ -17,6 +17,7 @@ import pytensor.tensor as pt import pytest import scipy.stats + from pymc.logprob.utils import ParameterValueError from pymc.testing import ( BaseTestDistributionRandom, diff --git a/tests/distributions/test_discrete_markov_chain.py b/tests/distributions/test_discrete_markov_chain.py index 0e55319de..b2b1d9796 100644 --- a/tests/distributions/test_discrete_markov_chain.py +++ b/tests/distributions/test_discrete_markov_chain.py @@ -4,6 +4,7 @@ # general imports import pytensor.tensor as pt import pytest + from pymc.distributions.shape_utils import change_dist_size from pymc.logprob.utils import ParameterValueError diff --git a/tests/model/test_marginal_model.py b/tests/model/test_marginal_model.py index fd1ce259c..7f97b15b5 100644 --- a/tests/model/test_marginal_model.py +++ b/tests/model/test_marginal_model.py @@ -1,4 +1,5 @@ import itertools + from contextlib import suppress as does_not_warn import numpy as np @@ -6,6 +7,7 @@ import pymc as pm import pytensor.tensor as pt import pytest + from arviz import InferenceData, dict_to_dataset from pymc.distributions import transforms from pymc.logprob.abstract import _logprob @@ -303,7 +305,7 @@ def test_recover_marginals_basic(): assert "k" in post assert "lp_k" in post assert post.k.shape == post.y.shape - assert post.lp_k.shape == post.k.shape + (len(p),) + assert post.lp_k.shape == (*post.k.shape, len(p)) def true_logp(y, sigma): y = y.repeat(len(p)).reshape(len(y), -1) @@ -375,7 +377,7 @@ def test_recover_batched_marginal(): assert "idx" in post assert "lp_idx" in post assert post.idx.shape == post.y.shape - assert post.lp_idx.shape == post.idx.shape + (2,) + assert post.lp_idx.shape == (*post.idx.shape, 2) @pytest.mark.xfail(reason="Still need to investigate") @@ -404,11 +406,11 @@ def test_nested_recover_marginals(): assert "idx" in post assert "lp_idx" in post assert post.idx.shape == post.y.shape - assert post.lp_idx.shape == post.idx.shape + (2,) + assert post.lp_idx.shape == (*post.idx.shape, 2) assert "sub_idx" in post assert "lp_sub_idx" in post assert post.sub_idx.shape == post.y.shape - assert post.lp_sub_idx.shape == post.sub_idx.shape + (2,) + assert post.lp_sub_idx.shape == (*post.sub_idx.shape, 2) def true_idx_logp(y): idx_0 = np.log(0.85 * 0.25 * norm.pdf(y, loc=0) + 0.15 * 0.25 * norm.pdf(y, loc=1)) @@ -638,7 +640,6 @@ def test_data_container(): @pytest.mark.parametrize("univariate", (True, False)) def test_vector_univariate_mixture(univariate): - with MarginalModel() as m: idx = pm.Bernoulli("idx", p=0.5, shape=(2,) if univariate else ()) @@ -704,7 +705,7 @@ def test_marginalized_hmm_normal_emission(batch_chain, batch_emission): if batch_emission: test_value = np.broadcast_to(test_value, (3, 4)) expected_logp *= 3 - np.testing.assert_allclose(logp_fn({f"emission": test_value}), expected_logp) + np.testing.assert_allclose(logp_fn({"emission": test_value}), expected_logp) @pytest.mark.parametrize( @@ -728,7 +729,7 @@ def test_marginalized_hmm_categorical_emission(categorical_emission): test_value = np.array([0, 0, 1]) expected_logp = np.log(0.1344) # Shown at the 10m22s mark in the video logp_fn = m.compile_logp() - np.testing.assert_allclose(logp_fn({f"emission": test_value}), expected_logp) + np.testing.assert_allclose(logp_fn({"emission": test_value}), expected_logp) @pytest.mark.parametrize("batch_emission1", (False, True)) @@ -764,7 +765,7 @@ def test_mutable_indexing_jax_backend(): from pymc.sampling.jax import get_jaxified_logp with MarginalModel() as model: - data = pm.Data(f"data", np.zeros(10)) + data = pm.Data("data", np.zeros(10)) cat_effect = pm.Normal("cat_effect", sigma=1, shape=5) cat_effect_idx = pm.Data("cat_effect_idx", np.array([0, 1] * 5)) diff --git a/tests/statespace/test_SARIMAX.py b/tests/statespace/test_SARIMAX.py index fe9d8435e..395182cb3 100644 --- a/tests/statespace/test_SARIMAX.py +++ b/tests/statespace/test_SARIMAX.py @@ -6,6 +6,7 @@ import pytensor.tensor as pt import pytest import statsmodels.api as sm + from numpy.testing import assert_allclose, assert_array_less from pymc_experimental.statespace import BayesianSARIMA @@ -218,7 +219,7 @@ def pymc_mod_interp(arima_mod_interp): @pytest.mark.parametrize( "p,d,q,P,D,Q,S,expected_names", - [order + (name,) for order, name in zip(test_orders, test_state_names)], + [(*order, name) for order, name in zip(test_orders, test_state_names)], ids=ids, ) def test_harvey_state_names(p, d, q, P, D, Q, S, expected_names): diff --git a/tests/statespace/test_VARMAX.py b/tests/statespace/test_VARMAX.py index 43faebe8e..ec37bc11d 100644 --- a/tests/statespace/test_VARMAX.py +++ b/tests/statespace/test_VARMAX.py @@ -7,6 +7,7 @@ import pytensor.tensor as pt import pytest import statsmodels.api as sm + from numpy.testing import assert_allclose, assert_array_less from pymc_experimental.statespace import BayesianVARMAX @@ -100,7 +101,7 @@ def test_VARMAX_update_matches_statsmodels(data, order, rng): sm_var = sm.tsa.VARMAX(data, order=(p, q)) - param_counts = [None] + np.cumsum(list(sm_var.parameters.values())).tolist() + param_counts = [None, *np.cumsum(list(sm_var.parameters.values())).tolist()] param_slices = [slice(a, b) for a, b in zip(param_counts[:-1], param_counts[1:])] param_lists = [trend, ar, ma, reg, state_cov, obs_cov] = [ sm_var.param_names[idx] for idx in param_slices diff --git a/tests/statespace/test_distributions.py b/tests/statespace/test_distributions.py index 1d049ae92..ab55eeba5 100644 --- a/tests/statespace/test_distributions.py +++ b/tests/statespace/test_distributions.py @@ -3,6 +3,7 @@ import pytensor import pytensor.tensor as pt import pytest + from numpy.testing import assert_allclose from scipy.stats import multivariate_normal @@ -174,7 +175,7 @@ def test_lgss_distribution_with_dims(output_name, ss_mod_me, pymc_model_2): steps=100, dims=[TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM], sequence_names=[], - k_endog=ss_mod_me.k_endog + k_endog=ss_mod_me.k_endog, ) # pylint: enable=unpacking-non-sequence idata = pm.sample_prior_predictive(draws=10) @@ -222,7 +223,7 @@ def test_lgss_with_time_varying_inputs(output_name, rng): *matrices, steps=9, sequence_names=["d", "Z"], - dims=[TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM] + dims=[TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM], ) # pylint: enable=unpacking-non-sequence idata = pm.sample_prior_predictive(draws=10) diff --git a/tests/statespace/test_kalman_filter.py b/tests/statespace/test_kalman_filter.py index 15d1effa5..a8582e2fe 100644 --- a/tests/statespace/test_kalman_filter.py +++ b/tests/statespace/test_kalman_filter.py @@ -2,6 +2,7 @@ import pytensor import pytensor.tensor as pt import pytest + from numpy.testing import assert_allclose, assert_array_less from pymc_experimental.statespace.filters import ( @@ -335,7 +336,7 @@ def test_kalman_filter_jax(filter): data = inputs.pop(0) data_specified = pt.specify_shape(data, (n, None)) data_specified.name = "data" - inputs = [data] + inputs + inputs = [data, *inputs] outputs = pytensor.graph.clone_replace(outputs, {data: data_specified}) diff --git a/tests/statespace/test_representation.py b/tests/statespace/test_representation.py index 10388d94d..4598a6f43 100644 --- a/tests/statespace/test_representation.py +++ b/tests/statespace/test_representation.py @@ -3,6 +3,7 @@ import numpy as np import pytensor import pytensor.tensor as pt + from numpy.testing import assert_allclose from pymc_experimental.statespace.core.representation import PytensorRepresentation diff --git a/tests/statespace/test_statespace.py b/tests/statespace/test_statespace.py index d93c66062..e0062933a 100644 --- a/tests/statespace/test_statespace.py +++ b/tests/statespace/test_statespace.py @@ -3,6 +3,7 @@ import pytensor import pytensor.tensor as pt import pytest + from numpy.testing import assert_allclose from pymc_experimental.statespace.core.statespace import FILTER_FACTORY, PyMCStateSpace diff --git a/tests/statespace/test_statespace_JAX.py b/tests/statespace/test_statespace_JAX.py index d9a0c4f96..53e0f40ad 100644 --- a/tests/statespace/test_statespace_JAX.py +++ b/tests/statespace/test_statespace_JAX.py @@ -5,6 +5,7 @@ import pytensor import pytensor.tensor as pt import pytest + from pymc.model.transform.optimization import freeze_dims_and_data from pymc_experimental.statespace.utils.constants import ( diff --git a/tests/statespace/test_structural.py b/tests/statespace/test_structural.py index 30a037811..6bffd36aa 100644 --- a/tests/statespace/test_structural.py +++ b/tests/statespace/test_structural.py @@ -1,5 +1,6 @@ import functools as ft import warnings + from collections import defaultdict from typing import Optional @@ -10,6 +11,7 @@ import pytensor.tensor as pt import pytest import statsmodels.api as sm + from numpy.testing import assert_allclose from scipy import linalg @@ -164,20 +166,20 @@ def _assert_params_info_correct(param_info, coords, param_dims): def create_structural_model_and_equivalent_statsmodel( rng, - level: Optional[bool] = False, - trend: Optional[bool] = False, - seasonal: Optional[int] = None, - freq_seasonal: Optional[list[dict]] = None, + level: bool | None = False, + trend: bool | None = False, + seasonal: int | None = None, + freq_seasonal: list[dict] | None = None, cycle: bool = False, - autoregressive: Optional[int] = None, - exog: Optional[np.ndarray] = None, - irregular: Optional[bool] = False, - stochastic_level: Optional[bool] = True, - stochastic_trend: Optional[bool] = False, - stochastic_seasonal: Optional[bool] = True, - stochastic_freq_seasonal: Optional[list[bool]] = None, - stochastic_cycle: Optional[bool] = False, - damped_cycle: Optional[bool] = False, + autoregressive: int | None = None, + exog: np.ndarray | None = None, + irregular: bool | None = False, + stochastic_level: bool | None = True, + stochastic_trend: bool | None = False, + stochastic_seasonal: bool | None = True, + stochastic_freq_seasonal: list[bool] | None = None, + stochastic_cycle: bool | None = False, + damped_cycle: bool | None = False, ): with warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -300,7 +302,7 @@ def create_structural_model_and_equivalent_statsmodel( sm_params["sigma2.level"] = sigma if stochastic_trend: sigma = sigma_level_value.pop(0) - sm_params[f"sigma2.trend"] = sigma + sm_params["sigma2.trend"] = sigma comp = st.LevelTrendComponent( name="level", order=level_trend_order, innovations_order=level_trend_innov_order @@ -624,7 +626,7 @@ def get_shift_factor(s): return 10 ** len(decimal) -@pytest.mark.parametrize("n", np.arange(1, 6, dtype="int").tolist() + [None]) +@pytest.mark.parametrize("n", [*np.arange(1, 6, dtype="int").tolist(), None]) @pytest.mark.parametrize("s", [5, 10, 25, 25.2]) def test_frequency_seasonality(n, s, rng): mod = st.FrequencySeasonality(season_length=s, n=n, name="season") diff --git a/tests/statespace/utilities/statsmodel_local_level.py b/tests/statespace/utilities/statsmodel_local_level.py index 9e59a9ed5..4422248e1 100644 --- a/tests/statespace/utilities/statsmodel_local_level.py +++ b/tests/statespace/utilities/statsmodel_local_level.py @@ -16,7 +16,7 @@ def __init__(self, endog, **kwargs): self.ssm["selection"] = np.eye(k_states) # Cache some indices - self._state_cov_idx = ("state_cov",) + np.diag_indices(k_posdef) + self._state_cov_idx = ("state_cov", *np.diag_indices(k_posdef)) @property def param_names(self): diff --git a/tests/statespace/utilities/test_helpers.py b/tests/statespace/utilities/test_helpers.py index 7f2183c14..4b4d07815 100644 --- a/tests/statespace/utilities/test_helpers.py +++ b/tests/statespace/utilities/test_helpers.py @@ -3,6 +3,7 @@ import pytensor import pytensor.tensor as pt import statsmodels.api as sm + from numpy.testing import assert_allclose from pymc import modelcontext diff --git a/tests/test_blackjax_smc.py b/tests/test_blackjax_smc.py index b669558fe..49db7de7f 100644 --- a/tests/test_blackjax_smc.py +++ b/tests/test_blackjax_smc.py @@ -17,6 +17,7 @@ import pytensor.tensor as pt import pytest import scipy + from numpy import dtype from xarray.core.utils import Frozen diff --git a/tests/test_laplace.py b/tests/test_laplace.py index 49e5614b2..3fefe3f73 100644 --- a/tests/test_laplace.py +++ b/tests/test_laplace.py @@ -25,7 +25,6 @@ + "To suppress this warning set `negate_output=False`:FutureWarning", ) def test_laplace(): - # Example originates from Bayesian Data Analyses, 3rd Edition # By Andrew Gelman, John Carlin, Hal Stern, David Dunson, # Aki Vehtari, and Donald Rubin. @@ -67,7 +66,6 @@ def test_laplace(): + "To suppress this warning set `negate_output=False`:FutureWarning", ) def test_laplace_only_fit(): - # Example originates from Bayesian Data Analyses, 3rd Edition # By Andrew Gelman, John Carlin, Hal Stern, David Dunson, # Aki Vehtari, and Donald Rubin. @@ -105,7 +103,6 @@ def test_laplace_only_fit(): + "To suppress this warning set `negate_output=False`:FutureWarning", ) def test_laplace_subset_of_rv(recwarn): - # Example originates from Bayesian Data Analyses, 3rd Edition # By Andrew Gelman, John Carlin, Hal Stern, David Dunson, # Aki Vehtari, and Donald Rubin. diff --git a/tests/test_model_builder.py b/tests/test_model_builder.py index 775f27302..d360166dc 100644 --- a/tests/test_model_builder.py +++ b/tests/test_model_builder.py @@ -16,7 +16,6 @@ import json import sys import tempfile -from typing import Dict, Union import numpy as np import pandas as pd @@ -140,7 +139,7 @@ def generate_and_preprocess_model_data(self, X: pd.DataFrame, y: pd.Series): self.y = y @staticmethod - def get_default_model_config() -> Dict: + def get_default_model_config() -> dict: return { "a": {"loc": 0, "scale": 10, "dims": ("numbers",)}, "b": {"loc": 0, "scale": 10}, @@ -148,13 +147,13 @@ def get_default_model_config() -> Dict: } def _generate_and_preprocess_model_data( - self, X: Union[pd.DataFrame, pd.Series], y: pd.Series + self, X: pd.DataFrame | pd.Series, y: pd.Series ) -> None: self.X = X self.y = y @staticmethod - def get_default_sampler_config() -> Dict: + def get_default_sampler_config() -> dict: return { "draws": 10, "tune": 10, diff --git a/tests/test_prior_from_trace.py b/tests/test_prior_from_trace.py index f6bcd3663..460c64f19 100644 --- a/tests/test_prior_from_trace.py +++ b/tests/test_prior_from_trace.py @@ -17,6 +17,7 @@ import numpy as np import pymc as pm import pytest + from pymc.distributions import transforms import pymc_experimental as pmx diff --git a/tests/test_splines.py b/tests/test_splines.py index d5eab9b50..a4afbd1ff 100644 --- a/tests/test_splines.py +++ b/tests/test_splines.py @@ -16,6 +16,7 @@ import numpy as np import pytensor.tensor as pt import pytest + from pytensor.sparse import SparseTensorType import pymc_experimental as pmx diff --git a/tests/utils.py b/tests/utils.py index 2d934bf6c..9576b4be6 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,4 +1,4 @@ -from typing import Sequence +from collections.abc import Sequence from pytensor.compile import SharedVariable from pytensor.graph import Constant, graph_inputs From e6bf0e38b9b89d135ec5cc89f8651c0c7ceb530e Mon Sep 17 00:00:00 2001 From: Jesse Grabowski Date: Sat, 27 Jul 2024 17:41:57 +0800 Subject: [PATCH 3/4] Additional fixes requested by ruff --- pymc_experimental/__init__.py | 19 +++++++++++++++---- .../distributions/multivariate/__init__.py | 2 ++ pymc_experimental/gp/__init__.py | 2 ++ pymc_experimental/inference/__init__.py | 2 ++ pymc_experimental/inference/fit.py | 7 +++---- pymc_experimental/model/marginal_model.py | 4 +--- .../model/transforms/autoreparam.py | 5 ++++- .../statespace/core/representation.py | 4 +--- .../statespace/filters/distributions.py | 4 ---- .../statespace/filters/kalman_filter.py | 17 ++++++++++------- .../statespace/models/structural.py | 3 ++- pymc_experimental/utils/linear_cg.py | 2 +- setup.py | 3 +-- tests/distributions/__init__.py | 2 ++ tests/statespace/test_VARMAX.py | 4 ++-- tests/test_model_builder.py | 2 +- 16 files changed, 49 insertions(+), 33 deletions(-) diff --git a/pymc_experimental/__init__.py b/pymc_experimental/__init__.py index 7b9cf7bb6..d519ff3db 100644 --- a/pymc_experimental/__init__.py +++ b/pymc_experimental/__init__.py @@ -13,6 +13,10 @@ # limitations under the License. import logging +from pymc_experimental import distributions, gp, statespace, utils +from pymc_experimental.inference.fit import fit +from pymc_experimental.model.marginal_model import MarginalModel +from pymc_experimental.model.model_api import as_model from pymc_experimental.version import __version__ _log = logging.getLogger("pmx") @@ -23,7 +27,14 @@ handler = logging.StreamHandler() _log.addHandler(handler) -from pymc_experimental import distributions, gp, statespace, utils -from pymc_experimental.inference.fit import fit -from pymc_experimental.model.marginal_model import MarginalModel -from pymc_experimental.model.model_api import as_model + +__all__ = [ + "distributions", + "gp", + "statespace", + "utils", + "fit", + "MarginalModel", + "as_model", + "__version__", +] diff --git a/pymc_experimental/distributions/multivariate/__init__.py b/pymc_experimental/distributions/multivariate/__init__.py index 64a79b248..12f6b493f 100644 --- a/pymc_experimental/distributions/multivariate/__init__.py +++ b/pymc_experimental/distributions/multivariate/__init__.py @@ -1 +1,3 @@ from pymc_experimental.distributions.multivariate.r2d2m2cp import R2D2M2CP + +__all__ = ["R2D2M2CP"] diff --git a/pymc_experimental/gp/__init__.py b/pymc_experimental/gp/__init__.py index c8804dd4e..ae827e947 100644 --- a/pymc_experimental/gp/__init__.py +++ b/pymc_experimental/gp/__init__.py @@ -14,3 +14,5 @@ from pymc_experimental.gp.latent_approx import KarhunenLoeveExpansion, ProjectedProcess + +__all__ = ["KarhunenLoeveExpansion", "ProjectedProcess"] diff --git a/pymc_experimental/inference/__init__.py b/pymc_experimental/inference/__init__.py index c74607bf5..8b5dbe189 100644 --- a/pymc_experimental/inference/__init__.py +++ b/pymc_experimental/inference/__init__.py @@ -14,3 +14,5 @@ from pymc_experimental.inference.fit import fit + +__all__ = ["fit"] diff --git a/pymc_experimental/inference/fit.py b/pymc_experimental/inference/fit.py index 9b1f891ed..f6c87d90d 100644 --- a/pymc_experimental/inference/fit.py +++ b/pymc_experimental/inference/fit.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from importlib.util import find_spec def fit(method, **kwargs): @@ -30,10 +31,8 @@ def fit(method, **kwargs): arviz.InferenceData """ if method == "pathfinder": - try: - import blackjax - except ImportError as exc: - raise RuntimeError("Need BlackJAX to use `pathfinder`") from exc + if find_spec("blackjax") is None: + raise RuntimeError("Need BlackJAX to use `pathfinder`") from pymc_experimental.inference.pathfinder import fit_pathfinder diff --git a/pymc_experimental/model/marginal_model.py b/pymc_experimental/model/marginal_model.py index 220e542f6..c594e8ac4 100644 --- a/pymc_experimental/model/marginal_model.py +++ b/pymc_experimental/model/marginal_model.py @@ -21,6 +21,7 @@ from pytensor import Mode, scan from pytensor.compile import SharedVariable from pytensor.graph import Constant, FunctionGraph, ancestors, clone_replace +from pytensor.graph.basic import graph_inputs from pytensor.graph.replace import graph_replace, vectorize_graph from pytensor.scan import map as scan_map from pytensor.tensor import TensorType, TensorVariable @@ -638,9 +639,6 @@ def is_elemwise_subgraph(rv_to_marginalize, other_input_rvs, output_rvs): return True -from pytensor.graph.basic import graph_inputs - - def collect_shared_vars(outputs, blockers): return [ inp for inp in graph_inputs(outputs, blockers=blockers) if isinstance(inp, SharedVariable) diff --git a/pymc_experimental/model/transforms/autoreparam.py b/pymc_experimental/model/transforms/autoreparam.py index a48cc6540..dc23010c2 100644 --- a/pymc_experimental/model/transforms/autoreparam.py +++ b/pymc_experimental/model/transforms/autoreparam.py @@ -419,6 +419,9 @@ def vip_reparametrize( lambda_names.append(lam.name) toposort_replace(fmodel, replacements, reverse=True) reparam_model = model_from_fgraph(fmodel) - model_lambdas = {n: reparam_model[l] for l, n in zip(lambda_names, var_names)} + model_lambdas = { + var_name: reparam_model[lambda_name] + for lambda_name, var_name in zip(lambda_names, var_names) + } vip = VIP(model_lambdas) return reparam_model, vip diff --git a/pymc_experimental/statespace/core/representation.py b/pymc_experimental/statespace/core/representation.py index c257378a3..60febea3b 100644 --- a/pymc_experimental/statespace/core/representation.py +++ b/pymc_experimental/statespace/core/representation.py @@ -1,7 +1,5 @@ import copy -from typing import Union - import numpy as np import pytensor import pytensor.tensor as pt @@ -12,7 +10,7 @@ ) floatX = pytensor.config.floatX -KeyLike = Union[tuple[str | int, ...], str] +KeyLike = tuple[str | int, ...] | str class PytensorRepresentation: diff --git a/pymc_experimental/statespace/filters/distributions.py b/pymc_experimental/statespace/filters/distributions.py index 0a2aa221f..1b4895081 100644 --- a/pymc_experimental/statespace/filters/distributions.py +++ b/pymc_experimental/statespace/filters/distributions.py @@ -96,8 +96,6 @@ def update(self, node: Node): class _LinearGaussianStateSpace(Continuous): - rv_op = LinearGaussianStateSpaceRV - def __new__( cls, name, @@ -360,8 +358,6 @@ def update(self, node: Node): class SequenceMvNormal(Continuous): - rv_op = KalmanFilterRV - def __new__(cls, *args, **kwargs): return super().__new__(cls, *args, **kwargs) diff --git a/pymc_experimental/statespace/filters/kalman_filter.py b/pymc_experimental/statespace/filters/kalman_filter.py index 887eef837..c2bfd2f39 100644 --- a/pymc_experimental/statespace/filters/kalman_filter.py +++ b/pymc_experimental/statespace/filters/kalman_filter.py @@ -351,11 +351,14 @@ def handle_missing_values( self, y, Z, H ) -> tuple[TensorVariable, TensorVariable, TensorVariable, float]: """ - This function handles missing values in the observation data `y` and adjusts the design matrix `Z` and the - observation noise covariance matrix `H` accordingly. Missing values are replaced with zeros to prevent - propagating NaNs through the computation. The function also returns a binary flag tensor `all_nan_flag`, - indicating if all values in the observation data are missing. This flag is used for numerical adjustments in - the update method. + Handle missing values in the observation data `y` + + Adjusts the design matrix `Z` and the observation noise covariance matrix `H` by removing rows and/or columns + associated with the data that is not observed at this iteration. Missing values are replaced with zeros to prevent + propagating NaNs through the computation. + + Return a binary flag tensor `all_nan_flag`,indicating if all values in the observation data are missing. This + flag is used for numerical adjustments in the update method. Parameters ---------- @@ -660,7 +663,7 @@ def update(self, a, P, y, c, d, Z, H, all_nan_flag): class CholeskyFilter(BaseFilter): - """ " + """ Kalman filter with Cholesky factorization Kalman filter implementation using a Cholesky factorization plus pt.solve_triangular to (attempt) to speed up @@ -712,7 +715,7 @@ class SingleTimeseriesFilter(BaseFilter): # TODO: This class should eventually be made irrelevant by pytensor re-writes. def check_params(self, data, a0, P0, c, d, T, Z, R, H, Q): - """ " + """ Wrap the data in an `Assert` `Op` to ensure there is only one observed state. """ data = assert_data_is_1d(data, pt.eq(data.shape[1], 1)) diff --git a/pymc_experimental/statespace/models/structural.py b/pymc_experimental/statespace/models/structural.py index 1e81f92ad..9df085659 100644 --- a/pymc_experimental/statespace/models/structural.py +++ b/pymc_experimental/statespace/models/structural.py @@ -3,6 +3,7 @@ from abc import ABC from collections.abc import Sequence +from itertools import pairwise from typing import Any import numpy as np @@ -198,7 +199,7 @@ def make_symbolic_graph(self) -> None: def _state_slices_from_info(self): info = self._component_info.copy() comp_states = np.cumsum([0] + [info["k_states"] for info in info.values()]) - state_slices = [slice(i, j) for i, j in zip(comp_states[:-1], comp_states[1:])] + state_slices = [slice(i, j) for i, j in pairwise(comp_states)] return state_slices diff --git a/pymc_experimental/utils/linear_cg.py b/pymc_experimental/utils/linear_cg.py index 93b64e577..4efaadb45 100644 --- a/pymc_experimental/utils/linear_cg.py +++ b/pymc_experimental/utils/linear_cg.py @@ -65,7 +65,7 @@ def linear_cg( initial_guess=None, preconditioner=None, terminate_cg_by_size=False, - use_eval_tolerange=False, + use_eval_tolerance=False, ): if initial_guess is None: initial_guess = np.zeros_like(rhs) diff --git a/setup.py b/setup.py index d37bff646..404e04110 100644 --- a/setup.py +++ b/setup.py @@ -13,6 +13,7 @@ # limitations under the License. import itertools +import os from codecs import open from os.path import dirname, join, realpath @@ -64,8 +65,6 @@ extras_require["complete"] = sorted(set(itertools.chain.from_iterable(extras_require.values()))) extras_require["dev"] = dev_install_reqs -import os - def read_version(): here = os.path.abspath(os.path.dirname(__file__)) diff --git a/tests/distributions/__init__.py b/tests/distributions/__init__.py index fa2a64480..d4b13bfcb 100644 --- a/tests/distributions/__init__.py +++ b/tests/distributions/__init__.py @@ -15,3 +15,5 @@ from pymc_experimental.distributions import histogram_utils from pymc_experimental.distributions.histogram_utils import histogram_approximation + +__all__ = ["histogram_utils", "histogram_approximation"] diff --git a/tests/statespace/test_VARMAX.py b/tests/statespace/test_VARMAX.py index ec37bc11d..0489a68f4 100644 --- a/tests/statespace/test_VARMAX.py +++ b/tests/statespace/test_VARMAX.py @@ -1,4 +1,4 @@ -from itertools import product +from itertools import pairwise, product import numpy as np import pandas as pd @@ -102,7 +102,7 @@ def test_VARMAX_update_matches_statsmodels(data, order, rng): sm_var = sm.tsa.VARMAX(data, order=(p, q)) param_counts = [None, *np.cumsum(list(sm_var.parameters.values())).tolist()] - param_slices = [slice(a, b) for a, b in zip(param_counts[:-1], param_counts[1:])] + param_slices = [slice(a, b) for a, b in pairwise(param_counts)] param_lists = [trend, ar, ma, reg, state_cov, obs_cov] = [ sm_var.param_names[idx] for idx in param_slices ] diff --git a/tests/test_model_builder.py b/tests/test_model_builder.py index d360166dc..8018dd443 100644 --- a/tests/test_model_builder.py +++ b/tests/test_model_builder.py @@ -225,7 +225,7 @@ def test_predict(fitted_model_instance): prediction_data = pd.DataFrame({"input": x_pred}) pred = fitted_model_instance.predict(prediction_data["input"]) # Perform elementwise comparison using numpy - assert type(pred) == np.ndarray + assert isinstance(pred, np.ndarray) assert len(pred) > 0 From 4e3d22dd18a2d9d83875eac8f96f98acaa5e69e6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 28 Jul 2024 00:29:50 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- pymc_experimental/model/transforms/autoreparam.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pymc_experimental/model/transforms/autoreparam.py b/pymc_experimental/model/transforms/autoreparam.py index 6c630b451..4fbf9ab93 100644 --- a/pymc_experimental/model/transforms/autoreparam.py +++ b/pymc_experimental/model/transforms/autoreparam.py @@ -1,5 +1,6 @@ -from collections.abc import Sequence import logging + +from collections.abc import Sequence from dataclasses import dataclass from functools import singledispatch