Merge pull request #699 from light-curve/embed-ztf-example

Add ZTF DR23 nearest-neighbour example to embed docs

Author: hombit

light-curve/README.md

@@ -828,13 +828,16 @@ See the benchmarks described in more detail in
 The `light_curve.embed` submodule provides pretrained neural network models that map raw photometric
 time series to dense vector embeddings suitable for downstream machine learning tasks such as
 classification, anomaly detection, and similarity search.
-Models are loaded directly from HuggingFace (weights are cached locally after the first download)
-and require `onnxruntime` and `huggingface-hub`:
+The examples below use `from_hf()` to download model weights from HuggingFace Hub
+(cached locally after the first download), which requires `onnxruntime` and `huggingface-hub`:
 
 ```
-pip install onnxruntime huggingface_hub
+pip install light-curve onnxruntime huggingface_hub
 ```
 
+Models can also be loaded directly from a local ONNX file via `from_onnx_file()`,
+without `huggingface-hub`.
+
 See the [onnxruntime install guide](https://onnxruntime.ai/docs/install/) for GPU and platform-specific
 packages (`onnxruntime-gpu`, `onnxruntime-directml`, etc.).
 Hardware-specific options such as execution providers can be passed via `ort_session_kwargs`
@@ -916,6 +919,62 @@ print(embedding.shape)  # (1, 1, 1, 384)
 Input fluxes should be in AB units. The default zero-point is 31.4 (LSST nJy); set `mag_zp=27.5`
 for ELAsTiCC / SNANA FITS data or `mag_zp=8.9` for Jy.
 
+### Example: nearest-neighbour search on ZTF DR23
+
+The following example reads a ZTF DR23 [HATS](https://hats.readthedocs.io) pixel directly from the public S3
+bucket using
+[`nested-pandas`](https://nested-pandas.readthedocs.io) and `s3fs`
+(install all dependencies for this example with
+`pip install huggingface_hub light-curve nested-pandas onnxruntime s3fs scipy`), embeds all well-observed
+light curves with Astromer2, and finds the closest neighbour to a given object by cosine distance.
+This particular pixel contains ~50k objects passing the quality cuts; embedding takes roughly
+12 minutes on M2 Pro (~15 ms per object).
+
+<!-- name: test_ztf_dr23_nn_example; mark: skip(reason="requires S3 access and ~12 minutes to run") -->
+
+```python
+import numpy as np
+import nested_pandas as npd
+from scipy.spatial.distance import cdist
+from upath import UPath
+from light_curve.embed import Astromer2
+
+TARGET_OID = 680213300009232  # ZTF r-band light curve
+MIN_OBS = 1000
+
+nf = npd.read_parquet(UPath(
+    "s3://ipac-irsa-ztf/contributed/dr23/lc/hats/ztf_dr23_lc-hats/dataset/Norder=5/Dir=0/Npix=2378/",
+    anon=True,
+))
+nf = nf.query("lightcurve.catflags == 0").query(f"lightcurve.list_lengths >= {MIN_OBS}")
+
+model = Astromer2.from_hf(reduction="beginning")
+
+
+def embed_row(row):
+    return {"embedding.values": model(row["lightcurve.hmjd"], row["lightcurve.mag"]).squeeze()}
+
+
+nf = nf.map_rows(embed_row, columns=["lightcurve.hmjd", "lightcurve.mag"], append_columns=True)
+
+oids = nf["objectid"].to_numpy()
+matrix = nf["embedding.values"].to_numpy().reshape(len(nf), -1)
+
+query_idx = np.where(oids == TARGET_OID)[0][0]
+distances = cdist(matrix[query_idx: query_idx + 1], matrix, metric="cosine")[0]
+distances[query_idx] = np.inf
+
+best_idx = np.argmin(distances)
+print(f"Nearest neighbour: OID {oids[best_idx]}, cosine distance {distances[best_idx]:.6f}")
+# Nearest neighbour: OID 680113300005170, cosine distance 0.000046
+```
+
+The nearest neighbour is OID `680113300005170` — the same physical
+object ([HZ Her / Her X-1](https://en.wikipedia.org/wiki/Hercules_X-1), an X-ray binary) observed in the *g*
+-band,
+recovered automatically from an *r*-band query through embedding similarity.
+See the light curve on [SNAD Viewer](https://ztf.snad.space/dr23/view/680113300005170).
+
 ## dm-dt map
 
 In addition to the feature extractors above, the package provides a separate dm–dt mapping tool.

light-curve/pyproject.toml

@@ -46,7 +46,7 @@ full = [
 # no linting tools, compatible with free-threading
 test-base = [
     "pytest",
-    "markdown-pytest",
+    "markdown-pytest>=0.6",
     "pytest-benchmark",
     "pytest-subtests>=0.10",
     "arro3-core",
@@ -60,7 +60,7 @@ test-base = [
 # Full testing environment (adds packages that don't support free-threading yet:
 # cesium, iminuit, nanoarrow, polars)
 test = [
-    {include-group = "test-base"},
+    { include-group = "test-base" },
     "numpy",
     "pyarrow",
     "pandas",
@@ -72,7 +72,7 @@ test = [
 ]
 # Development environment for free-threaded Python builds
 dev-free-threading = [
-    {include-group = "test-base"},
+    { include-group = "test-base" },
     "numpy",
     "pyarrow",
     "pandas",
@@ -81,7 +81,7 @@ dev-free-threading = [
 ]
 # Full development environment
 dev = [
-    {include-group = "dev-free-threading"},
+    { include-group = "dev-free-threading" },
     "iminuit>=2.21,<3",
     "nanoarrow",
     "polars",