codekingpro/portable-devtools

1from _typeshed import ConvertibleToInt
2from collections.abc import Callable, Iterable
3from typing import (
4    Any,
5    ClassVar,
6    Concatenate,
7    Final,
8    Literal as L,
9    Self,
10    TypeVar,
11    overload,
12)
13
14import numpy as np
15import numpy.typing as npt
16from numpy._typing import _IntLike_co
17
18from ._polybase import ABCPolyBase
19from ._polytypes import (
20    _Array1,
21    _Array2,
22    _CoefSeries,
23    _FuncBinOp,
24    _FuncCompanion,
25    _FuncDer,
26    _FuncFit,
27    _FuncFromRoots,
28    _FuncGauss,
29    _FuncInteg,
30    _FuncLine,
31    _FuncPoly2Ortho,
32    _FuncPow,
33    _FuncRoots,
34    _FuncUnOp,
35    _FuncVal,
36    _FuncVal2D,
37    _FuncVal3D,
38    _FuncVander,
39    _FuncVander2D,
40    _FuncVander3D,
41    _FuncWeight,
42    _Series,
43    _SeriesLikeCoef_co,
44)
45from .polyutils import trimcoef as chebtrim
46
47__all__ = [
48    "chebzero",
49    "chebone",
50    "chebx",
51    "chebdomain",
52    "chebline",
53    "chebadd",
54    "chebsub",
55    "chebmulx",
56    "chebmul",
57    "chebdiv",
58    "chebpow",
59    "chebval",
60    "chebder",
61    "chebint",
62    "cheb2poly",
63    "poly2cheb",
64    "chebfromroots",
65    "chebvander",
66    "chebfit",
67    "chebtrim",
68    "chebroots",
69    "chebpts1",
70    "chebpts2",
71    "Chebyshev",
72    "chebval2d",
73    "chebval3d",
74    "chebgrid2d",
75    "chebgrid3d",
76    "chebvander2d",
77    "chebvander3d",
78    "chebcompanion",
79    "chebgauss",
80    "chebweight",
81    "chebinterpolate",
82]
83
84_NumberOrObjectT = TypeVar("_NumberOrObjectT", bound=np.number | np.object_)
85_CoefScalarT = TypeVar("_CoefScalarT", bound=np.number | np.bool | np.object_)
86
87def _cseries_to_zseries(c: npt.NDArray[_NumberOrObjectT]) -> _Series[_NumberOrObjectT]: ...
88def _zseries_to_cseries(zs: npt.NDArray[_NumberOrObjectT]) -> _Series[_NumberOrObjectT]: ...
89def _zseries_mul(z1: npt.NDArray[_NumberOrObjectT], z2: npt.NDArray[_NumberOrObjectT]) -> _Series[_NumberOrObjectT]: ...
90def _zseries_div(z1: npt.NDArray[_NumberOrObjectT], z2: npt.NDArray[_NumberOrObjectT]) -> _Series[_NumberOrObjectT]: ...
91def _zseries_der(zs: npt.NDArray[_NumberOrObjectT]) -> _Series[_NumberOrObjectT]: ...
92def _zseries_int(zs: npt.NDArray[_NumberOrObjectT]) -> _Series[_NumberOrObjectT]: ...
93
94poly2cheb: Final[_FuncPoly2Ortho] = ...
95cheb2poly: Final[_FuncUnOp] = ...
96
97chebdomain: Final[_Array2[np.float64]] = ...
98chebzero: Final[_Array1[np.int_]] = ...
99chebone: Final[_Array1[np.int_]] = ...
100chebx: Final[_Array2[np.int_]] = ...
101
102chebline: Final[_FuncLine] = ...
103chebfromroots: Final[_FuncFromRoots] = ...
104chebadd: Final[_FuncBinOp] = ...
105chebsub: Final[_FuncBinOp] = ...
106chebmulx: Final[_FuncUnOp] = ...
107chebmul: Final[_FuncBinOp] = ...
108chebdiv: Final[_FuncBinOp] = ...
109chebpow: Final[_FuncPow] = ...
110chebder: Final[_FuncDer] = ...
111chebint: Final[_FuncInteg] = ...
112chebval: Final[_FuncVal] = ...
113chebval2d: Final[_FuncVal2D] = ...
114chebval3d: Final[_FuncVal3D] = ...
115chebgrid2d: Final[_FuncVal2D] = ...
116chebgrid3d: Final[_FuncVal3D] = ...
117chebvander: Final[_FuncVander] = ...
118chebvander2d: Final[_FuncVander2D] = ...
119chebvander3d: Final[_FuncVander3D] = ...
120chebfit: Final[_FuncFit] = ...
121chebcompanion: Final[_FuncCompanion] = ...
122chebroots: Final[_FuncRoots] = ...
123chebgauss: Final[_FuncGauss] = ...
124chebweight: Final[_FuncWeight] = ...
125def chebpts1(npts: ConvertibleToInt) -> np.ndarray[tuple[int], np.dtype[np.float64]]: ...
126def chebpts2(npts: ConvertibleToInt) -> np.ndarray[tuple[int], np.dtype[np.float64]]: ...
127
128# keep in sync with `Chebyshev.interpolate` (minus `domain` parameter)
129@overload
130def chebinterpolate(
131    func: np.ufunc,
132    deg: _IntLike_co,
133    args: tuple[()] = (),
134) -> npt.NDArray[np.float64 | np.complex128 | np.object_]: ...
135@overload
136def chebinterpolate(
137    func: Callable[[npt.NDArray[np.float64]], _CoefScalarT],
138    deg: _IntLike_co,
139    args: tuple[()] = (),
140) -> npt.NDArray[_CoefScalarT]: ...
141@overload
142def chebinterpolate(
143    func: Callable[Concatenate[npt.NDArray[np.float64], ...], _CoefScalarT],
144    deg: _IntLike_co,
145    args: Iterable[Any],
146) -> npt.NDArray[_CoefScalarT]: ...
147
148class Chebyshev(ABCPolyBase[L["T"]]):
149    basis_name: ClassVar[L["T"]] = "T"  # pyright: ignore[reportIncompatibleMethodOverride]
150    domain: _Array2[np.float64 | Any] = ...  # pyright: ignore[reportIncompatibleMethodOverride]
151    window: _Array2[np.float64 | Any] = ...  # pyright: ignore[reportIncompatibleMethodOverride]
152
153    @overload
154    @classmethod
155    def interpolate(
156        cls,
157        func: Callable[[npt.NDArray[np.float64]], _CoefSeries],
158        deg: _IntLike_co,
159        domain: _SeriesLikeCoef_co | None = None,
160        args: tuple[()] = (),
161    ) -> Self: ...
162    @overload
163    @classmethod
164    def interpolate(
165        cls,
166        func: Callable[Concatenate[npt.NDArray[np.float64], ...], _CoefSeries],
167        deg: _IntLike_co,
168        domain: _SeriesLikeCoef_co | None = None,
169        *,
170        args: Iterable[Any],
171    ) -> Self: ...
172    @overload
173    @classmethod
174    def interpolate(
175        cls,
176        func: Callable[Concatenate[npt.NDArray[np.float64], ...], _CoefSeries],
177        deg: _IntLike_co,
178        domain: _SeriesLikeCoef_co | None,
179        args: Iterable[Any],
180    ) -> Self: ...
181