From 378e341ece0eece4c46550c3fa04527c55b2de41 Mon Sep 17 00:00:00 2001 From: Sam Duffield Date: Thu, 18 Jun 2026 15:34:21 +0100 Subject: [PATCH] Move from pyright to ty --- .pre-commit-config.yaml | 19 ++++--------------- cuthbert/discrete/smoother.py | 3 ++- cuthbert/factorial/gaussian.py | 5 ++++- cuthbert/factorial/smc.py | 4 ++++ cuthbert/gaussian/kalman.py | 3 ++- cuthbert/gaussian/moments/smoother.py | 3 ++- cuthbert/gaussian/taylor/smoother.py | 3 ++- cuthbert/inference.py | 4 ++++ cuthbert/smc/backward_sampler.py | 3 ++- cuthbertlib/linearize/moments.py | 4 ++-- cuthbertlib/resampling/protocols.py | 3 +++ pkg/cuthbert/pyproject.toml | 2 +- pyproject.toml | 3 +++ 13 files changed, 35 insertions(+), 24 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1de731cb..ab8c92a9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -8,19 +8,8 @@ repos: args: [ --extend-select, I, --fix ] # Run the formatter. - id: ruff-format -- repo: https://github.com/RobertCraigie/pyright-python - rev: v1.1.405 +- repo: https://github.com/astral-sh/ty-pre-commit + # ty version. + rev: v0.0.50 hooks: - - id: pyright - additional_dependencies: [jax>=0.4.35, numba>=0.60.0] - exclude: ^tests/ - # Change this to automatically get from pyproject.toml or a requirements.txt? - # Seemingly there are good reasons why pre-commit cannot do this https://github.com/pre-commit/pre-commit/issues/730 - # So we have to duplicate the dependencies here, but suggestions are welcome on how to improve this pyright integration! - - - - - - - + - id: ty diff --git a/cuthbert/discrete/smoother.py b/cuthbert/discrete/smoother.py index 48bd984b..86bb4486 100644 --- a/cuthbert/discrete/smoother.py +++ b/cuthbert/discrete/smoother.py @@ -53,8 +53,9 @@ def build_smoother(get_trans_matrix: GetTransitionMatrix) -> Smoother: def smoother_prepare( filter_state: DiscreteFilterState, - get_trans_matrix: GetTransitionMatrix, model_inputs: ArrayTreeLike, + *, + get_trans_matrix: GetTransitionMatrix, key: KeyArray | None = None, ) -> DiscreteSmootherState: """Prepare a state for a smoother step. diff --git a/cuthbert/factorial/gaussian.py b/cuthbert/factorial/gaussian.py index 1ab2fd2d..ee7064ae 100644 --- a/cuthbert/factorial/gaussian.py +++ b/cuthbert/factorial/gaussian.py @@ -1,3 +1,6 @@ +# ty: ignore[invalid-return-type, unknown-argument] +# ty cannot unify NamedTuple `_replace`/`Self` with the constrained `KalmanState` +# TypeVar; the proper fix would be @overload instead of the TypeVar. """Factorial utilities for Kalman states.""" from typing import TypeVar @@ -66,7 +69,7 @@ def extract(factorial_state: KalmanState, factorial_inds: ArrayLike) -> KalmanSt new_elem = tree.map(lambda x: _extract_arr(x, factorial_inds), factorial_state.elem) new_state = factorial_state._replace(elem=new_elem) - if isinstance(factorial_state, LinearizedKalmanFilterState): + if isinstance(new_state, LinearizedKalmanFilterState): new_mean_prev = _extract_arr(factorial_state.mean_prev, factorial_inds) new_state = new_state._replace(mean_prev=new_mean_prev) diff --git a/cuthbert/factorial/smc.py b/cuthbert/factorial/smc.py index 921b400f..ee2d37df 100644 --- a/cuthbert/factorial/smc.py +++ b/cuthbert/factorial/smc.py @@ -1,3 +1,7 @@ +# ty: ignore[invalid-return-type, unknown-argument] +# ty cannot unify NamedTuple `_replace`/`Self` with the constrained +# `GeneralParticleFilterState` TypeVar; the proper fix would be @overload +# instead of the TypeVar. """Factorial utilities for SMC particle-filter states.""" from typing import TypeVar diff --git a/cuthbert/gaussian/kalman.py b/cuthbert/gaussian/kalman.py index 7ee2fd9c..3272b2fb 100644 --- a/cuthbert/gaussian/kalman.py +++ b/cuthbert/gaussian/kalman.py @@ -211,8 +211,9 @@ def filter_combine( def smoother_prepare( filter_state: KalmanFilterState, - get_dynamics_params: GetDynamicsParams, model_inputs: ArrayTreeLike, + *, + get_dynamics_params: GetDynamicsParams, store_gain: bool = False, store_chol_cov_given_next: bool = False, key: KeyArray | None = None, diff --git a/cuthbert/gaussian/moments/smoother.py b/cuthbert/gaussian/moments/smoother.py index 6ce868ac..e3414449 100644 --- a/cuthbert/gaussian/moments/smoother.py +++ b/cuthbert/gaussian/moments/smoother.py @@ -70,8 +70,9 @@ def build_smoother( def smoother_prepare( filter_state: LinearizedKalmanFilterState, - get_dynamics_params: GetDynamicsMoments, model_inputs: ArrayTreeLike, + *, + get_dynamics_params: GetDynamicsMoments, store_gain: bool = False, store_chol_cov_given_next: bool = False, key: KeyArray | None = None, diff --git a/cuthbert/gaussian/taylor/smoother.py b/cuthbert/gaussian/taylor/smoother.py index b9588dc1..cd5ee961 100644 --- a/cuthbert/gaussian/taylor/smoother.py +++ b/cuthbert/gaussian/taylor/smoother.py @@ -97,8 +97,9 @@ def build_smoother( def smoother_prepare( filter_state: LinearizedKalmanFilterState, - get_dynamics_log_density: GetDynamicsLogDensity, model_inputs: ArrayTreeLike, + *, + get_dynamics_log_density: GetDynamicsLogDensity, rtol: float | None = None, ignore_nan_dims: bool = False, store_gain: bool = False, diff --git a/cuthbert/inference.py b/cuthbert/inference.py index 2e1e2e95..33e77603 100644 --- a/cuthbert/inference.py +++ b/cuthbert/inference.py @@ -11,6 +11,7 @@ class InitPrepare(Protocol): def __call__( self, model_inputs: ArrayTreeLike, + *, key: KeyArray | None = None, ) -> ArrayTree: """Prepare the initial state for the inference. @@ -34,6 +35,7 @@ class FilterPrepare(Protocol): def __call__( self, model_inputs: ArrayTreeLike, + *, key: KeyArray | None = None, ) -> ArrayTree: """Prepare the state for the filter at the next time point. @@ -87,6 +89,7 @@ def __call__( self, filter_state: ArrayTreeLike, model_inputs: ArrayTreeLike, + *, key: KeyArray | None = None, ) -> ArrayTree: """Prepare the state for the smoother at the next time point. @@ -155,6 +158,7 @@ def __call__( self, filter_state: ArrayTreeLike, model_inputs: ArrayTreeLike | None = None, + *, key: KeyArray | None = None, ) -> ArrayTree: """Convert the filter state to a smoother state. diff --git a/cuthbert/smc/backward_sampler.py b/cuthbert/smc/backward_sampler.py index 9913404b..c16c64c0 100644 --- a/cuthbert/smc/backward_sampler.py +++ b/cuthbert/smc/backward_sampler.py @@ -74,9 +74,10 @@ def build_smoother( def convert_filter_to_smoother_state( filter_state: ParticleFilterState, + model_inputs: ArrayTreeLike | None = None, + *, resampling: Resampling, n_smoother_particles: int, - model_inputs: ArrayTreeLike | None = None, key: KeyArray | None = None, ) -> ParticleSmootherState: """Convert a particle filter state to a particle smoother state. diff --git a/cuthbertlib/linearize/moments.py b/cuthbertlib/linearize/moments.py index e2c737f0..4ba9e80a 100644 --- a/cuthbertlib/linearize/moments.py +++ b/cuthbertlib/linearize/moments.py @@ -63,14 +63,14 @@ def linearize_moments( - [sqrt-parallel-smoothers](https://github.com/EEA-sensors/sqrt-parallel-smoothers/blob/main/parsmooth/linearization/_extended.py) """ if has_aux: - mean_and_chol_cov_function = cast( + mean_and_chol_cov_function_with_aux = cast( MeanAndCholCovFuncAux, mean_and_chol_cov_function ) def mean_and_chol_cov_function_wrapper_aux( x: ArrayLike, ) -> tuple[Array, tuple[Array, Array, ArrayTree]]: - mean, chol_cov, aux = mean_and_chol_cov_function(x) + mean, chol_cov, aux = mean_and_chol_cov_function_with_aux(x) return mean, (mean, chol_cov, aux) F, (m, *extra) = jax.jacfwd( diff --git a/cuthbertlib/resampling/protocols.py b/cuthbertlib/resampling/protocols.py index 68312224..1a54c650 100644 --- a/cuthbertlib/resampling/protocols.py +++ b/cuthbertlib/resampling/protocols.py @@ -1,3 +1,6 @@ +# ty: ignore[invalid-return-type] +# f-strings are not supported docstrings and ty is confused about this +# See https://github.com/state-space-models/cuthbert/issues/261 """Shared protocols for resampling algorithms.""" from typing import Protocol, runtime_checkable diff --git a/pkg/cuthbert/pyproject.toml b/pkg/cuthbert/pyproject.toml index 45de84f0..08aa39e7 100644 --- a/pkg/cuthbert/pyproject.toml +++ b/pkg/cuthbert/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ [project.optional-dependencies] tests = [ "chex", - "pre-commit", "ruff", "pyright", + "pre-commit", "ruff", "ty", "pytest", "pytest-xdist", "entangled-cli; python_version >= '3.12'", ] diff --git a/pyproject.toml b/pyproject.toml index 33121a3f..e77f0f70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,9 @@ extend-select = ["D"] [tool.ruff.lint.pydocstyle] convention = "google" +[tool.ty.src] +exclude = ["tests"] + [tool.pytest] [tool.pytest.ini_options] markers = "examples: Run tangled example scripts as tests"