ran pre-commit on all files

Author: Michal-Novomestsky

.pre-commit-config.yaml

@@ -3,7 +3,7 @@ ci:
 
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.6.0
+  rev: v6.0.0
   hooks:
     -   id: check-merge-conflict
     -   id: check-toml
@@ -15,7 +15,7 @@ repos:
     -   id: trailing-whitespace
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.5.5
+  rev: v0.12.8
   hooks:
     - id: ruff
       args: [ --fix, --unsafe-fixes, --exit-non-zero-on-fix ]

notebooks/Making a Custom Statespace Model.ipynb

@@ -15,13 +15,13 @@
     }
    ],
    "source": [
-    "import numpy as np\n",
-    "import matplotlib.pyplot as plt\n",
     "import arviz as az\n",
-    "\n",
-    "from pymc_extras.statespace.core.statespace import PyMCStateSpace\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import pymc as pm\n",
     "import pytensor.tensor as pt\n",
-    "import pymc as pm"
+    "\n",
+    "from pymc_extras.statespace.core.statespace import PyMCStateSpace"
    ]
   },
   {
@@ -45,7 +45,7 @@
     "\n",
     "\n",
     "def print_model_ssm(mod, how=\"eval\"):\n",
-    "    nice_heading = f'{\"name\":<20}{\"__repr__\":<50}{\"shape\":<10}{\"value\":<20}'\n",
+    "    nice_heading = f\"{'name':<20}{'__repr__':<50}{'shape':<10}{'value':<20}\"\n",
     "    print(nice_heading)\n",
     "    print(\"=\" * len(nice_heading))\n",
     "    if how == \"eval\":\n",
@@ -1270,7 +1270,7 @@
    ],
    "source": [
     "az.plot_posterior(\n",
-    "    idata, var_names=[\"ar_params\", \"sigma_x\"], ref_val=true_ar.tolist() + [true_sigma_x]\n",
+    "    idata, var_names=[\"ar_params\", \"sigma_x\"], ref_val=[*true_ar.tolist(), true_sigma_x]\n",
     ");"
    ]
   },
@@ -1333,13 +1333,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from pymc_extras.statespace.models.utilities import make_default_coords\n",
     "from pymc_extras.statespace.utils.constants import (\n",
-    "    ALL_STATE_DIM,\n",
     "    ALL_STATE_AUX_DIM,\n",
-    "    OBS_STATE_DIM,\n",
+    "    ALL_STATE_DIM,\n",
     "    SHOCK_DIM,\n",
     ")\n",
-    "from pymc_extras.statespace.models.utilities import make_default_coords\n",
     "\n",
     "\n",
     "class AutoRegressiveThree(PyMCStateSpace):\n",

notebooks/discrete_markov_chain.ipynb

@@ -27,13 +27,13 @@
    "outputs": [],
    "source": [
     "import arviz as az\n",
+    "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import pandas as pd\n",
     "import pymc as pm\n",
     "import pytensor\n",
     "import pytensor.tensor as pt\n",
-    "import pandas as pd\n",
     "import statsmodels.api as sm\n",
-    "import matplotlib.pyplot as plt\n",
     "\n",
     "from matplotlib import ticker as mtick\n",
     "\n",
@@ -525,7 +525,7 @@
     "    \"dates\": dta_hamilton.index,\n",
     "    \"obs_dates\": dta_hamilton.index[order:],\n",
     "    \"states\": [\"State_1\", \"State_2\"],\n",
-    "    \"ar_params\": [f\"L{i+1}.phi\" for i in range(order)],\n",
+    "    \"ar_params\": [f\"L{i + 1}.phi\" for i in range(order)],\n",
     "}\n",
     "\n",
     "with pm.Model(coords=coords) as hmm:\n",

pymc_extras/printing.py

@@ -166,7 +166,7 @@ def model_table(
 
         for var in group:
             var_name = var.name
-            sep = f'[b]{" ~" if (var in model.basic_RVs) else " ="}[/b]'
+            sep = f"[b]{' ~' if (var in model.basic_RVs) else ' ='}[/b]"
             var_expr = variable_expression(model, var, truncate_deterministic)
             dims_expr = dims_expression(model, var)
             if dims_expr == "[]":

pymc_extras/statespace/core/statespace.py

@@ -60,7 +60,7 @@
 def _validate_filter_arg(filter_arg):
     if filter_arg.lower() not in FILTER_OUTPUT_TYPES:
         raise ValueError(
-            f'filter_output should be one of {", ".join(FILTER_OUTPUT_TYPES)}, received {filter_arg}'
+            f"filter_output should be one of {', '.join(FILTER_OUTPUT_TYPES)}, received {filter_arg}"
         )
 
 

pymc_extras/statespace/models/SARIMAX.py

@@ -224,7 +224,7 @@ def __init__(
         if state_structure not in SARIMAX_STATE_STRUCTURES:
             raise ValueError(
                 f"Got invalid argument {state_structure} for state structure, expected one of "
-                f'{", ".join(SARIMAX_STATE_STRUCTURES)}'
+                f"{', '.join(SARIMAX_STATE_STRUCTURES)}"
             )
 
         if state_structure == "interpretable" and (self.d + self.D) > 0:

pymc_extras/statespace/utils/data_tools.py

@@ -53,7 +53,7 @@ def _validate_data_shape(data_shape, n_obs, obs_coords=None, check_col_names=Fal
         if len(missing_cols) > 0:
             raise ValueError(
                 "Columns of DataFrame provided as data do not match state names. The following states were"
-                f'not found: {", ".join(missing_cols)}. This may result in unexpected results in complex'
+                f"not found: {', '.join(missing_cols)}. This may result in unexpected results in complex"
                 f"statespace models"
             )
 

tests/statespace/core/test_statespace.py

@@ -884,8 +884,7 @@ def test_invalid_scenarios():
     # Giving a list, tuple, or Series when a matrix of data is expected should always raise
     with pytest.raises(
         ValueError,
-        match="Scenario data for variable 'a' has the wrong number of columns. "
-        "Expected 2, got 1",
+        match="Scenario data for variable 'a' has the wrong number of columns. Expected 2, got 1",
     ):
         for data_type in [list, tuple, pd.Series]:
             ss_mod._validate_scenario_data(data_type(np.zeros(10)))
@@ -894,15 +893,14 @@ def test_invalid_scenarios():
     # Providing irrevelant data raises
     with pytest.raises(
         ValueError,
-        match="Scenario data provided for variable 'jk lol', which is not an exogenous " "variable",
+        match="Scenario data provided for variable 'jk lol', which is not an exogenous variable",
     ):
         ss_mod._validate_scenario_data({"jk lol": np.zeros(10)})
 
     # Incorrect 2nd dimension of a non-dataframe
     with pytest.raises(
         ValueError,
-        match="Scenario data for variable 'a' has the wrong number of columns. Expected "
-        "2, got 1",
+        match="Scenario data for variable 'a' has the wrong number of columns. Expected 2, got 1",
     ):
         scenario = np.zeros(10).tolist()
         ss_mod._validate_scenario_data(scenario)
@@ -1243,9 +1241,9 @@ def test_param_dims_coords(ss_mod_multi_component):
             assert dims is None
             continue
         for i, s in zip(shape, dims):
-            assert i == len(
-                ss_mod_multi_component.coords[s]
-            ), f"Mismatch between shape {i} and dimension {s}"
+            assert i == len(ss_mod_multi_component.coords[s]), (
+                f"Mismatch between shape {i} and dimension {s}"
+            )
 
 
 @pytest.mark.filterwarnings("ignore:Provided data contains missing values")

tests/statespace/filters/test_kalman_filter.py

@@ -73,9 +73,9 @@ def test_output_shapes_one_state_one_observed(filter_func, rng):
 
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
@@ -86,9 +86,9 @@ def test_output_shapes_when_all_states_are_stochastic(filter_func, rng):
     outputs = filter_func(*inputs)
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
@@ -99,9 +99,9 @@ def test_output_shapes_when_some_states_are_deterministic(filter_func, rng):
     outputs = filter_func(*inputs)
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.fixture
@@ -161,9 +161,9 @@ def test_output_shapes_with_time_varying_matrices(f_standard_nd, rng):
 
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)
@@ -175,9 +175,9 @@ def test_output_with_deterministic_observation_equation(filter_func, rng):
 
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize(
@@ -190,9 +190,9 @@ def test_output_with_multiple_observed(filter_func, filter_name, rng):
     outputs = filter_func(*inputs)
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize(
@@ -206,9 +206,9 @@ def test_missing_data(filter_func, filter_name, p, rng):
     outputs = filter_func(*inputs)
     for output_idx, name in enumerate(output_names):
         expected_output = get_expected_shape(name, p, m, r, n)
-        assert (
-            outputs[output_idx].shape == expected_output
-        ), f"Shape of {name} does not match expected"
+        assert outputs[output_idx].shape == expected_output, (
+            f"Shape of {name} does not match expected"
+        )
 
 
 @pytest.mark.parametrize("filter_func", filter_funcs, ids=filter_names)

tests/statespace/models/structural/test_against_statsmodels.py

@@ -67,16 +67,16 @@ def _assert_coord_shapes_match_matrices(mod, params):
     n_shocks = max(1, len(mod.coords[SHOCK_DIM]))
     n_obs = len(mod.coords[OBS_STATE_DIM])
 
-    assert x0.shape[-1:] == (
-        n_states,
-    ), f"x0 expected to have shape (n_states, ), found {x0.shape[-1:]}"
+    assert x0.shape[-1:] == (n_states,), (
+        f"x0 expected to have shape (n_states, ), found {x0.shape[-1:]}"
+    )
     assert P0.shape[-2:] == (
         n_states,
         n_states,
     ), f"P0 expected to have shape (n_states, n_states), found {P0.shape[-2:]}"
-    assert c.shape[-1:] == (
-        n_states,
-    ), f"c expected to have shape (n_states, ), found {c.shape[-1:]}"
+    assert c.shape[-1:] == (n_states,), (
+        f"c expected to have shape (n_states, ), found {c.shape[-1:]}"
+    )
     assert d.shape[-1:] == (n_obs,), f"d expected to have shape (n_obs, ), found {d.shape[-1:]}"
     assert T.shape[-2:] == (
         n_states,
@@ -107,9 +107,9 @@ def _assert_keys_match(test_dict, expected_dict):
     assert len(key_diff) == 0, f"{', '.join(key_diff)} were not found in the test_dict keys."
 
     key_diff = set(param_keys) - set(expected_keys)
-    assert (
-        len(key_diff) == 0
-    ), f"{', '.join(key_diff)} were keys of the tests_dict not in expected_dict."
+    assert len(key_diff) == 0, (
+        f"{', '.join(key_diff)} were keys of the tests_dict not in expected_dict."
+    )
 
 
 def _assert_param_dims_correct(param_dims, expected_dims):