Coverage for python / lsst / images / _geom.py: 40%
297 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-27 01:56 -0800
1# This file is part of lsst-images.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
12from __future__ import annotations
14__all__ = (
15 "XY",
16 "YX",
17 "Bounds",
18 "Box",
19 "Interval",
20 "SerializableBounds",
21)
23import math
24from collections.abc import Callable, Iterator, Sequence
25from typing import Any, ClassVar, NamedTuple, Protocol, Self, TypedDict, TypeVar, final, overload
27import numpy as np
28import pydantic
29import pydantic_core.core_schema as pcs
31# This pre-python-3.12 declaration is needed by Sphinx (probably the
32# autodoc-typehints plugin.
33T = TypeVar("T")
35# Interval and Box are defined as regular Python classes rather than
36# dataclasses or Pydantic models because we might want to implement them as
37# compiled-extension types in the future, and we want that to be transparent.
39# In a similar vein, we avoid declaring specific types for multidimensional
40# points or extents (other than ``tuple[int, ...]`` for numpy-compatible
41# shapes) in order to leave room for more fully-featured types to be added
42# upstream of this package in the future.
45class YX[T](NamedTuple):
46 """A pair of per-dimension objects, ordered ``(y, x)``.
48 Notes
49 -----
50 `YX` is used for slices, shapes, and other 2-d pairs when the most
51 natural ordering is ``(y, x)``. Because it is a `tuple`, however,
52 arithmetic operations behave as they would on a
53 `collections.abc.Sequence`, not a mathematical vector (e.g. adding
54 concatenates).
56 See Also
57 --------
58 XY
59 """
61 y: T
62 """The y / row object."""
64 x: T
65 """The x / column object."""
67 @property
68 def xy(self) -> XY:
69 """A tuple of the same objects in the opposite order."""
70 return XY(x=self.x, y=self.y)
72 def map[U](self, func: Callable[[T], U]) -> YX[U]:
73 """Apply a function to both objects."""
74 return YX(y=func(self.y), x=func(self.x))
77class XY[T](NamedTuple):
78 """A pair of per-dimension objects, ordered ``(x, y)``.
80 Notes
81 -----
82 `XY` is used for points and other 2-d pairs when the most natural ordering
83 is ``(x, y)``. Because it is a `tuple`, however, arithmetic operations
84 behave as they would on a `collections.abc.Sequence`, not a mathematical
85 vector (e.g. adding concatenates).
87 See Also
88 --------
89 YX
90 """
92 x: T
93 """The x / column object."""
95 y: T
96 """The y / row object."""
98 @property
99 def yx(self) -> YX:
100 """A tuple of the same objects in the opposite order."""
101 return YX(y=self.y, x=self.x)
103 def map[U](self, func: Callable[[T], U]) -> XY[U]:
104 """Apply a function to both objects."""
105 return XY(x=func(self.x), y=func(self.y))
108class _SerializedInterval(TypedDict):
109 start: int
110 stop: int
113@final
114class Interval:
115 """A 1-d integer interval with positive size.
117 Parameters
118 ----------
119 start
120 Inclusive minimum point in the interval.
121 stop
122 One past the maximum point in the interval.
124 Notes
125 -----
126 Adding or subtracting an `int` from an interval returns a shifted interval.
128 `Interval` implements the necessary hooks to be included directly in a
129 `pydantic.BaseModel`, even though it is neither a built-in type nor a
130 Pydantic model itself.
131 """
133 def __init__(self, start: int, stop: int):
134 # Coerce to be defensive against numpy int scalars.
135 self._start = int(start)
136 self._stop = int(stop)
137 if not (self._stop > self._start):
138 raise ValueError(f"Interval must have positive size; got [{self._start}, {self._stop})")
140 __slots__ = ("_start", "_stop")
142 factory: ClassVar[IntervalSliceFactory]
143 """A factory for creating intervals using slice syntax.
145 For example::
147 interval = Interval.factory[2:5]
148 """
150 @classmethod
151 def hull(cls, first: int | Interval, *args: int | Interval) -> Interval:
152 """Construct an interval that includes all of the given points and/or
153 intervals.
154 """
155 if type(first) is Interval:
156 rmin = first.min
157 rmax = first.max
158 else:
159 rmin = rmax = first
160 for arg in args:
161 if type(arg) is Interval:
162 rmin = min(rmin, arg.min)
163 rmax = max(rmax, arg.max)
164 else:
165 rmin = min(rmin, arg)
166 rmax = max(rmax, arg)
167 return Interval(start=rmin, stop=rmax + 1)
169 @classmethod
170 def from_size(cls, size: int, start: int = 0) -> Interval:
171 """Construct an interval from its size and optional start."""
172 return cls(start=start, stop=start + size)
174 @property
175 def start(self) -> int:
176 """Inclusive minimum point in the interval (`int`)."""
177 return self._start
179 @property
180 def stop(self) -> int:
181 """One past the maximum point in the interval (`int`)."""
182 return self._stop
184 @property
185 def min(self) -> int:
186 """Inclusive minimum point in the interval (`int`)."""
187 return self.start
189 @property
190 def max(self) -> int:
191 """Inclusive maximum point in the interval (`int`)."""
192 return self.stop - 1
194 @property
195 def size(self) -> int:
196 """Size of the interval (`int`)."""
197 return self.stop - self.start
199 @property
200 def range(self) -> __builtins__.range:
201 """An iterable over all values in the interval
202 (`__builtins__.range`).
203 """
204 return range(self.start, self.stop)
206 @property
207 def arange(self) -> np.ndarray:
208 """An array of all the values in the interval (`numpy.ndarray`).
210 Array values are integers.
211 """
212 return np.arange(self.start, self.stop)
214 def linspace(self, n: int | None = None, *, step: float | None = None) -> np.ndarray:
215 """Return an array of values that spans the interval.
217 Parameters
218 ----------
219 n
220 How many values to return. The default (if ``step`` is also not
221 provided) is the size of the interval, i.e. equivalent to the
222 `arange` property (but converted to `float`).
223 step
224 Set ``n`` such that the distance between points is equal to or
225 just less than this. Mutually exclusive with ``n``.
227 Returns
228 -------
229 numpy.ndarray
230 Array of `float` values.
232 See Also
233 --------
234 numpy.linspace
235 """
236 if n is None:
237 if step is None:
238 return self.arange.astype(np.float64)
239 n = math.ceil(self.size / step)
240 elif step is not None:
241 raise TypeError("'n' and 'step' cannot both be provided.")
242 return np.linspace(self.min, self.max, n, dtype=np.float64)
244 @property
245 def center(self) -> float:
246 """The center of the interval (`float`)."""
247 return 0.5 * (self.min + self.max)
249 def __str__(self) -> str:
250 return f"{self.start}:{self.stop}"
252 def __repr__(self) -> str:
253 return f"Interval(start={self.start}, stop={self.stop})"
255 def __eq__(self, other: object) -> bool:
256 if type(other) is Interval:
257 return self._start == other._start and self._stop == other._stop
258 return False
260 def __add__(self, other: int) -> Interval:
261 return Interval(start=self.start + other, stop=self.stop + other)
263 def __sub__(self, other: int) -> Interval:
264 return Interval(start=self.start - other, stop=self.stop - other)
266 def __contains__(self, x: int) -> bool:
267 return x >= self.start and x < self.stop
269 @overload
270 def contains(self, other: Interval | int) -> bool: ... 270 ↛ exitline 270 didn't return from function 'contains' because
272 @overload
273 def contains(self, other: np.ndarray) -> np.ndarray: ... 273 ↛ exitline 273 didn't return from function 'contains' because
275 def contains(self, other: Interval | int | float | np.ndarray) -> bool | np.ndarray:
276 """Test whether this interval fully contains another or one or more
277 points.
279 Parameters
280 ----------
281 other
282 Another interval to compare to, or one or more position values.
284 Returns
285 -------
286 `bool` | `numpy.ndarray`
287 If a single interval or value was passed, a single `bool`. If an
288 array was passed, an array with the same shape.
290 Notes
291 -----
292 In order to yield the desired behavior for floating-point arguments,
293 points are actually tested against an interval that is 0.5 larger on
294 both sides: this makes positions within the outer boundary of pixels
295 (but beyond the centers of those pixels, which have integer positions)
296 appear "on the image".
297 """
298 if isinstance(other, Interval):
299 return self.start <= other.start and self.stop >= other.stop
300 else:
301 result = np.logical_and(self.start - 0.5 <= other, other < self.stop + 0.5)
302 if not result.shape:
303 return bool(result)
304 return result
306 def intersection(self, other: Interval) -> Interval | None:
307 """Return an interval that is contained by both ``self`` and ``other``.
309 When there is no overlap between the intervals, `None` is returned.
310 """
311 new_start = max(self.start, other.start)
312 new_stop = min(self.stop, other.stop)
313 if new_start < new_stop:
314 return Interval(start=new_start, stop=new_stop)
315 return None
317 def dilated_by(self, padding: int) -> Interval:
318 """Return a new interval padded by the given amount on both sides."""
319 return Interval(start=self._start - padding, stop=self._stop + padding)
321 def slice_within(self, other: Interval) -> slice:
322 """Return the `slice` that corresponds to the values in this interval
323 when the items of the container being sliced correspond to ``other``.
325 This assumes ``other.contains(self)``.
326 """
327 return slice(self.start - other.start, self.stop - other.start)
329 @classmethod
330 def from_legacy(cls, legacy: Any) -> Interval:
331 """Convert from an `lsst.geom.IntervalI` instance."""
332 return cls(legacy.begin, legacy.end)
334 def to_legacy(self) -> Any:
335 """Convert to an `lsst.geom.IntervalI` instance."""
336 from lsst.geom import IntervalI
338 return IntervalI(min=self.min, max=self.max)
340 def __reduce__(self) -> tuple[type[Interval], tuple[int, int]]:
341 return (
342 Interval,
343 (
344 self._start,
345 self._stop,
346 ),
347 )
349 @classmethod
350 def __get_pydantic_core_schema__(
351 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
352 ) -> pcs.CoreSchema:
353 from_typed_dict = pcs.chain_schema(
354 [
355 handler(_SerializedInterval),
356 pcs.no_info_plain_validator_function(cls._validate),
357 ]
358 )
359 return pcs.json_or_python_schema(
360 json_schema=from_typed_dict,
361 python_schema=pcs.union_schema([pcs.is_instance_schema(Interval), from_typed_dict]),
362 serialization=pcs.plain_serializer_function_ser_schema(cls._serialize, info_arg=False),
363 )
365 @classmethod
366 def _validate(cls, data: _SerializedInterval) -> Interval:
367 return cls(**data)
369 def _serialize(self) -> _SerializedInterval:
370 return {"start": self._start, "stop": self._stop}
373class IntervalSliceFactory:
374 """A factory for `Interval` objects using array-slice syntax.
376 Notes
377 -----
378 When indexed with a single slice, this returns an `Interval`::
380 assert Interval.factory[3:6] == Interval(start=3, stop=6)
382 """
384 def __getitem__(self, s: slice) -> Interval:
385 if s.step is not None and s.step != 1:
386 raise ValueError(f"Slice {s} has non-unit step.")
387 return Interval(start=s.start or 0, stop=s.stop)
390Interval.factory = IntervalSliceFactory()
393class _SerializedBox(TypedDict):
394 y: _SerializedInterval
395 x: _SerializedInterval
398class Box:
399 """An axis-aligned 2-d rectangular region.
401 Parameters
402 ----------
403 y
404 Interval for the y dimension.
405 x
406 Interval for the x dimension.
408 Notes
409 -----
410 `Box` implements the necessary hooks to be included directly in a
411 `pydantic.BaseModel`, even though it is neither a built-in type nor a
412 Pydantic model itself.
413 """
415 def __init__(self, y: Interval, x: Interval):
416 self._intervals = YX(y, x)
418 __slots__ = ("_intervals",)
420 factory: ClassVar[BoxSliceFactory]
421 """A factory for creating boxes using slice syntax.
423 For example::
425 box = Box.factory[2:5, 3:9]
426 """
428 @classmethod
429 def from_shape(cls, shape: Sequence[int], start: Sequence[int] | None = None) -> Box:
430 """Construct a box from its shape and optional start.
432 Parameters
433 ----------
434 shape
435 Sequence of sizes, ordered ``(y, x)`` (except for `XY` instances).
436 start
437 Sequence of starts, ordered ``(y, x)`` (except for `XY` instances).
438 """
439 if start is None:
440 start = (0,) * len(shape)
441 match shape:
442 case XY(x=x_size, y=y_size):
443 pass
444 case [y_size, x_size]:
445 pass
446 case _:
447 raise ValueError(f"Invalid sequence for shape: {shape!r}.")
448 match start:
449 case XY(x=x_start, y=y_start):
450 pass
451 case [y_start, x_start]:
452 pass
453 case _:
454 raise ValueError(f"Invalid sequence for shape: {shape!r}.")
455 return Box(y=Interval.from_size(y_size, start=y_start), x=Interval.from_size(x_size, start=x_start))
457 @property
458 def start(self) -> YX[int]:
459 """Tuple holding the starts of the intervals ordered ``(y, x)``
460 (`YX` [`int`]).
461 """
462 return YX(self.y.start, self.x.start)
464 @property
465 def shape(self) -> YX[int]:
466 """Tuple holding the sizes of the intervals ordered ``(y, x)``
467 (`YX` [`int`]).
468 """
469 return YX(self.y.size, self.x.size)
471 @property
472 def x(self) -> Interval:
473 """Shortcut for the last dimension's interval (`int`)."""
474 return self._intervals[-1]
476 @property
477 def y(self) -> Interval:
478 """Shortcut for the second-to-last dimension's interval (`int`)."""
479 return self._intervals[-2]
481 def meshgrid(self, n: int | Sequence[int] | None = None, *, step: float | None = None) -> XY[np.ndarray]:
482 """Return a pair of 2-d arrays of the coordinate values of the box.
484 Parameters
485 ----------
486 n
487 Number of points in each dimension. If a sequence, points are
488 assumed to be ordered ``(x, y)`` unless a `YX` instance is
489 provided.
490 step
491 Set ``n`` such that the distance between points is equal to or
492 just less than this in each dimension. Mutually exclusive with
493 ``n``.
495 Returns
496 -------
497 `XY` [`numpy.ndarray`]
498 A pair of arrays, each of which is 2-d with floating-point values.
500 See Also
501 --------
502 numpy.meshgrid
503 """
504 match n:
505 case int():
506 ax = self.x.linspace(n)
507 ay = self.y.linspace(n)
508 case YX(y=ny, x=nx):
509 ax = self.x.linspace(nx)
510 ay = self.y.linspace(ny)
511 case [nx, ny]:
512 ax = self.x.linspace(nx)
513 ay = self.y.linspace(ny)
514 case None:
515 ax = self.x.linspace(step=step)
516 ay = self.y.linspace(step=step)
517 return XY(*np.meshgrid(ax, ay))
519 def __eq__(self, other: object) -> bool:
520 if type(other) is Box:
521 return self._intervals == other._intervals
522 return False
524 def __str__(self) -> str:
525 return f"[y={self.y}, x={self.x}]"
527 def __repr__(self) -> str:
528 return f"Box(y={self.y!r}, x={self.x!r})"
530 @overload
531 def contains(self, other: Box, /) -> bool: ... 531 ↛ exitline 531 didn't return from function 'contains' because
533 @overload
534 def contains(self, *, y: int, x: int) -> bool: ... 534 ↛ exitline 534 didn't return from function 'contains' because
536 @overload
537 def contains(self, *, y: np.ndarray, x: np.ndarray) -> np.ndarray: ... 537 ↛ exitline 537 didn't return from function 'contains' because
539 def contains(
540 self,
541 other: Box | None = None,
542 *,
543 y: int | np.ndarray | None = None,
544 x: int | np.ndarray | None = None,
545 ) -> bool | np.ndarray:
546 """Test whether this box fully contains another or one or more points.
548 Parameters
549 ----------
550 other
551 Another box to compare to. Not compatible with the ``y`` and ``x``
552 arguments.
553 y
554 One or more integer Y coordinates to test for containment.
555 If an array, an array of results will be returned.
556 x
557 One or more integer X coordinates to test for containment.
558 If an array, an array of results will be returned.
560 Returns
561 -------
562 `bool` | `numpy.ndarray`
563 If ``other`` was passed or ``x`` and ``y`` are both scalars, a
564 single `bool` value. If ``x`` and ``y`` are arrays, a boolean
565 array with their broadcasted shape.
567 Notes
568 -----
569 In order to yield the desired behavior for floating-point arguments,
570 points are actually tested against an interval that is 0.5 larger on
571 both sides: this makes positions within the outer boundary of pixels
572 (but beyond the centers of those pixels, which have integer positions)
573 appear "on the image".
574 """
575 if other is not None:
576 if x is not None or y is not None:
577 raise TypeError("Too many arguments to 'Box.contain'.")
578 return all(a.contains(b) for a, b in zip(self._intervals, other._intervals, strict=True))
579 elif x is None or y is None:
580 raise TypeError("Not enough arguments to 'Box.contain'.")
581 else:
582 result = np.logical_and(self.x.contains(x), self.y.contains(y))
583 if not result.shape:
584 return bool(result)
585 return result
587 def intersection(self, other: Box) -> Box | None:
588 """Return a box that is contained by both ``self`` and ``other``.
590 When there is no overlap between the boxes, `None` is returned.
591 """
592 intervals = []
593 for a, b in zip(self._intervals, other._intervals, strict=True):
594 if (r := a.intersection(b)) is None:
595 return None
596 intervals.append(r)
597 return Box(*intervals)
599 def dilated_by(self, padding: int) -> Box:
600 """Return a new box padded by the given amount on all sides."""
601 return Box(*[i.dilated_by(padding) for i in self._intervals])
603 def slice_within(self, other: Box) -> YX[slice]:
604 """Return a `tuple` of `slice` objects that correspond to the
605 positions in this box when the items of the container being sliced
606 correspond to ``other``.
608 This assumes ``other.contains(self)``.
609 """
610 return YX(self.y.slice_within(other.y), self.x.slice_within(other.x))
612 def boundary(self) -> Iterator[YX[int]]:
613 """Iterate over the corners of the box as ``(y, x)`` tuples."""
614 if len(self._intervals) != 2:
615 raise TypeError("Box is not 2-d.")
616 yield YX(self.y.min, self.x.min)
617 yield YX(self.y.min, self.x.max)
618 yield YX(self.y.max, self.x.max)
619 yield YX(self.y.max, self.x.min)
621 def __reduce__(self) -> tuple[type[Box], tuple[Interval, ...]]:
622 return (Box, self._intervals)
624 @classmethod
625 def from_legacy(cls, legacy: Any) -> Box:
626 """Convert from an `lsst.geom.Box2I` instance."""
627 return cls(y=Interval.from_legacy(legacy.y), x=Interval.from_legacy(legacy.x))
629 def to_legacy(self) -> Any:
630 """Convert to an `lsst.geom.BoxI` instance."""
631 from lsst.geom import Box2I
633 return Box2I(x=self.x.to_legacy(), y=self.y.to_legacy())
635 @classmethod
636 def __get_pydantic_core_schema__(
637 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler
638 ) -> pcs.CoreSchema:
639 from_typed_dict = pcs.chain_schema(
640 [
641 handler(_SerializedBox),
642 pcs.no_info_plain_validator_function(cls._validate),
643 ]
644 )
645 return pcs.json_or_python_schema(
646 json_schema=from_typed_dict,
647 python_schema=pcs.union_schema([pcs.is_instance_schema(Box), from_typed_dict]),
648 serialization=pcs.plain_serializer_function_ser_schema(cls._serialize, info_arg=False),
649 )
651 @classmethod
652 def _validate(cls, data: _SerializedBox) -> Box:
653 return cls(y=Interval._validate(data["y"]), x=Interval._validate(data["x"]))
655 def _serialize(self) -> _SerializedBox:
656 return {"y": self.y._serialize(), "x": self.x._serialize()}
658 def serialize(self) -> Box:
659 """Return a Pydantic-friendly representation of this object.
661 This method just returns the `Box` itself, since that already provides
662 Pydantic serialization hooks. It exists for compatibility with the
663 `Bounds` protocol.
664 """
665 return self
667 @classmethod
668 def deserialize(cls, serialized: SerializableBounds) -> Box:
669 """Deserialize a bounds object on the assumption it is a `Box`.
671 This method just returns the `Box` itself, since that already provides
672 Pydantic serialization hooks. It exists for compatibility with the
673 `Bounds` protocol.
674 """
675 assert isinstance(serialized, Box)
676 return serialized
679class BoxSliceFactory:
680 """A factory for `Box` objects using array-slice syntax.
682 Notes
683 -----
684 When indexed with one or more slices, this returns a `Box`:
686 assert (
687 Box.factory[3:6, -1:2]
688 == Box(x=Interval(start=-1, stop=2), y=Interval(start=3, stop=6)
689 )
690 """
692 def __getitem__(self, key: tuple[slice, slice]) -> Box:
693 match key:
694 case XY(x=x, y=y):
695 return Box(Interval.factory[y], Interval.factory[x])
696 case (y, x):
697 return Box(Interval.factory[y], Interval.factory[x])
698 case _:
699 raise TypeError("Expected exactly two slices.")
702Box.factory = BoxSliceFactory()
705# This is expected to become a union of concrete Bounds types that we can
706# serialize via pydantic. Right now that's only Box.
707type SerializableBounds = Box
710class Bounds(Protocol):
711 """A protocol for objects that represent the validity region for a function
712 defined in 2-d pixel coordinates.
714 Notes
715 -----
716 Most objects natively have a simple 2-d bounding box as their bounds
717 (typically the boundary of a sensor), and the `Box` class is hence the
718 most common bounds implementation. But sometimes a large chunk of that
719 box may be missing due to vignetting or bad amplifiers, and we may want to
720 transform from one coordinate system to another. The Bounds interface is
721 intended to handle both of these cases as well.
722 """
724 def boundary(self) -> Iterator[YX[int]]:
725 """Iterate over points on the boundary as ``(y, x)`` tuples."""
726 ...
728 @overload
729 def contains(self, *, x: int, y: int) -> bool: ... 729 ↛ exitline 729 didn't return from function 'contains' because
731 @overload
732 def contains(self, *, x: np.ndarray, y: np.ndarray) -> np.ndarray: ... 732 ↛ exitline 732 didn't return from function 'contains' because
734 def contains(self, *, x: int | np.ndarray, y: int | np.ndarray) -> bool | np.ndarray:
735 """Test whether this box fully contains another or one or more points.
737 Parameters
738 ----------
739 x
740 One or more integer X coordinates to test for containment.
741 If an array, an array of results will be returned.
742 y
743 One or more integer Y coordinates to test for containment.
744 If an array, an array of results will be returned.
746 Returns
747 -------
748 `bool` | `numpy.ndarray`
749 If ``x`` and ``y`` are both scalars, a single `bool` value. If
750 ``x`` and ``y`` are arrays, a boolean array with their broadcasted
751 shape.
752 """
753 ...
755 def serialize(self) -> SerializableBounds:
756 """Convert a bounds instance into a serializable object."""
757 ...
759 @classmethod
760 def deserialize(cls, serialized: SerializableBounds) -> Self:
761 """Convert a serialized bounds object into its in-memory form."""
762 match serialized:
763 case Box():
764 return serialized # type: ignore[return-value]
765 raise RuntimeError(f"Cannot deserialize {serialized!r}.")