DM-42685 Add a Tensor interface to analysis tools

python/lsst/analysis/tools/contexts/_baseContext.py

@@ -68,12 +68,10 @@ class ContextApplier:
     @overload
     def __get__(
         self, instance: AnalysisAction, klass: type[AnalysisAction] | None = None
-    ) -> Callable[[ContextType], None]:
-        ...
+    ) -> Callable[[ContextType], None]: ...
 
     @overload
-    def __get__(self, instance: None, klass: type[AnalysisAction] | None = None) -> ContextApplier:
-        ...
+    def __get__(self, instance: None, klass: type[AnalysisAction] | None = None) -> ContextApplier: ...
 
     def __get__(
         self, instance: AnalysisAction | None, klass: type[AnalysisAction] | None = None

python/lsst/analysis/tools/interfaces/_actions.py

@@ -45,7 +45,7 @@
 from lsst.pex.config.configurableActions import ConfigurableAction, ConfigurableActionField
 
 from ..contexts import ContextApplier
-from ._interfaces import KeyedData, KeyedDataSchema, MetricResultType, PlotResultType, Scalar, Vector
+from ._interfaces import KeyedData, KeyedDataSchema, MetricResultType, PlotResultType, Scalar, Tensor, Vector
 
 
 class AnalysisAction(ConfigurableAction):
@@ -155,6 +155,16 @@ def __call__(self, data: KeyedData, **kwargs) -> Vector:
         raise NotImplementedError("This is not implemented on the base class")
 
 
+class TensorAction(AnalysisAction):
+    """A `TensorAction` is an `AnalysisAction` that returns a `Tensor` when
+    called.
+    """
+
+    @abstractmethod
+    def __call__(self, data: KeyedData, **kwargs) -> Tensor:
+        raise NotImplementedError("This is not implemented on the base class")
+
+
 class ScalarAction(AnalysisAction):
     """A `ScalarAction` is an `AnalysisAction` that returns a `Scalar` when
     called.

python/lsst/analysis/tools/interfaces/_analysisTools.py

@@ -39,8 +39,7 @@
 
 @runtime_checkable
 class _HasOutputNames(Protocol):
-    def getOutputNames(self, config: pexConfig.Config | None = None) -> Iterable[str]:
-        ...
+    def getOutputNames(self, config: pexConfig.Config | None = None) -> Iterable[str]: ...
 
 
 def _finalizeWrapper(
@@ -116,6 +115,7 @@ class AnalysisTool(AnalysisAction):
     The stages themselves are also configurable, allowing control over various
     aspects of the individual `AnalysisAction`\ s.
     """
+
     prep = ConfigurableActionField[AnalysisAction](doc="Action to run to prepare inputs", default=BasePrep)
     process = ConfigurableActionField[AnalysisAction](
         doc="Action to process data into intended form", default=BaseProcess
@@ -177,9 +177,9 @@ def _call_single(self, data: KeyedData, **kwargs) -> KeyedResults:
         kwargs["metric_tags"] = list(self.metric_tags or ())
         prepped: KeyedData = self.prep(data, **kwargs)  # type: ignore
         processed: KeyedData = self.process(prepped, **kwargs)  # type: ignore
-        finalized: Mapping[str, PlotTypes] | PlotTypes | Mapping[
-            str, Measurement
-        ] | Measurement | JointResults = self.produce(
+        finalized: (
+            Mapping[str, PlotTypes] | PlotTypes | Mapping[str, Measurement] | Measurement | JointResults
+        ) = self.produce(
             processed, **kwargs
         )  # type: ignore
         return self._process_single_results(finalized)

python/lsst/analysis/tools/interfaces/_interfaces.py

@@ -22,6 +22,7 @@
 from __future__ import annotations
 
 __all__ = (
+    "Tensor",
     "Scalar",
     "ScalarType",
     "KeyedData",
@@ -34,7 +35,7 @@
 
 from abc import ABCMeta
 from numbers import Number
-from typing import Any, Iterable, Mapping, MutableMapping, TypeAlias
+from typing import Any, Iterable, Mapping, MutableMapping, Protocol, TypeAlias, runtime_checkable
 
 import numpy as np
 from healsparse import HealSparseMap
@@ -43,6 +44,36 @@
 from numpy.typing import NDArray
 
 
+@runtime_checkable
+class Tensor(Protocol):
+    r"""This is an interface only class and is intended to represent data that
+    is 2+ dimensions.
+
+    Technically one could use this for scalars or 1D arrays,
+    but for those the Scalar or Vector interface should be preferred.
+
+    `Tensor`\ s abstract around the idea of a multidimensional array, and work
+    with a variety of backends including Numpy, CuPy, Tensorflow, PyTorch,
+    MXNet, TVM, and mpi4py. This intentionally has a minimum interface to
+    comply with the industry standard dlpack which ensures each of these
+    backend native types will work.
+
+    To ensure that a `Tensor` is in a desired container (e.g. ndarray) one can
+    call the corresponding ``from_dlpack`` method. Whenever possible this will
+    be a zero copy action. For instance to work with a Tensor named
+    ``input_tensor`` as if it were a numpy object, one would do
+    ``image = np.from_dlpack(input_tensor)``.
+    """
+
+    ndim: int
+    shape: tuple[int, ...]
+    strides: tuple[int, ...]
+
+    def __dlpack__(self, /, *, stream: int | None = ...) -> Any: ...
+
+    def __dlpack_device__(self) -> tuple[int, int]: ...
+
+
 class ScalarMeta(ABCMeta):
     def __instancecheck__(cls: ABCMeta, instance: Any) -> Any:
         return isinstance(instance, tuple(cls.mro()[1:]))
@@ -72,18 +103,18 @@ def __init__(self) -> None:
 like an NDArray should be considered a Vector.
 """
 
-KeyedData = MutableMapping[str, Vector | Scalar | HealSparseMap]
+KeyedData = MutableMapping[str, Vector | Scalar | HealSparseMap | Tensor]
 """KeyedData is an interface where either a `Vector` or `Scalar` can be
 retrieved using a key which is of str type.
 """
 
-KeyedDataTypes = MutableMapping[str, type[Vector] | ScalarType | type[HealSparseMap]]
+KeyedDataTypes = MutableMapping[str, type[Vector] | ScalarType | type[HealSparseMap] | type[Tensor]]
 r"""A mapping of str keys to the Types which are valid in `KeyedData` objects.
 This is useful in conjunction with `AnalysisAction`\ 's ``getInputSchema`` and
 ``getOutputSchema`` methods.
 """
 
-KeyedDataSchema = Iterable[tuple[str, type[Vector] | ScalarType | type[HealSparseMap]]]
+KeyedDataSchema = Iterable[tuple[str, type[Vector] | ScalarType | type[HealSparseMap] | type[Tensor]]]
 r"""An interface that represents a type returned by `AnalysisAction`\ 's
 ``getInputSchema`` and ``getOutputSchema`` methods.
 """

python/lsst/analysis/tools/interfaces/_stages.py

@@ -26,6 +26,7 @@
 from typing import Any, cast
 
 import astropy.units as apu
+from healsparse import HealSparseMap
 from lsst.pex.config import ListField
 from lsst.pex.config.configurableActions import ConfigurableActionStructField
 from lsst.pex.config.dictField import DictField
@@ -38,6 +39,7 @@
     MetricAction,
     MetricResultType,
     NoPlot,
+    Tensor,
     VectorAction,
 )
 from ._interfaces import KeyedData, KeyedDataSchema, KeyedDataTypes, Scalar, Vector
@@ -46,19 +48,27 @@
 class BasePrep(KeyedDataAction):
     """Base class for actions which prepare data for processing."""
 
-    vectorKeys = ListField[str](doc="Keys to extract from KeyedData and return", default=[])
+    keysToLoad = ListField[str](doc="Keys to extract from KeyedData and return", default=[])
+
+    vectorKeys = ListField[str](doc="Keys from the input data which selectors will be applied", default=[])
 
     selectors = ConfigurableActionStructField[VectorAction](
         doc="Selectors for selecting rows, will be AND together",
     )
 
     def getInputSchema(self) -> KeyedDataSchema:
-        yield from ((column, Vector | Scalar) for column in self.vectorKeys)  # type: ignore
+        yield from (
+            (column, Vector | Scalar | HealSparseMap | Tensor)
+            for column in set(self.keysToLoad).union(self.vectorKeys)
+        )
         for action in self.selectors:
             yield from action.getInputSchema()
 
     def getOutputSchema(self) -> KeyedDataSchema:
-        return ((column, Vector | Scalar) for column in self.vectorKeys)  # type: ignore
+        return (
+            (column, Vector | Scalar | HealSparseMap | Tensor)
+            for column in set(self.keysToLoad).union(self.vectorKeys)
+        )
 
     def __call__(self, data: KeyedData, **kwargs) -> KeyedData:
         mask: Vector | None = None
@@ -69,16 +79,22 @@ def __call__(self, data: KeyedData, **kwargs) -> KeyedData:
             else:
                 mask *= subMask  # type: ignore
         result: dict[str, Any] = {}
-        for key in self.vectorKeys:
+        for key in set(self.keysToLoad).union(self.vectorKeys):
             formattedKey = key.format_map(kwargs)
             result[formattedKey] = cast(Vector, data[formattedKey])
         if mask is not None:
-            return {key: cast(Vector, col)[mask] for key, col in result.items()}
-        else:
-            return result
+            for key in self.vectorKeys:
+                # ignore type since there is not fully proper mypy support for
+                # vector type casting. In the future there will be, and this
+                # makes it clearer now what type things should be.
+                result[key] = cast(Vector, result[key])[mask]  # type: ignore
+        return result
 
     def addInputSchema(self, inputSchema: KeyedDataSchema) -> None:
-        self.vectorKeys = [name for name, _ in inputSchema]
+        existing = list(self.keysToLoad)
+        for name, _ in inputSchema:
+            existing.append(name)
+        self.keysToLoad = existing
 
 
 class BaseProcess(KeyedDataAction):

setup.cfg

@@ -1,7 +1,7 @@
 [flake8]
 max-line-length = 110
 max-doc-length = 79
-ignore = E133, E226, E228, N802, N803, N806, N812, N813, N815, N816, W503, E203
+ignore = E133, E226, E228, N802, N803, N806, N812, N813, N815, N816, W503, E203, E704
 exclude =
 	./bin,
 	doc/conf.py,