From 2f2e009506e20624472a25d2feaa1e7e784bd572 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 11:58:01 -0700 Subject: [PATCH 1/2] Add torch and lsdb dependencies and device helpers Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 2 ++ src/lsdb_ml/device.py | 77 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 src/lsdb_ml/device.py diff --git a/pyproject.toml b/pyproject.toml index 799fdb2..8f999d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,8 @@ classifiers = [ dynamic = ["version"] requires-python = ">=3.11" dependencies = [ + "lsdb>=0.11.0", # Brings nested-pandas, pyarrow and dask[complete] (incl. distributed) + "torch>=2.4", ] [project.urls] diff --git a/src/lsdb_ml/device.py b/src/lsdb_ml/device.py new file mode 100644 index 0000000..3dc0c55 --- /dev/null +++ b/src/lsdb_ml/device.py @@ -0,0 +1,77 @@ +"""Device resolution and batch movement helpers.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import torch + +DeviceLike = str | torch.device | None + + +def resolve_device(device: DeviceLike = "cpu") -> torch.device: + """Turn a device request into a concrete ``torch.device``. + + Parameters + ---------- + device : str, torch.device or None, default "cpu" + ``"cpu"`` (or ``None``) for the CPU, ``"auto"`` to pick CUDA if available, then + MPS, then CPU. Any other value (e.g. ``"cuda"``, ``"cuda:1"``, ``"mps"``, or a + ``torch.device``) is passed to ``torch.device`` and checked for availability. + + Returns + ------- + torch.device + The resolved device. + + Raises + ------ + RuntimeError + If an explicitly requested accelerator is not available. + + Examples + -------- + >>> resolve_device("cpu") + device(type='cpu') + >>> resolve_device(None) + device(type='cpu') + """ + if device is None: + return torch.device("cpu") + if isinstance(device, str) and device == "auto": + if torch.cuda.is_available(): + return torch.device("cuda") + if torch.backends.mps.is_available(): + return torch.device("mps") + return torch.device("cpu") + + resolved = torch.device(device) + if resolved.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"Device {resolved} was requested, but CUDA is not available.") + if resolved.type == "mps" and not torch.backends.mps.is_available(): + raise RuntimeError(f"Device {resolved} was requested, but MPS is not available.") + return resolved + + +def move_batch(batch: Mapping[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]: + """Move every tensor in a batch dict to ``device``. + + Non-tensor values are passed through unchanged. + + Parameters + ---------- + batch : Mapping[str, torch.Tensor] + A collated batch, as produced by :meth:`lsdb_ml.ResolvedSpec.collate`. + device : torch.device + The target device. + + Returns + ------- + dict[str, torch.Tensor] + A new dict with the tensors on ``device``. + """ + non_blocking = device.type == "cuda" + return { + key: value.to(device, non_blocking=non_blocking) if isinstance(value, torch.Tensor) else value + for key, value in batch.items() + } From e2668758c97080909f75ab7bdf782d6c14f6efc8 Mon Sep 17 00:00:00 2001 From: Doug Branton Date: Thu, 15 Oct 2026 11:58:01 -0700 Subject: [PATCH 2/2] Add ColumnSpec: meta validation and padded nested collation ColumnSpec maps flat and nested catalog columns to tensors. resolve() validates against catalog meta (nested sub-column dtypes come from NestedDtype.column_dtypes) without computing data, and reports every problem at once. Nested columns are padded to the batch maximum with a mask and lengths, fully vectorized via nest.len() and nest.to_flat(). Co-Authored-By: Claude Opus 5.5 --- src/lsdb_ml/__init__.py | 17 +- src/lsdb_ml/collate.py | 167 ++++++++++++ src/lsdb_ml/spec.py | 470 ++++++++++++++++++++++++++++++++++ tests/lsdb_ml/conftest.py | 57 +++++ tests/lsdb_ml/test_collate.py | 169 ++++++++++++ tests/lsdb_ml/test_device.py | 26 ++ tests/lsdb_ml/test_spec.py | 120 +++++++++ 7 files changed, 1025 insertions(+), 1 deletion(-) create mode 100644 src/lsdb_ml/collate.py create mode 100644 src/lsdb_ml/spec.py create mode 100644 tests/lsdb_ml/test_collate.py create mode 100644 tests/lsdb_ml/test_device.py create mode 100644 tests/lsdb_ml/test_spec.py diff --git a/src/lsdb_ml/__init__.py b/src/lsdb_ml/__init__.py index e42fde6..6502887 100644 --- a/src/lsdb_ml/__init__.py +++ b/src/lsdb_ml/__init__.py @@ -1,4 +1,19 @@ from ._version import __version__ +from .device import move_batch, resolve_device from .example_module import greetings, meaning +from .spec import ColumnSpec, Field, Flat, Nested, ResolvedSpec, SpecError, TensorInfo -__all__ = ["greetings", "meaning", "__version__"] +__all__ = [ + "ColumnSpec", + "Field", + "Flat", + "Nested", + "ResolvedSpec", + "SpecError", + "TensorInfo", + "move_batch", + "resolve_device", + "greetings", + "meaning", + "__version__", +] diff --git a/src/lsdb_ml/collate.py b/src/lsdb_ml/collate.py new file mode 100644 index 0000000..5c2b239 --- /dev/null +++ b/src/lsdb_ml/collate.py @@ -0,0 +1,167 @@ +"""Vectorized conversion of (nested) dataframe columns into padded numpy arrays. + +These helpers do the numeric work behind :class:`lsdb_ml.Flat` and :class:`lsdb_ml.Nested`. +Nothing here loops over rows in Python. +""" + +from __future__ import annotations + +from typing import Literal + +import numpy as np +import pandas as pd +import torch + +NullPolicy = Literal["raise"] | float +TruncatePolicy = Literal["raise", "head", "tail"] + + +def numpy_dtype(dtype: torch.dtype) -> np.dtype: + """Return the numpy dtype that corresponds to a torch dtype. + + Parameters + ---------- + dtype : torch.dtype + A torch dtype with a numpy equivalent. + + Returns + ------- + np.dtype + The matching numpy dtype. + + Raises + ------ + TypeError + If the torch dtype has no numpy equivalent (e.g. ``torch.bfloat16``). + + Examples + -------- + >>> numpy_dtype(torch.float32) + dtype('float32') + """ + try: + return torch.empty(0, dtype=dtype).numpy().dtype + except TypeError as err: + raise TypeError(f"torch dtype {dtype} has no numpy equivalent and cannot be collated.") from err + + +def frame_to_array(frame: pd.DataFrame, dtype: np.dtype, null: NullPolicy, what: str) -> np.ndarray: + """Convert the columns of ``frame`` into a writable ``(N, F)`` array. + + Parameters + ---------- + frame : pd.DataFrame + The columns to convert; may be arrow- or numpy-backed. + dtype : np.dtype + Output dtype. + null : "raise" or float + ``"raise"`` to fail on missing values, otherwise the value to fill them with. + what : str + Description of the data used in error messages. + + Returns + ------- + np.ndarray + A C-contiguous, writable array of shape ``(len(frame), len(frame.columns))``. + """ + missing = frame.isna().to_numpy() + if missing.any(): + if isinstance(null, str): + n_missing = int(missing.any(axis=1).sum()) + raise ValueError( + f"{what} contains {n_missing} row(s) with missing values. Pass null= " + "to fill them, or drop those rows before collation." + ) + frame = frame.fillna(null) + # np.array (not np.asarray) always copies: arrow-backed columns can convert to read-only + # views, which torch.from_numpy would warn about and share. + return np.array(frame.to_numpy(dtype=dtype), dtype=dtype, order="C", copy=True) + + +def pad_ragged( + values: np.ndarray, + lengths: np.ndarray, + *, + max_length: int | None = None, + truncate: TruncatePolicy = "raise", + pad_to: int | None = None, + pad_value: float = 0, + what: str = "nested column", +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Scatter flat ragged values into a padded ``(B, L, F)`` array with a mask. + + Parameters + ---------- + values : np.ndarray + Flat values of shape ``(sum(lengths), F)``, rows in order. + lengths : np.ndarray + Number of values in each of the ``B`` rows. + max_length : int or None, default None + Longest sequence to keep. Longer rows are handled by ``truncate``. + truncate : {"raise", "head", "tail"}, default "raise" + For rows over the limit: raise, keep the first values, or keep the last values. + pad_to : int or None, default None + Fixed padded length. When None, pad to the longest (kept) row in the batch. Rows + longer than ``pad_to`` are handled by ``truncate``. + pad_value : float, default 0 + Value written into padded positions. + what : str, default "nested column" + Description of the data used in error messages. + + Returns + ------- + padded : np.ndarray + Shape ``(B, L, F)``. + mask : np.ndarray + Boolean, shape ``(B, L)``; True where ``padded`` holds real data. + kept_lengths : np.ndarray + int64, shape ``(B,)``; number of real values per row after truncation. + + Examples + -------- + >>> padded, mask, kept = pad_ragged(np.arange(5.0)[:, None], np.array([2, 0, 3])) + >>> padded[..., 0] + array([[0., 1., 0.], + [0., 0., 0.], + [2., 3., 4.]]) + >>> mask.astype(int) + array([[1, 1, 0], + [0, 0, 0], + [1, 1, 1]]) + """ + lengths = np.asarray(lengths, dtype=np.int64) + n_rows = len(lengths) + n_features = values.shape[1] + if int(lengths.sum()) != len(values): + raise ValueError(f"{what}: lengths sum to {int(lengths.sum())} but there are {len(values)} values.") + + starts = np.cumsum(lengths) - lengths + limits = [limit for limit in (max_length, pad_to) if limit is not None] + limit = min(limits) if limits else None + + kept = lengths + source_starts = starts + if limit is not None and n_rows and int(lengths.max()) > limit: + if truncate == "raise": + raise ValueError( + f"{what} has rows of length up to {int(lengths.max())}, over the limit of {limit}. " + "Pass truncate='head' or truncate='tail' to cut them." + ) + kept = np.minimum(lengths, limit) + if truncate == "tail": + source_starts = starts + lengths - kept + elif truncate != "head": + raise ValueError(f"truncate must be 'raise', 'head' or 'tail', got {truncate!r}.") + + width = pad_to if pad_to is not None else (int(kept.max()) if n_rows else 0) + padded = np.full((n_rows, width, n_features), pad_value, dtype=values.dtype) + mask = np.zeros((n_rows, width), dtype=bool) + + total = int(kept.sum()) + if total: + row = np.repeat(np.arange(n_rows), kept) + position = np.arange(total) - np.repeat(np.cumsum(kept) - kept, kept) + source = np.repeat(source_starts, kept) + position + padded[row, position] = values[source] + mask[row, position] = True + return padded, mask, kept.astype(np.int64) diff --git a/src/lsdb_ml/spec.py b/src/lsdb_ml/spec.py new file mode 100644 index 0000000..fb76eda --- /dev/null +++ b/src/lsdb_ml/spec.py @@ -0,0 +1,470 @@ +"""Column specifications: map flat and nested catalog columns to tensors. + +A :class:`ColumnSpec` names the tensors a model needs and where each comes from. It is +checked against a catalog's ``meta`` (an empty ``NestedFrame`` that carries every dtype, +including nested sub-column dtypes) before any data is computed, and the resulting +:class:`ResolvedSpec` collates dataframe rows into a ``dict[str, torch.Tensor]``. + +>>> import lsdb_ml +>>> spec = lsdb_ml.ColumnSpec({"x": lsdb_ml.Flat(["ra", "dec"]), "lc": lsdb_ml.Nested("lc", ["flux"])}) +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass + +import numpy as np +import pandas as pd +import pyarrow as pa +import torch +from nested_pandas.series.dtype import NestedDtype + +from lsdb_ml.collate import NullPolicy, TruncatePolicy, frame_to_array, numpy_dtype, pad_ragged + + +class SpecError(ValueError): + """Raised when a column specification does not match a catalog's schema.""" + + +@dataclass(frozen=True) +class TensorInfo: + """Description of one tensor a resolved spec produces. + + Attributes + ---------- + dtype : torch.dtype + The tensor dtype. + shape : tuple + The shape; ``None`` marks a dimension that varies by batch (always the first one). + """ + + dtype: torch.dtype + shape: tuple[int | None, ...] + + +def _as_list(columns: str | Sequence[str]) -> list[str]: + return [columns] if isinstance(columns, str) else list(columns) + + +def _is_numeric(dtype) -> bool: + """True for pandas/numpy/pyarrow integer, float, and bool dtypes.""" + if isinstance(dtype, pa.DataType): + return pa.types.is_integer(dtype) or pa.types.is_floating(dtype) or pa.types.is_boolean(dtype) + if isinstance(dtype, pd.ArrowDtype): + return _is_numeric(dtype.pyarrow_dtype) + if isinstance(dtype, NestedDtype): + return False + try: + return pd.api.types.is_numeric_dtype(dtype) or pd.api.types.is_bool_dtype(dtype) + except TypeError: + return False + + +class Field(ABC): + """Base class for one entry of a :class:`ColumnSpec`. + + Subclasses implement :meth:`resolve`, which validates the field against catalog meta + and returns a :class:`ResolvedField` that does the actual collation. New column + representations (e.g. ragged tensors) are added as new ``Field`` types. + """ + + @abstractmethod + def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None: + """Validate against ``meta`` and build the collator. + + Parameters + ---------- + key : str + The output key this field was registered under. + meta : pd.DataFrame + The catalog meta (an empty frame with full dtypes). + problems : list[str] + Append a message for every problem found, instead of raising, so that all + problems in a spec are reported together. + + Returns + ------- + ResolvedField or None + None if any problem was found. + """ + + +class ResolvedField(ABC): + """A field validated against a schema, ready to collate frames.""" + + #: Top-level columns this field reads. + columns: list[str] + + @abstractmethod + def describe(self) -> dict[str, TensorInfo]: + """Return the tensors this field produces, keyed by output key.""" + + @abstractmethod + def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]: + """Convert the rows of ``frame`` to tensors.""" + + +@dataclass +class Flat(Field): + """One or more flat (non-nested) columns, stacked into a single tensor. + + Parameters + ---------- + columns : str or sequence of str + A single column name produces a ``(B,)`` tensor; a list produces ``(B, F)``. + dtype : torch.dtype, default torch.float32 + Output dtype. + null : "raise" or float, default "raise" + Missing values raise an error, or are replaced by this fill value. + transform : callable, optional + ``transform(frame[columns]) -> array-like`` of shape ``(B,)`` or ``(B, F)``. Use it + for non-numeric columns or custom encodings. When given, the dtype check is skipped + and the result is cast to ``dtype``. It must be picklable to be used from Dask + tasks or DataLoader workers. + """ + + columns: str | Sequence[str] + dtype: torch.dtype = torch.float32 + null: NullPolicy = "raise" + transform: Callable[[pd.DataFrame], np.ndarray] | None = None + + def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None: + """Validate the columns against ``meta``; see :meth:`Field.resolve`.""" + columns = _as_list(self.columns) + start = len(problems) + if not columns: + problems.append(f"'{key}': Flat needs at least one column.") + for column in columns: + if column not in meta.columns: + problems.append(f"'{key}': column '{column}' is not in the catalog.") + elif isinstance(meta.dtypes[column], NestedDtype): + problems.append(f"'{key}': column '{column}' is nested; use Nested('{column}', [...]).") + elif self.transform is None and not _is_numeric(meta.dtypes[column]): + problems.append( + f"'{key}': column '{column}' has non-numeric dtype {meta.dtypes[column]}; " + "pass transform= to encode it." + ) + try: + np_dtype = numpy_dtype(self.dtype) + except TypeError as err: + problems.append(f"'{key}': {err}") + if len(problems) > start: + return None + return _ResolvedFlat(key, columns, isinstance(self.columns, str), self, np_dtype) + + +class _ResolvedFlat(ResolvedField): + def __init__(self, key: str, columns: list[str], scalar: bool, spec: Flat, np_dtype: np.dtype): + self.key = key + self.columns = columns + self.scalar = scalar + self.spec = spec + self.np_dtype = np_dtype + + def describe(self) -> dict[str, TensorInfo]: + if self.spec.transform is not None: + shape: tuple[int | None, ...] = (None,) if self.scalar else (None, None) + else: + shape = (None,) if self.scalar else (None, len(self.columns)) + return {self.key: TensorInfo(self.spec.dtype, shape)} + + def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]: + what = f"'{self.key}' (columns {self.columns})" + if self.spec.transform is not None: + array = np.array(self.spec.transform(frame[self.columns]), dtype=self.np_dtype, copy=True) + if len(array) != len(frame): + raise ValueError(f"{what}: transform returned {len(array)} rows for {len(frame)} input rows.") + else: + array = frame_to_array(frame[self.columns], self.np_dtype, self.spec.null, what) + if self.scalar: + array = array[:, 0].copy() + return {self.key: torch.from_numpy(array)} + + +@dataclass +class Nested(Field): + """Sub-columns of a nested column, padded to a rectangular tensor with a mask. + + For output key ``k`` this produces: + + - ``k``: values, ``(B, L)`` for a single sub-column or ``(B, L, F)`` for a list; + - ``k_mask``: bool ``(B, L)``, True where a value is real rather than padding; + - ``k_lengths``: int64 ``(B,)``, the number of real values per row. + + ``L`` is the longest row in the batch (after truncation), or ``pad_to`` if set. Rows + whose nested value is null or empty have length 0. The padding value carries no + meaning; always use the mask. + + Parameters + ---------- + column : str + The nested column. + fields : str or sequence of str + Sub-column name(s) to read. + dtype : torch.dtype, default torch.float32 + Output dtype of the values. + max_length : int, optional + Longest sequence to keep; see ``truncate``. + truncate : {"raise", "head", "tail"}, default "raise" + What to do with rows longer than ``max_length`` (or ``pad_to``): raise, keep the + first values, or keep the last values. + pad_to : int, optional + Pad every batch to this fixed length instead of the batch maximum. + pad_value : float, default 0 + Value written into padded positions. + null : "raise" or float, default "raise" + Missing values *inside* a list raise an error, or are replaced by this value. + transform : callable, optional + ``transform(flat_values) -> array-like`` with one row per flat value, of shape + ``(N,)`` or ``(N, F)``, where ``flat_values`` is a DataFrame of the requested + sub-columns. When given, the dtype check is skipped. + """ + + column: str + fields: str | Sequence[str] + dtype: torch.dtype = torch.float32 + max_length: int | None = None + truncate: TruncatePolicy = "raise" + pad_to: int | None = None + pad_value: float = 0 + null: NullPolicy = "raise" + transform: Callable[[pd.DataFrame], np.ndarray] | None = None + + def resolve(self, key: str, meta: pd.DataFrame, problems: list[str]) -> ResolvedField | None: + """Validate the nested column and sub-columns against ``meta``; see :meth:`Field.resolve`.""" + subcolumns = _as_list(self.fields) + start = len(problems) + if not subcolumns: + problems.append(f"'{key}': Nested needs at least one sub-column.") + if self.truncate not in ("raise", "head", "tail"): + problems.append(f"'{key}': truncate must be 'raise', 'head' or 'tail', got {self.truncate!r}.") + for name, limit in (("max_length", self.max_length), ("pad_to", self.pad_to)): + if limit is not None and (isinstance(limit, bool) or not isinstance(limit, int) or limit < 1): + problems.append(f"'{key}': {name} must be a positive integer, got {limit!r}.") + if self.column not in meta.columns: + problems.append(f"'{key}': column '{self.column}' is not in the catalog.") + elif not isinstance(meta.dtypes[self.column], NestedDtype): + problems.append(f"'{key}': column '{self.column}' is not nested; use Flat('{self.column}').") + else: + available = meta.dtypes[self.column].column_dtypes + for sub in subcolumns: + if sub not in available: + problems.append( + f"'{key}': nested column '{self.column}' has no sub-column '{sub}' " + f"(available: {sorted(available)})." + ) + elif self.transform is None and not _is_numeric(available[sub]): + problems.append( + f"'{key}': sub-column '{self.column}.{sub}' has non-numeric type {available[sub]}; " + "pass transform= to encode it." + ) + try: + np_dtype = numpy_dtype(self.dtype) + except TypeError as err: + problems.append(f"'{key}': {err}") + if len(problems) > start: + return None + return _ResolvedNested(key, subcolumns, isinstance(self.fields, str), self, np_dtype) + + +class _ResolvedNested(ResolvedField): + def __init__(self, key: str, subcolumns: list[str], scalar: bool, spec: Nested, np_dtype: np.dtype): + self.key = key + self.columns = [spec.column] + self.subcolumns = subcolumns + self.scalar = scalar + self.spec = spec + self.np_dtype = np_dtype + + def describe(self) -> dict[str, TensorInfo]: + length = self.spec.pad_to + if self.scalar: + values_shape: tuple[int | None, ...] = (None, length) + else: + n_features = None if self.spec.transform is not None else len(self.subcolumns) + values_shape = (None, length, n_features) + return { + self.key: TensorInfo(self.spec.dtype, values_shape), + f"{self.key}_mask": TensorInfo(torch.bool, (None, length)), + f"{self.key}_lengths": TensorInfo(torch.int64, (None,)), + } + + def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]: + what = f"'{self.key}' (nested column '{self.spec.column}', fields {self.subcolumns})" + series = frame[self.spec.column] + # nest.len() counts null rows as 0, and to_flat() returns the values of all rows + # concatenated in row order; both are consistent across iloc slices and concat. + lengths = np.asarray(series.nest.len(), dtype=np.int64) + flat = series.nest.to_flat(self.subcolumns) + if self.spec.transform is not None: + values = np.array(self.spec.transform(flat), dtype=self.np_dtype, copy=True) + if len(values) != len(flat): + raise ValueError(f"{what}: transform returned {len(values)} rows for {len(flat)} values.") + if values.ndim == 1: + values = values[:, None] + else: + values = frame_to_array(flat, self.np_dtype, self.spec.null, what) + padded, mask, kept = pad_ragged( + values, + lengths, + max_length=self.spec.max_length, + truncate=self.spec.truncate, + pad_to=self.spec.pad_to, + pad_value=self.spec.pad_value, + what=what, + ) + if self.scalar and padded.shape[2] == 1: + padded = padded[..., 0] + return { + self.key: torch.from_numpy(np.ascontiguousarray(padded)), + f"{self.key}_mask": torch.from_numpy(mask), + f"{self.key}_lengths": torch.from_numpy(kept), + } + + +@dataclass +class ColumnSpec: + """Mapping from output keys to catalog columns. + + Parameters + ---------- + fields : Mapping[str, Field] + Output key to field, e.g. ``{"x": Flat(["ra", "dec"]), "lc": Nested("lc", "flux")}``. + index_key : str or None, default "_index" + Every batch also includes the frame index (``_healpix_29`` for LSDB catalogs) as an + int64 ``(B,)`` tensor under this key, so batches can be traced back to rows. None + disables it. + + Examples + -------- + >>> import nested_pandas as npd + >>> nf = npd.NestedFrame({"a": [1.0, 2.0]}, index=[10, 11]) + >>> nf = nf.join_nested(npd.NestedFrame({"f": [1.0, 2.0, 3.0]}, index=[10, 10, 11]), "lc") + >>> resolved = ColumnSpec({"a": Flat("a"), "lc": Nested("lc", "f")}).resolve(nf) + >>> batch = resolved.collate(nf) + >>> batch["lc"] + tensor([[1., 2.], + [3., 0.]]) + >>> batch["lc_mask"] + tensor([[ True, True], + [ True, False]]) + >>> batch["_index"] + tensor([10, 11]) + """ + + fields: Mapping[str, Field] + index_key: str | None = "_index" + + def resolve(self, schema) -> ResolvedSpec: + """Validate the spec against a catalog's schema without computing any data. + + Parameters + ---------- + schema : lsdb.Catalog or pd.DataFrame + A catalog (its ``meta`` is used) or any dataframe with the same columns and + dtypes, such as ``catalog.meta`` or a computed partition. + + Returns + ------- + ResolvedSpec + The validated spec, which can collate frames. + + Raises + ------ + SpecError + Listing every problem found. + """ + meta = schema.meta if hasattr(schema, "meta") and not isinstance(schema, pd.DataFrame) else schema + if not isinstance(meta, pd.DataFrame): + raise TypeError(f"Expected an lsdb catalog or a DataFrame, got {type(schema).__name__}.") + problems: list[str] = [] + if not self.fields: + problems.append("The spec has no fields.") + resolved: dict[str, ResolvedField] = {} + for key, spec_field in self.fields.items(): + if not isinstance(spec_field, Field): + problems.append(f"'{key}': expected a Field such as Flat or Nested, got {type(spec_field)}.") + continue + result = spec_field.resolve(key, meta, problems) + if result is not None: + resolved[key] = result + + seen: dict[str, str] = {} + if self.index_key is not None: + seen[self.index_key] = "the index" + if not pd.api.types.is_integer_dtype(meta.index.dtype): + problems.append( + f"index_key={self.index_key!r} needs an integer index, but the index has dtype " + f"{meta.index.dtype}; pass index_key=None to disable it." + ) + for key, result in resolved.items(): + for output_key in result.describe(): + if output_key in seen: + problems.append( + f"Output key '{output_key}' from '{key}' collides with {seen[output_key]}." + ) + seen[output_key] = f"'{key}'" + + if problems: + hint = "" + missing = [p for p in problems if "is not in the catalog" in p] + all_columns = getattr(schema, "all_columns", None) + if missing and all_columns is not None: + hint = ( + "\nSome columns may exist on disk but were not loaded; the catalog was opened " + f"with columns {list(meta.columns)}." + ) + raise SpecError("Column spec does not match the catalog:\n - " + "\n - ".join(problems) + hint) + return ResolvedSpec(resolved, self.index_key) + + +class ResolvedSpec: + """A :class:`ColumnSpec` validated against a schema; collates frames into tensors. + + Create it with :meth:`ColumnSpec.resolve`. It is picklable as long as any + ``transform`` callables are, so it can be shipped to Dask tasks and DataLoader workers. + """ + + def __init__(self, fields: dict[str, ResolvedField], index_key: str | None): + self.fields = fields + self.index_key = index_key + + @property + def columns(self) -> list[str]: + """Top-level catalog columns needed for collation, in first-use order.""" + columns: list[str] = [] + for resolved in self.fields.values(): + columns.extend(c for c in resolved.columns if c not in columns) + return columns + + def describe(self) -> dict[str, TensorInfo]: + """Return the dtype and shape of every tensor a batch contains, keyed by output key.""" + info: dict[str, TensorInfo] = {} + for resolved in self.fields.values(): + info.update(resolved.describe()) + if self.index_key is not None: + info[self.index_key] = TensorInfo(torch.int64, (None,)) + return info + + def collate(self, frame: pd.DataFrame) -> dict[str, torch.Tensor]: + """Convert the rows of ``frame`` to a batch of tensors. + + Parameters + ---------- + frame : pd.DataFrame + Rows to convert, e.g. a slice of a catalog partition. Row order is preserved: + row ``i`` of every output tensor comes from row ``i`` of ``frame``. + + Returns + ------- + dict[str, torch.Tensor] + One entry per key in :meth:`describe`. + """ + batch: dict[str, torch.Tensor] = {} + for resolved in self.fields.values(): + batch.update(resolved.collate(frame)) + if self.index_key is not None: + batch[self.index_key] = torch.from_numpy(np.array(frame.index.to_numpy(), dtype=np.int64)) + return batch diff --git a/tests/lsdb_ml/conftest.py b/tests/lsdb_ml/conftest.py index e69de29..ec295e2 100644 --- a/tests/lsdb_ml/conftest.py +++ b/tests/lsdb_ml/conftest.py @@ -0,0 +1,57 @@ +import lsdb +import nested_pandas as npd +import numpy as np +import pandas as pd +import pytest + + +def make_nested_frame(n_rows: int = 500, seed: int = 0, max_length: int = 6) -> npd.NestedFrame: + """Synthetic object table with a ragged nested light curve column. + + Light curve lengths range from 0 to ``max_length`` - 1, so some rows have empty + (null) nested values. + """ + rng = np.random.default_rng(seed) + base = pd.DataFrame( + { + "id": np.arange(n_rows, dtype=np.int64), + "ra": rng.uniform(0, 360, n_rows), + "dec": rng.uniform(-60, 60, n_rows), + "mag": rng.normal(20, 1, n_rows).astype(np.float32), + "flag": rng.integers(0, 2, n_rows).astype(bool), + "name": [f"obj{i}" for i in range(n_rows)], + } + ) + lengths = rng.integers(0, max_length, n_rows) + n_flat = int(lengths.sum()) + flat = pd.DataFrame( + { + "t": rng.uniform(0, 100, n_flat), + "flux": rng.normal(0, 1, n_flat).astype(np.float32), + "band": rng.choice(["g", "r"], n_flat), + }, + index=np.repeat(np.arange(n_rows), lengths), + ) + return npd.NestedFrame(base).join_nested(flat, "lc") + + +@pytest.fixture(scope="session") +def nested_frame() -> npd.NestedFrame: + """The in-memory frame behind ``nested_catalog``, indexed by row number.""" + return make_nested_frame() + + +@pytest.fixture(scope="session") +def nested_catalog(nested_frame) -> lsdb.Catalog: + """A lazy catalog with many small partitions and a ragged nested column.""" + return lsdb.from_dataframe(nested_frame, partition_rows=40, catalog_name="toy") + + +@pytest.fixture +def unloadable_catalog(nested_catalog) -> lsdb.Catalog: + """Same schema as ``nested_catalog``, but computing any partition raises.""" + + def fail(_): + raise AssertionError("A partition was computed, but no data should be read here.") + + return nested_catalog.map_partitions(fail, meta=nested_catalog.meta) diff --git a/tests/lsdb_ml/test_collate.py b/tests/lsdb_ml/test_collate.py new file mode 100644 index 0000000..2369eb2 --- /dev/null +++ b/tests/lsdb_ml/test_collate.py @@ -0,0 +1,169 @@ +import nested_pandas as npd +import numpy as np +import pandas as pd +import pytest +import torch + +from lsdb_ml import ColumnSpec, Flat, Nested +from lsdb_ml.collate import pad_ragged + + +def reference_rows(frame, column, fields): + """Per-row python reference: a list of (length, n_fields) arrays.""" + rows = [] + for value in frame[column].array: + if value is None or not isinstance(value, pd.DataFrame): + rows.append(np.zeros((0, len(fields)))) + else: + rows.append(value[fields].to_numpy(dtype=np.float64)) + return rows + + +def reference_pad(rows, width, keep="head"): + """Pad reference rows to ``width`` the slow way.""" + padded = np.zeros((len(rows), width, rows[0].shape[1] if rows else 0)) + mask = np.zeros((len(rows), width), dtype=bool) + for i, row in enumerate(rows): + row = row[:width] if keep == "head" else row[max(len(row) - width, 0) :] + padded[i, : len(row)] = row + mask[i, : len(row)] = True + return padded, mask + + +@pytest.fixture +def frame(nested_catalog): + """A frame mixing a sliced partition, a whole partition, and a concatenation.""" + # A computed partition, plus a slice and a concatenation, all exercise the + # flat-values/offsets bookkeeping on non-trivial arrow layouts. + parts = [nested_catalog.partitions[i].compute(progress_bar=False) for i in range(3)] + return pd.concat([parts[0].iloc[3:], parts[1], parts[2].iloc[:5]]) + + +def test_flat_columns(frame): + """Flat columns stack into (B,) or (B, F) tensors with the requested dtype.""" + resolved = ColumnSpec( + {"x": Flat(["ra", "dec", "mag"]), "mag": Flat("mag"), "flag": Flat("flag", dtype=torch.bool)} + ).resolve(frame) + batch = resolved.collate(frame) + assert batch["x"].shape == (len(frame), 3) and batch["x"].dtype == torch.float32 + np.testing.assert_allclose(batch["x"].numpy(), frame[["ra", "dec", "mag"]].to_numpy(dtype=np.float32)) + assert batch["mag"].shape == (len(frame),) + assert batch["flag"].dtype == torch.bool + np.testing.assert_array_equal(batch["flag"].numpy(), frame["flag"].to_numpy(dtype=bool)) + np.testing.assert_array_equal(batch["_index"].numpy(), frame.index.to_numpy()) + assert all(t.numpy().flags.writeable for t in batch.values()) + + +@pytest.mark.parametrize("fields", [["t", "flux"], ["flux"]]) +def test_nested_matches_reference(frame, fields): + """Vectorized nested padding matches a per-row reference, including empty rows.""" + resolved = ColumnSpec({"lc": Nested("lc", fields, dtype=torch.float64)}).resolve(frame) + batch = resolved.collate(frame) + rows = reference_rows(frame, "lc", fields) + lengths = np.array([len(r) for r in rows]) + assert (lengths == 0).any(), "fixture should include empty / null nested rows" + expected, expected_mask = reference_pad(rows, lengths.max()) + np.testing.assert_allclose(batch["lc"].numpy(), expected) + np.testing.assert_array_equal(batch["lc_mask"].numpy(), expected_mask) + np.testing.assert_array_equal(batch["lc_lengths"].numpy(), lengths) + + +def test_nested_single_field_is_2d(frame): + """A single nested sub-column produces (B, L) values.""" + batch = ColumnSpec({"f": Nested("lc", "flux")}).resolve(frame).collate(frame) + assert batch["f"].ndim == 2 + assert batch["f"].shape == batch["f_mask"].shape + + +@pytest.mark.parametrize("keep", ["head", "tail"]) +def test_truncation(frame, keep): + """Rows over max_length keep their first or last values.""" + resolved = ColumnSpec({"lc": Nested("lc", ["t"], max_length=2, truncate=keep, dtype=torch.float64)}) + batch = resolved.resolve(frame).collate(frame) + rows = reference_rows(frame, "lc", ["t"]) + expected, expected_mask = reference_pad(rows, 2, keep=keep) + np.testing.assert_allclose(batch["lc"].numpy(), expected) + np.testing.assert_array_equal(batch["lc_mask"].numpy(), expected_mask) + np.testing.assert_array_equal(batch["lc_lengths"].numpy(), np.minimum([len(r) for r in rows], 2)) + + +def test_truncation_raises_by_default(frame): + """Rows over max_length raise unless a truncate policy is given.""" + with pytest.raises(ValueError, match="over the limit of 2"): + ColumnSpec({"lc": Nested("lc", "t", max_length=2)}).resolve(frame).collate(frame) + + +def test_pad_to_fixed_width(frame): + """pad_to fixes the padded width and fills with pad_value.""" + batch = ColumnSpec({"lc": Nested("lc", "t", pad_to=9, pad_value=-1)}).resolve(frame).collate(frame) + assert batch["lc"].shape == (len(frame), 9) + assert (batch["lc"][~batch["lc_mask"]] == -1).all() + + +def test_transforms(frame): + """Transforms encode non-numeric flat and nested columns.""" + resolved = ColumnSpec( + { + "name_len": Flat("name", transform=lambda df: df["name"].str.len().to_numpy()), + "is_g": Nested("lc", "band", transform=lambda df: (df["band"] == "g").to_numpy()), + } + ).resolve(frame) + batch = resolved.collate(frame) + np.testing.assert_array_equal(batch["name_len"].numpy(), frame["name"].str.len().to_numpy()) + first = next(i for i, v in enumerate(frame["lc"].array) if isinstance(v, pd.DataFrame)) + expected = (frame["lc"].array[first]["band"] == "g").to_numpy() + np.testing.assert_array_equal(batch["is_g"][first, : len(expected)].numpy(), expected) + + +def test_empty_frame(frame): + """Zero-row frames collate to zero-row tensors.""" + resolved = ColumnSpec({"x": Flat(["ra", "dec"]), "lc": Nested("lc", ["t", "flux"])}).resolve(frame) + batch = resolved.collate(frame.iloc[:0]) + assert batch["x"].shape == (0, 2) + assert batch["lc"].shape == (0, 0, 2) + assert batch["lc_mask"].shape == (0, 0) + assert batch["_index"].shape == (0,) + + +def test_all_rows_empty_nested(frame): + """A batch where every nested row is empty pads to width 0.""" + lengths = np.asarray(frame["lc"].nest.len()) + empty_rows = frame[lengths == 0] + batch = ColumnSpec({"lc": Nested("lc", "flux")}).resolve(frame).collate(empty_rows) + assert batch["lc"].shape == (len(empty_rows), 0) + assert not batch["lc_mask"].any() + + +def test_null_values(): + """Missing values raise by default and are filled when a fill value is given.""" + nf = npd.NestedFrame({"a": [1.0, None, 3.0]}, index=[0, 1, 2]) + nf = nf.join_nested(npd.NestedFrame({"f": [1.0, None, 2.0]}, index=[0, 0, 2]), "lc") + with pytest.raises(ValueError, match="1 row\\(s\\) with missing values"): + ColumnSpec({"a": Flat("a")}).resolve(nf).collate(nf) + with pytest.raises(ValueError, match="missing values"): + ColumnSpec({"lc": Nested("lc", "f")}).resolve(nf).collate(nf) + batch = ColumnSpec({"a": Flat("a", null=-9.0), "lc": Nested("lc", "f", null=0.5)}).resolve(nf).collate(nf) + np.testing.assert_array_equal(batch["a"].numpy(), [1, -9, 3]) + np.testing.assert_array_equal(batch["lc"].numpy(), [[1, 0.5], [0, 0], [2, 0]]) + np.testing.assert_array_equal(batch["lc_lengths"].numpy(), [2, 0, 1]) + + +def test_transform_length_mismatch(frame): + """A transform that changes the row count is rejected.""" + resolved = ColumnSpec({"x": Flat("ra", transform=lambda df: [1.0])}).resolve(frame) + with pytest.raises(ValueError, match="transform returned 1 rows"): + resolved.collate(frame) + + +def test_pad_ragged_length_mismatch(): + """Inconsistent lengths and values are rejected.""" + with pytest.raises(ValueError, match="lengths sum to 3"): + pad_ragged(np.zeros((2, 1)), np.array([1, 2])) + + +def test_row_order_preserved(frame): + """Row i of every tensor comes from row i of the frame, even for a shuffled frame.""" + shuffled = frame.sample(frac=1, random_state=1) + batch = ColumnSpec({"id": Flat("id", dtype=torch.int64)}).resolve(shuffled).collate(shuffled) + np.testing.assert_array_equal(batch["id"].numpy(), shuffled["id"].to_numpy()) + np.testing.assert_array_equal(batch["_index"].numpy(), shuffled.index.to_numpy()) diff --git a/tests/lsdb_ml/test_device.py b/tests/lsdb_ml/test_device.py new file mode 100644 index 0000000..cb90019 --- /dev/null +++ b/tests/lsdb_ml/test_device.py @@ -0,0 +1,26 @@ +import pytest +import torch + +from lsdb_ml import move_batch, resolve_device + + +def test_resolve_device(): + """Known device requests resolve to torch devices.""" + assert resolve_device("cpu") == torch.device("cpu") + assert resolve_device(torch.device("cpu")) == torch.device("cpu") + assert resolve_device("auto").type in {"cpu", "cuda", "mps"} + + +@pytest.mark.skipif(torch.cuda.is_available(), reason="checks the error when CUDA is missing") +def test_resolve_missing_cuda(): + """Requesting CUDA without CUDA fails clearly.""" + with pytest.raises(RuntimeError, match="CUDA is not available"): + resolve_device("cuda") + + +def test_move_batch(): + """Tensors move to the device and other values pass through.""" + batch = {"x": torch.zeros(2), "meta": "keep"} + moved = move_batch(batch, torch.device("cpu")) + assert moved["x"].device.type == "cpu" + assert moved["meta"] == "keep" diff --git a/tests/lsdb_ml/test_spec.py b/tests/lsdb_ml/test_spec.py new file mode 100644 index 0000000..73e2896 --- /dev/null +++ b/tests/lsdb_ml/test_spec.py @@ -0,0 +1,120 @@ +import pytest +import torch + +from lsdb_ml import ColumnSpec, Flat, Nested, SpecError, TensorInfo + + +def test_resolve_valid_spec_from_catalog_without_compute(unloadable_catalog): + """A valid spec resolves against a catalog without computing any partition.""" + spec = ColumnSpec( + { + "x": Flat(["ra", "dec", "mag"]), + "flag": Flat("flag", dtype=torch.bool), + "lc": Nested("lc", ["t", "flux"], max_length=4, truncate="tail"), + "flux": Nested("lc", "flux", pad_to=8, truncate="head"), + } + ) + resolved = spec.resolve(unloadable_catalog) + assert resolved.columns == ["ra", "dec", "mag", "flag", "lc"] + assert resolved.describe() == { + "x": TensorInfo(torch.float32, (None, 3)), + "flag": TensorInfo(torch.bool, (None,)), + "lc": TensorInfo(torch.float32, (None, None, 2)), + "lc_mask": TensorInfo(torch.bool, (None, None)), + "lc_lengths": TensorInfo(torch.int64, (None,)), + "flux": TensorInfo(torch.float32, (None, 8)), + "flux_mask": TensorInfo(torch.bool, (None, 8)), + "flux_lengths": TensorInfo(torch.int64, (None,)), + "_index": TensorInfo(torch.int64, (None,)), + } + + +def test_unloadable_catalog_fixture_raises(unloadable_catalog): + """Guard: the no-compute fixture really fails when a partition is computed.""" + with pytest.raises(Exception, match="no data should be read"): + unloadable_catalog.partitions[0].compute(progress_bar=False) + + +def test_resolve_accepts_meta_frame(nested_catalog): + """resolve() also accepts the meta frame directly.""" + resolved = ColumnSpec({"x": Flat("ra")}).resolve(nested_catalog.meta) + assert resolved.columns == ["ra"] + + +def test_all_problems_reported_together(unloadable_catalog): + """Every problem in a spec is reported in one error.""" + spec = ColumnSpec( + { + "missing": Flat("nope"), + "nested_as_flat": Flat("lc"), + "string": Flat("name"), + "flat_as_nested": Nested("ra", "x"), + "missing_sub": Nested("lc", ["flux", "nope"]), + "string_sub": Nested("lc", "band"), + "bad_trunc": Nested("lc", "flux", truncate="middle"), + "bad_len": Nested("lc", "flux", max_length=0), + } + ) + with pytest.raises(SpecError) as err: + spec.resolve(unloadable_catalog) + message = str(err.value) + for expected in [ + "column 'nope' is not in the catalog", + "column 'lc' is nested", + "'name' has non-numeric dtype", + "column 'ra' is not nested", + "has no sub-column 'nope'", + "'lc.band' has non-numeric type", + "truncate must be", + "max_length must be a positive integer", + ]: + assert expected in message + assert "not loaded" in message # hint for missing columns on a catalog + + +def test_transform_allows_non_numeric(unloadable_catalog): + """A transform lifts the numeric dtype requirement.""" + spec = ColumnSpec( + { + "name_len": Flat("name", transform=lambda df: df["name"].str.len()), + "band": Nested("lc", "band", transform=lambda df: (df["band"] == "g").to_numpy()), + } + ) + spec.resolve(unloadable_catalog) + + +def test_key_collisions(unloadable_catalog): + """Output keys may not collide with each other or the index key.""" + with pytest.raises(SpecError, match="'lc_mask' from 'lc' collides"): + ColumnSpec({"lc_mask": Flat("ra"), "lc": Nested("lc", "flux")}).resolve(unloadable_catalog) + with pytest.raises(SpecError, match="'_index' from '_index' collides with the index"): + ColumnSpec({"_index": Flat("ra")}).resolve(unloadable_catalog) + ColumnSpec({"_index": Flat("ra")}, index_key=None).resolve(unloadable_catalog) + + +def test_non_field_and_empty_spec(unloadable_catalog): + """Non-Field values and empty specs are rejected.""" + with pytest.raises(SpecError, match="expected a Field"): + ColumnSpec({"x": "ra"}).resolve(unloadable_catalog) + with pytest.raises(SpecError, match="no fields"): + ColumnSpec({}).resolve(unloadable_catalog) + + +def test_unsupported_dtype(unloadable_catalog): + """Torch dtypes without a numpy equivalent are rejected.""" + with pytest.raises(SpecError, match="no numpy equivalent"): + ColumnSpec({"x": Flat("ra", dtype=torch.bfloat16)}).resolve(unloadable_catalog) + + +def test_non_integer_index(nested_frame): + """The index key requires an integer index, and can be disabled.""" + frame = nested_frame.reset_index(drop=True).set_index("name") + with pytest.raises(SpecError, match="needs an integer index"): + ColumnSpec({"x": Flat("ra")}).resolve(frame) + ColumnSpec({"x": Flat("ra")}, index_key=None).resolve(frame) + + +def test_bad_schema_type(): + """resolve() rejects objects that are neither catalogs nor frames.""" + with pytest.raises(TypeError, match="Expected an lsdb catalog"): + ColumnSpec({"x": Flat("ra")}).resolve(42)