diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1de731c..ab8c92a 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 48bd984..86bb448 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 1ab2fd2..ee7064a 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 921b400..ee2d37d 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 7ee2fd9..3272b2f 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 6ce868a..e341444 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 b9588dc..cd5ee96 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 2e1e2e9..33e7760 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 9913404..c16c64c 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 e2c737f..4ba9e80 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 6831222..1a54c65 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 45de84f..08aa39e 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 33121a3..e77f0f7 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"