Additional fixes requested by ruff

Author: jessegrabowski

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__",
+]

pymc_experimental/distributions/multivariate/__init__.py

@@ -1 +1,3 @@
 from pymc_experimental.distributions.multivariate.r2d2m2cp import R2D2M2CP
+
+__all__ = ["R2D2M2CP"]

pymc_experimental/gp/__init__.py

@@ -14,3 +14,5 @@
 
 
 from pymc_experimental.gp.latent_approx import KarhunenLoeveExpansion, ProjectedProcess
+
+__all__ = ["KarhunenLoeveExpansion", "ProjectedProcess"]

pymc_experimental/inference/__init__.py

@@ -14,3 +14,5 @@
 
 
 from pymc_experimental.inference.fit import fit
+
+__all__ = ["fit"]

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
 

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)

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

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:

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)
 

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))

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
 

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)

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__))

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"]

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
     ]

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