Hyrax: GPU-Enabled Astronomy ML Framework

• Hyrax is an open‐source, modular, GPU-enabled Python framework designed for comprehensive astronomy machine learning workflows.
• It integrates data acquisition, preprocessing, training, inference, vector indexing, and interactive visualization to support both supervised and unsupervised discovery.
• Its scalable architecture leverages GPU support, multi-GPU execution, and built-in experiment tracking for efficient, reproducible ML experimentation.

Hyrax is an open-source, modular, GPU-enabled Python framework for astronomy machine learning designed to support the full ML lifecycle rather than only model definition. In the formulation given by "Hyrax: An Extensible Framework for Rapid ML Experimentation and Unsupervised Discovery in the Era of Rubin, Roman, and Euclid," it spans data acquisition, dataset abstraction, training, inference, experiment comparison, vector-database indexing, and interactive latent-space exploration, with an explicit emphasis on multimodal survey data and on unsupervised discovery workflows in which scalar loss values are not sufficient for scientific assessment (Ghosh et al., 18 May 2026).

1. Motivation and conceptual scope

Hyrax was introduced in the context of Rubin, Roman, Euclid, CSST, and related survey programs, where the limiting factor in astronomy ML is described as shifting from model architecture design to infrastructure. The motivating pressures are concrete: Rubin alone is expected to produce over 3 petabytes of images annually, while modern survey analysis increasingly requires survey-native data access, scalable handling close to the data, multimodal learning over images, spectra, light curves, and metadata, GPU and multi-GPU execution, and mechanisms for experiment management, reproducibility, clustering, similarity search, and human-in-the-loop discovery (Ghosh et al., 18 May 2026).

The framework is explicitly characterized as application-agnostic, modular and extensible, astronomy-aware, built on existing open-source tools where possible, custom where astronomy requires it, usable from laptop to multi-GPU HPC, and supportive of unsupervised discovery rather than only supervised benchmarking. Its recurring workflow stages are exposed through a small set of high-level verbs: hyrax.download(), hyrax.train(), hyrax.infer(), and hyrax.visualize() (Ghosh et al., 18 May 2026).

At the level of scientific positioning, Hyrax addresses three broad problems identified by its authors: fragmented astronomy ML infrastructure, scalability and deployment barriers, and the lack of astronomy-native support for unsupervised discovery. A central implication is that Hyrax is not presented as a single model family or benchmark package; it is presented as reusable astronomy ML infrastructure for the survey era (Ghosh et al., 18 May 2026).

2. System architecture and data abstractions

Hyrax organizes the end-to-end workflow as a linked sequence of modules: data acquisition, dataset organization, training, inference, optional vector-database insertion, and interactive visualization. That sequence is intended to accommodate both supervised and unsupervised pipelines without forcing a single scientific use case or data modality (Ghosh et al., 18 May 2026).

A key architectural distinction is the separation between data acquisition and model-facing data exposure. Modality-specific access is handled through DatasetClass abstractions for images, spectra, and time series or light curves. A configurable DataProvider sits above these classes and assembles one or more of them into structured model inputs according to request configuration, enabling both single-modality and multimodal pipelines. The DataProvider can also preload or cache items in memory; the authors state that this often sped up training epochs by an order of magnitude (Ghosh et al., 18 May 2026).

For survey ingestion, the paper reports support at the time of writing for Rubin-LSST and HSC-SSP, while also allowing users to bypass built-in downloaders and provide custom data directly. The LSST-specific dataset class integrates with the Rubin Data Butler and is optimized for large imaging collections: survey patches are retrieved via Butler, cached locally at the patch level, cutouts are generated on the fly from cached larger units, multi-threaded requests are used, and outputs are streamed directly to GPU training. This design is meant to reduce repeated remote queries and avoid bulk download of very large image sets. If local storage permits, generated cutouts may also be persisted for reuse across experiments (Ghosh et al., 18 May 2026).

The acquisition layer includes usability features that are operational rather than algorithmic: automatic recovery from failed connections, progress and ETA monitoring, tracking of partially or unsuccessfully acquired files, and parallelized retrieval. Preprocessing is intentionally split between framework-level orchestration and task-specific implementation. The paper lists, as application examples, magnitude-to-flux conversion, temporal merging of light-curve points, feature normalization, spectral rest-frame correction and resampling, image cutout generation, flips, rotations, shifts, and synthetic source injection (Ghosh et al., 18 May 2026).

3. Training, inference, and experiment management

Training and inference are centralized through hyrax.train() and hyrax.infer(). Users specify models, datasets, and training logic for a training or validation loop, while Hyrax handles orchestration, configuration capture, output management, and experiment bookkeeping. The framework accepts any PyTorch-based model, supports custom dataset classes, and also includes prewritten models for common use cases (Ghosh et al., 18 May 2026).

GPU support is a first-class design feature. Hyrax is described as scaling from laptops to HPC systems, with automatic multi-GPU support and efficient GPU feeding via caching and prefetching. Hyperparameter optimization is integrated through Optuna, with the paper explicitly naming learning rate, batch size, architectural choices, and latent dimensionality as examples of tunable parameters, although no framework-specific optimization objective is imposed (Ghosh et al., 18 May 2026).

Experiment tracking is built directly into the workflow via MLflow and TensorBoard. All runs automatically log hyperparameters, evaluation metrics, run metadata, metric histories, configuration details, and run states, with additional support for custom task-specific quantities. This is framed as operationally important because, at survey scale, comparison across many runs becomes a central bottleneck. Inference supports both supervised and representation-learning use cases; in the latter, embeddings may be forwarded into vector databases and interactive exploration interfaces (Ghosh et al., 18 May 2026).

The paper does not introduce a single Hyrax-specific loss function or training objective. Instead, it treats Hyrax as infrastructure that can host diverse objectives. Its demonstrations include reconstruction-based autoencoders, SimCLR, BYOL, supervised classifiers, multimodal transient models, and CNN-based detection pipelines, with the framework managing orchestration and provenance rather than constraining scientific model design (Ghosh et al., 18 May 2026).

4. Vector databases, similarity search, and latent-space analysis

One of Hyrax’s distinctive components is native support for vector databases as part of an astronomy workflow for learned embeddings. The motivation is explicit: scientific use of representation learning requires nearest-neighbor retrieval, anomaly surfacing, clustering, and interactive browsing at scales where naive similarity search becomes impractical. The paper states the relevant asymptotics as $\mathcal{O}(n)$ for a single query against all vectors, $\mathcal{O}(n^2)$ for exhaustive pairwise comparisons, and approximately $\mathcal{O}(\log n)$ for HNSW-based approximate nearest-neighbor search (Ghosh et al., 18 May 2026).

At the time of publication, supported vector-database backends were LanceDB, ChromaDB, and Qdrant, exposed through a plugin-based architecture. All supported backends use HNSW with a user-defined distance metric such as cosine similarity. The paper reports practical runtimes of milliseconds for databases with tens of millions of records, and retrieval is exposed through search_by_vector() after embeddings are inserted during inference (Ghosh et al., 18 May 2026).

Hyrax complements vector retrieval with integrated 2D and 3D latent-space exploration. Dimensionality reduction is performed with UMAP, with planned support for t-SNE. The 2D interface, available through hyrax.visualize() in Jupyter, is built on HoloViews and Datashader, supporting scalable rendering of very large point clouds, coloring by arbitrary scalar quantities, interactive region selection, object-ID retrieval, metadata display in an interactive table, display of a random subset of associated images, resampling through a selection, export of selected subsets as a Pandas DataFrame, and handoff of selections to clustering or similarity search. Because the visualizer is coupled to the DataProvider, it can also display spectra or light curves rather than only images (Ghosh et al., 18 May 2026).

The 3D explorer is browser-based and implemented with three.js, supporting rotate, pan, zoom, interactive region selection, retrieval of metadata and associated imaging, and GPU-accelerated browser interaction. The paper treats the 2D and 3D interfaces as complementary: the former is more tightly integrated with notebook workflows and survey platforms, while the latter gives a more intuitive view of complex embedding structure (Ghosh et al., 18 May 2026).

A major methodological point is that training or validation loss is not treated as an adequate proxy for scientific utility in unsupervised learning. The DP1 galaxy demonstration is used to argue this directly: different latent spaces with different validation losses produced materially different organization of mergers and artifacts. This suggests that Hyrax’s combination of embedding generation, vector indexing, visualization, and iterative inspection is not an auxiliary convenience but part of its core scientific logic (Ghosh et al., 18 May 2026).

5. Demonstrated applications across astronomy

The framework is demonstrated on five representative applications spanning unsupervised discovery, clustering-assisted candidate generation, multimodal transient classification, false-positive rejection in moving-object searches, and supervised detection of semi-resolved dwarf galaxies. These are presented as evidence of breadth rather than as a single benchmark suite (Ghosh et al., 18 May 2026).

Application	Data	Reported outcome
Unsupervised DP1 galaxy analysis	$\sim 4\times 10^5$ Rubin LSST DP1 galaxies in ECDFS	surfaced artifact families, merger candidates, and low-surface-brightness candidates without labels
Cluster-lens candidate workflow	DP1 galaxies embedded from RA, Dec, and $z_{\rm phot}$	identified a candidate arc in an overdensity at $z_{\rm phot}\sim 0.65$
Early-time transient classification	ZTF alerts, $\sim 30{,}000$ real alerts and 18,245 objects	adding spectra reduced SN Ia/core-collapse SN and TDE/AGN confusion
KBMOD false-positive filtering	DEEP coadded stamps, $21\times 21$ pixels	filtered out 72.2% of the total dataset while recovering 95.2% of inserted synthetics expected to pass
Semi-resolved dwarf detection	HSC and LSST-like imaging with synthetic injections	achieved a false-positive rate of 0.2%

In the Rubin DP1 galaxy experiment, Hyrax was used with six-band $ugrizy$ imaging cutouts from ECDFS, a final sample of $4\times 10^5$ objects, and unsupervised or self-supervised models including a convolutional autoencoder, SimCLR, and BYOL. The workflow used UMAP projection, vector-database retrieval, and 2D latent-space exploration. For the $\mathcal{O}(n^2)$0 sample, artifacts separated into an initial island associated with filter-specific defects; after removing these and retraining, a second generation of outliers isolated diffraction spikes, satellite trails, and cosmic-ray hits. The same latent-space analysis surfaced merger candidates relative to the Euclid Quick Release 1 visual morphology catalog, using the merging_merger vote-fraction threshold of 70%, and low-surface-brightness candidates relative to the DES Y3 Gold LSB catalog. In both cases, the paper emphasizes that many nearest neighbors of known or newly identified candidates were not present in the reference catalogs (Ghosh et al., 18 May 2026).

The cluster-lens workflow used UMAP on RA, Dec, and photometric redshift from LePHARE, followed by DBSCAN to identify overdensities interpreted as likely groups or clusters in joint RA-Dec-$\mathcal{O}(n^2)$1 space. Image cutouts of 10" and 20" were then inspected visually. The reported result is a candidate arc in a galaxy cluster field in ECDFS corresponding to a DES overdensity at $\mathcal{O}(n^2)$2, with Euclid VIS morphology consistent with a sheared lensed arc and a second candidate blended in Rubin imaging. The paper explicitly notes the absence of spectroscopic confirmation and does not provide purity or recall metrics for this demonstration (Ghosh et al., 18 May 2026).

For multimodal early-time transient classification, Hyrax hosted the AppleCiDEr pipeline on ZTF alerts. The paper reports approximately 30,000 real alerts and 18,245 objects across five classes: SN I, SN II, CV, AGN, and TDE. Inputs included multiband $\mathcal{O}(n^2)$3 photometric light curves, image triplets of science, reference, and difference cutouts, 17 selected metadata features, and spectra when available. The comparison between three-modal and four-modal models showed that adding spectra reduced class degeneracies, especially SN Ia versus core-collapse SN and TDE versus AGN. TDE performance remained weak because the sample contained only 54 confirmed TDEs (Ghosh et al., 18 May 2026).

The DEEP moving-object application used Hyrax for supervised filtering after KBMOD candidate generation. Coadded result stamps of size 21 by 21 pixels were constructed by following the KBMOD tracklet trajectory through the image stack and averaging the selected pixels. Labels were derived from inserted synthetic objects for the true class and from trajectory parameters 85 degrees away from the ecliptic for a false-positive-rich set. After augmentation, the true set grew from about 9,000 to about 100,000 and the false set from about 50,000 to about 200,000. In first DEEP processing runs, about 300,000 tracklets remained after initial filtering, the ML model removed 72.2% of the total dataset, and 87,490 “good” results remained; 95.2% of the inserted synthetic objects expected to pass were recovered (Ghosh et al., 18 May 2026).

In the dwarf-galaxy application, Hyrax wrapped Rubin source_injection on HSC $\mathcal{O}(n^2)$4 coadds processed with a modified version of the LSST Science Pipelines. Synthetic semi-resolved dwarf galaxies were injected directly into pixel-level data, the resulting images were split into 50" postage stamps, and a ResNet18 CNN was trained with hyrax.train(). The scientific regime targeted was low-mass dwarf galaxies with $\mathcal{O}(n^2)$5, Local Volume $\mathcal{O}(n^2)$6 Mpc, and low surface brightness $\mathcal{O}(n^2)$7 mag arcsec$\mathcal{O}(n^2)$8. The quantitative headline was a false-positive rate of 0.2%, which the authors interpret as making visual inspection tractable at LSST scale (Ghosh et al., 18 May 2026).

6. Software stack, reproducibility, and deployment

Hyrax is implemented in a broad scientific Python and ML ecosystem. The paper lists NumPy, SciPy, Astropy, Pandas, and Matplotlib as general scientific components; PyTorch, TorchVision, PyTorch Ignite, Optuna, MLflow, and TensorBoard for ML and training; HuggingFace Datasets, Pydantic, PyArrow or Apache Arrow, nested-pandas, ONNX, ONNX Runtime, and ONNXScript for data and workflow; pooch, schwimmbad, tqdm, psutil, more-itertools, and colorama as survey-access and utility tools; HoloViews, Datashader, Bokeh, Plotly, and three.js for visualization; Cython for performance; and LanceDB, ChromaDB, and Qdrant for vector-database support. The framework is built on the LINCC Frameworks Python Project Template (Ghosh et al., 18 May 2026).

Its workflow organization is strongly configuration-driven. Users specify dataset choices, model backbone, train or validation logic, hyperparameter space, and inference behavior, and Hyrax records configuration and outputs in a structured form. This is presented as a mechanism for quick local prototyping, later scale-up to HPC or cloud, reproducibility, and systematic comparison across many runs (Ghosh et al., 18 May 2026).

Reproducibility and software engineering receive explicit attention. The paper states that Hyrax is installable via pip install hyrax, that source code is available on GitHub, and that documentation and tutorials are available. It also reports a unit test suite over core modules, CI on every pull request across multiple Python versions, code coverage via Codecov, style enforcement via ruff and pre-commit, semantic versioning for a stable public API, and a frozen release corresponding to the paper at version v0.8.1 (Ghosh et al., 18 May 2026).

Deployment targets include laptops, HPC clusters, survey platforms such as the Rubin Science Platform and Roman Research Nexus, and browser and notebook environments. The scalability strategy combines patch-level survey caching, on-the-fly cutout generation, multi-threaded data preparation, memory preloading in the DataProvider, GPU and multi-GPU support, and approximate nearest-neighbor indexing for embeddings (Ghosh et al., 18 May 2026).

7. Scientific significance, limitations, and adjacent terminology

The main contribution attributed to Hyrax is not a novel loss function or a new model architecture, but a reusable framework that integrates acquisition, dataset handling, multimodal composition, training, inference, experiment tracking, unsupervised exploration, vector search, and visualization into one astronomy-aware system. The paper’s strongest conceptual claim is that unsupervised discovery should be treated as an operational workflow rather than as model fitting alone, because scalar metrics can fail to predict scientific usefulness of learned representations (Ghosh et al., 18 May 2026).

The framework’s strengths, as presented, are end-to-end scope, astronomy-aware design, strong support for unsupervised discovery through vector databases and 2D or 3D latent visualization, modularity with support for arbitrary PyTorch-based models and custom datasets, scalability through GPU and multi-GPU execution and ANN indexing, and explicit reproducibility mechanisms. The diversity of the five demonstrations is used to support the claim that Hyrax is not tied to one subfield: the use cases span galaxies, gravitational lenses, transients, solar-system objects, and dwarf galaxies (Ghosh et al., 18 May 2026).

The limitations are equally explicit. Hyrax is primarily an infrastructure paper, so several demonstrations are described as science previews rather than final benchmark studies. Exact hyperparameters and full training details are often omitted. Quantitative comparisons against alternative frameworks are sparse. Some capabilities, such as broader latent-space reduction methods and richer multimodal visualization beyond images in the 3D interface, are noted as planned rather than fully realized. Future priorities identified in the paper include alert-broker integration, support for community foundation models, better templates and tooling for external contributions, expanded survey support beyond Rubin and HSC, additional latent-space reduction methods such as t-SNE, and richer visualization of spectra and time series (Ghosh et al., 18 May 2026).

In astronomical nomenclature, Hyrax should be distinguished from HIRAX, the Hydrogen Intensity and Real-time Analysis eXperiment, which is a South African 21 cm radio interferometer for intensity mapping, FRB science, and cross-correlation cosmology rather than a machine-learning software framework (Crichton et al., 2021).