themisim 0.1.0__py3-none-any.whl

themisim/__init__.py

@@ -0,0 +1,73 @@
+"""THEMIS all-sky imager similarity search.
+
+A standalone, pip-installable extraction of the THEMIS similarity-search
+engine: download the archive, build a FAISS index from the pretrained SimCLR
+encoder, and query it by ``(site, datetime, frame)``.
+
+Quick start (query an existing index)::
+
+    from themisim import query
+    df = query("fsmi", "2015-03-18T06", 412, artifacts="data/artifacts")
+
+Build an index from scratch::
+
+    from themisim import download_archive, build_index
+    download_archive("data/cdf", sites=["fsmi"], start="2015-03", end="2015-03")
+    build_index("data/cdf", "data/artifacts", "weights/<checkpoint>.tar")
+"""
+from __future__ import annotations
+
+# Establish the torch-before-faiss import order before anything pulls in faiss
+# (the query stack imports faiss). torch and faiss ship conflicting CUDA libs
+# and break each other's loader if faiss is imported first; importing torch
+# here (when present) is a no-op safety net. Absent torch (query-only installs)
+# this is simply skipped and faiss loads alone, which is fine.
+try:  # pragma: no cover - environment dependent
+    import torch as _torch  # noqa: F401
+except Exception:  # pragma: no cover
+    pass
+
+from themisim.export import (  # noqa: E402
+    hits_to_dataframe,
+    results_csv_filename,
+)
+from themisim.query import query, resolve_global_id  # noqa: E402
+from themisim.search import Hit, SearchEngine  # noqa: E402
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "query",
+    "resolve_global_id",
+    "SearchEngine",
+    "Hit",
+    "hits_to_dataframe",
+    "results_csv_filename",
+    "download_archive",
+    "build_index",
+    "fetch_weights",
+    "visualize_results",
+    "__version__",
+]
+
+
+def __getattr__(name: str):
+    """Lazily expose the build pipeline + weights helper.
+
+    Kept out of the eager import path so ``import themisim`` stays
+    cheap (and torch-free) for query-only use; ``build_index`` /
+    ``download_archive`` pull in the embed/index stack only when first used.
+    """
+    if name in ("build_index", "download_archive"):
+        from themisim import pipeline
+
+        return getattr(pipeline, name)
+    if name == "fetch_weights":
+        from themisim.weights import fetch_weights
+
+        return fetch_weights
+    if name == "visualize_results":
+        from themisim.viz import visualize_results
+
+        return visualize_results
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

themisim/catalog.py

@@ -0,0 +1,220 @@
+"""Discover CDFs on the public THEMIS archive and build a download work-list.
+
+Distilled from the original project's ``completeness.py`` (which also did
+multi-USB local diffing). The public library only needs to (1) crawl the
+Berkeley Apache autoindex to find which CDFs exist for a set of sites / dates,
+and (2) turn that listing into the ``(url, target_path, size_bytes)`` work-list
+that :mod:`themisim.download` consumes.
+
+Crawling is seeded per site (and per year when a date range is given) so a
+filtered run does not walk the entire ~1M-file tree.
+"""
+from __future__ import annotations
+
+import re
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from pathlib import Path
+from typing import Callable, List, Optional, Sequence, Set
+
+import pandas as pd
+
+from themisim.config import ARCHIVE_BASE_URL
+
+# Strict on `thg_l1_asf_<site>_<YYYYMMDDHH>_v01.cdf`. The tree also carries
+# `thg_l1_ast_*` thumbnails; ignore those (the index is built from asf only).
+#
+# v01 only, deliberately: the rest of the pipeline is v01-throughout — the
+# inventory walk globs `*_v01.cdf` (inventory.py), the embed filename parser
+# requires `_v01` (embed.py), and the CDF download URL is reconstructed with a
+# hardcoded `_v01.cdf` suffix (export.py). Matching any `_v<NN>` here would let
+# the crawler queue a reprocessed `_v02.cdf` for download that every downstream
+# stage then silently drops. Keeping the crawl v01-only makes the whole pipeline
+# consistent; supporting reprocessed versions would require making the shard
+# naming and URL reconstruction version-aware end to end.
+CDF_NAME_RE = re.compile(r"^thg_l1_asf_([a-z]{4})_(\d{10})_v01\.cdf$")
+HREF_RE = re.compile(rb'href="([^"?#][^"]*)"', re.I)
+
+WORKLIST_COLUMNS = ["site", "datetime", "filename", "url", "target_path", "size_bytes"]
+
+
+def _opener(user_agent: str) -> urllib.request.OpenerDirector:
+    opener = urllib.request.build_opener()
+    opener.addheaders = [("User-Agent", user_agent)]
+    return opener
+
+
+def _list_directory(
+    url: str, opener: urllib.request.OpenerDirector, timeout: int = 30
+) -> List[str]:
+    """All href values on an Apache autoindex page (no filtering)."""
+    with opener.open(url, timeout=timeout) as r:
+        body = r.read()
+    return [m.decode("utf-8", "ignore") for m in HREF_RE.findall(body)]
+
+
+def _seed_urls(
+    base_url: str,
+    sites: Optional[Sequence[str]],
+    start: Optional[str],
+    end: Optional[str],
+) -> List[str]:
+    """Starting directories for the crawl.
+
+    With no site filter we crawl from ``base_url`` (the whole archive). With
+    sites we crawl each ``<base>/<site>/``; if a year range is derivable from
+    ``start``/``end`` we descend to ``<base>/<site>/<YYYY>/`` so a narrow date
+    request touches only the relevant years.
+    """
+    if not base_url.endswith("/"):
+        base_url += "/"
+    if not sites:
+        return [base_url]
+
+    years: Optional[List[int]] = None
+    if start or end:
+        y0 = int(start[:4]) if start else 1990
+        y1 = int(end[:4]) if end else 2100
+        years = list(range(y0, y1 + 1))
+
+    seeds: List[str] = []
+    for site in sites:
+        site_base = urllib.parse.urljoin(base_url, f"{site}/")
+        if years is None:
+            seeds.append(site_base)
+        else:
+            for y in years:
+                seeds.append(urllib.parse.urljoin(site_base, f"{y}/"))
+    return seeds
+
+
+def _in_range(dtstr: str, start: Optional[str], end: Optional[str]) -> bool:
+    """Is the ``YYYYMMDDHH`` stamp within ``[start, end]``?
+
+    ``start``/``end`` are compared as left-anchored prefixes of ``dtstr``
+    (``YYYY``, ``YYYY-MM`` or ``YYYYMM`` etc.), so ``--start 2015-03`` keeps
+    everything from March 2015 onward.
+    """
+    def norm(s: str) -> str:
+        return s.replace("-", "")
+
+    if start and dtstr < norm(start).ljust(len(dtstr), "0"):
+        return False
+    if end:
+        e = norm(end)
+        # Inclusive upper bound: pad with '9' so a YYYYMM end keeps that month.
+        if dtstr > e.ljust(len(dtstr), "9"):
+            return False
+    return True
+
+
+def crawl(
+    sites: Optional[Sequence[str]] = None,
+    *,
+    start: Optional[str] = None,
+    end: Optional[str] = None,
+    base_url: str = ARCHIVE_BASE_URL,
+    rate_limit_s: float = 0.3,
+    user_agent: str = "themis-asi-search/0.1",
+    cache_path: Optional[Path] = None,
+    on_dir: Optional[Callable[[int, int, int, str], None]] = None,
+) -> pd.DataFrame:
+    """BFS the archive autoindex and return discovered CDFs.
+
+    Returns a DataFrame with columns ``site, datetime, filename, remote_url``.
+    If ``cache_path`` is given the listing is checkpointed there and resumed on
+    a re-run (an interrupted crawl loses at most a few directory listings).
+    """
+    opener = _opener(user_agent)
+    rows: List[dict] = []
+    seen: Set[str] = set()
+
+    if cache_path is not None and Path(cache_path).exists():
+        prev = pd.read_parquet(cache_path)
+        rows.extend(prev.to_dict("records"))
+        seen.update(prev["remote_url"].tolist())
+
+    queue: List[str] = _seed_urls(base_url, sites, start, end)
+    dirs_done = 0
+    last_ckpt = 0
+
+    while queue:
+        url = queue.pop(0)
+        time.sleep(rate_limit_s)
+        try:
+            hrefs = _list_directory(url, opener)
+        except (urllib.error.URLError, TimeoutError) as exc:
+            if on_dir is not None:
+                on_dir(dirs_done, len(queue), len(rows), f"WARN {url}: {exc}")
+            continue
+
+        for h in hrefs:
+            if h in ("../", "..") or h.startswith("?") or h.startswith("/"):
+                continue
+            child = urllib.parse.urljoin(url, h)
+            if h.endswith("/"):
+                queue.append(child)
+                continue
+            m = CDF_NAME_RE.match(h)
+            if not m:
+                continue
+            if not _in_range(m.group(2), start, end):
+                continue
+            if child in seen:
+                continue
+            seen.add(child)
+            rows.append(
+                {
+                    "site": m.group(1),
+                    "datetime": m.group(2),
+                    "filename": h,
+                    "remote_url": child,
+                }
+            )
+
+        dirs_done += 1
+        if on_dir is not None:
+            on_dir(dirs_done, len(queue), len(rows), url)
+        if cache_path is not None and dirs_done - last_ckpt >= 50:
+            _write(rows, Path(cache_path))
+            last_ckpt = dirs_done
+
+    if cache_path is not None:
+        _write(rows, Path(cache_path))
+    return pd.DataFrame(rows, columns=["site", "datetime", "filename", "remote_url"])
+
+
+def _write(rows: List[dict], out_path: Path) -> None:
+    if not rows:
+        return
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    pd.DataFrame(rows).drop_duplicates("remote_url").reset_index(drop=True).to_parquet(
+        out_path, engine="pyarrow", index=False
+    )
+
+
+def build_worklist(crawl_df: pd.DataFrame, data_root: Path) -> pd.DataFrame:
+    """Turn a crawl listing into a download work-list.
+
+    Each CDF's local target mirrors the archive layout
+    ``<data_root>/<site>/<YYYY>/<MM>/<filename>``. ``size_bytes`` is left at 0;
+    the downloader HEADs each URL for the authoritative size.
+    """
+    data_root = Path(data_root)
+    rows = []
+    for r in crawl_df.itertuples(index=False):
+        yyyy, mm = r.datetime[:4], r.datetime[4:6]
+        target = data_root / r.site / yyyy / mm / r.filename
+        rows.append(
+            {
+                "site": r.site,
+                "datetime": r.datetime,
+                "filename": r.filename,
+                "url": r.remote_url,
+                "target_path": str(target),
+                "size_bytes": 0,
+            }
+        )
+    return pd.DataFrame(rows, columns=WORKLIST_COLUMNS)

themisim/cli/__init__.py

@@ -0,0 +1 @@
+"""Command-line entry points for themis-asi-search."""

themisim/cli/build_index.py

@@ -0,0 +1,96 @@
+"""``themis-build-index`` — embed downloaded CDFs and build the FAISS index.
+
+Runs the full build: inventory -> embed (SimCLR encoder) -> concatenate ->
+train + add to an OPQ-IVF-PQ index. Uses the GPU when available and falls back
+to CPU otherwise. Resumable — already-embedded hours are skipped.
+
+Example::
+
+    themis-build-index --data-root ./data/cdf --artifacts ./data/artifacts \\
+        --checkpoint ./weights/aurora-fm-no-finetune.tar
+"""
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+from themisim.config import default_artifacts_root, default_data_root
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(description=__doc__.split("\n\n", 1)[0])
+    p.add_argument("--data-root", type=Path, default=default_data_root())
+    p.add_argument("--artifacts", type=Path, default=default_artifacts_root())
+    p.add_argument(
+        "--checkpoint",
+        type=Path,
+        default=None,
+        help="SimCLR checkpoint .tar (default: fetch/verify into the weights dir)",
+    )
+    p.add_argument(
+        "--nlist",
+        default="auto",
+        help="IVF cells; 'auto' scales with the vector count (default: auto)",
+    )
+    p.add_argument(
+        "--device",
+        default="auto",
+        help="'auto' (GPU if present, else CPU), 'cuda', or 'cpu'",
+    )
+    p.add_argument("--batch-size", type=int, default=128)
+    p.add_argument("--num-workers", type=int, default=4, help="DataLoader workers")
+    p.add_argument(
+        "--no-gpu-quantizer",
+        action="store_true",
+        help="train the IVF k-means on CPU even if GPUs are present",
+    )
+    p.add_argument(
+        "--no-streaming",
+        action="store_true",
+        help="use the eager per-CDF embed loop instead of the faster "
+        "persistent-worker streaming path",
+    )
+    p.add_argument(
+        "--no-compile",
+        action="store_true",
+        help="skip torch.compile on the encoder during streaming embed "
+        "(avoids its warmup; only affects GPU runs)",
+    )
+    args = p.parse_args(argv)
+
+    from themisim.pipeline import build_index
+
+    checkpoint = args.checkpoint
+    if checkpoint is None:
+        from themisim.weights import fetch_weights
+
+        print("No --checkpoint given; fetching/verifying weights ...", flush=True)
+        checkpoint = fetch_weights()
+
+    t0 = time.time()
+
+    def progress(stage: str, info: object) -> None:
+        print(f"  [{stage}] {info}  (+{time.time() - t0:.0f}s)", flush=True)
+
+    nlist = args.nlist if args.nlist == "auto" else int(args.nlist)
+    out = build_index(
+        args.data_root,
+        args.artifacts,
+        checkpoint,
+        nlist=nlist,
+        device=args.device,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        use_gpu_quantizer=not args.no_gpu_quantizer,
+        streaming=not args.no_streaming,
+        use_compile=not args.no_compile,
+        on_progress=progress,
+    )
+    size = out.stat().st_size
+    print(f"\nWrote {out} ({size / 1e6:.1f} MB)")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

themisim/cli/download.py

@@ -0,0 +1,86 @@
+"""``themis-download`` — fetch THEMIS ASI CDFs from the Berkeley archive.
+
+Examples
+--------
+Full archive (every site, every date)::
+
+    themis-download --data-root ./data/cdf
+
+A single site and month::
+
+    themis-download --data-root ./data/cdf --sites fsmi --start 2015-03 --end 2015-03
+"""
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+from themisim.config import default_data_root
+
+
+def _split_sites(value: str | None) -> list[str] | None:
+    if not value:
+        return None
+    return [s.strip().lower() for s in value.split(",") if s.strip()]
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(description=__doc__.split("\n\n", 1)[0])
+    p.add_argument(
+        "--data-root",
+        type=Path,
+        default=default_data_root(),
+        help="where to store CDFs (default: $THEMIS_ASI_DATA_ROOT or ./data/cdf)",
+    )
+    p.add_argument(
+        "--sites",
+        default=None,
+        help="comma-separated 4-letter site codes (default: all sites)",
+    )
+    p.add_argument("--start", default=None, help="earliest date, e.g. 2015 or 2015-03")
+    p.add_argument("--end", default=None, help="latest date, e.g. 2016 or 2016-12")
+    p.add_argument("--workers", type=int, default=6, help="parallel downloads")
+    p.add_argument("--limit", type=int, default=None, help="cap files (smoke tests)")
+    args = p.parse_args(argv)
+
+    # Imported here so `themis-download --help` doesn't pay the import cost.
+    from themisim.pipeline import download_archive
+
+    state = {"t0": time.time()}
+
+    def on_dir(done: int, queued: int, found: int, msg: str) -> None:
+        if done % 25 == 0 or msg.startswith("WARN"):
+            print(f"[crawl] dirs={done} queued={queued} found={found:,}  {msg}", flush=True)
+
+    def on_event(res, done: int, total: int) -> None:
+        if res.status in ("failed", "stopped"):
+            print(f"  [{res.status}] {res.url}  {res.error}", flush=True)
+        if done % 50 == 0 or done == total:
+            dt = max(time.time() - state["t0"], 1e-6)
+            print(f"  {done:,}/{total:,} files  ({done / dt:.1f}/s)", flush=True)
+
+    print(f"Crawling archive (sites={args.sites or 'ALL'}, "
+          f"start={args.start}, end={args.end}) ...", flush=True)
+    res = download_archive(
+        args.data_root,
+        sites=_split_sites(args.sites),
+        start=args.start,
+        end=args.end,
+        workers=args.workers,
+        limit=args.limit,
+        on_dir=on_dir,
+        on_event=on_event,
+    )
+
+    print("\n=== summary ===")
+    if len(res):
+        print(res.groupby("status").size().to_string())
+        print(f"bytes downloaded: {res['bytes_written'].sum() / 1e9:.2f} GB")
+    else:
+        print("(no files matched)")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

themisim/cli/query.py

@@ -0,0 +1,93 @@
+"""``themis-query`` — similarity search by ``(site, datetime, frame)``.
+
+Writes (or prints) a CSV identical to the dashboard's export:
+``site, datetime, score, source_cdf``.
+
+Example::
+
+    themis-query --site fsmi --datetime 2015-03-18T06 --frame 412 \\
+        --artifacts ./data/artifacts --output results.csv
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from themisim.config import default_artifacts_root
+from themisim.search import (
+    DEFAULT_DIVERSIFY_SECONDS,
+    DEFAULT_K,
+    DEFAULT_NPROBE,
+    DEFAULT_PREFILTER,
+)
+
+
+def main(argv: list[str] | None = None) -> int:
+    p = argparse.ArgumentParser(description=__doc__.split("\n\n", 1)[0])
+    p.add_argument("--site", required=True, help="4-letter site code, e.g. fsmi")
+    p.add_argument(
+        "--datetime",
+        required=True,
+        help="hourly CDF, e.g. 2015-03-18T06 (sub-hour fields ignored)",
+    )
+    p.add_argument("--frame", type=int, required=True, help="0-based frame in the hour")
+    p.add_argument("--artifacts", type=Path, default=default_artifacts_root())
+    p.add_argument(
+        "--results",
+        type=int,
+        default=DEFAULT_K,
+        help="count mode: number of top matches to return (ignored with --min-score)",
+    )
+    p.add_argument(
+        "--min-score",
+        type=float,
+        default=None,
+        help="threshold mode: return every match with cosine similarity >= this "
+        "cutoff (0..1), best first, instead of a fixed --results count",
+    )
+    p.add_argument("--prefilter", type=int, default=DEFAULT_PREFILTER)
+    p.add_argument("--nprobe", type=int, default=DEFAULT_NPROBE)
+    p.add_argument("--diversify-seconds", type=int, default=DEFAULT_DIVERSIFY_SECONDS)
+    p.add_argument(
+        "--output",
+        type=Path,
+        default=None,
+        help="CSV path. Default: auto-named in the cwd; '-' prints to stdout",
+    )
+    args = p.parse_args(argv)
+
+    from themisim.export import results_csv_filename
+    from themisim.query import get_engine, query, resolve_global_id
+
+    engine = get_engine(args.artifacts)
+    # Resolve once so we can name the output file the way the dashboard does.
+    gid = resolve_global_id(engine, args.site, args.datetime, args.frame)
+    df = query(
+        args.site,
+        args.datetime,
+        args.frame,
+        results=args.results,
+        min_score=args.min_score,
+        prefilter=args.prefilter,
+        nprobe=args.nprobe,
+        diversify_seconds=args.diversify_seconds,
+        engine=engine,
+    )
+
+    if str(args.output) == "-":
+        print(df.to_csv(index=False))
+        return 0
+
+    if args.output is None:
+        out = Path.cwd() / results_csv_filename(
+            args.site, int(engine.time_ns_of(gid))
+        )
+    else:
+        out = args.output
+    df.to_csv(out, index=False)
+    print(f"Wrote {len(df)} rows -> {out}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())