From e50ffbdb6513887799436d9449d9f4515afbae09 Mon Sep 17 00:00:00 2001
From: jorenham <jhammudoglu@gmail.com>
Date: Sat, 7 Dec 2024 21:02:18 +0100
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20complete=20and=20improve=20`interpo?=
 =?UTF-8?q?late.interpnd`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 scipy-stubs/interpolate/_ndgriddata.pyi | 74 +++++++++++++-----
 scipy-stubs/interpolate/interpnd.pyi    | 99 +++++++++++++++++++------
 2 files changed, 134 insertions(+), 39 deletions(-)

diff --git a/scipy-stubs/interpolate/_ndgriddata.pyi b/scipy-stubs/interpolate/_ndgriddata.pyi
index c0df48ce..7c316262 100644
--- a/scipy-stubs/interpolate/_ndgriddata.pyi
+++ b/scipy-stubs/interpolate/_ndgriddata.pyi
@@ -1,33 +1,71 @@
-from typing import Any
-from typing_extensions import override
+from typing import Generic, Literal, TypeAlias, TypedDict, overload
+from typing_extensions import TypeVar, Unpack, override
 
 import numpy as np
 import optype.numpy as onp
-from scipy._typing import Untyped
+from scipy.spatial._ckdtree import cKDTree
 from .interpnd import CloughTocher2DInterpolator, LinearNDInterpolator, NDInterpolatorBase
 
 __all__ = ["CloughTocher2DInterpolator", "LinearNDInterpolator", "NearestNDInterpolator", "griddata"]
 
-class NearestNDInterpolator(NDInterpolatorBase):
-    tree: Untyped
-    values: Untyped
-    @override
+_SCT_co = TypeVar("_SCT_co", bound=np.float64 | np.complex128, default=np.float64 | np.complex128, covariant=True)
+
+class _TreeOptions(TypedDict, total=False):
+    leafsize: onp.ToJustInt
+    compact_nodes: onp.ToBool
+    copy_data: onp.ToBool
+    balanced_tree: onp.ToBool
+    boxsize: onp.ToFloatND | None
+
+class _QueryOptions(TypedDict, total=False):
+    eps: onp.ToFloat
+    p: onp.ToFloat
+    distance_upper_bound: onp.ToFloat
+    workers: int
+
+_Method: TypeAlias = Literal["nearest", "linear", "cubic"]
+
+###
+
+class NearestNDInterpolator(NDInterpolatorBase[_SCT_co], Generic[_SCT_co]):
+    tree: cKDTree
+
+    @overload
+    def __init__(
+        self: NearestNDInterpolator[np.float64],
+        /,
+        x: onp.ToFloat2D,
+        y: onp.ToFloat1D,
+        rescale: bool = False,
+        tree_options: _TreeOptions | None = None,
+    ) -> None: ...
+    @overload
     def __init__(
         self,
         /,
-        x: onp.ToFloatND,
-        y: onp.ToFloatND,
+        x: onp.ToFloat2D,
+        y: onp.ToComplex1D,
         rescale: bool = False,
-        tree_options: Untyped | None = None,
+        tree_options: _TreeOptions | None = None,
     ) -> None: ...
     @override
-    def __call__(self, /, *args: onp.ToFloatND, **query_options: Untyped) -> onp.ArrayND[np.floating[Any]]: ...
+    def __call__(self, /, *args: onp.ToFloatND, **query_options: Unpack[_QueryOptions]) -> onp.Array[onp.AtLeast1D, _SCT_co]: ...
 
+@overload
+def griddata(
+    points: onp.ToFloat1D | onp.ToFloat2D,
+    values: onp.ToFloat1D,
+    xi: onp.ToFloat2D | tuple[onp.ToFloat1D | onp.ToFloat2D, ...],
+    method: _Method = "linear",
+    fill_value: onp.ToFloat = ...,  # np.nan
+    rescale: onp.ToBool = False,
+) -> onp.Array[onp.AtLeast1D, np.float64]: ...
+@overload
 def griddata(
-    points: onp.ToFloatND,
-    values: onp.ToFloatND,
-    xi: Untyped,
-    method: str = "linear",
-    fill_value: float = ...,
-    rescale: bool = False,
-) -> Untyped: ...
+    points: onp.ToFloat1D | onp.ToFloat2D,
+    values: onp.ToComplex1D,
+    xi: onp.ToFloat2D | tuple[onp.ToFloat1D | onp.ToFloat2D, ...],
+    method: _Method = "linear",
+    fill_value: onp.ToComplex = ...,  # np.nan
+    rescale: onp.ToBool = False,
+) -> onp.Array[onp.AtLeast1D, np.float64 | np.complex128]: ...
diff --git a/scipy-stubs/interpolate/interpnd.pyi b/scipy-stubs/interpolate/interpnd.pyi
index 637aead0..23e1fa06 100644
--- a/scipy-stubs/interpolate/interpnd.pyi
+++ b/scipy-stubs/interpolate/interpnd.pyi
@@ -1,49 +1,106 @@
-from typing import Any
-from typing_extensions import override
+# scipy/interpolate/interpnd.pyx
+
+from typing import Generic, overload
+from typing_extensions import TypeVar
 
 import numpy as np
-import numpy.typing as npt
 import optype.numpy as onp
-from scipy._typing import Untyped
+from scipy.spatial._qhull import Delaunay, DelaunayInfo_t
+
+_SCT_co = TypeVar("_SCT_co", bound=np.float64 | np.complex128, default=np.float64 | np.complex128, covariant=True)
+
+###
 
 class GradientEstimationWarning(Warning): ...
 
-class NDInterpolatorBase:
+class NDInterpolatorBase(Generic[_SCT_co]):
+    points: onp.Array2D[np.float64]
+    values: onp.ArrayND[_SCT_co] | None
+    is_complex: bool
+    scale: onp.Array1D[np.float64] | None
+    offset: onp.Array1D[np.float64]  # only if rescale=True
+
+    @overload
     def __init__(
-        self,
+        self: NDInterpolatorBase[np.float64],
         /,
-        points: onp.ToFloatND,
+        points: onp.ToFloat2D | Delaunay,
         values: onp.ToFloatND,
-        fill_value: float = ...,
+        fill_value: onp.ToFloat = ...,  # np.nan
         ndim: int | None = None,
         rescale: bool = False,
         need_contiguous: bool = True,
         need_values: bool = True,
     ) -> None: ...
-    def __call__(self, /, *args: onp.ToFloatND) -> onp.ArrayND[np.floating[Any]]: ...
-
-class LinearNDInterpolator(NDInterpolatorBase):
-    @override
+    @overload
     def __init__(
         self,
         /,
-        points: onp.ToFloatND,
-        values: onp.ToFloatND,
-        fill_value: float = ...,
+        points: onp.ToFloat2D | Delaunay,
+        values: onp.ToComplexND | None,
+        fill_value: onp.ToComplex = ...,  # np.nan
+        ndim: int | None = None,
         rescale: bool = False,
+        need_contiguous: bool = True,
+        need_values: bool = True,
     ) -> None: ...
+    def __call__(self, /, *args: onp.ToFloatND) -> onp.Array[onp.AtLeast1D, _SCT_co]: ...
 
-class CloughTocher2DInterpolator(NDInterpolatorBase):
-    @override
+class LinearNDInterpolator(NDInterpolatorBase[_SCT_co], Generic[_SCT_co]):
+    @overload
+    def __init__(
+        self: LinearNDInterpolator[np.float64],
+        /,
+        points: onp.ToFloat2D | Delaunay,
+        values: onp.ToFloatND,
+        fill_value: onp.ToFloat = ...,  # np.nan
+        rescale: bool = False,
+    ) -> None: ...
+    @overload
     def __init__(
         self,
         /,
-        points: onp.ToFloatND,
+        points: onp.ToFloat2D | Delaunay,
+        values: onp.ToComplexND,
+        fill_value: onp.ToComplex = ...,  # np.nan
+        rescale: bool = False,
+    ) -> None: ...
+
+class CloughTocher2DInterpolator(NDInterpolatorBase[_SCT_co], Generic[_SCT_co]):
+    @overload
+    def __init__(
+        self: CloughTocher2DInterpolator[np.float64],
+        /,
+        points: onp.ToFloat2D | Delaunay,
         values: onp.ToFloatND,
-        fill_value: float = ...,
-        tol: float = 1e-06,
+        fill_value: onp.ToFloat = ...,  # np.nan
+        tol: onp.ToFloat = 1e-06,
+        maxiter: int = 400,
+        rescale: bool = False,
+    ) -> None: ...
+    @overload
+    def __init__(
+        self,
+        /,
+        points: onp.ToFloat2D | Delaunay,
+        values: onp.ToComplexND,
+        fill_value: onp.ToComplex = ...,  # np.nan
+        tol: onp.ToFloat = 1e-06,
         maxiter: int = 400,
         rescale: bool = False,
     ) -> None: ...
 
-def estimate_gradients_2d_global(tri: Untyped, y: npt.ArrayLike, maxiter: int = 400, tol: float = 1e-6) -> Untyped: ...
+@overload
+def estimate_gradients_2d_global(
+    tri: DelaunayInfo_t,
+    y: onp.ToFloat1D | onp.ToFloat2D,
+    maxiter: onp.ToJustInt = 400,
+    tol: float = 1e-6,
+) -> onp.Array3D[np.float64]: ...
+@overload
+def estimate_gradients_2d_global(
+    tri: DelaunayInfo_t,
+    y: onp.ToComplex1D | onp.ToComplex2D,
+    maxiter: onp.ToJustInt = 400,
+    tol: float = 1e-6,
+) -> onp.Array3D[np.float64] | onp.Array3D[np.complex128]: ...