Allow floats

[nstarman]

Signed-off-by: nstarman [email protected]

Description

Allow float inputs to redshift-dependent methods.

PR Checklist

• Give a detailed description of the PR above.
• Document changes in the CHANGES.rst file. See existing changelog for examples.
• Add tests, if applicable, to ensure code coverage never decreases.
• Make sure the docs are up to date, if applicable, particularly the docstrings and RST files in docs folder.
• Ensure linear history by rebasing, when requested by the maintainer.

[codecov-commenter]

Codecov Report

Merging #27 (527cc67) into main (0ac84d0) will not change coverage.
The diff coverage is 100.00%.

📣 This organization is not using Codecov’s GitHub App Integration. We recommend you install it so Codecov can continue to function properly for your repositories. Learn more

@@           Coverage Diff           @@
##             main      #27   +/-   ##
=======================================
  Coverage   72.37%   72.37%           
=======================================
  Files          13       13           
  Lines         286      286           
=======================================
  Hits          207      207           
  Misses         79       79           

Impacted Files	Coverage Δ
src/cosmology/api/background.py	72.50% <100.00%> (ø)
src/cosmology/api/components.py	69.81% <100.00%> (ø)
src/cosmology/api/standard.py	73.68% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[nstarman]

I think this covers all the inputs.

[ntessore]

I wonder if it's worth adding an abstraction layer to our input type definition here. We could add a new ArrayLike generic alias that is only float | ArrayT for the time being, but can be expanded later without having to change all function signatures.

ArrayLike = Union[float | ArrayT]

...

    def scale_factor(self, z: ArrayLike[ArrayT], /) -> ArrayT:
        ...

[ntessore]

Thinking about this more, we could even target PEP 696 and define:

ArrayT = TypeVar('ArrayT', bound=Array)
ArrayLikeT = TypeVar('ArrayLikeT', default=float | ArrayT)

class Cosmology(Generic[ArrayT, ArrayLikeT], Protocol): ...

Cosmologies using numpy arrays could then accept numpy.typing.ArrayLike inputs, and e.g. cosmo.comoving_distance([z1, z2]) would remain legal, which is a big plus in my mind.

[nstarman]

I agree it may be worth expanding the Generic type to be float | Array. I want to add mypy tests to the test suite to confirm this.
However this will not enable lists as inputs because they are neither float nor Array. I agree lists are convenient, but they are very different structures from a numerical library's Array. Permitting lists means Cosmology libraries would have to process all input! That's an unfortunate slowdown. If a user has a list, I don't think it's too onerous to call xp.asarray. And if a user wants speed, then working with arrays from the start will guarantee the best possible performance.

[ntessore]

However this will not enable lists as inputs because they are neither float nor Array.

Yes, that's exactly the intention. The default behaviour is to have ArrayLike be floats and arrays, which are always supported. But if a cosmology accepts e.g. all numpy array_like input, which includes sequences of floats, then it can set the ArrayLikeT to something more permissive than float | Array, like the full np.typing.ArrayLike.

[nstarman]

I see. Sorry for the confusion. Yes, this would be a good addition when that PEP is rolled out!
I agree we can make the Cosmology be

class Cosmology(Protocol[InputType=ArrayLike, Array]):

So that inputs to functions can be changed as you said, but output types remain Array. This can be supported by variadic generics so that only 1 or both types can be given.

[ntessore]

Fair enough. I am just a little wary of restricting functionality that currently is universally supported.

[nstarman]

We might be able to support it now by going for a variadic generic implementation. This requires a Mypy plugin, which might be quite simple. I'll take a look.

[ntessore]

We might be able to support it now by going for a variadic generic implementation. This requires a Mypy plugin, which might be quite simple. I'll take a look.

This already works right now if you use TypeVar from typing_extensions to gain the default= parameter. We don't need variadic generics, the ArrayLikeT will use its default if not given:

from typing_extensions import TypeVar

ArrayT = TypeVar('ArrayT', bound=Array)
ArrayLikeT = TypeVar('ArrayLikeT', default=float | ArrayT)

class Cosmology(Generic[ArrayT, ArrayLikeT], Protocol): ...

reveal_type(Cosmology[Array])  # Cosmology[Array, float | Array]

[nstarman]

I did a quick experiment using the typing_extensions TypeVar but I can't get mypy to recognize it as a TypeVar. I don't think it's supported yet. So we'll have to wait a bit for a default this way. I think we can still accomplish this with variadic generics and a mypy plugin. I'll take a closer look.

[ntessore]

Ok cool. It works well with pyright so I assumed mypy would be the same. We can also leave this for now if it's blocking.

[nstarman]

Pyright has a faster release exhedule. I do think this is a blocker for just floats, which we do want to support. I'll make a reminder issue to return to this, because it's definitely a valuable addition.

[nstarman]

I added a short table to the docs on which array libraries allow the addition of an array and a float.

[nstarman]

The review got dismissed because I pushed to the branch. The branch protection rules are set up so that an approval only applies to the most recent version of the PR.

[nstarman]

Needed to add a little bit of docs to get the toctree working.