Typing against Cosmology

[ntessore]

As far as I can tell, the definition

class CosmologyAPI(Protocol[ArrayT_co]):
    ...

means that I need to explicitly pass the array type when specifying the type hint for a function. I find that quite awkward, particularly if it means I have to declare and pass another type var to support arbitrary cosmology classes:

# I don't know yet what cosmology I am given, so I don't know its array type
T = TypeVar('T')

def myfunction(cosmo: Cosmology[T]):
    ...

However, if we let ArrayT be a type var in the individual functions instead, we don't really need the Generic class there:

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

Having the type var at the method level also results in correct behaviour when floats are passed to the function. That's good because, realistically, no one will explicitly convert their floats to 0-dim arrays before calling these functions, so no one will ever strictly adhere to the API as it stands.

I understand that the generic class definition is necessary for hinting the properties of the cosmological constants, but that is such a minor part of the interface that I think they could simply be float, or maybe Array if we want support for arrays of values.

[nstarman]

As far as I can tell, the definition

class CosmologyAPI(Protocol[ArrayT_co]):
    ...

means that I need to explicitly pass the array type when specifying the type hint for a function.

It currently kind-of means that. There's four things to note:

1. Depending how tools like mypy are configured (https://mypy.readthedocs.io/en/stable/config_file.html#confval-disallow_any_generics) skipping the array type will be treated as Any
2. If mypy is on strict=True and one doesn't care about the array type, Any can be given explicitly: CosmologyAPI[Any]
3. https://peps.python.org/pep-0696/ is coming to Python, so in python 3.12/13(?) we can set a default and obviate this whole discussion. Perhaps there will be a __future__ import allowing this in earlier python versions.
4. Type Aliases can be used to pre-specify the type hint.

from typing import TypeAlias  # or from typing_extensions for older python
from cosmology.api import BackgroundCosmology
from dask import Array

DaskCosmo:  TypeAlias = BackgroundCosmology[Array]

...

I find that quite awkward, particularly if it means I have to declare and pass another type var to support arbitrary cosmology classes:

# I don't know yet what cosmology I am given, so I don't know its array type
T = TypeVar('T')

def myfunction(cosmo: Cosmology[T]):
    ...

Just a note that the above discussion of Any and default types is relevant here.

But more generally, I think this function is actually a good example why we DO want type vars. This function is missing a return type annotation; filling this in...

T = TypeVar('T')
def myfunction(cosmo: Cosmology[T]) -> T:
    ...

If Cosmology wasn't generic then this function couldn't be generic over the array type.

However, if we let ArrayT be a type var in the individual functions instead, we don't really need the Generic class there:

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

Having the type var at the method level also results in correct behaviour when floats are passed to the function.

This is true, but we couldn't use this function as intended inside of myfunction since we can't trace the type of the Cosmology through the function.

dask_cosmo = SomeCosmology(H0=dask.Array(70))

def myfunction(cosmo: Cosmology[NO TYPE], z: T) -> T:
    return mycosmo.H * cosmo.scale_factor(z)

myfunction(dask_cosmo, np.array([1, 2, 3])  # THIS IS WRONG! does it return a dask or numpy array?

That's good because, realistically, no one will explicitly convert their floats to 0-dim arrays before calling these functions, so no one will ever strictly adhere to the API as it stands.

I think this is the real crux of the issue. We need to expand the API to permit float inputs! I think Array outputs is fine since that's up to the library implementor, but yes, float inputs should be allowed.

I understand that the generic class definition is necessary for hinting the properties of the cosmological constants, but that is such a minor part of the interface that I think they could simply be float, or maybe Array if we want support for arrays of values.

It's not just the constants like c and G, it's also H0. These are used everywhere through the functions, so their type is important. float isn't bad, we might / probably should expand the interface to allow float. The problem with Array is that it only works for single variables. While any individual variable can satisfy isinstance(variable, Array) since Array is a Protocol, the protocol can't describe interactions between two different arrays. What happens when a dask array is added to a jax array? We can't resolve this, so if a Cosmology can have any Array constant it can have array constants of two different types and that breaks everything.

This didn't fit into the above discussion, but I think it's very important to make the Cosmology API work well in ML and high precision contexts. We should be able to have Cosmology[32bit] and Cosmology[256bit]. Doing this will require that the parameters like H0 support these precisions. I agree we should have significantly improved support for float. In the next few days I'll try drafting a PR enabling this.

[ntessore]

Ok, if I understand it correctly, the interface is saying that a cosmology object of type Cosmology[MyArray] works exclusively with arrays of type MyArray. I thought the interface would take any Array API compatible input and have to work out how to process it at the method level!

Taking CAMB as an example, which computes everything in Fortran via an interface using numpy arrays, it then makes sense that the CAMBCosmology interface is Cosmology[NDArray[np.floating[Any]]]. But that means no CAMB cosmology will ever work with dask arrays, nor any of the handful of other Boltzmann codes which all work on numpy arrays as well.

But where before I thought they would all have to individually make their inputs compatible, actually the current interface makes our life easier, because we would only have to write e.g. a single "dask adapter" cosmology that takes dask chunks and feeds them to a Cosmology[NDArray[np.floating[Any]]] instance as numpy arrays, and all existing Boltzmann codes would work on dask arrays at once. Do I understand that correctly?

For the constants, I agree that having them be 0D arrays is a good idea, if there is scope for supporting arrays of parameters. Running a whole MCMC chain via dask parameter arrays is definitely something that should be written down under the "vision" of the interface!

Lastly, for the scalar vs 0d array input, would it not suffice to define all functions like this:

from numbers import Number
from array_api import Array
...

NumericT = TypeVar('NumericT', bound=Number|Array)

class Cosmology(...):
    ...

    def scale_factor(self: Self, z: NumericT) -> NumericT:
        ...

[nstarman]

But where before I thought they would all have to individually make their inputs compatible, actually the current interface makes our life easier, because we would only have to write e.g. a single "dask adapter" cosmology that takes dask chunks and feeds them to a Cosmology[NDArray[np.floating[Any]]] instance as numpy arrays, and all existing Boltzmann codes would work on dask arrays at once. Do I understand that correctly?

Yes as you said, our CAMBCosmology wrapper is Cosmology[NDArray[np.floating[Any]]] because CAMB only works with numpy arrays. So yes, a dask version could take our CAMBCosmology and define a dask adapter function / wrapper, and voila.

Something along these lines is all it would take:

from dataclasses import make_dataclass

DaskCAMBCosmology = make_dataclass(
    "DaskCAMBCosmology",
    [],
    bases=(CAMBCosmology,),
    namespace={k: dask_wrapper_func(v) for k, v in ...}   # introspect the methods in CAMBCosmology
)

In future if CAMB natively allows for dask arrays then we can generalize the wrapper and the user can then specify whether they are working with numpy or dask.

[nstarman]

Lastly, for the scalar vs 0d array input, would it not suffice to define all functions like this:
...

Yes! I have a branch for a PR following cosmology-api/cosmology.api#25 that does this. It starts with just float. Number has a lot of undesirable weirdness that I haven't sorted out yet. We can expand from just float — now or in a future release.

[ntessore]

Lastly, for the scalar vs 0d array input, would it not suffice to define all functions like this:

...

Yes! I have a branch for a PR following cosmology-api/cosmology.api#25 that does this. It starts with just float. Number has a lot of undesirable weirdness that I haven't sorted out yet. We can expand from just float — now or in a future release.

Oh yeah I remember now, it should have been Real to get arithmetic operations, but IIRC even then many numpy functions were not typed correctly. Using float is probably fine even if, strictly, it doesn't allow numpy.floating inputs.

[nstarman]

Yeah. It would be nice to figure out how to type scalars. For most of the array libraries, e.g. pytorch and jax, float is special-cased. I don't thing numpy.floating will work generally. But for libraries that work with numpy arrays, I think they should work with numpy.floating if they work for float.