Clean up type hints

Use built-in type hints in place of `typing.List` and `typing.Dict`

Author: jessegrabowski

pymc_experimental/statespace/core/representation.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple, Type, Union
+from typing import Optional, Type, Union
 
 import numpy as np
 import pytensor
@@ -10,7 +10,7 @@
 )
 
 floatX = pytensor.config.floatX
-KeyLike = Union[Tuple[Union[str, int]], str]
+KeyLike = Union[tuple[Union[str, int]], str]
 
 
 class PytensorRepresentation:
@@ -228,8 +228,8 @@ def _update_shape(self, key: KeyLike, value: Union[np.ndarray, pt.TensorType]) -
         self.shapes[key] = shape
 
     def _add_time_dim_to_slice(
-        self, name: str, slice_: Union[List[int], Tuple[int]], n_dim: int
-    ) -> Tuple[int]:
+        self, name: str, slice_: Union[list[int], tuple[int]], n_dim: int
+    ) -> tuple[int]:
         # Case 1: There is never a time dim. No changes needed.
         if name in NEVER_TIME_VARYING:
             return slice_

pymc_experimental/statespace/core/statespace.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Optional, Sequence, Union
 
 import numpy as np
 import pandas as pd
@@ -219,8 +219,8 @@ def __init__(
         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_coords: Optional[dict[str, Sequence[str]]] = None
+        self._fit_dims: Optional[dict[str, Sequence[str]]] = None
         self._fit_data: Optional[pt.TensorVariable] = None
 
         self._needs_exog_data = False
@@ -237,7 +237,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: Optional[list[pt.TensorVariable]] = None
 
         if filter_type.lower() not in FILTER_FACTORY.keys():
             raise NotImplementedError(
@@ -296,7 +296,7 @@ def _print_data_requirements(self) -> None:
             f"{out}"
         )
 
-    def _unpack_statespace_with_placeholders(self) -> Tuple:
+    def _unpack_statespace_with_placeholders(self) -> tuple[pt.TensorVariable, ...]:
         """
         Helper function to quickly obtain all statespace matrices in the standard order. Matrices returned by this
         method will include pytensor placeholders.
@@ -314,7 +314,7 @@ def _unpack_statespace_with_placeholders(self) -> Tuple:
 
         return a0, P0, c, d, T, Z, R, H, Q
 
-    def unpack_statespace(self) -> list[pt.TensorVariable, ...]:
+    def unpack_statespace(self) -> list[pt.TensorVariable]:
         """
         Helper function to quickly obtain all statespace matrices in the standard order.
         """
@@ -329,7 +329,7 @@ def unpack_statespace(self) -> list[pt.TensorVariable, ...]:
         return self.subbed_ssm
 
     @property
-    def param_names(self) -> List[str]:
+    def param_names(self) -> list[str]:
         """
         Names of model parameters
 
@@ -339,7 +339,7 @@ def param_names(self) -> List[str]:
         raise NotImplementedError("The param_names property has not been implemented!")
 
     @property
-    def data_names(self) -> List[str]:
+    def data_names(self) -> list[str]:
         """
         Names of data variables expected by the model.
 
@@ -349,7 +349,7 @@ def data_names(self) -> List[str]:
         raise NotImplementedError("The data_names property has not been implemented!")
 
     @property
-    def param_info(self) -> Dict[str, Dict[str, Any]]:
+    def param_info(self) -> dict[str, dict[str, Any]]:
         """
         Information about parameters needed to declare priors
 
@@ -377,7 +377,7 @@ def data_info(self) -> dict[str, dict[str, Any]]:
         raise NotImplementedError("The data_info property has not been implemented!")
 
     @property
-    def state_names(self) -> List[str]:
+    def state_names(self) -> list[str]:
         """
         A k_states length list of strings, associated with the model's hidden states
 
@@ -386,22 +386,22 @@ def state_names(self) -> List[str]:
         raise NotImplementedError("The state_names property has not been implemented!")
 
     @property
-    def observed_states(self) -> List[str]:
+    def observed_states(self) -> list[str]:
         """
         A k_endog length list of strings, associated with the model's observed states
         """
         raise NotImplementedError("The observed_states property has not been implemented!")
 
     @property
-    def shock_names(self) -> List[str]:
+    def shock_names(self) -> list[str]:
         """
         A k_posdef length list of strings, associated with the model's shock processes
 
         """
         raise NotImplementedError("The shock_names property has not been implemented!")
 
     @property
-    def default_priors(self) -> Dict[str, Callable]:
+    def default_priors(self) -> dict[str, Callable]:
         """
         Dictionary of parameter names and callable functions to construct default priors for the model
 
@@ -411,7 +411,7 @@ def default_priors(self) -> Dict[str, Callable]:
         raise NotImplementedError("The default_priors property has not been implemented!")
 
     @property
-    def coords(self) -> Dict[str, Sequence[str]]:
+    def coords(self) -> dict[str, Sequence[str]]:
         """
         PyMC model coordinates
 
@@ -422,7 +422,7 @@ def coords(self) -> Dict[str, Sequence[str]]:
         raise NotImplementedError("The coords property has not been implemented!")
 
     @property
-    def param_dims(self) -> Dict[str, Sequence[str]]:
+    def param_dims(self) -> dict[str, Sequence[str]]:
         """
         Dictionary of named dimensions for each model parameter
 
@@ -483,7 +483,9 @@ def make_and_register_variable(self, name, shape, dtype=floatX) -> Variable:
         self._name_to_variable[name] = placeholder
         return placeholder
 
-    def make_and_register_data(self, name, shape, dtype=floatX) -> Variable:
+    def make_and_register_data(
+        self, name: str, shape: Union[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
 
@@ -577,7 +579,9 @@ 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[Tuple, Tuple]:
+    def _get_matrix_shape_and_dims(
+        self, name: str
+    ) -> tuple[Optional[tuple[int]], Optional[tuple[str]]]:
         """
         Get the shape and dimensions of a matrix associated with the specified name.
 
@@ -614,7 +618,11 @@ def _get_matrix_shape_and_dims(self, name: str) -> Tuple[Tuple, Tuple]:
 
         return shape, dims
 
-    def _get_output_shape_and_dims(self, idata: InferenceData, filter_output: str) -> Tuple:
+    def _get_output_shape_and_dims(
+        self, idata: InferenceData, filter_output: str
+    ) -> tuple[
+        Optional[tuple[int]], Optional[tuple[int]], Optional[tuple[str]], Optional[tuple[str]]
+    ]:
         """
         Get the shapes and dimensions of the output variables from the provided InferenceData.
 
@@ -756,15 +764,15 @@ def _insert_data_variables(self):
         replacement_dict = {data: pymc_model[name] for name, data in self._name_to_data.items()}
         self.subbed_ssm = graph_replace(self.subbed_ssm, replace=replacement_dict, strict=True)
 
-    def _register_matrices_with_pymc_model(self) -> List[pt.TensorVariable]:
+    def _register_matrices_with_pymc_model(self) -> list[pt.TensorVariable]:
         """
         Add all statespace matrices to the PyMC model currently on the context stack as pm.Deterministic nodes, and
         adds named dimensions if they are found.
 
         Returns
         -------
         registered_matrices: list of pt.TensorVariable
-            List of statespace matrices, wrapped in pm.Deterministic
+            list of statespace matrices, wrapped in pm.Deterministic
         """
 
         pm_mod = modelcontext(None)
@@ -788,9 +796,7 @@ def _register_matrices_with_pymc_model(self) -> List[pt.TensorVariable]:
         return registered_matrices
 
     @staticmethod
-    def _register_kalman_filter_outputs_with_pymc_model(
-        outputs: tuple[pt.TensorVariable, ...]
-    ) -> None:
+    def _register_kalman_filter_outputs_with_pymc_model(outputs: tuple[pt.TensorVariable]) -> None:
         mod = modelcontext(None)
         states, covs = outputs[:4], outputs[4:]
 
@@ -1014,7 +1020,7 @@ def _build_dummy_graph(self) -> None:
 
         Returns
         -------
-        List[pm.Flat]
+        list[pm.Flat]
             A list of pm.Flat variables representing all parameters estimated by the model.
         """
         for name in self.param_names:
@@ -1026,7 +1032,7 @@ def _build_dummy_graph(self) -> None:
 
     def _kalman_filter_outputs_from_dummy_graph(
         self,
-    ) -> tuple[list[pt.TensorVariable, ...], tuple[pt.TensorVariable, pt.TensorVariable]]:
+    ) -> tuple[list[pt.TensorVariable], list[tuple[pt.TensorVariable, pt.TensorVariable]]]:
         """
         Builds a Kalman filter graph using "dummy" pm.Flat distributions for the model variables and sorts the returns
         into (mean, covariance) pairs for each of filtered, predicted, and smoothed output.
@@ -1379,9 +1385,9 @@ def sample_unconditional_prior(
 
     def sample_unconditional_posterior(
         self,
-        idata,
-        steps=None,
-        use_data_time_dim=False,
+        idata: InferenceData,
+        steps: Optional[int] = None,
+        use_data_time_dim: bool = False,
         random_seed: Optional[RandomState] = None,
         **kwargs,
     ) -> InferenceData:
@@ -1775,7 +1781,8 @@ def _sort_obs_inputs_by_time_varying(self, d, Z):
 
         return seqs, non_seqs
 
-    def _sort_obs_scan_args(self, args):
+    @staticmethod
+    def _sort_obs_scan_args(args):
         args = list(args)
 
         # If a matrix is time-varying, pytensor will put a [t] on the name