From 6360c1ede82fdbe02d84bfd2e8bb80576521bf49 Mon Sep 17 00:00:00 2001 From: Roy Smart Date: Sat, 25 Jul 2026 20:24:49 -0600 Subject: [PATCH] Add `ctis.arange`, a centered variant of `named_arrays.arange` `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 (rather than starting at `start` and leaving a gap on the right like `named_arrays.arange`). This is convenient for building a coordinate grid with a fixed pitch centered on a field of view. It is built on `named_arrays.linspace`, so it is unit-safe (unlike `named_arrays.arange`, which wraps `numpy.arange` and misbehaves with `astropy.units.Quantity`). When `start`, `stop`, and `step` are vectors, each component is centered independently along its own axis, producing an outer-product grid. Requires named-arrays >=2.2 for the vector-of-`ScalarArray`s case (`linspace` accepting a scalar-array `num`). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01CtdKmedevWkDab6BWupXqQ --- ctis/__init__.py | 2 + ctis/_arange.py | 101 +++++++++++++++++++++++++++++++++++++++++++ ctis/_arange_test.py | 80 ++++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- 4 files changed, 184 insertions(+), 1 deletion(-) create mode 100644 ctis/_arange.py create mode 100644 ctis/_arange_test.py diff --git a/ctis/__init__.py b/ctis/__init__.py index ef8c7cf..5e829ac 100644 --- a/ctis/__init__.py +++ b/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", diff --git a/ctis/_arange.py b/ctis/_arange.py new file mode 100644 index 0000000..5732a28 --- /dev/null +++ b/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) diff --git a/ctis/_arange_test.py b/ctis/_arange_test.py new file mode 100644 index 0000000..7dda3bf --- /dev/null +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 10919f6..9874bce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ classifiers = [ ] dependencies = [ "astropy", - "named-arrays~=2.0", + "named-arrays~=2.2", ] dynamic = ["version"]