Depreciate duplicate functions in statespace (#308)

* Delete `pytensor_scipy.py`

* Delete `block_diagonal`, use `pt.linalg.block_diag`

* Relax tolerance in loglikelihood tests

* Use updated pandas frequency string

Author: jessegrabowski

pymc_experimental/statespace/filters/kalman_filter.py

@@ -9,15 +9,14 @@
 from pytensor.raise_op import Assert
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.nlinalg import matrix_dot
-from pytensor.tensor.slinalg import solve_triangular
+from pytensor.tensor.slinalg import solve_discrete_are, solve_triangular
 
 from pymc_experimental.statespace.filters.utilities import (
     quad_form_sym,
     split_vars_into_seq_and_nonseq,
     stabilize,
 )
 from pymc_experimental.statespace.utils.constants import JITTER_DEFAULT, MISSING_FILL
-from pymc_experimental.statespace.utils.pytensor_scipy import solve_discrete_are
 
 MVN_CONST = pt.log(2 * pt.constant(np.pi, dtype="float64"))
 PARAM_NAMES = ["c", "d", "T", "Z", "R", "H", "Q"]

pymc_experimental/statespace/models/SARIMAX.py

@@ -2,6 +2,7 @@
 
 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
 from pymc_experimental.statespace.models.utilities import (
@@ -19,7 +20,6 @@
     SEASONAL_AR_PARAM_DIM,
     SEASONAL_MA_PARAM_DIM,
 )
-from pymc_experimental.statespace.utils.pytensor_scipy import solve_discrete_lyapunov
 
 
 def _verify_order(p, d, q, P, D, Q, S):

pymc_experimental/statespace/models/VARMAX.py

@@ -3,6 +3,7 @@
 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
 from pymc_experimental.statespace.models.utilities import make_default_coords
@@ -16,7 +17,6 @@
     SHOCK_AUX_DIM,
     SHOCK_DIM,
 )
-from pymc_experimental.statespace.utils.pytensor_scipy import solve_discrete_lyapunov
 
 floatX = pytensor.config.floatX
 

pymc_experimental/statespace/models/structural.py

@@ -44,23 +44,6 @@ def _frequency_transition_block(s, j):
     return pt.stack([[pt.cos(lam), pt.sin(lam)], [-pt.sin(lam), pt.cos(lam)]]).squeeze()
 
 
-def block_diagonal(matrices: list[pt.matrix]):
-    rows = [x.shape[0] for x in matrices]
-    cols = [x.shape[1] for x in matrices]
-    out = pt.zeros((sum(rows), sum(cols)))
-    row_cursor = 0
-    col_cursor = 0
-
-    for row, col, mat in zip(rows, cols, matrices):
-        row_slice = slice(row_cursor, row_cursor + row)
-        col_slice = slice(col_cursor, col_cursor + col)
-        row_cursor += row
-        col_cursor += col
-
-        out = pt.set_subtensor(out[row_slice, col_slice], mat)
-    return out
-
-
 class StructuralTimeSeries(PyMCStateSpace):
     r"""
     Structural Time Series Model
@@ -527,7 +510,7 @@ def make_slice(name, x, o_x):
         initial_state = pt.concatenate(conform_time_varying_and_time_invariant_matrices(x0, o_x0))
         initial_state.name = x0.name
 
-        initial_state_cov = block_diagonal([P0, o_P0])
+        initial_state_cov = pt.linalg.block_diag(P0, o_P0)
         initial_state_cov.name = P0.name
 
         state_intercept = pt.concatenate(conform_time_varying_and_time_invariant_matrices(c, o_c))
@@ -536,19 +519,19 @@ def make_slice(name, x, o_x):
         obs_intercept = d + o_d
         obs_intercept.name = d.name
 
-        transition = block_diagonal([T, o_T])
+        transition = pt.linalg.block_diag(T, o_T)
         transition.name = T.name
 
         design = pt.concatenate(conform_time_varying_and_time_invariant_matrices(Z, o_Z), axis=-1)
         design.name = Z.name
 
-        selection = block_diagonal([R, o_R])
+        selection = pt.linalg.block_diag(R, o_R)
         selection.name = R.name
 
         obs_cov = H + o_H
         obs_cov.name = H.name
 
-        state_cov = block_diagonal([Q, o_Q])
+        state_cov = pt.linalg.block_diag(Q, o_Q)
         state_cov.name = Q.name
 
         new_ssm = PytensorRepresentation(
@@ -1326,7 +1309,7 @@ def make_symbolic_graph(self) -> None:
         self.ssm["initial_state", init_state_idx] = init_state
 
         T_mats = [_frequency_transition_block(self.season_length, j + 1) for j in range(self.n)]
-        T = block_diagonal(T_mats)
+        T = pt.linalg.block_diag(*T_mats)
         self.ssm["transition", :, :] = T
 
         if self.innovations:

pymc_experimental/statespace/utils/pytensor_scipy.py

@@ -26,8 +26,8 @@
 
 # TODO: These are pretty loose because of all the stabilizing of covariance matrices that is done inside the kalman
 #  filters. When that is improved, this should be tightened.
-ATOL = 1e-6 if floatX.endswith("64") else 1e-4
-RTOL = 1e-6 if floatX.endswith("64") else 1e-4
+ATOL = 1e-5 if floatX.endswith("64") else 1e-4
+RTOL = 1e-5 if floatX.endswith("64") else 1e-4
 
 filter_names = [
     "standard",

pymc_experimental/tests/statespace/test_distributions.py

@@ -19,8 +19,15 @@
 
 
 def load_nile_test_data():
+    from importlib.metadata import version
+
     nile = pd.read_csv("pymc_experimental/tests/statespace/test_data/nile.csv", dtype={"x": floatX})
-    nile.index = pd.date_range(start="1871-01-01", end="1970-01-01", freq="AS-Jan")
+    major, minor, rev = map(int, version("pandas").split("."))
+    if major >= 2 and minor >= 2 and rev >= 0:
+        freq_str = "YS-JAN"
+    else:
+        freq_str = "AS-JAN"
+    nile.index = pd.date_range(start="1871-01-01", end="1970-01-01", freq=freq_str)
     nile.rename(columns={"x": "height"}, inplace=True)
     nile = (nile - nile.mean()) / nile.std()
     nile = nile.astype(floatX)