Repository: yt-project/unyt (language: python)

Issue:
NEP 18 progress tracker
We want to implement NEP 18 progressively on main. #200 will be the first step towards that goal. This issue will be used to keep track of what's left to do in this project.

The full list of functions that need re-implementation is found here https://github.com/numpy/numpy/issues/15544#issuecomment-584453408

Here it is, converted to a TODO list (the `numpy.` namespace is implicit everywhere)

<details> <summary> TODO </summary>

- [ ] allclose
- [ ] amax
- [ ] amin
- [ ] angle
- [ ] any
- [ ] append
- [ ] apply_along_axis
- [ ] apply_over_axes
- [ ] argmax
- [ ] argmin
- [ ] argpartition
- [ ] argsort
- [ ] argwhere
- [ ] around
- [ ] array_equal
- [ ] array_equiv
- [ ] array_repr
- [ ] array_split
- [ ] array_str
- [ ] asfarray
- [ ] asscalar
- [ ] atleast_1d
- [ ] atleast_2d
- [ ] atleast_3d
- [ ] average
- [ ] bincount
- [ ] block
- [ ] broadcast_arrays
- [ ] broadcast_to
- [ ] busday_count
- [ ] busday_offset
- [ ] can_cast
- [ ] choose
- [ ] clip
- [ ] column_stack
- [ ] common_type
- [ ] compress
- [ ] convolve
- [ ] copy
- [ ] copyto
- [ ] corrcoef
- [ ] correlate
- [ ] count_nonzero
- [ ] cov
- [ ] cumprod
- [ ] cumproduct
- [ ] cumsum
- [ ] datetime_as_string
- [ ] delete
- [ ] diag
- [ ] diag_indices_from
- [ ] diagflat
- [ ] diagonal
- [ ] diff
- [ ] digitize
- [ ] dsplit
- [ ] dstack
- [ ] ediff1d
- [ ] einsum
- [ ] einsum_path
- [ ] empty_like
- [ ] expand_dims
- [ ] extract
- [ ] fft.fft
- [ ] fft.fft2
- [ ] fft.fftn
- [ ] fft.fftshift
- [ ] fft.hfft
- [ ] fft.ifft
- [ ] fft.ifft2
- [ ] fft.ifftn
- [ ] fft.ifftshift
- [ ] fft.ihfft
- [ ] fft.irfft
- [ ] fft.irfft2
- [ ] fft.irfftn
- [ ] fft.rfft
- [ ] fft.rfft2
- [ ] fft.rfftn
- [ ] fill_diagonal
- [ ] fix
- [ ] flatnonzero
- [ ] flip
- [ ] fliplr
- [ ] flipud
- [ ] full_like
- [ ] fv
- [ ] geomspace
- [ ] gradient
- [ ] histogram_bin_edges
- [ ] hsplit
- [ ] i0
- [ ] imag
- [ ] in1d
- [ ] insert
- [ ] interp
- [ ] ipmt
- [ ] irr
- [ ] is_busday
- [ ] isclose
- [ ] iscomplex
- [ ] iscomplexobj
- [ ] isin
- [ ] isneginf
- [ ] isposinf
- [ ] isreal
- [ ] isrealobj
- [ ] ix_
- [ ] lexsort
- [ ] lib.scimath.arccos
- [ ] lib.scimath.arcsin
- [ ] lib.scimath.arctanh
- [ ] lib.scimath.log
- [ ] lib.scimath.log10
- [ ] lib.scimath.log2
- [ ] lib.scimath.logn
- [ ] lib.scimath.power
- [ ] lib.scimath.sqrt
- [ ] linalg.cholesky
- [ ] linalg.cond
- [ ] linalg.det
- [ ] linalg.eig
- [ ] linalg.eigh
- [ ] linalg.eigvals
- [ ] linalg.eigvalsh
- [ ] linalg.lstsq
- [ ] linalg.matrix_power
- [ ] linalg.matrix_rank
- [ ] linalg.multi_dot
- [ ] linalg.qr
- [ ] linalg.slogdet
- [ ] linalg.solve
- [ ] linalg.svd
- [ ] linalg.tensorsolve
- [ ] linspace
- [ ] logspace
- [ ] may_share_memory
- [ ] mean
- [ ] median
- [ ] meshgrid
- [ ] min_scalar_type
- [ ] mirr
- [ ] moveaxis
- [ ] msort
- [ ] nan_to_num
- [ ] nanargmax
- [ ] nanargmin
- [ ] nancumprod
- [ ] nancumsum
- [ ] nanmax
- [ ] nanmean
- [ ] nanmedian
- [ ] nanmin
- [ ] nanpercentile
- [ ] nanprod
- [ ] nanquantile
- [ ] nanstd
- [ ] nansum
- [ ] nanvar
- [ ] ndim
- [ ] nonzero
- [ ] nper
- [ ] npv
- [ ] ones_like
- [ ] packbits
- [ ] pad
- [ ] partition
- [ ] percentile
- [ ] piecewise
- [ ] place
- [ ] pmt
- [ ] poly
- [ ] polyadd
- [ ] polyder
- [ ] polydiv
- [ ] polyfit
- [ ] polyint
- [ ] polymul
- [ ] polysub
- [ ] polyval
- [ ] ppmt
- [ ] prod
- [ ] product
- [ ] ptp
- [ ] put
- [ ] put_along_axis
- [ ] putmask
- [ ] pv
- [ ] quantile
- [ ] rate
- [ ] ravel
- [ ] ravel_multi_index
- [ ] real
- [ ] real_if_close
- [ ] repeat
- [ ] reshape
- [ ] resize
- [ ] result_type
- [ ] roll
- [ ] rollaxis
- [ ] roots
- [ ] rot90
- [ ] round_
- [ ] save
- [ ] savetxt
- [ ] savez
- [ ] savez_compressed
- [ ] searchsorted
- [ ] select
- [ ] setdiff1d
- [ ] setxor1d
- [ ] shape
- [ ] shares_memory
- [ ] sinc
- [ ] size
- [ ] sometrue
- [ ] sort
- [ ] sort_complex
- [ ] split
- [ ] squeeze
- [ ] std
- [ ] sum
- [ ] swapaxes
- [ ] take
- [ ] take_along_axis
- [ ] tensordot
- [ ] tile
- [ ] trace
- [ ] transpose
- [ ] trapz
- [ ] tril
- [ ] tril_indices_from
- [ ] trim_zeros
- [ ] triu
- [ ] triu_indices_from
- [ ] unique
- [ ] unpackbits
- [ ] unravel_index
- [ ] unwrap
- [ ] vander
- [ ] var
- [ ] vsplit
- [ ] where
- [ ] zeros_like

</details>

### Invalid (nothing to be done)

#### deprecated API
- alen
- alltrue

#### API that applies to non-numeric types

- all
- `numpy.char` namespace
<details> <summary> full list (numpy.char) </summary>

-  char.add
-  char.capitalize
-  char.center
-  char.count
-  char.decode
-  char.encode
-  char.endswith
-  char.equal
-  char.expandtabs
-  char.find
-  char.greater
-  char.greater_equal
-  char.index
-  char.isalnum
-  char.isalpha
-  char.isdecimal
-  char.isdigit
-  char.islower
-  char.isnumeric
-  char.isspace
-  char.istitle
-  char.isupper
-  char.join
-  char.less
-  char.less_equal
-  char.ljust
-  char.lower
-  char.lstrip
-  char.mod
-  char.multiply
-  char.not_equal
-  char.partition
-  char.replace
-  char.rfind
-  char.rindex
-  char.rjust
-  char.rpartition
-  char.rsplit
-  char.rstrip
-  char.split
-  char.splitlines
-  char.startswith
-  char.str_len
-  char.strip
-  char.swapcase
-  char.title
-  char.translate
-  char.upper
-  char.zfill

</details>


### In progress

### DONE

<details> <summary> #200 </summary>

- [x] array2string
- [x] dot
- [x] histogram
- [x] histogram2d
- [x] histogramdd
- [x] inner
- [x] kron
- [x] linalg.inv
- [x] linalg.pinv
- [x] linalg.tensorinv
- [x] outer
- [x] vdot

</details>

<details> <summary> #293 </summary>

- [x] concatenate
- [x] cross
- [x] hstack
- [x] intersect1d
- [x] linalg.norm
- [x] stack
- [x] union1d
- [x] vstack

</details>


Discussion / hints:
@ngoldbaum I generated the TODO list.
I propose we dispatch the load on a voluntary basis and announce what we're going to tackle here. For instance, in my next PR, I want to deal with every function on the list that has a `u*` unit preserving equivalent.
I note that there's an inconsistency in how I handled the optional `out` argument in my first batch #200 and my second one #293.
My next PR will be about harmonising these before it gets out of control
edit: this is #295
My next PR will aim at implementing all functions 28 with names starting with "a".
For the tests you're making as xfail, are you reporting bugs upstream to numpy? I'd prefer not to accumulate lots of tests for numpy bugs in unyt that never get fixed upstream.
I don't recall reporting those upstream no. I think I would have linked the issues in my tests otherwise. I agree that we should report, I think I just didn't anticipate how much larger the full implementation would be when I wrote #200
OK, I'd appreciate it if you'd do that so we don't accumulate a bunch of xfail tests forever.
Will do !
I'd like to change the rules you came up with for how to handle `out`. Right now if you pass in an `ndarray`, `__array_function__` is bypassed. This behavior for wrapped functions is subtly different from how unyt handles ufuncs, and I think we need to make sure that they behave identically, following what we already do for ufuncs.

I'm still working on this, but I think we at least need to make this change:

```
diff --git a/unyt/array.py b/unyt/array.py
index 47c6acb..6795397 100644
--- a/unyt/array.py
+++ b/unyt/array.py
@@ -1958,7 +1958,7 @@ class unyt_array(np.ndarray):
             return func._implementation(*args, **kwargs)
         # Note: this allows subclasses that don't override
         # __array_function__ to handle unyt_array objects
-        if not all(issubclass(t, unyt_array) for t in types):
+        if not all(issubclass(t, unyt_array) or t is np.ndarray for t in types):
             return NotImplemented
         return _HANDLED_FUNCTIONS[func](*args, **kwargs)
```

Which allows the  `__array_function__` handlers in unyt to mix `ndarray` and `unyt_array`. If someone passes in an `ndarray` for any argument that's not `out`, we treat it as dimensionless. If an ndarray is passed in as `out`, we need to fill the `out` ndarray with the data from the result, and return the result as a `unyt_array`. This means that the result and `out` will *not* be the same array. I personally think that's fine. It also follows how we handle this for ufuncs.

Why not raise an error? Practicality mostly. I don't think it's very useful to raise an error when we pass in an ndarray and there are valid reasons why you'd want to supply an ndarray (e.g. passing it to other code that doesn't use unyt, or maybe the array is a graphics buffer). Also I think it's worth making sure ufuncs behave identically.
To make that a little more concerete, here's how this distinction between ufuncs and wrapped functions works currently in practice:

```
import numpy as np
import unyt as u

a = [[1, 2, 3], [1, 2, 3], [1, 2, 3]]*u.g
b = [[1, 2, 3], [1, 2, 3], [1, 2, 3]]*u.kg
c = np.zeros((3, 3), dtype='int64')
# matmul is a ufunc and works according to the rules of unyt's __array_ufunc__
res = np.matmul(a, b, out=c)
print(res)
print(c)
print(type(res), type(c))

c = np.zeros((3, 3), dtype='int64')
# dot is a wrapper and works according to the rules of unyt's __array_function__
res = np.dot(a, b, out=c)
print(res)
print(c)
print(type(res), type(c))
```

Currently this prints:

```
[[ 6. 12. 18.]
 [ 6. 12. 18.]
 [ 6. 12. 18.]] g*kg
[[ 6. 12. 18.]
 [ 6. 12. 18.]
 [ 6. 12. 18.]]
<class 'unyt.array.unyt_array'> <class 'numpy.ndarray'>
[[ 6 12 18]
 [ 6 12 18]
 [ 6 12 18]]
[[ 6 12 18]
 [ 6 12 18]
 [ 6 12 18]]
<class 'numpy.ndarray'> <class 'numpy.ndarray'>

```

> This means that the result and out will not be the same array. I personally think that's fine. It also follows how we handle this for ufuncs.

FTR This was the reason why I originally didn't allow the `ndarray` type for `out`, but it makes a lot of sense to harmonise rules with ufuncs, so I agree.

> Why not raise an error? Practicality mostly. I don't think it's very useful to raise an error when we pass in an ndarray and there are valid reasons why you'd want to supply an ndarray (e.g. passing it to other code that doesn't use unyt, or maybe the array is a graphics buffer).

strong +1 on this
@ngoldbaum Correction, I'm now working on a list of "noop" functions that for various reasons do not need special treatment, so we can treat the _rest_ of `NOT_HANDLED_FUNCTIONS` as a TODO list.
I'll resume actual implementations when I'm done with this.
This is #299 
FTR, I don't think I'll continue maintaining this issue as a TODO list now that we have a test embedded version of it.
@ngoldbaum FYI I think I'll soon reach the point where only the functions for which implementation is *not* obvious to me remain, for instance:
- `np.savetxt` -> should we do anything like add the unit information in a header comment ? `np.loadtxt` isn't part of NEP 18 so I guess we would need to add a `unyt.loadtxt` function too ?
- `np.array_equal` -> how should an array with units compare to a pure ndarray ?

Generally, functions that modify an array (in place of nor) and for which it _could_ make sense to support melding data from pure ndarrays (like `np.where`) require a discussion, IMO.

I don't have strong opinions of what *should* be done, I just think we need a careful decision that can be applied uniformly.
We already have `unyt.savetxt` and `unyt.loadtxt`. If `loadtxt` isn't wrappable then I think the best thing to do is leave the numpy implementations alone? Kinda seems like an oversight on the numpy side that `loadtxt` isn't overridable.

I think `array_equal` should check units. A pure `ndarray` should be treated the same as a dimensionless `unyt_array`.

Not quite sure what you mean about `np.where`, maybe an example illustrating where your confusion is coming from would help me?
> Not quite sure what you mean about np.where, maybe an example illustrating where your confusion is coming from would help me?

nevermind, I don't think that I took a good example (`np.where` is the one function I never get right in numpy...)
I think a better example of what I wanted to illustrate is `np.isin` what should be the behaviour for, e.g
```python
>>> np.isin([1, 2, 3]*cm, [1, 2, 3])
```
?
I'm not totally sure for `isin`. I guess probably ignoring the units if one of the operands is an ndarray, but if both operands have different explicit units that's probably an error?

In general, I think if one supplies an ndarray and it wouldn't be useful to raise an error because the units don't match, probably best to not raise an error. An example of where it is useful is `numpy.array_equal`, since there the user is explicitly asking if the arrays are equal and ambiguity might be a problem.

Unfortunately it's hard to make these decisions without knowing how they'll play out in practice, we'll probably end up making some mistakes.
> Unfortunately it's hard to make these decisions without knowing how they'll play out in practice, we'll probably end up making some mistakes.

Maybe we can try going for a release candidate first this time around ?
If that would be useful, sure? I don't think we have many active downstream users besides yt who would know that we did that.
kind of a chicken+egg problem the way I see it: we *probably* don't have many of them but we could nontheless try to reach some of them through yt's user mailing list ? If we don't receive any feedback for like a month after a rc, we could probably go for the actual release ?
>  If loadtxt isn't wrappable then I think the best thing to do is leave the numpy implementations alone? Kinda seems like an oversight on the numpy side that loadtxt isn't overridable.

I was thinking again about this and it dawn on me that `loadtxt` cannot use the mechanisms from NEP 18 since it cannot know what array class to instantiate from its arguments.
It's still a bit of a mystery to me that `np.savetxt` *is* included in NEP 18, because there doesn't seem to be much point to it. Maybe we should wrap it just so it raises awareness of its limitations ? something like
```
>>> np.savetxt(arr)
warning: np.savetxt does not preserve units, and will only save the raw numerical data from the unyt_array object.
If this is the intended behaviour, you can call np.savetxt(arr.d) instead to silence this warning.
If you want to preserve units, use `unyt.savetxt` (and `unyt.loadtxt`) instead.
```
WDYT ?
SGTM, although I think you meant `np.savetxt(arr.d)` in the error message.
absolutely ! (edited)
@ngoldbaum numpy.msort is about to be deprecated (marked as such in numpy 1.24.0rc1)
see https://github.com/numpy/numpy/releases/tag/v1.24.0rc1

It's currently not implemented, so I'd vote for moving it to `removed_functions`, but there are currently no example of a function that's only _deprecated_ there. What would you advise ?

No strong opinion. Not implementing it seems fine.
I made an attempt at handling `numpy.poly*` functions, however:
- to me, it doesn't make much sense to me to special case `unyt_arrays` there (or maybe just to add errors stating that all non-dimensionless units are invalid inputs)
- I discovered that most of these function aren't actually trappable (probably a bug in numpy)
- astropy is explicitly ignoring them:

https://github.com/astropy/astropy/blob/2b74b87e24b23ef217cf4afe3b6cf1e382676844/astropy/units/quantity_helper/function_helpers.py#L117-L122

So I'm not sure what to do from here. I would probably report upstream and mark them as not doable in unyt for the time being.
That all sounds good to me, I agree that the `Polynomial` class is probably not worth worrying about at all. Please report any bugs you find upstream to numpy, although if you can try to make the reproducer in the issue you report not make use of unyt as that will make it easier for upstream to understand.
it seems I can't actually reproduce the problem I was referring to with what I assumed constituted a minimal example, and taking a look at numpy's source code I doubt that there's actually a problem there so I have some suspicion that it only happens when other, seemingly unrelated components of unyt are present (other wrapped functions, maybe). Any way I'm not convinced it's worth the effort of invegating this further for the time being.
Here a set of functions that I think call for a single uniform API decision.
- `in1d`
- `searchsorted`
- `setdiff1d`
- `setxor1d`
- `where`
- `piecewise`
- `place`
- `put`
- `putmask`
- `put_along_axis`
- `select`
- `choose`
- `compress` (fits the theme, not the problem)
- `clip`

What they have in common is that they assemble an array from 2 or more sources (arrays or scalars), and so we need to decide how pure ndarrays are to be treated. I suppose that it should be the same as everywhere else: pure numbers are treated as dimensionless, so errors should be raised if ones tries to combine such sources with arrays that have actual units, yet I'm having doubts about that, so I'd rather ask for feedback before I implement them.

That sounds good to me. If it turns out to be annoying to deal with this choice in yt we can revisit it.
We are down to only 8 functions (!!!) on the TODO list:

```python
    np.compress,
    np.i0,
    np.in1d,
    np.interp,
    np.ix_,
    np.linalg.svd,
    np.take_along_axis,
    np.unwrap,
```

here's where I'm at:

- I think `np.i0` is a noop ? It doesn't seem sensical to pass unitful data to it.
- `np.compress`, `np.take_along_axis`, `np.interp`, and `np.unwrap` seem relatively straightforward
- `np.in1d` warrants discussion : what should happen if one or the other array is dimensionless and the other isn't ? My guess is the following, assuming `a` is dimensionless and `b` isn't:
  - `np.in1d(a, b)` -> ok (treat `a` as if having the same units as `b`)
  - `np.in1d(b, a)` -> error
- `np.ix_` and `np.linalg.svd` are complete mysteries to me, and I would like a hand (or guidance) on what should be done with them.

I left `np.interp` out of #397 because I now think that it should be discussed alongside `np.in1d`
For `in1d`, I think we should be consistent with `isin`, where we do `_validate_units_consistency((element, test_elements))`. Your proposed solution also seems fine, but let's just keep things consistent, so if you'd prefer that then please change `isin` as well. I suspect that there may be code in the wild that uses `isin` or `in1d` and ignores units, so maybe best to avoid breaking that code?

Astropy has `np.ix_` unconditionally raising a `TypeError`. We could probably also just bail on it. For `svd`, maybe also look at `astropy.units.test.test_quantity_non_ufuncs::test_svd`?

> For in1d, I think we should be consistent with isin, where we do _validate_units_consistency((element, test_elements))

Right, I forgot about that one. I agree it should be kept consistent, and my proposed solution is both more complicated and more breaking, so I'm going with the simpler approach in #398 
Is that all of them? If so let’s close this and open a new one to discuss todos for the unyt 3.0 release.
almost there ! `np.ix_` and `np.linalg.svd` are still not done (`np.ix_` should go to `_UNSUPPORTED_FUNCTIONS`, as discussed)

Fix patch:
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index 06232e47..ae2e8116 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -37,6 +37,7 @@
     np.piecewise,  # astropy.units doens't have a simple implementation either
     np.packbits,
     np.unpackbits,
+    np.ix_,
 }
 
 _HANDLED_FUNCTIONS = {}
@@ -119,6 +120,19 @@ def linalg_pinv(a, *args, **kwargs):
     return np.linalg.pinv._implementation(a, *args, **kwargs).view(np.ndarray) / a.units
 
 
+@implements(np.linalg.svd)
+def linalg_svd(a, full_matrices=True, compute_uv=True, *args, **kwargs):
+    ret_units = a.units
+    retv = np.linalg.svd._implementation(
+        a.view(np.ndarray), full_matrices, compute_uv, *args, **kwargs
+    )
+    if compute_uv:
+        u, s, vh = retv
+        return (u, s * ret_units, vh)
+    else:
+        return retv * ret_units
+
+
 def _sanitize_range(_range, units):
     # helper function to histogram* functions
     ndim = len(units)


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index b75e6a0b..c2f1a12f 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -163,13 +163,6 @@
 }
 
 
-# this set represents all functions that need inspection, tests, or both
-# it is always possible that some of its elements belong in NOOP_FUNCTIONS
-TODO_FUNCTIONS = {
-    np.ix_,
-    np.linalg.svd,
-}
-
 DEPRECATED_FUNCTIONS = {
     "alen",  # deprecated in numpy 1.18, removed in 1.22
     "asscalar",  # deprecated in numpy 1.18, removed in 1.22
@@ -189,9 +182,7 @@
     "cumproduct",  # deprecated in numpy 1.25
 }
 
-NOT_HANDLED_FUNCTIONS = (
-    NOOP_FUNCTIONS | TODO_FUNCTIONS | UNSUPPORTED_FUNCTIONS | IGNORED_FUNCTIONS
-)
+NOT_HANDLED_FUNCTIONS = NOOP_FUNCTIONS | UNSUPPORTED_FUNCTIONS | IGNORED_FUNCTIONS
 
 for func in DEPRECATED_FUNCTIONS:
     if hasattr(np, func):
@@ -1166,6 +1157,19 @@ def test_eigvals(func):
     assert w.units == cm
 
 
+def test_linalg_svd():
+    a = (np.random.randn(9, 6) + 1j * np.random.randn(9, 6)) * cm
+    u, s, vh = np.linalg.svd(a)
+    assert type(u) is np.ndarray
+    assert type(vh) is np.ndarray
+    assert type(s) is unyt_array
+    assert s.units == cm
+
+    s = np.linalg.svd(a, compute_uv=False)
+    assert type(s) is unyt_array
+    assert s.units == cm
+
+
 def test_savetxt(tmp_path):
     a = [1, 2, 3] * cm
     with pytest.raises(