themisim 0.1.0__tar.gz

themisim-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jeremiah Johnson
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

themisim-0.1.0/PKG-INFO

@@ -0,0 +1,196 @@
+Metadata-Version: 2.4
+Name: themisim
+Version: 0.1.0
+Summary: Similarity search over THEMIS all-sky imager imagery using SimCLR encoder features and a FAISS index
+Author: Jeremiah Johnson
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/jwjohnson314/themisim
+Project-URL: Documentation, https://jwjohnson314.github.io/themisim/
+Project-URL: Source, https://github.com/jwjohnson314/themisim
+Project-URL: Issues, https://github.com/jwjohnson314/themisim/issues
+Keywords: aurora,THEMIS,all-sky-imager,similarity-search,faiss,self-supervised
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Classifier: Topic :: Scientific/Engineering :: Image Processing
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Requires-Dist: torchvision>=0.15
+Requires-Dist: numpy<2,>=1.24
+Requires-Dist: pandas>=2.0
+Requires-Dist: pyarrow>=11
+Requires-Dist: Pillow>=9.4
+Requires-Dist: cdflib>=1.2
+Requires-Dist: faiss-cpu>=1.8
+Provides-Extra: dev
+Requires-Dist: pytest>=7.4; extra == "dev"
+Requires-Dist: pytest-cov>=4; extra == "dev"
+Provides-Extra: notebook
+Requires-Dist: jupyter; extra == "notebook"
+Requires-Dist: matplotlib; extra == "notebook"
+Provides-Extra: docs
+Requires-Dist: sphinx>=7; extra == "docs"
+Requires-Dist: furo>=2024.1; extra == "docs"
+Dynamic: license-file
+
+# THEMISim
+
+Similarity search over [THEMIS](https://themis.ssl.berkeley.edu) all-sky imager
+(ASI) auroral imagery. A SimCLR-trained ResNet-18 encoder turns each 256×256
+frame into a 512-D feature vector where proximity as measured by cosine similarity corresponds to visual and morphological similarity rather than pixel-level identity; ~1 billion of these are indexed with FAISS (OPQ-IVF-PQ). Given any indexed frame, the library quickly and efficiently returns the most visually similar frames across the whole archive.
+
+This software does three things:
+
+1. **Download** THEMIS ASI CDFs from the Berkeley archive and encoder weights from HuggingFace.
+2. **Build** a FAISS index from the pretrained encoder weights.
+3. **Query** the index by `(site, datetime, frame)`, returning a tidy
+   `DataFrame`/CSV.
+
+## Install
+
+```bash
+pip install themisim
+```
+
+`faiss-cpu` is pulled in as a dependency (the query path is CPU-only). If you
+already have a GPU `faiss` from conda, install with `--no-deps` to keep it.
+A CUDA-enabled PyTorch makes index building much faster but is optional — every
+stage falls back to CPU.
+
+## Quick start — querying an existing index
+
+```python
+from themisim import query
+
+df = query("fsmi", "2015-03-18T06", 412, artifacts="data/artifacts")
+df.head()
+#    site             datetime     score                                         source_cdf
+# 0  fsmi  2015-03-18 06:20:36  1.000000  http://themis.ssl.berkeley.edu/.../thg_l1_asf_fsmi_2015031806_v01.cdf
+# ...
+```
+
+- `datetime` names the **hourly** CDF (`YYYY-MM-DDTHH`); sub-hour fields are
+  ignored. `2015-03-18 06`, `2015031806`, and a `datetime` object also work.
+- `frame` is the 0-based frame index within that hour (THEMIS runs at a
+  3-second cadence, ~1200 frames/hour).
+- The top result is normally the query frame itself (`score ≈ 1.0`).
+
+Search parameters:
+
+| arg | default | meaning |
+|-----|---------|---------|
+| `results` | 24 | count mode: number of results to return |
+| `min_score` | `None` | threshold mode: return *every* match scoring ≥ this cutoff (overrides `results`) |
+| `prefilter` | 500 | FAISS candidates pulled before exact cosine rerank |
+| `nprobe` | 64 | IVF cells inspected per query |
+| `diversify_seconds` | 30 | drop near-duplicate frames within ±N s at the same site (0 disables) |
+
+### Bounding the result set: count vs. threshold
+
+There are two ways to decide how many results come back:
+
+- **Count** (default) — return the top `results` matches.
+- **Threshold** — pass `min_score` (a cosine similarity in `[0, 1]`) and the
+  search ignores `results`, returning *every* (temporally diversified) match
+  scoring at or above the cutoff, best first. The result set is capped at 2000
+  for safety. Because the candidate pool is the top `prefilter` FAISS matches,
+  raise `prefilter` to surface more low-cutoff hits.
+
+```python
+# every frame at least 0.9 cosine-similar to the query, not just the top 24
+df = query("fsmi", "2015-03-18T06", 412, artifacts="data/artifacts",
+           min_score=0.9)
+```
+
+### From the command line
+
+```bash
+themis-query --site fsmi --datetime 2015-03-18T06 --frame 412 \
+    --artifacts data/artifacts --output results.csv
+
+# threshold mode: every match scoring >= 0.9, instead of a fixed --results count
+themis-query --site fsmi --datetime 2015-03-18T06 --frame 412 \
+    --artifacts data/artifacts --min-score 0.9 --output results.csv
+```
+
+Output columns are `site, datetime, score, source_cdf` — identical to the
+dashboard's CSV export.
+
+## Build an index from scratch
+
+Note that the full THEMIS ASI archive is ~1M CDFs (over 100 TB) and building the index takes on the order of GPU-days; the
+prebuilt index (about 72 GB) is **not** currently downloadable. If you are interested in obtaining it please email me to discuss.
+
+
+```bash
+# 1. fetch the encoder weights (44 MB)
+python -c "from themisim import fetch_weights; print(fetch_weights())"
+
+# 2. one-shot pipeline (download + build)
+SITES=fsmi START=2015-03 END=2015-03 ./scripts/run_pipeline.sh
+```
+
+or step by step:
+
+```bash
+themis-download   --data-root data/cdf --sites fsmi --start 2015-03 --end 2015-03
+themis-build-index --data-root data/cdf --artifacts data/artifacts \
+                   --checkpoint weights/aurora-fm-no-finetune.tar
+```
+
+The equivalent Python API:
+
+```python
+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/aurora-fm-no-finetune.tar")
+```
+
+`build_index` runs inventory → embed → concat → train/build, resuming cleanly
+if interrupted (already-embedded hours are skipped). `--nlist auto` scales the
+IVF cell count to the dataset size.
+
+## Model weights
+
+`fetch_weights()` downloads and SHA-256-verifies the ~44 MB SimCLR checkpoint
+into `./weights` (override with `$THEMIS_ASI_WEIGHTS`). It is pulled from the
+public Hugging Face repo
+[`Jwjohnson314/Aurora-FM`](https://huggingface.co/Jwjohnson314/Aurora-FM); override
+the URL via `$THEMIS_ASI_WEIGHTS_URL` or pass `url=`. If you already have the
+`.tar`, drop it in the weights directory and it will be verified and reused
+without a network call.
+
+## Paths / configuration
+
+Resolved from explicit arguments, then environment variables, then defaults:
+
+| what | env var | default |
+|------|---------|---------|
+| downloaded CDFs | `THEMIS_ASI_DATA_ROOT` | `./data/cdf` |
+| index artifacts | `THEMIS_ASI_ARTIFACTS` | `./data/artifacts` |
+| model weights | `THEMIS_ASI_WEIGHTS` | `./weights` |
+
+An artifacts directory contains `index.faiss`, `manifest.parquet`,
+`vectors.f16.dat`, and `vectors.meta.json`.
+
+## License, data & citation
+
+The **code** in this repository is MIT-licensed — see [LICENSE](LICENSE).
+
+The assets this tool downloads carry their own terms, which you must honor when
+publishing results:
+
+- **Model weights** (`Jwjohnson314/Aurora-FM`) are released under
+  **CC-BY-4.0** — attribution required.
+- **THEMIS ASI data** is provided by the THEMIS mission (UC Berkeley / NASA) and
+  is subject to the
+  [THEMIS data use & citation policy](https://themis.ssl.berkeley.edu/roadrules.shtml).
+  Acknowledge the mission and instrument teams in any publication.

themisim-0.1.0/README.md

@@ -0,0 +1,155 @@
+# THEMISim
+
+Similarity search over [THEMIS](https://themis.ssl.berkeley.edu) all-sky imager
+(ASI) auroral imagery. A SimCLR-trained ResNet-18 encoder turns each 256×256
+frame into a 512-D feature vector where proximity as measured by cosine similarity corresponds to visual and morphological similarity rather than pixel-level identity; ~1 billion of these are indexed with FAISS (OPQ-IVF-PQ). Given any indexed frame, the library quickly and efficiently returns the most visually similar frames across the whole archive.
+
+This software does three things:
+
+1. **Download** THEMIS ASI CDFs from the Berkeley archive and encoder weights from HuggingFace.
+2. **Build** a FAISS index from the pretrained encoder weights.
+3. **Query** the index by `(site, datetime, frame)`, returning a tidy
+   `DataFrame`/CSV.
+
+## Install
+
+```bash
+pip install themisim
+```
+
+`faiss-cpu` is pulled in as a dependency (the query path is CPU-only). If you
+already have a GPU `faiss` from conda, install with `--no-deps` to keep it.
+A CUDA-enabled PyTorch makes index building much faster but is optional — every
+stage falls back to CPU.
+
+## Quick start — querying an existing index
+
+```python
+from themisim import query
+
+df = query("fsmi", "2015-03-18T06", 412, artifacts="data/artifacts")
+df.head()
+#    site             datetime     score                                         source_cdf
+# 0  fsmi  2015-03-18 06:20:36  1.000000  http://themis.ssl.berkeley.edu/.../thg_l1_asf_fsmi_2015031806_v01.cdf
+# ...
+```
+
+- `datetime` names the **hourly** CDF (`YYYY-MM-DDTHH`); sub-hour fields are
+  ignored. `2015-03-18 06`, `2015031806`, and a `datetime` object also work.
+- `frame` is the 0-based frame index within that hour (THEMIS runs at a
+  3-second cadence, ~1200 frames/hour).
+- The top result is normally the query frame itself (`score ≈ 1.0`).
+
+Search parameters:
+
+| arg | default | meaning |
+|-----|---------|---------|
+| `results` | 24 | count mode: number of results to return |
+| `min_score` | `None` | threshold mode: return *every* match scoring ≥ this cutoff (overrides `results`) |
+| `prefilter` | 500 | FAISS candidates pulled before exact cosine rerank |
+| `nprobe` | 64 | IVF cells inspected per query |
+| `diversify_seconds` | 30 | drop near-duplicate frames within ±N s at the same site (0 disables) |
+
+### Bounding the result set: count vs. threshold
+
+There are two ways to decide how many results come back:
+
+- **Count** (default) — return the top `results` matches.
+- **Threshold** — pass `min_score` (a cosine similarity in `[0, 1]`) and the
+  search ignores `results`, returning *every* (temporally diversified) match
+  scoring at or above the cutoff, best first. The result set is capped at 2000
+  for safety. Because the candidate pool is the top `prefilter` FAISS matches,
+  raise `prefilter` to surface more low-cutoff hits.
+
+```python
+# every frame at least 0.9 cosine-similar to the query, not just the top 24
+df = query("fsmi", "2015-03-18T06", 412, artifacts="data/artifacts",
+           min_score=0.9)
+```
+
+### From the command line
+
+```bash
+themis-query --site fsmi --datetime 2015-03-18T06 --frame 412 \
+    --artifacts data/artifacts --output results.csv
+
+# threshold mode: every match scoring >= 0.9, instead of a fixed --results count
+themis-query --site fsmi --datetime 2015-03-18T06 --frame 412 \
+    --artifacts data/artifacts --min-score 0.9 --output results.csv
+```
+
+Output columns are `site, datetime, score, source_cdf` — identical to the
+dashboard's CSV export.
+
+## Build an index from scratch
+
+Note that the full THEMIS ASI archive is ~1M CDFs (over 100 TB) and building the index takes on the order of GPU-days; the
+prebuilt index (about 72 GB) is **not** currently downloadable. If you are interested in obtaining it please email me to discuss.
+
+
+```bash
+# 1. fetch the encoder weights (44 MB)
+python -c "from themisim import fetch_weights; print(fetch_weights())"
+
+# 2. one-shot pipeline (download + build)
+SITES=fsmi START=2015-03 END=2015-03 ./scripts/run_pipeline.sh
+```
+
+or step by step:
+
+```bash
+themis-download   --data-root data/cdf --sites fsmi --start 2015-03 --end 2015-03
+themis-build-index --data-root data/cdf --artifacts data/artifacts \
+                   --checkpoint weights/aurora-fm-no-finetune.tar
+```
+
+The equivalent Python API:
+
+```python
+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/aurora-fm-no-finetune.tar")
+```
+
+`build_index` runs inventory → embed → concat → train/build, resuming cleanly
+if interrupted (already-embedded hours are skipped). `--nlist auto` scales the
+IVF cell count to the dataset size.
+
+## Model weights
+
+`fetch_weights()` downloads and SHA-256-verifies the ~44 MB SimCLR checkpoint
+into `./weights` (override with `$THEMIS_ASI_WEIGHTS`). It is pulled from the
+public Hugging Face repo
+[`Jwjohnson314/Aurora-FM`](https://huggingface.co/Jwjohnson314/Aurora-FM); override
+the URL via `$THEMIS_ASI_WEIGHTS_URL` or pass `url=`. If you already have the
+`.tar`, drop it in the weights directory and it will be verified and reused
+without a network call.
+
+## Paths / configuration
+
+Resolved from explicit arguments, then environment variables, then defaults:
+
+| what | env var | default |
+|------|---------|---------|
+| downloaded CDFs | `THEMIS_ASI_DATA_ROOT` | `./data/cdf` |
+| index artifacts | `THEMIS_ASI_ARTIFACTS` | `./data/artifacts` |
+| model weights | `THEMIS_ASI_WEIGHTS` | `./weights` |
+
+An artifacts directory contains `index.faiss`, `manifest.parquet`,
+`vectors.f16.dat`, and `vectors.meta.json`.
+
+## License, data & citation
+
+The **code** in this repository is MIT-licensed — see [LICENSE](LICENSE).
+
+The assets this tool downloads carry their own terms, which you must honor when
+publishing results:
+
+- **Model weights** (`Jwjohnson314/Aurora-FM`) are released under
+  **CC-BY-4.0** — attribution required.
+- **THEMIS ASI data** is provided by the THEMIS mission (UC Berkeley / NASA) and
+  is subject to the
+  [THEMIS data use & citation policy](https://themis.ssl.berkeley.edu/roadrules.shtml).
+  Acknowledge the mission and instrument teams in any publication.

themisim-0.1.0/pyproject.toml

@@ -0,0 +1,62 @@
+[build-system]
+requires = ["setuptools>=77", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "themisim"
+version = "0.1.0"
+description = "Similarity search over THEMIS all-sky imager imagery using SimCLR encoder features and a FAISS index"
+readme = "README.md"
+requires-python = ">=3.11"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [{ name = "Jeremiah Johnson" }]
+keywords = ["aurora", "THEMIS", "all-sky-imager", "similarity-search", "faiss", "self-supervised"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Astronomy",
+    "Topic :: Scientific/Engineering :: Image Processing",
+]
+dependencies = [
+    "torch>=2.0",
+    "torchvision>=0.15",
+    "numpy>=1.24,<2",
+    "pandas>=2.0",
+    "pyarrow>=11",
+    "Pillow>=9.4",
+    "cdflib>=1.2",
+    # CPU FAISS is sufficient: the query path is CPU-only (the index is
+    # memory-mapped and disk-paged). GPU users who already have a conda
+    # faiss build can install with --no-deps to keep it.
+    "faiss-cpu>=1.8",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=7.4", "pytest-cov>=4"]
+notebook = ["jupyter", "matplotlib"]
+docs = ["sphinx>=7", "furo>=2024.1"]
+
+[project.scripts]
+themis-download = "themisim.cli.download:main"
+themis-build-index = "themisim.cli.build_index:main"
+themis-query = "themisim.cli.query:main"
+
+[project.urls]
+Homepage = "https://github.com/jwjohnson314/themisim"
+Documentation = "https://jwjohnson314.github.io/themisim/"
+Source = "https://github.com/jwjohnson314/themisim"
+Issues = "https://github.com/jwjohnson314/themisim/issues"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+addopts = "-ra"
+pythonpath = ["src"]

themisim-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

themisim-0.1.0/src/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}")