Repository: yt-project/unyt

Issue:
BUG: `logspace` incorrectly accepts units
* unyt version: 3.0.3
* Python version: 3.12
* Operating System: Linux (Fedora 40)

### Description

The `logspace` function accepts `start` and `stop` such that the returned array spans between `base**start` and `base**stop`. Currently you can pass `start, stop = 1*cm, 2*cm` and get values between 1 and 100 cm, which is mathematically incorrect. (But dimensionless unyt arrays should be accepted.) At the same time, it should be possible to pass the keyword argument `base` with units and have these applied to the returned array.

### What I Did

```
>>> np.logspace(1 * u.cm, 2 * u.cm, num=5)
unyt_array([ 10.        ,  17.7827941 ,  31.6227766 ,  56.23413252,
       100.        ], 'cm')
>>> np.logspace(1, 2, base=10*u.cm)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/txwx36/code/unyt/unyt/array.py", line 2045, in __array_function__
    return _HANDLED_FUNCTIONS[func](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/txwx36/code/unyt/unyt/_array_functions.py", line 511, in logspace
    * start.units
      ^^^^^^^^^^^
AttributeError: 'int' object has no attribute 'units'
```


Discussion:
kyleaoman: pre-commit.ci autofix
kyleaoman: pre-commit.ci autofix
kyleaoman: pre-commit.ci autofix
neutrinoceros: Thank you so much for catching and fixing so many (of my) mistakes ! There's a merge conflict now that prevents merging, can you please update your branch with a rebase ? 
neutrinoceros: pre-commit.ci autofix
kyleaoman: pre-commit.ci autofix
kyleaoman: Incidentally, would be great to have a release of `3.0.4` before I get to the point of merging SWIFTSIM/swiftsimio#219. I'm getting close to handling all of the functions but expecting some refactoring after that, plus adding a few more tests and updating docs, so hoping to merge sometime in January - any chance of a release around that time?
neutrinoceros: We're hoping to get unyt 3.0.4 out before 2025, though I can't make any promises.
Do you expect to have more patches for unyt after this one ?

Patch:
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index 3bd811b..c8e92f9 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -621,26 +621,94 @@ def array_equiv(a1, a2, *args, **kwargs) -> bool:
     )
 
 
-@implements(np.linspace)
-def linspace(start, stop, *args, **kwargs):
+def _linspace(
+    start,
+    stop,
+    num=50,
+    endpoint=True,
+    retstep=False,
+    dtype=None,
+    axis=0,
+    *,
+    device=None,
+):
     _validate_units_consistency((start, stop))
-    return (
-        np.linspace._implementation(
-            np.asarray(start), np.asarray(stop), *args, **kwargs
+    kwargs = {
+        "num": num,
+        "endpoint": endpoint,
+        "retstep": retstep,
+        "dtype": dtype,
+        "axis": axis,
+    }
+    if NUMPY_VERSION >= Version("2.1.0.dev0"):
+        kwargs["device"] = device
+    result = np.linspace._implementation(np.asarray(start), np.asarray(stop), **kwargs)
+    if retstep:
+        return result[0] * start.units, result[1] * start.units
+    else:
+        return result * start.units
+
+
+if NUMPY_VERSION >= Version("2.0.0.dev0"):
+
+    @implements(np.linspace)
+    def linspace(
+        start,
+        stop,
+        num=50,
+        endpoint=True,
+        retstep=False,
+        dtype=None,
+        axis=0,
+        *,
+        device=None,
+    ):
+        return _linspace(
+            start,
+            stop,
+            num=50,
+            endpoint=endpoint,
+            retstep=retstep,
+            dtype=dtype,
+            axis=axis,
+            device=device,
+        )
+
+else:
+
+    @implements(np.linspace)
+    def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0):
+        return _linspace(
+            start,
+            stop,
+            num=50,
+            endpoint=endpoint,
+            retstep=retstep,
+            dtype=dtype,
+            axis=axis,
         )
-        * start.units
-    )
 
 
 @implements(np.logspace)
-def logspace(start, stop, *args, **kwargs):
-    _validate_units_consistency((start, stop))
-    return (
-        np.logspace._implementation(
-            np.asarray(start), np.asarray(stop), *args, **kwargs
+def logspace(start, stop, num=50, endpoint=True, base=10, dtype=None, axis=0):
+    startu = getattr(start, "units", NULL_UNIT)
+    stopu = getattr(stop, "units", NULL_UNIT)
+    if startu != NULL_UNIT or stopu != NULL_UNIT:
+        raise TypeError(
+            "The first two arguments to numpy.logspace must be dimensionless, "
+            f"got units={startu} (arg1) and units={stopu} (arg2). If output with"
+            " units is desired, apply the units to the `base` kwarg."
         )
-        * start.units
+    result = np.logspace._implementation(
+        np.asarray(start),
+        np.asarray(stop),
+        num=num,
+        endpoint=endpoint,
+        base=np.asarray(base),
+        dtype=dtype,
+        axis=axis,
     )
+    return result * getattr(base, "units", NULL_UNIT)
 
 
 @implements(np.geomspace)


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index 291db7d..fbabac1 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -9,7 +9,7 @@ import pytest
 from numpy.testing import assert_allclose
 from packaging.version import Version
 
-from unyt import A, K, Msun, cm, degC, delta_degC, g, km, rad, s
+from unyt import A, K, Msun, cm, degC, delta_degC, dimensionless, g, km, rad, s
 from unyt._array_functions import (
     _HANDLED_FUNCTIONS as HANDLED_FUNCTIONS,
     _UNSUPPORTED_FUNCTIONS as UNSUPPORTED_FUNCTIONS,
@@ -1233,16 +1233,42 @@ def test_isclose_error():
         np.isclose(x, y)
 
 
-@pytest.mark.parametrize(
-    "func",
-    [
-        np.linspace,
-        np.logspace,
-        np.geomspace,
-    ],
-)
-def test_xspace(func):
-    res = func(1 * cm, 11 * cm, 10)
+def test_linspace():
+    res = np.linspace(1 * cm, 11 * cm, 10)
+    assert type(res) is unyt_array
+    assert res.units == cm
+
+
+def test_linspace_with_retstep():
+    res, step = np.linspace(1 * cm, 11 * cm, 10, retstep=True)
+    assert type(res) is unyt_array
+    assert res.units == cm
+    assert type(step) is unyt_quantity
+    assert step.units == cm
+
+
+def test_logspace_with_units_raises():
+    with pytest.raises(
+        TypeError,
+        match="The first two arguments to numpy.logspace must be dimensionless",
+    ):
+        np.logspace(1 * cm, 2 * cm)
+
+
+def test_logspace_with_base():
+    res = np.logspace(1, 2, base=10 * cm)
+    assert type(res) is unyt_array
+    assert res.units == cm
+
+
+def test_logspace_with_dimless_input():
+    res = np.logspace(1 * dimensionless, 2 * dimensionless)
+    assert type(res) is unyt_array
+    assert res.units == dimensionless
+
+
+def test_geomspace():
+    res = np.geomspace(1 * cm, 11 * cm, 10)
     assert type(res) is unyt_array
     assert res.units == cm