[Claude-led] Feature/inference

docs/examples/predict_on_catalog.py

@@ -0,0 +1,80 @@
+"""Run a PyTorch model over every row of an LSDB catalog on Dask workers.
+
+Run with ``python docs/examples/predict_on_catalog.py``. It uses a synthetic catalog;
+for a real one, replace ``make_catalog()`` with ``lsdb.open_catalog(...)``.
+
+The model is built once per Dask worker (by a worker plugin) and the result is a lazy
+catalog with prediction columns, aligned row by row with the input.
+"""
+
+from pathlib import Path
+
+import lsdb
+import nested_pandas as npd
+import numpy as np
+import pandas as pd
+import torch
+from distributed import Client, LocalCluster
+from torch import nn
+
+from lsdb_ml import ColumnSpec, Flat, ModelDeployment, ModelSpec, Nested
+
+
+def make_catalog(n_objects=5000, seed=0):
+    """Objects with a magnitude and a ragged light curve."""
+    rng = np.random.default_rng(seed)
+    lengths = rng.integers(0, 40, n_objects)
+    objects = pd.DataFrame(
+        {
+            "object_id": np.arange(n_objects),
+            "ra": rng.uniform(0, 360, n_objects),
+            "dec": rng.uniform(-60, 60, n_objects),
+            "mag": rng.normal(20, 1, n_objects).astype(np.float32),
+        }
+    )
+    light_curves = pd.DataFrame(
+        {"mjd": rng.uniform(0, 100, lengths.sum()), "flux": rng.normal(1, 0.1, lengths.sum())},
+        index=np.repeat(np.arange(n_objects), lengths),
+    )
+    return lsdb.from_dataframe(npd.NestedFrame(objects).join_nested(light_curves, "lc"), partition_rows=500)
+
+
+class Scorer(nn.Module):
+    """Toy model combining a flat column with a masked summary of a nested column."""
+
+    def __init__(self):
+        super().__init__()
+        self.head = nn.Linear(2, 2)
+
+    def forward(self, mag, flux, mask):
+        """Return two scores per object from the magnitude and the mean light-curve flux."""
+        mean_flux = (flux * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+        return self.head(torch.stack([mag, mean_flux], dim=1))
+
+
+def main(weights_path):
+    """Save some weights, deploy the model and predict on the catalog."""
+    torch.manual_seed(0)
+    torch.save(Scorer().state_dict(), weights_path)  # stands in for a trained model
+
+    catalog = make_catalog().query("mag < 21")  # any lazy lsdb pipeline works here
+    spec = ColumnSpec({"mag": Flat("mag"), "flux": Nested("lc", "flux", max_length=32, truncate="tail")})
+
+    # Workers build Scorer() and load the weights themselves; only the spec is shipped.
+    with ModelDeployment(ModelSpec(Scorer, weights=weights_path, device="cpu")) as deployed:
+        scored = deployed.predict(
+            catalog,
+            spec,
+            outputs={"score_a": "float32", "score_b": "float32"},
+            inputs=["mag", "flux", "flux_mask"],
+            batch_size=1024,
+        )
+        print(scored.dtypes[["object_id", "mag", "score_a", "score_b"]])  # meta, nothing computed yet
+        result = scored.compute()
+    print(result[["object_id", "mag", "score_a", "score_b"]].head())
+    print(f"{len(result)} rows scored")
+
+
+if __name__ == "__main__":
+    with LocalCluster(n_workers=2, threads_per_worker=2, dashboard_address=None) as cluster, Client(cluster):
+        main(Path("scorer_weights.pt").resolve())

docs/examples/train_from_catalog.py

@@ -0,0 +1,87 @@
+"""Train a small model on an LSDB catalog with a nested light-curve column.
+
+Run with ``python docs/examples/train_from_catalog.py``. It builds a synthetic catalog,
+so no data download is needed. To use a real catalog, replace ``make_catalog()`` with
+``lsdb.open_catalog(...)`` and adjust the spec's column names.
+"""
+
+import lsdb
+import nested_pandas as npd
+import numpy as np
+import pandas as pd
+import torch
+from distributed import Client
+from torch import nn
+from torch.utils.data import DataLoader
+
+from lsdb_ml import CatalogIterableDataset, ColumnSpec, Flat, Nested, SupervisedStep, TrainSession
+
+
+def make_catalog(n_objects=2000, seed=0):
+    """Objects whose label is the mean flux of their (ragged) light curve, plus noise."""
+    rng = np.random.default_rng(seed)
+    lengths = rng.integers(1, 30, n_objects)
+    flux = rng.normal(0, 1, lengths.sum()).astype(np.float32)
+    owner = np.repeat(np.arange(n_objects), lengths)
+    mean_flux = np.bincount(owner, weights=flux) / lengths
+    objects = pd.DataFrame(
+        {
+            "ra": rng.uniform(0, 360, n_objects),
+            "dec": rng.uniform(-60, 60, n_objects),
+            "label": (mean_flux + rng.normal(0, 0.01, n_objects)).astype(np.float32),
+        }
+    )
+    light_curves = pd.DataFrame({"mjd": rng.uniform(0, 100, lengths.sum()), "flux": flux}, index=owner)
+    frame = npd.NestedFrame(objects).join_nested(light_curves, "lc")
+    return lsdb.from_dataframe(frame, partition_rows=200)
+
+
+class MaskedMeanModel(nn.Module):
+    """Predict the label from the masked mean of a learned per-observation embedding."""
+
+    def __init__(self):
+        super().__init__()
+        self.embed = nn.Linear(1, 8)
+        self.head = nn.Linear(8, 1)
+
+    def forward(self, flux, mask):
+        """Pool padded observations using the mask, so padding never contributes."""
+        hidden = self.embed(flux.unsqueeze(-1)) * mask.unsqueeze(-1)
+        pooled = hidden.sum(dim=1) / mask.sum(dim=1, keepdim=True).clamp(min=1)
+        return self.head(pooled).squeeze(-1)
+
+
+def main():
+    """Stream the catalog for a few epochs and train until the loss is small."""
+    catalog = make_catalog()
+    spec = ColumnSpec({"flux": Nested("lc", "flux"), "label": Flat("label")})
+
+    torch.manual_seed(0)
+    model = MaskedMeanModel()
+    step = SupervisedStep(
+        loss_fn=nn.MSELoss(),
+        optimizer=torch.optim.Adam(model.parameters(), lr=1e-2),
+        inputs=["flux", "flux_mask"],
+        target="label",
+    )
+    dataset = CatalogIterableDataset(catalog, spec, batch_size=64, seed=0, epochs=20, partitions_per_chunk=2)
+    loader = DataLoader(dataset, batch_size=None, num_workers=0)
+
+    recent = []
+    with TrainSession(model, step, loader, checkpoint_dir="checkpoints", checkpoint_every=100) as session:
+        for _, metrics in session:
+            if metrics is None:  # skipped (empty) batch
+                continue
+            recent = (recent + [metrics["loss"]])[-10:]
+            if session.batch_count % 10 == 0:
+                print(f"batch {session.batch_count}: mean loss of last 10 batches {np.mean(recent):.5f}")
+            if len(recent) == 10 and np.mean(recent) < 1e-3:
+                print(f"converged after {session.batch_count} batches")
+                session.stop()
+    print("final checkpoint: checkpoints/checkpoint-final.pt")
+
+
+if __name__ == "__main__":
+    # A distributed client lets the stream compute the next chunk while the model trains.
+    with Client(processes=False, n_workers=2, threads_per_worker=1, dashboard_address=None):
+        main()

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]

src/lsdb_ml/__init__.py

@@ -1,4 +1,30 @@
 from ._version import __version__
+from .data import CatalogIterableDataset
+from .device import move_batch, resolve_device
 from .example_module import greetings, meaning
+from .inference import ModelDeployment, ModelSpec, OutputSpec, predict
+from .spec import ColumnSpec, Field, Flat, Nested, ResolvedSpec, SpecError, TensorInfo
+from .training import SupervisedStep, TrainSession, TrainStep
 
-__all__ = ["greetings", "meaning", "__version__"]
+__all__ = [
+    "CatalogIterableDataset",
+    "SupervisedStep",
+    "TrainSession",
+    "TrainStep",
+    "ModelDeployment",
+    "ModelSpec",
+    "OutputSpec",
+    "predict",
+    "ColumnSpec",
+    "Field",
+    "Flat",
+    "Nested",
+    "ResolvedSpec",
+    "SpecError",
+    "TensorInfo",
+    "move_batch",
+    "resolve_device",
+    "greetings",
+    "meaning",
+    "__version__",
+]

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=<fill value> "
+                "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)

src/lsdb_ml/data/__init__.py

@@ -0,0 +1,14 @@
+"""Streaming LSDB catalogs into PyTorch."""
+
+from .dataset import CatalogIterableDataset, iter_slices
+from .sharding import ShardInfo, assign_partitions
+from .source import CatalogStreamSource, PartitionSource
+
+__all__ = [
+    "CatalogIterableDataset",
+    "CatalogStreamSource",
+    "PartitionSource",
+    "ShardInfo",
+    "assign_partitions",
+    "iter_slices",
+]