Quantity basics

[mhvk]

Support for numpy, array-api-strict, dask, and jax, with basic tests for all , and ufuncs for numpy only. Pushing now in part to check whether CI works properly.

EDIT: I pushed a version that does not yet include numpy ufuncs, but for which the basics are now tested more thoroughly.

[nstarman]

Looking good! some quick comments.

[nstarman]

I think eventually, but not urgently.

[mhvk]

@nstarman, @adrn, @jeffjennings - here's the start of Quantity 2.0. This PR just sets the stage, showing that operations work as expected for ndarray, dask, and jax (and array_api_strict as a sanity check).

I have a separate branch that implements ufuncs, but want to clean up the tests for that before pushing it.

[nstarman]

Should this be "stuff that implements the Array API".
Unfortunately the Array API doesn't vendor a Protocol themself, but based on the docs the minimal protocol would be

class SupportsArrayAPI(Protocol):
    def __array_namespace__(self, *, api_version: str | None = None) -> Any: ...

Suggested change

	value: Any
	value: SupportsArrayAPI

[mhvk]

The problem is that apart from numpy, no-one actually defines __array_namespace__ (well, at least dask and jax do not). Also, I do want to add support for python numbers. Overall, it seems easier to just punt on this for now.

[nstarman]

Jax should define __array_namespace__ (link.
Hopefully dask will soon as well.
IMO supporting python numbers should be handled by a wrapper class that provides the Array API.

import math

class PythonNumber:
    def __init__(self, value): ...
    def __array_namespace__(self, ...): return math

If we stick to the Array API it will massively reduce compatibility logic in every function..

[mhvk]

Ah, you're right, JAX actually does define it already.

For python numbers, I was thinking along similar lines. That may be where your __value_namespace__ can come in handy. Although perhaps wrapping array_api_compat.array_namespace may be better.

Anyway, how about we do a typing PR after this one? I do think it is a good idea to start with this as well-typed as possible, but fear I'll need more than a little help!

[nstarman]

Although perhaps wrapping array_api_compat.array_namespace may be better.

Agreed. If we can avoid __value_namespace__ entirely, that would be nice.

[nstarman]

One of the major annoyances of dataclasses is that they chose not to implement the converter argument to the fields. https://peps.python.org/pep-0712/. I think this was an atrocious decision that makes everyone's lives more difficult. For us, this means it's much harder to help users when value is not an Array API-compatible object, like converting a float to a PythonNumber[float]-wrapped value.

Thankfully there is a well-defined way around this - the dataclass_transform decorator. I'd be happy to contribute this in a followup PR. Essentially we produce our own dataclass decorator which calls python's built-in dataclass but provides its own __init__.

def __init__(self, *args, **kwargs):
    fs = fields(self)
    # Iterate through the fields applying the Field.metadata["converter"] then setting the attribute.
    object.__setattr__(self, f.name, value)

[mhvk]

Sounds like a plan!

[mhvk]

@nstarman - sorry for the many pushes; clearly too late to get things straight in one go. But now really ready for review (already a lot better given your earlier comments!).

[nstarman]

Nitpicks, and suggestions for followup.
It's great to see how much the Array API and dataclasses can simplify our code structure.

[nstarman]

Agreed. Something like

@runtime_checkable
class QuantityAPI(Protocol):
    value: Array
    unit: Unit
    
    def to(self, ...): ...

Then this function becomes

def get_value_and_unit(arg, default_unit=None):
    return (arg.value, arg.unit) if isinstance(arg, QuantityAPI) else (arg, default_unit)

[mhvk]

I would like to actually raise an error on having a .unit attribute and not .value (e.g., how do we deal with Column?)

[mhvk]

p.s. There's part of me that still would like to avoid Quantity having anything to do with unit conversion - that can be up to the unit. E.g., q.to_value(unit) could become q // unit or q >> unit or so (somewhat partial to // myself, since it is close - q1 // q2 already returns a dimensionless quantity, so all one really needs is to avoid acting on the "make integer" part).

[nstarman]

As in not have .to() and .to_value()?

[nstarman]

Like in https://unxt.readthedocs.io/en/latest/#converting-units

(based on the Julia paradigm)

[mhvk]

OK to postpone this until we have a clearer idea what exactly the Protocol will be?

[nstarman]

Yes. Should we open an Issue?

[mhvk]

See #3

[nstarman]

As followup, we should get mypy set up to do static testing to make sure it's happy with this magic.

[mhvk]

Indeed, maybe typing the lot should be next so that we can start with a smaller base to test - then the (u)functions after that!

[mhvk]

@nstarman - I updated the minimum versions. Shall we go ahead and get this in, so we can move to the next step (typing?).

[nstarman]

Yes! Are you ok to do squash merge commits, or do you want to clean up the commit history?

[mhvk]

I think this is fairly clean - or do you want me to go back and add version numbers when I introduce the dependencies?

[nstarman]

LGTM.

[nstarman]

🎉