diff --git a/pyproject.toml b/pyproject.toml index 95a2dc6..1b9ce17 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -109,3 +109,7 @@ ignore_missing_imports=true [[tool.mypy.overrides]] module="kalman.*" ignore_missing_imports=true + +[[tool.mypy.overrides]] +module="scipy.*" +ignore_missing_imports=true diff --git a/src/gen_experiments/data.py b/src/gen_experiments/data.py index a3091ad..0d651f2 100644 --- a/src/gen_experiments/data.py +++ b/src/gen_experiments/data.py @@ -1,31 +1,31 @@ from math import ceil from pathlib import Path -from typing import Callable +from typing import Callable, Optional, cast from warnings import warn import mitosis import numpy as np import scipy -from gen_experiments.utils import GridsearchResultDetails +from gen_experiments.utils import Float1D, Float2D, GridsearchResultDetails INTEGRATOR_KEYWORDS = {"rtol": 1e-12, "method": "LSODA", "atol": 1e-12} TRIALS_FOLDER = Path(__file__).parent.absolute() / "trials" def gen_data( - rhs_func, - n_coord, - seed=None, - n_trajectories=1, - x0_center=None, - ic_stdev=3, - noise_abs=None, - noise_rel=None, - nonnegative=False, - dt=0.01, - t_end=10, -): + rhs_func: Callable, + n_coord: int, + seed: Optional[int] = None, + n_trajectories: int = 1, + x0_center: Optional[Float1D] = None, + ic_stdev: float = 3, + noise_abs: Optional[float] = None, + noise_rel: Optional[float] = None, + nonnegative: bool = False, + dt: float = 0.01, + t_end: float = 10, +) -> tuple[float, Float1D, Float2D, Float2D, Float2D, Float2D]: """Generate random training and test data Note that test data has no noise. @@ -57,7 +57,7 @@ def gen_data( rng = np.random.default_rng(seed) if x0_center is None: x0_center = np.zeros((n_coord)) - t_train = np.arange(0, t_end, dt) + t_train = np.arange(0, t_end, dt, dtype=np.float_) t_train_span = (t_train[0], t_train[-1]) if nonnegative: shape = ((x0_center + 1) / ic_stdev) ** 2 @@ -123,10 +123,11 @@ def _alert_short(arr): x_train_true = np.copy(x_train) if noise_rel is not None: noise_abs = np.sqrt(_signal_avg_power(x_test) * noise_rel) - x_train = x_train + noise_abs * rng.standard_normal(x_train.shape) + x_train = x_train + cast(float, noise_abs) * rng.standard_normal(x_train.shape) x_train = list(x_train) x_test = list(x_test) x_dot_test = list(x_dot_test) + x_train_true = list(x_train_true) return dt, t_train, x_train, x_test, x_dot_test, x_train_true @@ -211,7 +212,7 @@ def gen_pde_data( x_train_true = np.copy(x_train) if noise_rel is not None: noise_abs = _max_amplitude(x_test) * noise_rel - x_train = x_train + noise_abs * rng.standard_normal(x_train.shape) + x_train = x_train + cast(float, noise_abs) * rng.standard_normal(x_train.shape) x_train = [np.moveaxis(x_train, 0, -2)] x_train_true = np.moveaxis(x_train_true, 0, -2) x_test = [np.moveaxis(x_test, [0, 1], [-1, -2])] diff --git a/src/gen_experiments/debug.py b/src/gen_experiments/debug.py new file mode 100644 index 0000000..a1193ca --- /dev/null +++ b/src/gen_experiments/debug.py @@ -0,0 +1,25 @@ +from typing import Annotated, Generic, TypedDict, TypeVar + +import numpy as np +from numpy.typing import DTypeLike, NBitBase, NDArray + +# T = TypeVar("T") + +# class Foo[T]: +# items: list[T] + +# def __init__(self, thing: T): +# self.items = [thing, thing] + +# Bar = + + +T = TypeVar("T", bound=np.generic) +Foo = NDArray[T] +Bar = Annotated[NDArray, "foobar"] + +lil_foo = NDArray[np.void] + + +def baz(qux: Foo[np.void]): + pass diff --git a/src/gen_experiments/utils.py b/src/gen_experiments/utils.py index 7841f4d..65580ee 100644 --- a/src/gen_experiments/utils.py +++ b/src/gen_experiments/utils.py @@ -63,12 +63,12 @@ class SavedData(TypedDict): T = TypeVar("T", bound=np.generic) -GridsearchResult = Annotated[NDArray[T], "(n_metrics, n_plot_axis)"] # type: ignore +GridsearchResult = Annotated[NDArray[T], "(n_metrics, n_plot_axis)"] SeriesData = Annotated[ list[ tuple[ Annotated[GridsearchResult, "metrics"], - Annotated[GridsearchResult, "arg_opts"], + Annotated[GridsearchResult[np.void], "arg_opts"], ] ], "len=n_grid_axes",