pymc-extras 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

pymc_extras/statespace/filters/distributions.py

@@ -6,11 +6,9 @@ import pytensor.tensor as pt
 from pymc import intX
 from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import Continuous, SymbolicRandomVariable
-from pymc.distributions.multivariate import MvNormal
 from pymc.distributions.shape_utils import get_support_shape_1d
 from pymc.logprob.abstract import _logprob
 from pytensor.graph.basic import Node
-from pytensor.tensor.random.basic import MvNormalRV
 
 floatX = pytensor.config.floatX
 COV_ZERO_TOL = 0
@@ -49,44 +47,6 @@ def make_signature(sequence_names):
     return f"{signature},[rng]->[rng],({time},{state_and_obs})"
 
 
-class MvNormalSVDRV(MvNormalRV):
-    name = "multivariate_normal"
-    signature = "(n),(n,n)->(n)"
-    dtype = "floatX"
-    _print_name = ("MultivariateNormal", "\\operatorname{MultivariateNormal}")
-
-
-class MvNormalSVD(MvNormal):
-    """Dummy distribution intended to be rewritten into a JAX multivariate_normal with method="svd".
-
-    A JAX MvNormal robust to low-rank covariance matrices
-    """
-
-    rv_op = MvNormalSVDRV()
-
-
-try:
-    import jax.random
-
-    from pytensor.link.jax.dispatch.random import jax_sample_fn
-
-    @jax_sample_fn.register(MvNormalSVDRV)
-    def jax_sample_fn_mvnormal_svd(op, node):
-        def sample_fn(rng, size, dtype, *parameters):
-            rng_key = rng["jax_state"]
-            rng_key, sampling_key = jax.random.split(rng_key, 2)
-            sample = jax.random.multivariate_normal(
-                sampling_key, *parameters, shape=size, dtype=dtype, method="svd"
-            )
-            rng["jax_state"] = rng_key
-            return (rng, sample)
-
-        return sample_fn
-
-except ImportError:
-    pass
-
-
 class LinearGaussianStateSpaceRV(SymbolicRandomVariable):
     default_output = 1
     _print_name = ("LinearGuassianStateSpace", "\\operatorname{LinearGuassianStateSpace}")
@@ -244,8 +204,12 @@ class _LinearGaussianStateSpace(Continuous):
             k = T.shape[0]
             a = state[:k]
 
-            middle_rng, a_innovation = MvNormalSVD.dist(mu=0, cov=Q, rng=rng).owner.outputs
-            next_rng, y_innovation = MvNormalSVD.dist(mu=0, cov=H, rng=middle_rng).owner.outputs
+            middle_rng, a_innovation = pm.MvNormal.dist(
+                mu=0, cov=Q, rng=rng, method="svd"
+            ).owner.outputs
+            next_rng, y_innovation = pm.MvNormal.dist(
+                mu=0, cov=H, rng=middle_rng, method="svd"
+            ).owner.outputs
 
             a_mu = c + T @ a
             a_next = a_mu + R @ a_innovation
@@ -260,8 +224,8 @@ class _LinearGaussianStateSpace(Continuous):
         Z_init = Z_ if Z_ in non_sequences else Z_[0]
         H_init = H_ if H_ in non_sequences else H_[0]
 
-        init_x_ = MvNormalSVD.dist(a0_, P0_, rng=rng)
-        init_y_ = MvNormalSVD.dist(Z_init @ init_x_, H_init, rng=rng)
+        init_x_ = pm.MvNormal.dist(a0_, P0_, rng=rng, method="svd")
+        init_y_ = pm.MvNormal.dist(Z_init @ init_x_, H_init, rng=rng, method="svd")
 
         init_dist_ = pt.concatenate([init_x_, init_y_], axis=0)
 
@@ -421,7 +385,7 @@ class SequenceMvNormal(Continuous):
         rng = pytensor.shared(np.random.default_rng())
 
         def step(mu, cov, rng):
-            new_rng, mvn = MvNormalSVD.dist(mu=mu, cov=cov, rng=rng).owner.outputs
+            new_rng, mvn = pm.MvNormal.dist(mu=mu, cov=cov, rng=rng, method="svd").owner.outputs
             return mvn, {rng: new_rng}
 
         mvn_seq, updates = pytensor.scan(

pymc_extras/statespace/utils/data_tools.py

@@ -87,12 +87,7 @@ def preprocess_pandas_data(data, n_obs, obs_coords=None, check_column_names=Fals
     col_names = data.columns
     _validate_data_shape(data.shape, n_obs, obs_coords, check_column_names, col_names)
 
-    if isinstance(data.index, pd.RangeIndex):
-        if obs_coords is not None:
-            warnings.warn(NO_TIME_INDEX_WARNING)
-        return preprocess_numpy_data(data.values, n_obs, obs_coords)
-
-    elif isinstance(data.index, pd.DatetimeIndex):
+    if isinstance(data.index, pd.DatetimeIndex):
         if data.index.freq is None:
             warnings.warn(NO_FREQ_INFO_WARNING)
             data.index.freq = data.index.inferred_freq
@@ -100,10 +95,30 @@ def preprocess_pandas_data(data, n_obs, obs_coords=None, check_column_names=Fals
         index = data.index
         return data.values, index
 
+    elif isinstance(data.index, pd.RangeIndex):
+        if obs_coords is not None:
+            warnings.warn(NO_TIME_INDEX_WARNING)
+        return preprocess_numpy_data(data.values, n_obs, obs_coords)
+
+    elif isinstance(data.index, pd.MultiIndex):
+        if obs_coords is not None:
+            warnings.warn(NO_TIME_INDEX_WARNING)
+
+        raise NotImplementedError("MultiIndex panel data is not currently supported.")
+
     else:
-        raise IndexError(
-            f"Expected pd.DatetimeIndex or pd.RangeIndex on data, found {type(data.index)}"
-        )
+        if obs_coords is not None:
+            warnings.warn(NO_TIME_INDEX_WARNING)
+
+        index = data.index
+        if not np.issubdtype(index.dtype, np.integer):
+            raise IndexError("Provided index is not an integer index.")
+
+        index_diff = index.to_series().diff().dropna().values
+        if not (index_diff == 1).all():
+            raise IndexError("Provided index is not monotonic increasing.")
+
+        return preprocess_numpy_data(data.values, n_obs, obs_coords)
 
 
 def add_data_to_active_model(values, index, data_dims=None):

pymc_extras/version.txt

@@ -1 +1 @@
-0.2.3
+0.2.5

{pymc_extras-0.2.3.dist-info → pymc_extras-0.2.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: pymc-extras
-Version: 0.2.3
+Version: 0.2.5
 Summary: A home for new additions to PyMC, which may include unusual probability distribitions, advanced model fitting algorithms, or any code that may be inappropriate to include in the pymc repository, but may want to be made available to users.
 Home-page: http://github.com/pymc-devs/pymc-extras
 Maintainer: PyMC Developers
@@ -12,6 +12,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Intended Audience :: Science/Research
 Classifier: Topic :: Scientific/Engineering
@@ -20,7 +21,7 @@ Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: pymc>=5.20
+Requires-Dist: pymc>=5.21.1
 Requires-Dist: scikit-learn
 Requires-Dist: better-optimize
 Provides-Extra: dask-histogram
@@ -40,6 +41,7 @@ Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: license
+Dynamic: license-file
 Dynamic: maintainer
 Dynamic: maintainer-email
 Dynamic: provides-extra

{pymc_extras-0.2.3.dist-info → pymc_extras-0.2.5.dist-info}/RECORD

@@ -1,26 +1,28 @@
-pymc_extras/__init__.py,sha256=IFIEZdPX_Ugq57Bu7jlyrJLpKng-P0FBAAAzl2pFXLE,1266
+pymc_extras/__init__.py,sha256=lYGf9TcwUHROIElkX7Epnb7-IppcmiSEYuxdtRzqS3s,1195
 pymc_extras/linearmodel.py,sha256=6eitl15Ec15mSZu7zoHZ7Wwy4U1DPwqfAgwEt6ILeIc,3920
 pymc_extras/model_builder.py,sha256=sAw77fxdiy046BvDPjocuMlbJ0Efj-CDAGtmcwYmoG0,26361
 pymc_extras/printing.py,sha256=G8mj9dRd6i0PcsbcEWZm56ek6V8mmil78RI4MUhywBs,6506
 pymc_extras/version.py,sha256=VxPGCBzhtSegu-Jp5cjzn0n4DGU0wuPUh-KyZKB6uPM,240
-pymc_extras/version.txt,sha256=OrlMBNJJhvOvKIuhzaLAu928Wonf8JcYKAX1RXjh6nU,6
-pymc_extras/distributions/__init__.py,sha256=gTX7tvX8NcgP7V72URV7GeqF1aAEjGVbuW8LMxhXceY,1295
+pymc_extras/version.txt,sha256=6Vn3UOktu3YUriislvCjcnLK7YHYu7dYeRr3v7thBqA,6
+pymc_extras/distributions/__init__.py,sha256=fDbrBt9mxEVp2CDPwnyCW3oiutzZ0PduB8EUH3fUrjI,1377
 pymc_extras/distributions/continuous.py,sha256=z-nvQgGncYISdRY8cWsa-56V0bQGq70jYwU-i8VZ0Uk,11253
 pymc_extras/distributions/discrete.py,sha256=vrARNuiQAEXrs7yQgImV1PO8AV1uyEC_LBhr6F9IcOg,13032
 pymc_extras/distributions/histogram_utils.py,sha256=5RTvlGCUrp2qzshrchmPyWxjhs6RIYL62SMikjDM1TU,5814
 pymc_extras/distributions/timeseries.py,sha256=M5MZ-nik_tgkaoZ1hdUGEZ9g04DQyVLwszVJqSKwNcY,12719
 pymc_extras/distributions/multivariate/__init__.py,sha256=E8OeLW9tTotCbrUjEo4um76-_WQD56PehsPzkKmhfyA,93
 pymc_extras/distributions/multivariate/r2d2m2cp.py,sha256=bUj9bB-hQi6CpaJfvJjgNPi727uTbvAdxl9fm1zNBqY,16005
+pymc_extras/distributions/transforms/__init__.py,sha256=FUp2vyRE6_2eUcQ_FVt5Dn0-vy5I-puV-Kz13-QtLNc,104
+pymc_extras/distributions/transforms/partial_order.py,sha256=oEZlc9WgnGR46uFEjLzKEUxlhzIo2vrUUbBE3vYrsfQ,8404
 pymc_extras/gp/__init__.py,sha256=sFHw2y3lEl5tG_FDQHZUonQ_k0DF1JRf0Rp8dpHmge0,745
 pymc_extras/gp/latent_approx.py,sha256=cDEMM6H1BL2qyKg7BZU-ISrKn2HJe7hDaM4Y8GgQDf4,6682
-pymc_extras/inference/__init__.py,sha256=5cXpaQQnW0mJJ3x8wSxmYu63l--Xab5D_gMtjA6Q3uU,666
+pymc_extras/inference/__init__.py,sha256=UH6S0bGfQKKyTSuqf7yezdy9PeE2bDU8U1v4eIRv4ZI,887
 pymc_extras/inference/find_map.py,sha256=vl5l0ei48PnX-uTuHVTr-9QpCEHc8xog-KK6sOnJ8LU,16513
-pymc_extras/inference/fit.py,sha256=S9R48dh74s6K0MC9Iys4NAwVjP6rVRfx6SF-kPiR70E,1165
-pymc_extras/inference/laplace.py,sha256=uOZGp8ssQuhvCHV_Y_v3icsr4rhcYgr_qlr9dS7pcSM,21761
+pymc_extras/inference/fit.py,sha256=oe20RAajImZ-VD9Ucbzri8Bof4Y2KHNhNRG19v9O3lI,1336
+pymc_extras/inference/laplace.py,sha256=cqarAdbFaOH74AkPUF4c7c_Hswa5mqmhgHpsgrkebHY,21860
 pymc_extras/inference/pathfinder/__init__.py,sha256=FhAYrCWNx_dCrynEdjg2CZ9tIinvcVLBm67pNx_Y3kA,101
-pymc_extras/inference/pathfinder/importance_sampling.py,sha256=VvmuaE3aw_Mo3tMwswfF0rqe19mnhOCpzIScaJzjA1Y,6159
-pymc_extras/inference/pathfinder/lbfgs.py,sha256=P0UIOVtspdLzDU6alK-y91qzVAzXjYAXPuGmZ1nRqMo,5715
-pymc_extras/inference/pathfinder/pathfinder.py,sha256=fomZ5voVcWxvhWpeIZV7IHGIJCasT1g0ivC4dC3-0GM,63694
+pymc_extras/inference/pathfinder/importance_sampling.py,sha256=NwxepXOFit3cA5zEebniKdlnJ1rZWg56aMlH4MEOcG4,6264
+pymc_extras/inference/pathfinder/lbfgs.py,sha256=GOoJBil5Kft_iFwGNUGKSeqzI5x_shA4KQWDwgGuQtQ,7110
+pymc_extras/inference/pathfinder/pathfinder.py,sha256=GW04HQurj_3Nlo1C6_K2tEIeigo8x0buV3FqDLA88PQ,64439
 pymc_extras/inference/smc/__init__.py,sha256=wyaT4NJl1YsSQRLiDy-i0Jq3CbJZ2BQd4nnCk-dIngY,603
 pymc_extras/inference/smc/sampling.py,sha256=AYwmKqGoV6pBtKnh9SUbBKbN7VcoFgb3MmNWV7SivMA,15365
 pymc_extras/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -37,9 +39,9 @@ pymc_extras/statespace/__init__.py,sha256=0MtZj7yT6jcyERvITnn-nkhyY8fO6Za4_vV53C
 pymc_extras/statespace/core/__init__.py,sha256=huHEiXAm8zV2MZyZ8GBHp6q7_fnWqveM7lC6ilpb3iE,309
 pymc_extras/statespace/core/compile.py,sha256=9FZfE8Bi3VfElxujfOIKRVvmyL9M5R0WfNEqPc5kbVQ,1603
 pymc_extras/statespace/core/representation.py,sha256=DwNIun6wdeEA20oWBx5M4govyWTf5JI87aGQ_E6Mb4U,18956
-pymc_extras/statespace/core/statespace.py,sha256=K_WVnWKlI6sR2kgriq9sctQVvwXCeAirm14TthDpmRM,96860
+pymc_extras/statespace/core/statespace.py,sha256=Tx-821UNNLqsZgHzRmwaQ6s-agp_OthqSsbfwDpA1o0,96927
 pymc_extras/statespace/filters/__init__.py,sha256=N9Q4D0gAq_ZtT-GtrqiX1HkSg6Orv7o1TbrWUtnbTJE,420
-pymc_extras/statespace/filters/distributions.py,sha256=-9j__vRqL5hKyYFnQr5HKHA5kEFzwiuSccH4mslTOuQ,12900
+pymc_extras/statespace/filters/distributions.py,sha256=ejimTFLgBFZMEznxY5zh6u4Vrqij60i0k2_sxdPcZ3A,11878
 pymc_extras/statespace/filters/kalman_filter.py,sha256=HELC3aK4k8EdWlUAk5_F7y7YkIz-Xi_0j2AwRgAXgcc,31949
 pymc_extras/statespace/filters/kalman_smoother.py,sha256=ypH9K_88nfJ5K2Cq737aWL3p8v4UfI7MxnYs54WPdDs,4329
 pymc_extras/statespace/filters/utilities.py,sha256=iwdaYnO1cO06t_XUjLLRmqb8vwzzVH6Nx1iyZcbJL2k,1584
@@ -52,21 +54,22 @@ pymc_extras/statespace/models/utilities.py,sha256=G9GuHKsghmIYOlfkPtvxBWF-FZY5-5
 pymc_extras/statespace/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pymc_extras/statespace/utils/constants.py,sha256=Kf6j75ABaDQeRODxKQ76wTUQV4F5sTjn1KBcZgCQx20,2403
 pymc_extras/statespace/utils/coord_tools.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-pymc_extras/statespace/utils/data_tools.py,sha256=caanvrxDu9g-dEKff2bbmaVTs6-71kkSoYIiiSUXhw4,5985
+pymc_extras/statespace/utils/data_tools.py,sha256=01sz6XDtLYK9I5xghxYpD-PuDzGXv9D-wFGfTV6FGEw,6566
 pymc_extras/utils/__init__.py,sha256=yxI9cJ7fCtVQS0GFw0y6mDGZIQZiK53vm3UNKqIuGSk,758
 pymc_extras/utils/linear_cg.py,sha256=KkXhuimFsrKtNd_0By2ApxQQQNm5FdBtmDQJOVbLYkA,10056
 pymc_extras/utils/model_equivalence.py,sha256=8QIftID2HDxD659i0RXHazQ-l2Q5YegCRLcDqb2p9Pc,2187
 pymc_extras/utils/pivoted_cholesky.py,sha256=QtnjP0pAl9b77fLAu-semwT4_9dcoiqx3dz1xKGBjMk,1871
 pymc_extras/utils/prior.py,sha256=QlWVr7uKIK9VncBw7Fz3YgaASKGDfqpORZHc-vz_9gQ,6841
 pymc_extras/utils/spline.py,sha256=qGq0gcoMG5dpdazKFzG0RXkkCWP8ADPPXN-653-oFn4,4820
+pymc_extras-0.2.5.dist-info/licenses/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
 tests/__init__.py,sha256=-ree9OWVCyTeXLR944OWjrQX2os15HXrRNkhJ7QdRjc,603
 tests/test_blackjax_smc.py,sha256=jcNgcMBxaKyPg9UvHnWQtwoL79LXlSpZfALe3RGEZnQ,7233
 tests/test_find_map.py,sha256=B8ThnXNyfTQeem24QaLoTitFrsxKoq2VQINUdOwzna0,3379
 tests/test_histogram_approximation.py,sha256=w-xb2Rr0Qft6sm6F3BTmXXnpuqyefC1SUL6YxzqA5X4,4674
-tests/test_laplace.py,sha256=u4o-0y4v1emaTMYr_rOyL_EKY_bQIz0DUXFuwuDbfNg,9314
+tests/test_laplace.py,sha256=fArHjwMR7x98K-gZLvrvb3AwNZ7_fo7E0A4SJyt4EGU,9843
 tests/test_linearmodel.py,sha256=iB8ApNqIX9_nUHoo-Tm51xuPdrva5t4VLLut6qXB5Ao,6906
 tests/test_model_builder.py,sha256=QiINEihBR9rx8xM4Nqlg4urZKoyo58aTKDtxl9SJF1s,11249
-tests/test_pathfinder.py,sha256=GnSbZJ9QuFW9UVbkWaVgMVqQZTCttOyz_rSflxhQ-EA,4955
+tests/test_pathfinder.py,sha256=vlMI1p2Ja5X4QIaSV4h6U41I303rEppfO0JqE3xe1Rs,10023
 tests/test_pivoted_cholesky.py,sha256=PuMdMSCzO4KdQWpUF4SEBeuH_qsINCIH8TYtmmJ1NKo,692
 tests/test_printing.py,sha256=HnvwwjrjBuxXFAJdyU0K_lvKGLgh4nzHAnhsIUpenbY,5211
 tests/test_prior_from_trace.py,sha256=HOzR3l98pl7TEJquo_kSugED4wBTgHo4-8lgnpmacs8,5516
@@ -77,6 +80,7 @@ tests/distributions/test_continuous.py,sha256=1-bu-IP6RgLUJnuPYpOD8ZS1ahYbKtsJ9o
 tests/distributions/test_discrete.py,sha256=CjjaUpppsvQ6zLzV15ZsbwNOKrDmEdz4VWcleoCXUi0,7776
 tests/distributions/test_discrete_markov_chain.py,sha256=8RCHZXSB8IWjniuKaGGlM_iTWGmdrcOqginxmrAeEJg,9212
 tests/distributions/test_multivariate.py,sha256=LBvBuoT_3rzi8rR38b8L441Y-9Ff0cIXeRBKiEn6kjs,10452
+tests/distributions/test_transform.py,sha256=QM9sSQ5eSbuT2pM76nUMWqb-tQa7DGZbT9uwFDqIRUk,2672
 tests/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/model/test_model_api.py,sha256=FJvMTmexovRELZOUcUyk-6Vwk9qSiH7hIFoiArgl5mk,1040
 tests/model/marginal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -87,7 +91,7 @@ tests/statespace/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 tests/statespace/test_ETS.py,sha256=IPg3uQ7xEGqDMEHu993vtUTV7r-uNAxmw23sr5MVGfQ,15582
 tests/statespace/test_SARIMAX.py,sha256=1BYNOm9aSHnpn-qbpe3YsQVH8m-mXcp_gvKgWhWn1W4,12948
 tests/statespace/test_VARMAX.py,sha256=rJnea9_WEGo9I0iv2eaSbwwFQv0tlIjpN7KAE0eQewU,6336
-tests/statespace/test_coord_assignment.py,sha256=2GBm46-0eI4QNh4bvp3D7az58stcA5Zo6VgOo_JkCig,3821
+tests/statespace/test_coord_assignment.py,sha256=2Mo5196ibkBTwscE7kqQoUsgQphdaagVkOccDi7D4RI,5980
 tests/statespace/test_distributions.py,sha256=WQ_ROyd-PL3cimXTyEtyVaMEVtS7Hue2Z0lN7UnGDyo,9122
 tests/statespace/test_kalman_filter.py,sha256=s2n62FzXl9elU_uqaMNaEaexUfq3SXe3_YvQ2lM6hiQ,11600
 tests/statespace/test_representation.py,sha256=1KAJY4ZaVhb1WdAJLx2UYSXuVYsMNWX98gEDF7P0B4s,6210
@@ -98,8 +102,7 @@ tests/statespace/utilities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJ
 tests/statespace/utilities/shared_fixtures.py,sha256=SNw8Bvj1Yw11TxAW6n20Bq0B8oaYtVTiFFEVNH_wnp4,164
 tests/statespace/utilities/statsmodel_local_level.py,sha256=SQAzaYaSDwiVhUQ1iWjt4MgfAd54RuzVtnslIs3xdS8,1225
 tests/statespace/utilities/test_helpers.py,sha256=oH24a6Q45NFFFI3Kx9mhKbxsCvo9ErCorKFoTjDB3-4,9159
-pymc_extras-0.2.3.dist-info/LICENSE,sha256=WjiLhUKEysJvy5e9jk6WwFv9tmAPtnov1uJ6gcH1kIs,11720
-pymc_extras-0.2.3.dist-info/METADATA,sha256=ZTiMM7hvVRF3O_liRu4Aea_EuxJc4vHfTD2CbRRQrcU,5152
-pymc_extras-0.2.3.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-pymc_extras-0.2.3.dist-info/top_level.txt,sha256=D6RkgBiXiZCel0nvsYg_zYEoT_VuwocyIY98EMaulj0,18
-pymc_extras-0.2.3.dist-info/RECORD,,
+pymc_extras-0.2.5.dist-info/METADATA,sha256=02v5liTQQ55sV8xeFl5EjFpwbOKSYKG6g5lE_4htpBo,5227
+pymc_extras-0.2.5.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+pymc_extras-0.2.5.dist-info/top_level.txt,sha256=D6RkgBiXiZCel0nvsYg_zYEoT_VuwocyIY98EMaulj0,18
+pymc_extras-0.2.5.dist-info/RECORD,,

{pymc_extras-0.2.3.dist-info → pymc_extras-0.2.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.0)
+Generator: setuptools (78.1.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tests/distributions/test_transform.py

@@ -0,0 +1,77 @@
+#   Copyright 2025 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   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.
+import numpy as np
+import pymc as pm
+
+from pymc_extras.distributions.transforms import PartialOrder
+
+
+class TestPartialOrder:
+    adj_mats = np.array(
+        [
+            # 0 < {1, 2} < 3
+            [[0, 1, 1, 0], [0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 0, 0]],
+            # 1 < 0 < 3 < 2
+            [[0, 0, 0, 1], [1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0]],
+        ]
+    )
+
+    valid_values = np.array([[0, 2, 1, 3], [1, 0, 3, 2]], dtype=float)
+
+    # Test that forward and backward are inverses of eachother
+    # And that it works when extra dimensions are added in data
+    def test_forward_backward_dimensionality(self):
+        po = PartialOrder(self.adj_mats)
+        po0 = PartialOrder(self.adj_mats[0])
+        vv = self.valid_values
+        vv0 = self.valid_values[0]
+
+        testsets = [
+            (vv, po),
+            (po.initvals(), po),
+            (vv0, po0),
+            (po0.initvals(), po0),
+            (np.tile(vv0, (2, 1)), po0),
+            (np.tile(vv0, (2, 3, 2, 1)), po0),
+            (np.tile(vv, (2, 3, 2, 1, 1)), po),
+        ]
+
+        for vv, po in testsets:
+            fw = po.forward(vv)
+            bw = po.backward(fw)
+            np.testing.assert_allclose(bw.eval(), vv)
+
+    def test_sample_model(self):
+        po = PartialOrder(self.adj_mats)
+        with pm.Model() as model:
+            x = pm.Normal(
+                "x",
+                size=(3, 2, 4),
+                transform=po,
+                initval=po.initvals(shape=(3, 2, 4), lower=-1, upper=1),
+            )
+            idata = pm.sample()
+
+        # Check that the order constraints are satisfied
+        # Move chain, draw and "3" dimensions to the back
+        xvs = idata.posterior.x.values.transpose(3, 4, 0, 1, 2)
+        x0 = xvs[0]  # 0 < {1, 2} < 3
+        assert (
+            (x0[0] < x0[1]).all()
+            and (x0[0] < x0[2]).all()
+            and (x0[1] < x0[3]).all()
+            and (x0[2] < x0[3]).all()
+        )
+        x1 = xvs[1]  # 1 < 0 < 3 < 2
+        assert (x1[1] < x1[0]).all() and (x1[0] < x1[3]).all() and (x1[3] < x1[2]).all()

tests/statespace/test_coord_assignment.py

@@ -8,6 +8,7 @@ import pytensor.tensor as pt
 import pytest
 
 from pymc_extras.statespace.models import structural
+from pymc_extras.statespace.models.structural import LevelTrendComponent
 from pymc_extras.statespace.utils.constants import (
     FILTER_OUTPUT_DIMS,
     FILTER_OUTPUT_NAMES,
@@ -114,3 +115,67 @@ def test_data_index_is_coord(f, warning, create_model):
     with warning:
         pymc_model = create_model(f)
     assert TIME_DIM in pymc_model.coords
+
+
+def make_model(index):
+    n = len(index)
+    a = pd.DataFrame(index=index, columns=["A", "B", "C", "D"], data=np.arange(n * 4).reshape(n, 4))
+
+    mod = LevelTrendComponent(order=2, innovations_order=[0, 1])
+    ss_mod = mod.build(name="a", verbose=False)
+
+    initial_trend_dims, sigma_trend_dims, P0_dims = ss_mod.param_dims.values()
+    coords = ss_mod.coords
+
+    with pm.Model(coords=coords) as model:
+        P0_diag = pm.Gamma("P0_diag", alpha=5, beta=5)
+        P0 = pm.Deterministic("P0", pt.eye(ss_mod.k_states) * P0_diag, dims=P0_dims)
+
+        initial_trend = pm.Normal("initial_trend", dims=initial_trend_dims)
+        sigma_trend = pm.Gamma("sigma_trend", alpha=2, beta=50, dims=sigma_trend_dims)
+
+        with pytest.warns(UserWarning, match="No time index found on the supplied data"):
+            ss_mod.build_statespace_graph(
+                a["A"],
+                mode="JAX",
+            )
+    return model
+
+
+def test_integer_index():
+    index = np.arange(8).astype(int)
+    model = make_model(index)
+    assert TIME_DIM in model.coords
+    np.testing.assert_allclose(model.coords[TIME_DIM], index)
+
+
+def test_float_index_raises():
+    index = np.linspace(0, 1, 8)
+
+    with pytest.raises(IndexError, match="Provided index is not an integer index"):
+        make_model(index)
+
+
+def test_non_strictly_monotone_index_raises():
+    # Decreases
+    index = [0, 1, 2, 1, 2, 3]
+    with pytest.raises(IndexError, match="Provided index is not monotonic increasing"):
+        make_model(index)
+
+    # Has gaps
+    index = [0, 1, 2, 3, 5, 6]
+    with pytest.raises(IndexError, match="Provided index is not monotonic increasing"):
+        make_model(index)
+
+    # Has duplicates
+    index = [0, 1, 1, 2, 3, 4]
+    with pytest.raises(IndexError, match="Provided index is not monotonic increasing"):
+        make_model(index)
+
+
+def test_multiindex_raises():
+    index = pd.MultiIndex.from_tuples([(0, 0), (1, 1), (2, 2), (3, 3)])
+    with pytest.raises(
+        NotImplementedError, match="MultiIndex panel data is not currently supported"
+    ):
+        make_model(index)

tests/test_laplace.py

@@ -263,3 +263,19 @@ def test_fit_laplace(fit_in_unconstrained_space, mode, gradient_backend: Gradien
         else:
             assert idata.fit.rows.values.tolist() == ["mu", "sigma"]
             np.testing.assert_allclose(idata.fit.mean_vector.values, np.array([3.0, 1.5]), atol=0.1)
+
+
+def test_laplace_scalar():
+    # Example model from Statistical Rethinking
+    data = np.array([0, 0, 0, 1, 1, 1, 1, 1, 1])
+
+    with pm.Model():
+        p = pm.Uniform("p", 0, 1)
+        w = pm.Binomial("w", n=len(data), p=p, observed=data.sum())
+
+        idata_laplace = pmx.fit_laplace(progressbar=False)
+
+    assert idata_laplace.fit.mean_vector.shape == (1,)
+    assert idata_laplace.fit.covariance_matrix.shape == (1, 1)
+
+    np.testing.assert_allclose(idata_laplace.fit.mean_vector.values.item(), data.mean(), atol=0.1)