Add ctis.arange, a centered variant of named_arrays.arange

ctis/__init__.py

@@ -3,12 +3,14 @@
 spectrograph.
 """
 
+from ._arange import arange
 from ._regrid import regrid
 from . import scenes
 from . import instruments
 from . import inverters
 
 __all__ = [
+    "arange",
     "regrid",
     "scenes",
     "instruments",

ctis/_arange.py

@@ -0,0 +1,101 @@
+"""
+Evenly-spaced sequences for building CTIS coordinate grids.
+"""
+
+import numpy as np
+import astropy.units as u
+import named_arrays as na
+
+__all__ = [
+    "arange",
+]
+
+
+def arange(
+    start: "float | u.Quantity | na.AbstractVectorArray",
+    stop: "float | u.Quantity | na.AbstractVectorArray",
+    axis: "str | na.AbstractVectorArray",
+    step: "float | u.Quantity | na.AbstractVectorArray" = 1,
+) -> "na.AbstractArray":
+    """
+    Return evenly-spaced values over a range, centered within it.
+
+    This is the centered counterpart of :func:`named_arrays.arange`, which
+    starts exactly at `start` and leaves the remainder of the range as a gap on
+    the right. :func:`ctis.arange` fits as many samples spaced by `step` as
+    possible into ``[start, stop]`` and centers them, splitting the leftover
+    evenly between the two ends. This is convenient for building a coordinate
+    grid with a fixed pitch that is centered on a field of view rather than
+    biased toward one edge.
+
+    Unlike :func:`named_arrays.arange`, this is built on
+    :func:`named_arrays.linspace`, so it works correctly when the arguments are
+    :class:`astropy.units.Quantity` instances.
+
+    If `start`, `stop`, and `step` are instances of
+    :class:`named_arrays.AbstractVectorArray`, each component is centered
+    independently along its own axis (given by the matching component of
+    `axis`), producing an outer-product grid.
+
+    Parameters
+    ----------
+    start
+        The lower bound of the range.
+    stop
+        The upper bound of the range. The samples are centered within
+        ``[start, stop]``, so `start` and `stop` are samples only when `step`
+        divides the range evenly.
+    axis
+        The name of the new logical axis of the result. If `start`, `stop`, and
+        `step` are vectors, this should be a vector of axis names, one per
+        component.
+    step
+        The spacing between adjacent samples.
+
+    Examples
+    --------
+
+    :func:`named_arrays.arange` starts at `start`, leaving the remainder of the
+    range as a gap on the right.
+
+    .. jupyter-execute::
+
+        import named_arrays as na
+        import ctis
+
+        na.arange(0, 10, axis="x", step=3)
+
+    :func:`ctis.arange` uses the same samples and step, but splits that
+    remainder evenly between the two ends of the range.
+
+    .. jupyter-execute::
+
+        ctis.arange(0, 10, axis="x", step=3)
+
+    Vector arguments produce a centered grid, with each component sampled
+    independently along its own axis.
+
+    .. jupyter-execute::
+
+        import astropy.units as u
+
+        ctis.arange(
+            start=na.Cartesian2dVectorArray(-10, -8) * u.arcsec,
+            stop=na.Cartesian2dVectorArray(10, 8) * u.arcsec,
+            axis=na.Cartesian2dVectorArray("x", "y"),
+            step=na.Cartesian2dVectorArray(6, 5) * u.arcsec,
+        )
+    """
+    # the number of samples spaced by `step` that fit within the range,
+    # computed per-component so vector arguments produce a grid.
+    num = na.as_named_array((stop - start) / step)
+    num = (np.floor(num + 1e-10) + 1).astype(int)
+    if not isinstance(num, na.AbstractVectorArray):
+        # na.linspace requires a Python int for the count of a scalar axis
+        num = int(num.ndarray)
+
+    # center the samples by splitting the leftover evenly between both ends
+    span = (num - 1) * step
+    offset = (stop - start - span) / 2
+
+    return na.linspace(start + offset, stop - offset, axis=axis, num=num)

ctis/_arange_test.py

@@ -0,0 +1,80 @@
+import pytest
+import numpy as np
+import astropy.units as u
+import named_arrays as na
+import ctis
+
+
+@pytest.mark.parametrize(
+    argnames="start,stop,step",
+    argvalues=[
+        (0, 10, 3),  # range not divisible by step
+        (0, 9, 3),  # range divisible by step
+        (-10 * u.arcsec, 10 * u.arcsec, 6 * u.arcsec),  # astropy Quantity
+        (500 * u.nm, 600 * u.nm, 7 * u.nm),  # astropy Quantity
+    ],
+)
+def test_arange(
+    start: float | u.Quantity,
+    stop: float | u.Quantity,
+    step: float | u.Quantity,
+):
+    result = ctis.arange(start, stop, axis="x", step=step)
+
+    assert isinstance(result, na.AbstractScalarArray)
+    assert result.axes == ("x",)
+    assert na.unit_normalized(result) == na.unit_normalized(start)
+
+    # the samples are spaced by `step`
+    diff = np.diff(result, axis="x")
+    assert np.allclose(diff, step)
+
+    # the samples are centered within the range: the gaps at the two ends are
+    # equal, and they are smaller than a full step
+    gap_lower = result.min("x") - start
+    gap_upper = stop - result.max("x")
+    assert np.allclose(gap_lower, gap_upper)
+    assert np.all(gap_lower >= 0 * step)
+    assert np.all(gap_lower < step)
+
+    # the samples fit within the range
+    assert np.all(result >= start)
+    assert np.all(result <= stop)
+
+
+@pytest.mark.parametrize(
+    argnames="wrap",
+    argvalues=[
+        lambda x: x,  # plain Quantity components
+        lambda x: na.ScalarArray(x),  # ScalarArray components
+    ],
+    ids=["quantity", "scalararray"],
+)
+def test_arange_vector(wrap):
+    start = na.Cartesian2dVectorArray(x=wrap(-10 * u.arcsec), y=wrap(-8 * u.arcsec))
+    stop = na.Cartesian2dVectorArray(x=wrap(10 * u.arcsec), y=wrap(8 * u.arcsec))
+    step = na.Cartesian2dVectorArray(x=wrap(6 * u.arcsec), y=wrap(5 * u.arcsec))
+    axis = na.Cartesian2dVectorArray("x", "y")
+
+    result = ctis.arange(start, stop, axis=axis, step=step)
+
+    assert isinstance(result, na.Cartesian2dVectorArray)
+
+    # the components form an outer-product grid, one axis per component
+    assert set(na.shape(result)) == {"x", "y"}
+
+    # each component is spaced by its own step and centered in its own range
+    for c in ["x", "y"]:
+        component = getattr(result, c)
+        ax = getattr(axis, c)
+        lower = getattr(start, c)
+        upper = getattr(stop, c)
+        pitch = getattr(step, c)
+
+        assert np.allclose(np.diff(component, axis=ax), pitch)
+
+        gap_lower = component.min(ax) - lower
+        gap_upper = upper - component.max(ax)
+        assert np.allclose(gap_lower, gap_upper)
+        assert np.all(gap_lower >= 0 * pitch)
+        assert np.all(gap_lower < pitch)

pyproject.toml

@@ -17,7 +17,7 @@ classifiers = [
 ]
 dependencies = [
     "astropy",
-    "named-arrays~=2.0",
+    "named-arrays~=2.2",
 ]
 dynamic = ["version"]