fixing linting

Signed-off-by: Nathaniel <[email protected]>

Author: NathanielF

causalpy/pymc_models.py

@@ -677,7 +677,7 @@ def build_model(  # type: ignore
         Z: np.ndarray,
         y: np.ndarray,
         t: np.ndarray,
-        coords: Dict[str, Any],
+        coords: dict[str, Any],
         priors,
         vs_prior_type=None,
         vs_hyperparams=None,
@@ -721,7 +721,8 @@ def build_model(  # type: ignore
                 warnings.warn(
                     "Variable selection priors specified. "
                     "The 'mus' and 'sigmas' in the priors dict will be ignored "
-                    "for beta coefficients. Only 'eta' and 'lkj_sd' will be used."
+                    "for beta coefficients. Only 'eta' and 'lkj_sd' will be used.",
+                    stacklevel=2,
                 )
 
             # Create coefficient priors

causalpy/variable_selection_priors.py

@@ -20,7 +20,7 @@
 top of the pymc-extras Prior infrastructure.
 """
 
-from typing import Any, Dict, Optional, Union
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -30,7 +30,7 @@
 
 
 def _relaxed_bernoulli_transform(
-    p: Union[float, pt.TensorVariable], temperature: float = 0.1
+    p: float | pt.TensorVariable, temperature: float = 0.1
 ):
     """
     Transform function for relaxed (continuous) Bernoulli distribution.
@@ -100,7 +100,7 @@ def __init__(
         pi_beta: float = 2,
         slab_sigma: float = 2,
         temperature: float = 0.1,
-        dims: Optional[Union[str, tuple]] = None,
+        dims: str | tuple | None = None,
     ):
         self.pi_alpha = pi_alpha
         self.pi_beta = pi_beta
@@ -176,11 +176,11 @@ class HorseshoePrior:
 
     def __init__(
         self,
-        tau0: Optional[float] = None,
+        tau0: float | None = None,
         nu: float = 3,
         c2_alpha: float = 2,
         c2_beta: float = 2,
-        dims: Optional[Union[str, tuple]] = None,
+        dims: str | tuple | None = None,
     ):
         self.tau0 = tau0
         self.nu = nu
@@ -277,7 +277,7 @@ class VariableSelectionPrior:
     ...     beta = vs_prior.create_prior(name="beta", n_params=5, dims="features")
     """
 
-    def __init__(self, prior_type: str, hyperparams: Optional[Dict[str, Any]] = None):
+    def __init__(self, prior_type: str, hyperparams: dict[str, Any] | None = None):
         """Initialize the variable selection prior factory."""
         self.prior_type = prior_type.lower()
         self.hyperparams = hyperparams or {}
@@ -292,8 +292,8 @@ def __init__(self, prior_type: str, hyperparams: Optional[Dict[str, Any]] = None
         self._prior_instance = None
 
     def _get_default_hyperparams(
-        self, n_params: int, X: Optional[np.ndarray] = None
-    ) -> Dict[str, Any]:
+        self, n_params: int, X: np.ndarray | None = None
+    ) -> dict[str, Any]:
         """
         Get default hyperparameters for the chosen prior type.
 
@@ -346,10 +346,10 @@ def create_prior(
         self,
         name: str,
         n_params: int,
-        dims: Optional[Union[str, tuple]] = None,
-        X: Optional[np.ndarray] = None,
-        hyperparams: Optional[Dict[str, Any]] = None,
-    ) -> Union[pm.Deterministic, pm.Distribution]:
+        dims: str | tuple | None = None,
+        X: np.ndarray | None = None,
+        hyperparams: dict[str, Any] | None = None,
+    ) -> pm.Deterministic | pm.Distribution:
         """
         Create the specified prior on a coefficient vector.
 
@@ -559,10 +559,10 @@ def create_variable_selection_prior(
     prior_type: str,
     name: str,
     n_params: int,
-    dims: Optional[Union[str, tuple]] = None,
-    X: Optional[np.ndarray] = None,
-    hyperparams: Optional[Dict[str, Any]] = None,
-) -> Union[pm.Deterministic, pm.Distribution]:
+    dims: str | tuple | None = None,
+    X: np.ndarray | None = None,
+    hyperparams: dict[str, Any] | None = None,
+) -> pm.Deterministic | pm.Distribution:
     """
     Convenience function to create a variable selection prior in one call.