Simplified API

[ntessore]

Over the past couple of weeks, I have been trying to work more with the Cosmology API, for a number of projects that we want to integrate into GLASS.

I have come to think the Cosmology API would be massively simplified if we

1. remove all higher-level protocols;
2. rename all low-level protocols to match the function they contain.

This is a radical change, so let me explain by way of example.

# cosmology definitions for this module
class Cosmology:
    import cosmology.api as cosmo

    class distance_weight(
        cosmo.H_over_H0[Array],
        Protocol,
    ): ...

    class volume_weight(
        cosmo.transverse_comoving_distance[Array],
        cosmo.H_over_H0[Array],
        Protocol,
    ): ...


def distance_weight(
    z: Array,
    cosmo: Cosmology.distance_weight[Array],
) -> Array:
    return 1 / cosmo.H_over_H0(z)


def volume_weight(
    z: Array,
    cosmo: Cosmology.volume_weight[Array],
) -> Array:
    return cosmo.transverse_comoving_distance(z) ** 2 / cosmo.H_over_H0(z)

This is the only always-correct, fully-consistent, and flexible solution I have found. It supports all standard and non-standard cosmologies, as well as perturbations, and whatever else we might come up with in the future.

And if type intersections ever make it into Python, then it won't even be necessary to define the helper protocols.

---

Edit: There actually is a mypy extension that supports protocol intersections: https://github.com/klausweiss/typing-protocol-intersection

Assuming that in cosmology.api we might do the suggested

from typing_protocol_intersection import ProtocolIntersection as Has

then the above example could be typed as

import cosmology.api as cosmo


def distance_weight(
    z: Array,
    cosmo: cosmo.Has[
        cosmo.H_over_H0[Array],
    ],
) -> Array:
    return 1 / cosmo.H_over_H0(z)


def volume_weight(
    z: Array,
    cosmo: cosmo.Has[
        cosmo.transverse_comoving_distance[Array],
        cosmo.H_over_H0[Array],
    ],
) -> Array:
    return cosmo.transverse_comoving_distance(z) ** 2 / cosmo.H_over_H0(z)

which to me seems like a fantastic way to do this.

[nstarman]

Have you read https://github.com/jorenham/optype?
It's a great library that works through basically all the strict type definitions in Python. With this you can precisely describe basically any Python object. IMO this is what we want for the Cosmology API, but for cosmology methods, not base Python ones.
Importantly, this library is very strict about following Python's typing conventions and approaches, so I think it's also a good guide about how we should think about structure. (On a separate note I think we should make this library a dependency and use it to improve the type annotations)

I want to focus in on the guide: https://github.com/jorenham/optype?tab=readme-ov-file#optype-1

There are five flavors of things that live within optype,

• The optype.Just[T] and its optype.Just{Int,Float,Complex} subtypes only accept instances of the type itself, while rejecting instances of strict subtypes. This can be be used to e.g. work around the float and complex [type promotions][BAD], annotating object() sentinels with Just[object], rejecting bool in functions that accept int, etc.
• optype.Can{} types describe what can be done with it. For instance, any CanAbs[T] type can be used as argument to the abs() builtin function with return type T. Most Can{} implement a single special method, whose name directly matched that of the type. CanAbs implements __abs__, CanAdd implements __add__, etc.
• optype.Has{} is the analogue of Can{}, but for special attributes. HasName has a __name__ attribute, HasDict has a __dict__, etc.
  ...

Pretty much everything in the Cosmology API would fall under the CanX or HasX categories. We've named everything HasX which we might consider changing 🤷. But back to your point, this is the Pythonic approach. Here's why I'm resistant to aspects of the suggested refactor.

1. removing all higher-level protocols.

These are just intersection types that significantly simplify typing a high-level function.
I like the look of "a mypy extension that supports protocol intersections: https://github.com/klausweiss/typing-protocol-intersection" but this is a no-go for 2 reasons:

• we can't force everyone to use a mypy extension. Realistically it just won't happen. Especially for inexperienced users
  -. mypy extensions only work with mypy not Pyright. BasedMypy. BasedPyright. etc. Many (all?) of these type checkers are better than mypy and shouldn't be rendered incompatible.

def fancy_func(cosmo: HasScaleFactor0 & HasOmegMatter0 & ..., z: Array) -> Array: ...

vs

def fancy_func(cosmo: StandardCosmology, z: Array) -> Array: ...

If a user wants to granularly type a function then we provide both intermediate intersection types and the nitty-gritty low-level HasX types. If someone wants to use the mypy extension in their own library to make it easy to create custom intersection types from the low-level CosmologyAPI types then by all means!

2. rename all low-level protocols to match the function they contain.

Aren't they already?

@runtime_checkable
class HasOmegaTot0(Protocol[Array]):
    r"""The object has a total density attribute -- :math:`Omega_{\rm tot}`."""

    @property
    def Omega_tot0(self) -> Array:
        r"""Omega total; the total density/critical density at z=0."""

This is CamelCase. Any Python linter will be unhappy with if the class is snake_cased, like omega_tot0.
The consistent regex rule is "Has" + re.sub(r'(^|_)([a-z0-9])', lambda m: m.group(2).upper(), name) (we should check all classes satisfy this)

3. Classes w/in classes

class Cosmology:
    import cosmology.api as cosmo

    class distance_weight(
        cosmo.H_over_H0[Array],
        Protocol,
    ): ...

🤮 . Oh no!

Summary.

I think we were already careful to write this Pythonically and ended up with something very close to how optype does things, where the author thought longer and more carefully about typing and typed all of Python.

If in GLASS you don't want to use the high-level protocols, then don't? 🧐
If in GLASS you wan't to make fast intersection types with https://github.com/klausweiss/typing-protocol-intersection, then ! I'm struggling to understand why that necessitates changes here.
Plz no snake-casing classes.

[ntessore]

Thanks for the great points!

I guess I should have been clearer about the main point: the medium & high-level protocols are not adding any value. I can confidently state that most cosmology-using functions only use a handful of methods, if not fewer. The value of the cosmology API is to define the syntax and semantics of the individual methods. Which of these methods a cosmo object is required to have varies so much, it should be specified on a function-by-function case. The high-level protocols do not help; if anything, they are a net negative because library writers are tempted to just go with StandardCosmology instead of figuring out the real requirements of their code.

The same is even more true for an API for perturbations: here it becomes impossible to even define a sensible set of high-level groupings, since provider codes are so different. We should acknowledge that fact, and only define the syntax and semantics of individual methods, not functional groups.

As for the classes-within-classes, this was only to come up with a convention for where to stick the many individual subclasses that protocol intersections currently require. A ProtocolIntersection removes the need for explicitly defining such subclasses. For type-checking, this is only available as a mypy plugin right now (but only required for type-checking, at runtime it all works perfectly fine).

Finally, renaming the low-level protocols to match the functions they contain is a massive productivity booster. I know the astropy.cosmology function names, but I have to figure out the CamelCase variant each and every time. Besides, our method names aren't exactly following any sort of convention, either...

[nstarman]

I guess I should have been clearer about the main point: the medium & high-level protocols are not adding any value... if anything, they are a net negative because library writers are tempted to just go with StandardCosmology instead of figuring out the real requirements of their code.

I can see that being the case. It's certainly a balance we discussed when writing the Protocols. How about we move them to a less exposed namespace, e.g. a submodule called groupings? so it would be cosmology.api.groupings.DistanceMeasures. We can also add an admonition to the docstring.

The same is even more true for an API for perturbations: here it becomes impossible to even define a sensible set of high-level groupings, since provider codes are so different. We should acknowledge that fact, and only define the syntax and semantics of individual methods, not functional groups.

Agreed.

As for the classes-within-classes, this was only to come up with a convention for where to stick the many individual subclasses that protocol intersections currently require.

The intermediate and high-level classes are type intersections! Use inheritance, not nested classes 😆 See https://peps.python.org/pep-0544/#unions-and-intersections-of-protocols.

Finally, renaming the low-level protocols to match the functions they contain is a massive productivity booster. I know the astropy.cosmology function names, but I have to figure out the CamelCase variant each and every time.

What IDE are you using? VSCode suggests for autocompletion all the relevant methods regardless of capitalization. So using the https://pep8.org correct naming convention shouldn't be a slow-down. I also just tried Sublime Text and Nova which also work, but they don't have / I haven't setup ML prediction, which takes cosmology.api as context and does even better prediction.

Another problem with the suggested rename is that it implies it's the type of of the function, not the type of an object that has that method. See https://peps.python.org/pep-0544/#callback-protocols.

class HasFunc(Protocol):  # protocol for an object with the method `func`
    def func(self, arg1: Type1, arg2: Type2) -> Type3: ...

class Func(Protocol):  # protocol for func itself
    def __call__(self, arg1: Type1, arg2: Type2) -> Type3: ...

Besides, our method names aren't exactly following any sort of convention

PEP8 for casing.
And then it's a re-think on the Astropy, CAMB, CLASS names for the functions, trying to standardize e.g. where the underscores and z=0 suffix endings are.

[ntessore]

I use vi 😄 The confusion between the protocol and the method was pretty much intentional here, so that people who aren't familiar with protocols don't even see there are two things in play here. But ok, I understand how this is slightly uncomfortable.

Moving the existing intersections to be the second thing that people see sounds like a good compromise, if you don't want to deprecate them outright.

All in all, what do you then think of an interface like this:

# No implementation, just a description of the interface!
import cosmology.api


def flat_angular_diameter_distance(
    z: Array,
    cosmo: cosmology.api.Cosmology[  # name for the intersection, until there's native support
        cosmology.api.HasOmegaK0[Array],
        cosmology.api.HasComovingDistance[Array],
    ],
) -> Array:
    # Do some cosmology with any object that implements the API
    if cosmo.Omega_k0 != 0:
        raise ValueError("This function only works for flat cosmologies")
    return cosmo.comoving_distance(z) / (1 + z)

[nstarman]

I think that looks good! Can you explain how cosmology.api.Cosmology[ would work? Is this https://github.com/klausweiss/typing-protocol-intersection ?