Add type stubs to scipy/signal/_peak_finding.pyi

[pavyamsiri]

I have added type stubs to scipy/signal/_peak_finding.pyi although I think the type annotations could be improved a bit.

In particular:

• I am using npt.NDArray[np.generic] to denote that the argument must be np.ndarray but the inner type is not relevant. Is there a better type for that?
• I am using a lot array likes as some of the functions cast the arguments using np.asarray but some also require that two arraylike arguments have the same shape. Is it possible to write this requirement as a type?
• I am using a lot of concrete dtypes for arrays which I found by inspecting the source. I'm not sure how subtyping works with Cython because some of these arrays are passed straight into a cython function which has a concrete dtype.
• I am using TypedDict but because the presence of the keys are dependent on whether the input arguments to the function are None or not, all values must be NotRequired but I feel like this makes the typed dict a bit useless. Making overloads is possible I guess but that would be combinatorically bad.

Open to any improvements and suggestions.

[jorenham]

Thanks for this!
In general this looks like a solid improvement.
I left a couple of minor suggestions that could help make this even better.

[jorenham]

if every key is NotRequired, then you could also use _FindPeaksResultsDict(TypedDict, total=False): ... instead. See https://docs.python.org/3/library/typing.html#typing.TypedDict for details

[pavyamsiri]

Thanks for the tip. Haven't used TypedDict before.

[jorenham]

I usually tend to put a trailing comma in cases like this, which causes ruff format to place each parameters on a new line, improving readability. But I'm also fine if you leave it like this :)

[jorenham]

I believe that peaks must always be an array-like of integers. So even though this is perfectly fine as-is, you could also consider using numpy._typing._ArrayLikeInt_co for this (no need to worry about the API stability here; at the moment I'm the only active numpy maintainer that works on typing ;) )

[pavyamsiri]

That's nice to know! I feel a bit weird using ArrayLike when the source assumes the dtype but it is nice to know that there are types like that in numpy although it seems to be behind a private module.

[jorenham]

The axis argument accepts everything that implement __index__.

So for example, instances of this Index type are also allowed.

>>> class Index:
...     def __init__(self, i: int, /) -> None:
...         self._i = int(i)  # upcast `bool` or other `int` subtypes
...     def __index__(self, /) -> int:
...         return self._i

You could use the typing.SupportsIndex or optype.CanIndex protocols for this. (the latter is optionally generic on the return type of __index__ which allows using it with e.g. Literal).

Note that there are some scipy functions that also require __lt__ or __add__ as well. But in any case, it's better to have overly-broad parameter types than overly-narrow ones.

[jorenham]

stubgen was wrong in this case: rel_height also accepts e.g. np.float16(0.5) and np.bool_(True). The scipy._typing.AnyReal type alias could be used in this case.

[pavyamsiri]

This argument is one of the ones that gets passed straight into a cython function, here is the signature

def _peak_widths(const np.float64_t[::1] x not None,
                 const np.intp_t[::1] peaks not None,
                 np.float64_t rel_height,
                 const np.float64_t[::1] prominences not None,
                 const np.intp_t[::1] left_bases not None,
                 const np.intp_t[::1] right_bases not None):

I'm not familiar with cython so I didn't know how it handles type conversion from python into cython. I think the cython docs say something like python objects are converted to their c types automatically for def functions, so I think rel_height can be any type that can get converted into np.float64_t which I think is just an alias for double. Not sure which types are allowed in that case.

The AnyReal type alias is a nice suggestion though.

[jorenham]

I'm guessing that (at least) float is also allowed here at runtime

[pavyamsiri]

This is the same as rel_height where it is passed straight into a cython function which I didn't know before did automatic type conversion. This means using AnyReal should be fine I think.

[jorenham]

Does this always have to return an array of float64 dtype, or can it also be e.g. float16?

[pavyamsiri]

I had it as a placeholder but wavelet is the same wavelet that gets passed into _cwt

def _cwt(data, wavelet, widths, dtype=None, **kwargs):
    # Determine output type
    if dtype is None:
        if np.asarray(wavelet(1, widths[0], **kwargs)).dtype.char in 'FDG':
            dtype = np.complex128
        else:
            dtype = np.float64

    output = np.empty((len(widths), len(data)), dtype=dtype)
    for ind, width in enumerate(widths):
        N = np.min([10 * width, len(data)])
        wavelet_data = np.conj(wavelet(N, width, **kwargs)[::-1])
        output[ind] = convolve(data, wavelet_data, mode='same')
    return output

which means it can just be an ArrayLike?

[jorenham]

I know it's a bit weird, but an ndarray isn't assignable to a Sequence, which the docs say that this should accept

[pavyamsiri]

Oh thats a bit weird/annoying. The reason why I chose Sequence is because I didn't look at the docs for find_peaks_cwt (because it is a bit long) but instead for the function that max_distances gets passed to which is called _identify_ridge_lines which has says that max_distances is a 1-D sequence.

The source code for that function only calls len and indexes max_distances which fits the definition of Sequence. The values from max_distances gets compared to np.intp as well.

Would the type then be like onpt.Array[tuple[int], AnyInt]? That might still be too narrow considering that Sequence is a relatively weak type constraint.

[jorenham]

The problem with Sequence is that apart from __getitem__ it also has several other methods like .count() and .index() that aren't present in np.ndarray 🤷🏻

And yea I agree that onpt.Array (or some other direct np.ndarray alias) would be indeed too narrow. So you could use e.g. _ArrayLikeInt_co, which is compatible with both Sequence[int] and npt.NDArray[integer[Any]]. And as far as I'm concerned, the fact that it would also allow integer scalars isn't much of a problem, as we're talking about the input of a function.

[jorenham]

If you're feeling up to it, you could narrow this to a 1-D array as e.g. optype.numpy.Array[tuple[int], np.intp] (but it's also perfectly fine if you leave it like this).

[jorenham]

I've seen quite a lot of npt.NDArray[np.int64] and npt.NDArray[np.inpt] being used in this module, so perhaps it could help to introduce type aliases for these two as well

[pavyamsiri]

Thanks for the feedback! The PR wasn't really complete because I'm not too well versed in all of the typing available, so working on this repo is really helping me learn about the possibilities.

I'll try to fix some of the issues although I think there are still a few hangups:

• Not sure on wavelet regarding the output type and also the input types as well
• Not sure on max_distances where the docs of two functions disagree a bit on the type constraints.
• The dtype of NDArray for the find_peaks function arguments are too narrow.

[jorenham]

The PR wasn't really complete because I'm not too well versed in all of the typing available, so working on this repo is really helping me learn about the possibilities.

That's completely understandable; it took me quite a while to figure out how to properly type abstract things like an "array-like" and a "dtype-like".
In such cases it'll certainly help a lot to know about the available typing tools.
So maybe it'll help to give a quick summary of the ones that I have used in these cases:

• There's scipy._typing, where I've put several type aliases that are usually related to numpy or scipy in general
• Similarly, there are the package-specific scipy.{package}._typing types. These are only used within that package.
• The numpy._typing module is not part of the public API, but as a numpy maintainer I know that these aren't gonna change anytime soon. Especially the dtype-specific array-like aliases such as _ArrayLikeFloat_co and _ArrayLike[T: ] are very useful in scipy-stubs. These are widely used within numpy's bundled stubs, and have withstood the test of time. But there aren't documented, so you'll need to learn about them by looking at the code itself, e.g. https://github.com/numpy/numpy/blob/v2.2.0.dev0/numpy/_typing/_array_like.py (I'm not saying you should do this or anything, I just wanted to mention it).
• optype is another library of mine that I've used quite a bit here, often as alternative to collections.abc. For instance, the single-method protocols (those that start with optype.Can{}) allow for writing very precise annotations.
• optype.typing is a a bit more high-level, and has very useful aliases, such as optype.typing.AnyInt that describes exactly everything that the builtins.int constructor can accept.
• Finally there's optype.numpy, which I plan to improve in the coming weeks, so that it can be used instead of numpy._typing. Within scipy-stubs I always import this as onpt (matching the conventional npt alias for numpy.typing). One of its most useful types is onpt.Array, which is the shape-typed analogue of npt.NDArray, and the type params are optional, so that onpt.Array is the same as npt.NDArray[np.generic].

[jorenham]

• Not sure on wavelet regarding the output type and also the input types as well

After experimenting a bit with it, it looks like it is allowed to return both ndarray and list, as long as they're 1-D, and the length must be equal to the first argument, so the return type should be a (1-d) "array-like".
It always gets passed 2 arguments as the docs say, but their types turn out to be rather weird

1. The first argument is a builtins.int on the first call, but np.int_ or np.float64 (huh?) after that, depending on whether widths is an array-like of integers or floats.
2. The second argument is (always) a numpy scalar type SCT: np.generic, and matches the scalar-type/dtype of the widths array-like.

So find_peaks_cwt(x, np.arange(1, 10, dtype=np.uint8), f) first calls f as (int, np.uint8), and then as (np.int64, np.uint8) (np.int_ is an alias for np.int64 on my machine).
And find_peaks_cwt(x, np.arange(1, 10, dtype=np.float16), f) will first call f as (int, np.float16), and then as (np.float64, np.float16). The second time the f was called, its *args were (np.float64(10.0), np.float16(1.0)) (wasn't the first argument is supposed to be a discrete size?).

tldr;

I guess the signature is best described as an overloaded one, that looks a bit like

• (int | np.int_, SCT) -> array_like in case of widths: ArrayLike[SCT: np.integer[Any]]
• (int | np.float64, SCT) in case of widths: ArrayLike[SCT: np.floating[Any]]

(the ArrayLike alias is fictitious unfortunately)

[jorenham]

• Not sure on max_distances where the docs of two functions disagree a bit on the type constraints.

You could try it out for a bit in a (i)python console or something, or you could play it safe and use npt.ArrayLike (which also covers Sequence)

---

• The dtype of NDArray for the find_peaks function arguments are too narrow.

As these can be scalar, ndarray or a sequence, it's basically an "array-like".
And in this case, we probably only should allow reals.
So you could use numpy._typing._ArrayLikeFloat_co for this, which also allows you to get rid of the tuple (although for documentation purposes it might help if it stays).

[jorenham]

This could also be a float, according to the docs

[jorenham]

The problem with Sequence is that apart from __getitem__ it also has several other methods like .count() and .index() that aren't present in np.ndarray 🤷🏻

And yea I agree that onpt.Array (or some other direct np.ndarray alias) would be indeed too narrow. So you could use e.g. _ArrayLikeInt_co, which is compatible with both Sequence[int] and npt.NDArray[integer[Any]]. And as far as I'm concerned, the fact that it would also allow integer scalars isn't much of a problem, as we're talking about the input of a function.

[jorenham]

See my other comment on this signature

[jorenham]

totally optional; but you could use a TypeVar(..., bound=np.generic) instead of the current np.generics

[pavyamsiri]

Wouldn't that require the comparator function to have two arguments of the same type? I think that in the ideal case the function signature makes sense but what if someone were to pass a function that takes in a float and an int? They would still be able to do the comparison despite the types being different.

Assuming that using TypeVar enforces that the two argument types to be the same that is.

[jorenham]

Wouldn't that require the comparator function to have two arguments of the same type?

Yes exactly.

---

I think that in the ideal case the function signature makes sense but what if someone were to pass a function that takes in a float and an int?

If someone would pass a function that accepts arrays whose scalar types are incompatible, then it would indeed not be allowed.
And that's a good thing, because it could prevent some potentially type-unsafe situations.

---

Assuming that using TypeVar enforces that the two argument types to be the same that is.

Well, it kinda depends on what you mean with "the same".
Because in this case, we also need to take the relationship of NDArray and its scalar type parameter into account: It's covariant.
So that means that you're allowed to pass something like (NDArray[np.int32], NDArray[np.integer[Any]]) -> NDArray[np.bool_] as function if you also pass data: NDarray[np.int32].

For a demonstration of this, see https://basedpyright.com/?pythonVersion=3.13&typeCheckingMode=all&reportUnusedExpression=true&code=GYJw9gtgBAxmA28CmMAuBLMA7AzgOgEMAjGKdCABzBFSgGUkBHAVySxiQChRIpUBPCuiwBzMpWq0AwgUTFkAGlgEcqTpwDEULGFRI%2BACwK0CMOCAAmwsajB9BKAygDWSEDiUADYak9kcUCo46CJY8vq2UJ7A8GDGnuoCFPoAgiAgBPwA2gAqALpQALz0TKzsSLkFUFpwAG4EIOgEWLQAFADu1M44AJSa9slQaRnZ%2BUUlLGwclVDVZFj1jc1tFmBIOFgA5LSdIM59nBZIwFDAla0WxgQAXEPpmZVKcJQNxtS3MnJEyFlZww-5JT-UZ5PJA%2B7ZIhgBCgnpQAC0AD4oFCENdOLNZiAkKhmCAsIFEK1nhRXrYQBcrllrvCAIxgqCXVAELK0655HoHI4nHwAfWQvJicVQrRudxGWR8DKIt2BWSFxgZAHo4UjxQ9UfA8ujMVBsbj8VAsgR0FAADzFIim4DUQLoJRW%2BZQABe6AooodPTy6kOV3GMBUIrlUqU8HQqlaGVESFaAFZOX1w7ycJIkBZxsBKcylHyBQrUAcgA

[pavyamsiri]

Ah got it. Just misunderstood the semantics. I can add this then.

[jorenham]

I think that besides int, you could also pass some np.integer[Any] and np.bool_ to order. And coincidentally, scipy._typing.AnyInt is an alias for exactly that :)

There's also some other orders like this below

[pavyamsiri]

• Not sure on wavelet regarding the output type and also the input types as well

After experimenting a bit with it, it looks like it is allowed to return both ndarray and list, as long as they're 1-D, and the length must be equal to the first argument, so the return type should be a (1-d) "array-like". It always gets passed 2 arguments as the docs say, but their types turn out to be rather weird

1. The first argument is a builtins.int on the first call, but np.int_ or np.float64 (huh?) after that, depending on whether widths is an array-like of integers or floats.
2. The second argument is (always) a numpy scalar type SCT: np.generic, and matches the scalar-type/dtype of the widths array-like.

So find_peaks_cwt(x, np.arange(1, 10, dtype=np.uint8), f) first calls f as (int, np.uint8), and then as (np.int64, np.uint8) (np.int_ is an alias for np.int64 on my machine). And find_peaks_cwt(x, np.arange(1, 10, dtype=np.float16), f) will first call f as (int, np.float16), and then as (np.float64, np.float16). The second time the f was called, its *args were (np.float64(10.0), np.float16(1.0)) (wasn't the first argument is supposed to be a discrete size?).

tldr;

I guess the signature is best described as an overloaded one, that looks a bit like

• (int | np.int_, SCT) -> array_like in case of widths: ArrayLike[SCT: np.integer[Any]]
• (int | np.float64, SCT) in case of widths: ArrayLike[SCT: np.floating[Any]]

(the ArrayLike alias is fictitious unfortunately)

I finally got around to understanding this signature and I agree that the overloaded ones are good. I'm not sure how to make a type alias like the one you suggested though.

I think a simple way to annotate without having to use too many generics is widths: _ArrayLikeInt_co implies that wavelet: Callable[Concatenate[int | np.int_, AnyInt, ...], npt.ArrayLike] and then likewise for the float case.

This doesn't enforce that the wavelet second argument is of the same type as the scalar type of the array like but it should be broader right?

[pavyamsiri]

Sorry was a bit busy the last couple of days so couldn't finish off the PR.

[pavyamsiri]

Actually on the topic of wavelet I don't think npt.ArrayLike is the correct output type. The second time it is called the code reverses the output using [::-1]

wavelet_data = np.conj(wavelet(N, width, **kwargs)[::-1])

so I think it expects that wavelet returns some non-scalar arraylike. For example np.float16(0.3) is an array like I believe but you can't reverse it like this. I would opt to put it as NDArray[np.generic] as the output type but I'm open to any improvements to make it less broad.

[pavyamsiri]

Actually on the topic of wavelet I don't think npt.ArrayLike is the correct output type. The second time it is called the code reverses the output using [::-1]

wavelet_data = np.conj(wavelet(N, width, **kwargs)[::-1])

so I think it expects that wavelet returns some non-scalar arraylike. For example np.float16(0.3) is an array like I believe but you can't reverse it like this. I would opt to put it as NDArray[np.generic] as the output type but I'm open to any improvements to make it less broad.

I also don't think we need to write overloads for this as well although it maybe more accurate. I think the spirit of the wavelet function is that it should be able to handle float values as well because it represents a continuous mathematical function.

Following this interpretation as well, the first argument which corresponds to the number of points should be something like AnyInt although the source code does not respect that and passes in np.float64 which kinds of ruins things. So maybe the type annotation should be like

_WaveletFunction: TypeAlias = Callable[Concatenate[AnyInt, AnyReal, ...], npt.NDArray[np.generic]]

with the only caveats being that the source actually passes in np.float64 as the first argument sometimes (I guess we can concede and add as part of type union) and that the output is allowed to be any non-scalar array likes like tuple[float, ...] is fine as well.

[jorenham]

I think a simple way to annotate without having to use too many generics is widths: _ArrayLikeInt_co implies that wavelet: Callable[Concatenate[int | np.int_, AnyInt, ...], npt.ArrayLike] and then likewise for the float case.

Yes that seems about right. Just keep in mind that _ArrayLikeFloat_co is more general than _ArrayLikeInt_co (the latter is assignable to the former).

This doesn't enforce that the wavelet second argument is of the same type as the scalar type of the array like but it should be broader right?

Yes; the parameters of a Callable behave as if contravariant

[jorenham]

Actually on the topic of wavelet I don't think npt.ArrayLike is the correct output type. The second time it is called the code reverses the output using [::-1]

Good catch! I then just got lucky when I tried it with a list as return value. So I guess that anything can be returned that implements __getattr__: (Self, slice) -> npt.ArrayLike, which includes, Sequence andnpt.NDArray. With optype you could write that as optype.CanGetitem[slice, npt.ArrayLike] btw.

[pavyamsiri]

I added the following changes:

1. Use _ArrayLikeFloat_co for widths parameter
   This is because the documents specify float or sequence so the
   dtype is implied to be float.
2. Change type annotation for max_distances from _ArrayLikeInt_co to
   _ArrayLikeFloat_co. In case of None the array is given by the
   widths parameter divided by 4.0 so it is has to be compatible with
   float dtypes.
3. Change the output for _WaveletFunction to be more accurate. The
   output has to be sliceable with the type of the sliced access being
   an array like.
4. Change the type annotation for the vector parameter from ArrayLike
   to NDArray[generic]. The docs specify the input as ndarray not an
   array like and as there is no conversion in the source code, we can't use ArrayLike.

[jorenham]

Great! Can I merge this?

[pavyamsiri]

Great! Can I merge this?

sorry just found a mistake. I think max_distances is supposed to be an ndarray according to the docs so not an array like. Let me change that first. Would NDArray[np.floating[Any]] be good or is a better annotation?

[jorenham]

Great! Can I merge this?

sorry just found a mistake. I think max_distances is supposed to be an ndarray according to the docs so not an array like. Let me change that first. Would NDArray[np.floating[Any]] be good or is a better annotation?

The docs aren't always accurate I have found, so be sure to check that first. But in case they are, then I guess that NDArray[np.floating[Any] | np.integer[Any] | np.bool_] would work here, so that all coercible dtypes are allowed.

[pavyamsiri]

Yeah I actually was a bit confused with the docs because it says max_distances is an ndarray but its default value is widths / 4.0. I thought this conflicted because widths can be a scalar but the function firsts casts widths to an ndarray first so I think the docs are right in this case.

[pavyamsiri]

Should be good to merge now after CI!

[jorenham]

Thanks @pavyamsiri !