Abstract↩︎

1 Introduction↩︎

Machine learning (ML) provides powerful tools for exploring, organizing, and analyzing scientific data. Applications range from automated data curation [1] and feature identification [2], [3] to dimensionality reduction and visualization of high-dimensional datasets [4], [5]. In ultracold atom experiments, many ML tasks follow a common workflow involving data preprocessing, feature extraction, classification, and statistical interpretation. However, analysis pipelines are often developed independently for individual studies or adapted from pretrained architectures in ways that are tightly coupled to specific datasets or experimental systems [6]. As a result, substantial effort is frequently required to reproduce, modify, or extend existing implementations for new applications. This presents an opportunity to streamline the process by providing a unified framework that simplifies the development and deployment of these pipelines.

Here, we present Quantum Gas Analysis and Inference (Q-GAIN), a Python framework for developing and deploying ML- and statistics-based analysis tools for scientific imaging data, with a specific focus on cold-atom quantum gas experiments. Q-GAIN provides an object-oriented, modular interface that supports reusable analysis workflows while encouraging the adaptation, replacement, and extension of its built-in components. The framework integrates tools for retrieving and preprocessing experimental data, maintaining compatibility with existing analysis pipelines, ML-based object classification and tracking, and physics-informed statistical analysis. Its modular architecture allows existing analysis modules to be applied directly to new labeled or unlabeled datasets, or adapted to systems with different physical parameters and imaging conditions.

Although Q-GAIN was originally developed for ultracold atom systems, the framework is not restricted to this domain. New analysis environments can be supported by implementing specialized modules while retaining the layered architecture for data handling, model deployment, and statistical analysis. This design simplifies the development of new analysis workflows by separating domain-specific functionality from shared infrastructure. The framework supports workflows spanning data preprocessing, ML-assisted feature tagging and transformation, downstream statistical analysis, and export of processed results into multiple file formats for dataset generation and curation.

We first present a high-level overview of the Q-GAIN framework using a standard ML task: the classification of handwritten digits from the MNIST dataset [7] (in Section 3). This example demonstrates how to construct new specialized classes and how to repurpose existing models and tools. In Section 4.1, we discuss a reimplementation within Q-GAIN of an existing Python library, SolDet [8]. Finally, in Section 4.2, we present a new specialized [1.0] [1.0]VortexDetector class for identifying and counting vortices in images of ring-shaped Bose-Einstein condensates (BECs), using ML architectures inspired by Ref. [9].

Q-GAIN is available for download from the Q-GAIN GitHub organization, with source code distributed under the GNU Lesser General Public License in the Q-GAIN GitHub repository [10]. Comprehensive API documentation is provided via docstrings embedded in the source code and as HTML pages hosted on GitHub Pages. The library includes type hints for all code to support clarity, maintainability, and extensibility. Unit tests are also distributed with the Q-GAIN repository to facilitate validation and performance evaluation.

2 Framework overview↩︎

The core functionality of Q-GAIN is built around a base [1.0]Detector class, which provides a unified interface for constructing reusable analysis workflows. Instances of the [1.0]Detector class, or any of its subclasses, are referred to as [1.0]detector_objects. These objects provide access to the framework’s primary capabilities, including data handling, ML model training and inference, statistical analysis, and visualization. Their behavior can be customized through configurable modules and user-defined function calls supplied during initialization. A simplified overview of the Q-GAIN class structure is shown in Fig. 1 (a).

Q-GAIN provides flexible data-handling functionality for importing and preprocessing experimental data, loading it into a detector, and exporting analysis results in various formats. By default, [1.0]detector_objects support importing absorption images from ultracold atom experiments stored in HDF files generated by Labscript [11] and provide preprocessing routines, including image cropping, masking, and optical-depth calculation. Support for additional data modalities, including alternative image formats, audio, or text data, can be added by supplying custom preprocessing routines during initialization. Metadata and labels associated with labeled datasets can also be incorporated during the import stage.

Specialized analysis functionality within Q-GAIN is implemented via modular components called “tools.” These tools are managed through subclasses of the [1.0]Controller class, which provides a shared interface for deploying analysis routines. The [1.0]Controller subclasses are specialized for distinct categories of analysis tasks while maintaining a consistent code infrastructure. Users can extend the framework by implementing new controllers or integrating additional tools into existing ones. At present, Q-GAIN includes three primary controller subclasses: [1.0]MLControl, [1.0]StatControl, and [1.0]PlotControl.

The [1.0]MLControl class provides ML-related functionality, including training and inference of PyTorch-compatible models. Currently, Q-GAIN supports classification and object-detection workflows through built-in ML tools. Additional ML models can be directly integrated into [1.0]MLControl using the [1.0]add_new_tool() method. Users can configure standard training parameters such as batch size, number of training epochs, learning rates, optimizers, early stopping criteria, and model checkpointing. Saved model weights can subsequently be reloaded to initialize pretrained models for inference or continued training. The framework also supports configurable validation metrics that are tracked and stored during training.

The statistical analysis capabilities are provided by the [1.0]StatControl class, which manages deployable statistical analysis tools, referred to as [1.0]stat_tools. A [1.0]stat_tool implements downstream analysis routines that operate on ML-generated or preprocessed data products. New statistical analysis methods can be integrated into [1.0]StatControl using the same modular interface as [1.0]MLControl. Configuration parameters and fitted model parameters can optionally be saved and reloaded to support reproducible analysis workflows.

Visualization functionality is implemented through the [1.0]PlotControl class, which manages plotting-oriented tools referred to as [1.0]plot_tools. By default, Q-GAIN includes plotting utilities for visualization tasks such as object-detection scatter plots and classifier confusion matrices. Additional plotting routines can be integrated through the [1.0]add_new_tool() method. The framework supports rendering through Matplotlib [12], including optional use of style sheets and automated export of generated figures.

A typical Q-GAIN workflow consists of importing and preprocessing experimental data, applying ML-based feature identification or classification, and performing downstream statistical analysis. Many ML and statistical analysis tools require an initial training or fitting stage before deployment on new systems or datasets. Once trained, these tools can be reused to automate analysis of new data, with results exported in multiple file formats or visualized directly via the framework’s plotting utilities.

3 Minimal example: MNIST dataset↩︎

To demonstrate the modular design and extensibility of Q-GAIN, we implement a handwritten digit classifier using the MNIST dataset [7]. MNIST is a standard ML benchmark rather than a physics application, and it serves as a compact example of how Q-GAIN separates data management, model execution, statistical analysis, and visualization into reusable, interchangeable components.

The architecture of Q-GAIN is designed to reduce coupling between analysis stages while enabling workflows developed for one application to be adapted to new datasets and experimental systems with minimal task-specific code. The MNIST example demonstrates this design philosophy by combining custom preprocessing routines with the ML and statistical analysis infrastructure already available within the framework.

In this example, a specialized [1.0]MNISTDetector subclass is constructed by extending the base [1.0]Detector interface with dataset-specific preprocessing routines and lightweight wrappers around existing analysis tools. The workflow consists of importing and preprocessing raw data, loading processed samples into the detector, training and deploying ML models, and applying downstream statistical-analysis routines.

3.1 Data management and preprocessing↩︎

Q-GAIN organizes analysis tasks using a directory-based experiment structure that separates raw and processed data from trained models, statistical analysis objects, and exported results. Each experiment directory contains the data products associated with a specific analysis workflow, along with the metadata required to reproduce the workflow. A schematic overview of the resulting data organization within Q-GAIN is shown in Fig. 2.

Within an experiment directory, Q-GAIN maintains a roster file that tracks the locations and descriptors of all imported data samples. Each entry in the roster corresponds to a single data point and contains metadata used to organize and retrieve the associated files. These descriptors are represented using user-defined tags that can encode classification labels, dataset partitions, or other application-specific identifiers. The roster system allows labeled and unlabeled datasets to be managed using a common interface while enabling data subsets to be selectively loaded during later analysis stages.

Data are imported into an experiment through the detector’s [1.0]import_data() method. During import, a user-defined preprocessing function transforms raw input data into a format compatible with Q-GAIN. This preprocessing stage also provides a mechanism for attaching metadata and labels to individual samples. To support extensibility across different data modalities, preprocessing functions are supplied during detector initialization through the [1.0]process_fn argument.

Q-GAIN expects preprocessing functions to return a list-like object containing dictionaries with, at minimum, the keys [1.0]data, [1.0]tag, and [1.0]sub_dir. The [1.0]data field stores the target data used during analysis, while [1.0]sub_dir specifies the experiment subdirectory in which the sample should be stored. The [1.0]tag field provides a descriptive identifier that can subsequently be used for dataset filtering, partitioning, or as a supervised-learning target. Additional metadata can be included through arbitrary user-defined entries. If no preprocessing function is supplied, Q-GAIN defaults to built-in routines for importing absorption images from Labscript [11] experiment files.

class MNISTDetector(Detector):
    def __init__(self) -> None:
        super().__init__(process_fn=self.mnist_process_fn)

    def mnist_process_fn(self, input_data: MNIST, *, train: bool = True) -> list[dict]:
        data_samples = []
        data_set = MNIST(input_data, train=train, download=True)
        for (image, label) in tqdm(data_set, desc="Processing data"):
            sample = {}
            sample["data"] = image
            sample["tag"] = "training" if train else "testing"
            sample["sub_dir"] = sample["tag"]
            sample["label"] = label
            data_samples += [sample]
        return data_samples

mnist_det = MNISTDetector()
mnist_det.import_data(exp_path, **{"train": True})
mnist_det.import_data(exp_path, **{"train": False})
mnist_det.load_data(tags="training", minmax=[0, 255], scale=True)

During import, processed samples and associated metadata are serialized into hierarchical data format (HDF) files. A simplified representation of the roster and data-file organization is shown in Fig. 2 (b) and Fig. 2 (c), respectively. The framework stores datasets along with user-provided metadata, ensuring that experimental parameters, labels, and auxiliary analysis quantities remain associated with individual data samples throughout the workflow.

Listing [lst:mnist_import] demonstrates the preprocessing pipeline for the MNIST example. A custom preprocessing function downloads the MNIST dataset and packages the resulting arrays together with dataset-specific metadata. The [1.0]tag and [1.0]sub_dir entries are used to separate training and testing data into distinct experiment subdirectories, while the handwritten digit labels are stored through an additional [1.0]label metadata entry. The preprocessing function is supplied during initialization of the [1.0]MNISTDetector, after which the data can be imported with calls to [1.0]import_data().

Once imported, datasets can be reloaded into a detector object repeatedly without additional preprocessing. The [1.0]load_data() method retrieves samples associated with selected tags and optionally applies normalization operations to array-like data, as shown in line 21 in Listing [lst:mnist_import]. This design separates the computationally expensive import and preprocessing stages from later training and inference workflows, enabling datasets to be reused across multiple analysis tasks.

3.2 ML controller and model integration↩︎

ML functionality within Q-GAIN is managed through the [1.0]MLControl controller class, which provides a unified interface for training, inference, and deployment of PyTorch-compatible models. The controller encapsulates individual ML analysis routines within modular components referred to as [1.0]MLTools. This abstraction allows heterogeneous models and datasets to share a common workflow interface while remaining independent of specific ML architectures or data representations.

[1.0]MLTools are typically initialized during detector construction, but can also be added dynamically through the [1.0]add_new_tool() method. Each tool is defined through three primary components: a PyTorch model, a dataset interface compatible with PyTorch [1.0]DataLoaders, and a loss function. Because Q-GAIN imposes no constraints beyond standard PyTorch interfaces, existing neural-network architectures and dataset classes can be integrated directly into the framework with minimal modification. This design enables previously developed models to be reused across new analysis workflows while retaining a consistent training and inference interface.

The MNIST example demonstrates this interoperability by specializing one of the built-in Q-GAIN architectures, rather than introducing a new network model. This generalized classification model operates by employing convolutional layers and residual connections [13]. To adapt the classifier to handwritten digit recognition, the number of output classes is modified through the initialization parameter [1.0]num_classes. The number of convolution stacks and their kernel sizes are configured with the [1.0]stack_size and [1.0]kernel_size parameters, respectively. The number of layers is set by specifying a list of filter sizes with the parameter [1.0]filter_list. The remainder of the training and inference workflow is handled directly by Q-GAIN. A lightweight dataset wrapper is additionally implemented to expose the imported MNIST samples through the standard PyTorch dataset interface. The resulting implementation is shown in Listing [lst:mnist_ML].

]
class MNISTDetector(Detector):
    def __init__(self) -> None:
        super().__init__(cl_model=qgain.models.MLST2021CNNmodern,
                         cl_kwargs={"filter_list": [8, 16, 32, 64], "stack_size": 1, "kernel_size": 4, "num_classes": 10},
                         cl_dataset_fn=self.MNISTDataset,
                         cl_aug=None,
                         cl_loss_fn=torch.nn.NLLLoss)

    class MNISTDataset(torch.utils.data.Dataset):
        def __init__(self, data: list[dict] | tuple[dict]) -> None:
            self.data = list(data)

        def __len__(self) -> int:
            return len(self.data)

        def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
            img = self.data[idx]["data"]
            img = torch.from_numpy(img[np.newaxis, np.newaxis, :]).float()
            label = int(self.data[idx]["label"]) if "label" in self.data[idx] else None
            return img, label

    def use_models(self, model_paths: list, model_list: list | tuple,
                   data: list | dict | None = None) -> None:
        super().use_models(model_list=model_list, model_paths=model_paths, data=data)
        target_data = self.data if data is None else data
        for item in target_data:
            if "CL_pred" in item:
                item["CL_pred"] = np.argmax(item["CL_pred"])
        if data is None:
            self.data = target_data

mnist_det = MNISTDetector()
mnist_det.load_data(tags="training", data_frac=0.9, minmax=[0, 255], scale=True)
mnist_det.train_nn(epochs=100, patience=10)
mnist_det.load_data(tags=["testing"], minmax=[0, 255], scale=True, keep=False)
mnist_det.use_models(model_list=["CL"], model_paths=["20260223_190725_CL.pt"])

Q-GAIN supports configurable training workflows with the detector method [1.0]train_nn(), which exposes common optimization parameters including learning rate, batch size, optimizer selection, training epochs, and early stopping criteria; see Table 1. Models can be trained using either internally loaded detector data or externally supplied datasets. Trained model weights may optionally be saved and later reloaded for inference or continued optimization.

Inference is performed using the [1.0]use_models() method, which applies trained [1.0]MLTools to loaded data samples. For the MNIST example, the inherited [1.0]use_models() method is lightly overridden to convert the probability outputs of the classifier into discrete digit labels. This modification illustrates how detector subclasses can adapt existing analysis models to application-specific output formats while preserving the underlying framework workflow.

class MNISTDetector(Detector):
    def __init__(self) -> None:
        super().__init__(process_fn=self.mnist_process_fn,
                         cl_model=qgain.models.MLST2021CNNmodern,
                         cl_kwargs={"filter_list": [8, 16, 32, 64], "stack_size": 1, "kernel_size": 4, "num_classes": 10},
                         cl_dataset_fn=self.MNISTDataset,
                         cl_aug=None,
                         cl_loss_fn=torch.nn.NLLLoss)
        
        self.controllers["ML Controller"].add_new_tool(model=DummyModel, name="DummyModel", dataset_fn=DummyDataset, loss_fn=torch.nn.SmoothL1Loss, device=0, metrics=[{"name": "SmoothL1", "metric": torch.nn.SmoothL1Loss()}])

        self.controllers["ML Controller"].get_tool("CL").device = "cuda:1"
        self.controllers["ML Controller"].get_tool("CL").metrics += [{"name": "GaussianNL", "metric": torch.nn.GaussianNLLLoss()}]   

Beyond model execution, [1.0]MLControl supports additional runtime configuration, including device selection and user-defined validation metrics. Models can be deployed on CPUs or CUDA-enabled GPUs, and arbitrary validation metrics may be incorporated alongside built-in loss tracking during training. These capabilities are exposed through configurable [1.0]MLTool attributes and are illustrated in Listing [lst:ML_params]. Together, these components provide a reusable infrastructure for integrating ML models into scientific-analysis workflows while separating model-specific implementation details from workflow orchestration and data management.

3.3 Statistical analysis tools↩︎

In addition to ML workflows, Q-GAIN provides infrastructure for integrating downstream statistical analysis routines via the [1.0]StatControl controller class. Statistical analysis methods are encapsulated in modular components, referred to as [1.0]stat_tools, and are managed using the same controller-based architecture as ML workflows. This design allows statistical inference routines to operate directly on imported, preprocessed, or ML-generated data products while remaining integrated within a unified analysis pipeline.

New [1.0]stat_tools can be attached to a detector during initialization or dynamically added through the [1.0]add_new_tool() method. Each tool is implemented as a lightweight wrapper around a statistical-analysis algorithm and is expected to expose a [1.0]transform() method that operates on collections of detector data entries. Optionally, [1.0]stat_tools may implement a [1.0]fit() method for training or parameter estimation before deployment. This interface allows external analysis packages to be incorporated into Q-GAIN while preserving a consistent workflow abstraction.

The MNIST example demonstrates this capability using a K-Means clustering implementation based on Scikit-learn [14]. A custom [1.0]KMeansTool wrapper is introduced to adapt the Scikit-learn implementation to the Q-GAIN data interface. The wrapper’s [1.0]fit() method extracts image arrays from detector data entries, performs clustering, and constructs a mapping between cluster assignments and handwritten digit labels. The corresponding [1.0]transform() method subsequently applies the trained clustering model to new samples and converts the predicted cluster indices into digit classifications. The resulting implementation is shown in Listing [lst:mnist_stat].

class MNISTDetector(Detector):
    def __init__(self) -> None:
        super().__init__(stat_tools=[{"name": "KMeans", "tool": self.KMeansTool},], stats_kwargs=[{"n_clusters": 10,}])
    
    class KMeansTool(KMeans):
        def __init__(self, **kwargs: dict):
            super().__init__(**kwargs)
        
        def fit(self, data: list[dict]):
            target_data = []
            self.clust_to_labels = np.zeros((10))
            gmatrix = np.zeros((10,10))
            for sample in data:
                target_data += [sample["data"].flatten()]
            print("Fitting to data.")
            super().fit(np.array(target_data))
            for ground, pred in zip(data, self.labels_):
                gmatrix[int(pred), int(ground["label"])] += 1
            for cluster in range(10):
                self.clust_to_labels[cluster] = gmatrix[cluster, :].argmax()
            del gmatrix
            del target_data

        def transform(self, data: list[dict]):
            target_data = []
            predictions = []
            for sample in data:
                target_data += [sample["data"].flatten()]
            print("Determining predictions.")
            result = self.predict(np.array(target_data))
            for res in result:
                predictions += [self.clust_to_labels[res]]
            del result
            del target_data
            return predictions

mnist_det = MNISTDetector()
mnist_det.load_data(tags=["training"], minmax=[0, 255], scale=True)
mnist_det.define_stat(save=True)
mnist_det.load_data(tags=["testing"], minmax=[0, 255], scale=True, keep=False)
mnist_det.use_models(model_list=["KMeans"], model_paths=["20260223_202454_KMeans.pkl"])

Statistical analysis workflows are deployed via the same detector-level interface used for ML inference. The [1.0]define_stat() method initializes and fits statistical-analysis tools when required, while [1.0]use_models() applies one or more fitted tools to loaded detector data. Because fitted analysis objects can optionally be serialized and reloaded, computationally expensive fitting procedures need not be repeated for subsequent analysis tasks.

Together, the [1.0]StatControl infrastructure and the [1.0]stat_tool abstraction enable classical statistical analysis methods and ML workflows to coexist within a common scientific analysis framework. This allows downstream inference stages to be composed from heterogeneous analysis methods while maintaining a unified data-management and execution interface.

3.4 Export and reproducibility infrastructure↩︎

Q-GAIN includes infrastructure for saving and reusing analysis results, trained models, fitted statistical-analysis tools, and processed datasets to support reproducible workflows. Because imported data, preprocessing outputs, trained model weights, and fitted analysis objects are stored within the experiment directory structure, analysis pipelines can be resumed or reapplied without repeating earlier computational stages.

Processed detector data and analysis results can be exported through the detector method [1.0]export() using multiple output formats, including comma-separated values, HDF, HTML, pickle, and NumPy object arrays. Exported quantities are selected by specifying the corresponding detector-data keys, allowing subsets of analysis outputs to be shared or integrated into external workflows. This functionality enables Q-GAIN to interoperate with downstream analysis environments while maintaining a consistent internal data representation.

The framework also supports generating reusable datasets via the [1.0]generate_samples() method. This method creates portable copies of experiment data structures containing user-selected data fields and metadata entries. The resulting datasets can subsequently be distributed, re-imported into independent Q-GAIN experiments, or reused in alternative training and analysis workflows without requiring repeated preprocessing of the original raw data.

The reproducibility infrastructure of Q-GAIN extends beyond dataset management to include the persistence of trained ML models and fitted statistical analysis tools. Saved model weights can be reloaded for inference or continued training, while fitted [1.0]stat_tool objects can be serialized and restored without repeating computationally expensive fitting procedures. Together, these capabilities allow complete analysis pipelines to be reconstructed from stored experiment directories, facilitating iterative development and reproducible scientific analysis.

4 Scientific applications↩︎

The primary goal of Q-GAIN is to provide a reusable framework for combining ML and statistical analysis workflows within a consistent scientific analysis environment. In this section, we present two applications that demonstrate how this architecture can support both the implementation of existing analysis pipelines and the development of new domain-specific workflows.

Section 4.1 presents the integration of the SolDet soliton detector analysis framework into Q-GAIN, illustrating how an established scientific ML workflow can be migrated into the framework while preserving domain-specific functionality and reusing shared infrastructure. Section 4.2 then demonstrates the extensibility of the framework through the implementation of a specialized vortex-detection workflow for ring-shaped BECs. Together, these examples illustrate how Q-GAIN separates reusable workflow infrastructure from application-specific scientific analysis logic.

4.1 Soliton detection and classification↩︎

The SolDet sub-package illustrates how an existing scientific analysis workflow can be incorporated into Q-GAIN. Originally developed for automated identification and characterization of solitons in BECs, SolDet combines ML-based excitation detection with physically informed statistical analysis methods. Its integration into Q-GAIN demonstrates how specialized scientific workflows can be modernized, generalized, and extended while reusing shared infrastructure for data management, model execution, and reproducible analysis.

4.1.1 Integration of SolDet into Q-GAIN↩︎

SolDet was originally developed as an ML framework for automated localization and identification of solitonic excitations in BECs [15]. The original implementation combined ML-based detection and classification with physically informed statistical analysis to characterize identified excitations. However, as the project evolved, maintaining the standalone implementation became increasingly challenging due to growing software-library incompatibilities and workflow constraints tied to experiment-specific data formats.

To improve maintainability, extensibility, and interoperability, SolDet’s codebase was updated and integrated into Q-GAIN as a dedicated sub-package. Within this framework, SolDet subclasses the base [1.0]Detector class to construct the specialized [1.0]SolitonDetector class. This migration preserved the original scientific analysis workflow while replacing custom infrastructure for data management, preprocessing, model execution, and statistical analysis orchestration with reusable components provided by Q-GAIN.

Integration into Q-GAIN also generalized SolDet’s data-management workflow. The original implementation relied on a highly specific HDF file organization scheme associated with the experimental system on which it was initially developed. By adopting the modular preprocessing and experiment infrastructure of Q-GAIN, the modernized SolDet implementation can now support alternative HDF organization schemes and experimental datasets through configurable preprocessing routines.

Finally, the migration enabled a PyTorch-based ML infrastructure. The original neural network implementations were replaced with specialized classification and object-detection architectures provided by Q-GAIN. This enabled improved interoperability with modern ML tooling while preserving compatibility with the broader Q-GAIN workflow. This redesign additionally introduced architectural refinements, including residual connections and deeper convolutional structures, as well as improvements to the physically informed statistical analysis routines to better support images containing multiple excitations.

The resulting [1.0]SolitonDetector demonstrates how an existing scientific-analysis workflow can be incorporated into Q-GAIN while preserving domain-specific functionality and reusing shared framework infrastructure.

4.1.2 ML-based soliton detection and classification↩︎

The [1.0]SolitonDetector integrates two coordinated ML models for automated analysis of BEC images: an object detector for localization of solitonic excitations and a classifier for identifying excitation classes. Together, these models provide the initial analysis stage of the SolDet workflow before downstream physics-informed statistical analysis is applied.

The classifier assigns the image into one of three classes defined based on the number of potential excitations ( [1.0]class-0: no excitations; [1.0]class-1: a lone excitation; and [1.0]class-2: ambiguous or multiple excitations), while the object detector estimates the excitation positions within the BEC image. To support localization across images with different spatial resolutions and feature sizes, the object detector operates in a compressed cell-space representation, where each cell encodes both the probability that a soliton is present and its approximate position in the original image. The size of this representation can be configured during detector initialization to accommodate different imaging conditions and experimental systems.

Both ML models are integrated into Q-GAIN through the standard [1.0]MLControl interface described in Section 3.2. The [1.0]SolitonDetector supplies custom dataset wrappers for both the classifier and object detector, enabling detector data entries to be converted into PyTorch-compatible datasets while preserving associated metadata and labels. The dataset wrappers additionally support optional data augmentation through horizontal and vertical reflections and \(180^\circ\) rotations.

The classifier model uses a negative log-likelihood loss function, while the object detector employs a custom loss function based on a one-dimensional adaptation of the localization framework introduced in Ref. [9]. This loss combines a cross-entropy term that determines the presence of solitons within individual cells with a mean-squared-error term that refines their predicted positions. Additional validation metrics, including classification and detection accuracy, are incorporated through the configurable metric infrastructure provided by Q-GAIN.

class SolitonDetector(Detector):
    def __init__(self, od_kwargs: dict | None = None, cl_kwargs: dict | None = None,
                 *, augment: bool = True) -> None:
        od_kwargs = {} if od_kwargs is None else od_kwargs
        cl_kwargs = {"num_classes": 3} if cl_kwargs is None else cl_kwargs
        super().__init__(process_fn=process_data,
                         od_model=ObjectDetector, od_dataset_fn=SolitonODDataset,
                         od_loss_fn=MetzLoss, cl_model=SolDetClassifier,
                         cl_dataset_fn=SolitonClassDataset, cl_loss_fn=NLLLoss,
                         cl_aug=augment, od_aug=augment,
                         od_kwargs=od_kwargs, cl_kwargs=cl_kwargs,
                         stat_tools=[{"name": "pie classifier", "tool": PIEClassifier},
                                     {"name": "quality estimator", "tool": QE}],
                         stats_kwargs=[{"func": "modern", 
                                        "transformer": PowerTransformer},
                                       {"func": "modern",
                                        "transformer": PowerTransformer}])

        self.controllers["ML Controller"].get_tool("OD").metrics += [{"name": "Accuracy", "metric": od_accu_metric}]
        self.controllers["ML Controller"].get_tool("CL").metrics += [{"name": "Accuracy", "metric": cl_accu_metric}]
        self.controllers["Plot Controller"].add_new_tool(plot_tools=[{"name": "pie classifier", "tool": self.__plot_pie}])
        self.controllers["Plot Controller"].add_new_tool(plot_tools=[{"name": "quality estimator", "tool": self.__plot_qe}])

    def use_models(self, model_paths: list,
                   model_list: list | tuple = ("classifier", "object detector", "pie classifier", "quality estimator"),
                   data: list | dict | None = None) -> None:
        super().use_models(model_list=model_list, model_paths=model_paths, data=data)

        target_data = self.data if data is None else data
        for item in target_data:
            if "OD_pred" in item and type(item["OD_pred"]) is np.ndarray:
                item["OD_pred"] = pos_41labels_conversion(item["OD_pred"][0])
            if "CL_pred" in item and item["CL_pred"].flatten().shape[0] > 1:
                item["CL_pred"] = np.argmax(item["CL_pred"])
        if data is None:
            self.data = target_data

A simplified overview of the detector configuration is shown in Listing [lst:soldet]. The detector overrides the inherited [1.0]use_models() workflow in order to convert raw network outputs into application-specific predictions. Classifier probability vectors are transformed into discrete excitation labels, while object-detector outputs are converted from compressed cell-space coordinates into image-space soliton positions. The resulting predictions are stored directly in the detector data entries via the [1.0]CL_pred and [1.0]OD_pred fields.

Once initialized, the [1.0]SolitonDetector can classify and track solitons in images of degenerate gases, such as the one shown in Fig. 3 (a). To evaluate the performance of the updated SolDet implementation, the public Dark solitons in BECs dataset 2.0 [16] was partitioned into training, validation, and test sets using the same procedures as in the original SolDet studies [15]. The classifier was trained and tested using the [1.0]class-0, [1.0]class-1, [1.0]class-2, and [1.0]class-8¹ data. The object detector was trained using only the [1.0]class-0 and [1.0]class-1 data due to the lack of reliable position labels for [1.0]class-2, though [1.0]class-2 data was used during testing. Both models employed a 10-fold cross-validation with a 90:10 training-validation split.

Prior to testing, the test set was further filtered based on the initial label agreement among human annotators to facilitate a consistent comparison of performance with the legacy version of SolDet. The test results were used to calculate weighted \(F_1\)² scores for both the classifier and the object detector. For the latter, the \(F_1\) score was calculated based on the agreement between the number of identified excitations and the class label. Confusion matrices for the ML classifier and object detector models with the highest \(F_1\) score are shown in Fig. 3 (b) and Fig. 3 (c), respectively. Additionally, \(R^2\) scores of the excitation positions for the object detector were also calculated for each fold. To handle object detector failure modes, such as false positives (detecting a soliton without a corresponding ground-truth position) or false negatives (failing to detect a soliton at a known position), missing position labels were manually assigned a value of zero.

After approximately 500 training epochs, the SolDet models achieved close agreement with the human-generated labels. The classifier produced a mean weighted \(F_1\) score of \(91(1)~\%\), with the best-performing fold reaching \(92.3~\%\). This represents a marginal improvement over the original SolDet classifier reported in Ref. [15], which achieved a best-model \(F_1\) score of \(91.6~\%\) on the corresponding consensus-labeled test subset, though this increase in performance might be in part due to the random initializations during each round of cross-validation.

The object detector achieved a mean weighted \(F_1\) score of \(89.0(8)~\%\), with the best-performing fold reaching \(90.2~\%\). This same fold reached an \(R^2\) score of \(99.95\%\) for the [1.0]class-1 subset of test data. As shown in Fig. 3 (c), the object detector reliably identifies the absence or presence of single solitons in [1.0]class-0 and [1.0]class-1. Performance on [1.0]class-2 data remains more variable due to the presence of multiple excitations and ambiguous non-solitonic structures. These complex excitation patterns motivate the downstream physics-informed analysis stages described in the following subsection.

4.1.3 Physics-informed statistical analysis↩︎

Beyond ML-based localization and classification, the [1.0]SolitonDetector incorporates downstream physics-informed statistical analysis for characterizing identified excitations. These tools operate on the outputs of the ML workflow and provide physically interpretable descriptions of detected solitonic excitations.

The SolDet implementation includes two statistical analysis tools: the physically-informed excitation (PIE) classifier and the quality estimator (QE). The PIE classifier categorizes detected solitonic excitations into one of six physically-motivated excitation subclasses summarized in Table 2, while the QE estimates the degree to which a detected excitation resembles a canonical longitudinal soliton. Additional details regarding the underlying physically informed analysis methods are provided in Ref. [18].

Both analysis stages are implemented within Q-GAIN using the [1.0]stat_tool interface described in Section 3.3. Each tool exposes [1.0]fit() and [1.0]transform() methods that integrate the corresponding analysis routines into the shared detector workflow. This allows the statistical-analysis stages to operate directly on detector data entries produced by the ML pipeline while remaining compatible with the reusable Q-GAIN infrastructure.

The physically informed analysis routines are based on fitting procedures that estimate parameters describing the structure and geometry of detected solitons. The implementation supports multiple fitting strategies, including Gaussian approximations, the original SolDet fitting routine, and a fitting implementation designed to better support images containing multiple excitations. The resulting fitting parameters can additionally be transformed using configurable statistical transformations, with the default implementation employing a Yeo-Johnson transformation to improve the Gaussianity of the parameter distributions; see lines 12–17 in Listing [lst:soldet].

The [1.0]SolitonDetector extends the inherited detector workflow by incorporating the PIE and QE tools directly into the default analysis pipeline. Following ML-based classification and localization, the detector applies the corresponding statistical analysis routines and stores their outputs in the detector data entries via the fields [1.0]pie classifier_pred and [1.0]quality [1.0]estimator_pred. An example of the resulting combined workflow output is shown in the legend of Fig. 3, while representative PIE-classification and QE-analysis results are shown in Fig. 3 (d) and Fig. 3 (e), respectively.

To evaluate the physically informed analysis stages, a fitting subset was constructed using \(70~\%\) of each PIE subclass (A–F) of [1.0]class-1 data; see Table 2. The PIE classifier was fitted using all PIE subclasses, while the QE was fitted only using subclass A corresponding to longitudinal solitons. The remaining \(30~\%\) of [1.0]class-1 data was reserved for testing.

The PIE classifier achieved a weighted \(F_1\) score of \(94.8~\%\) relative to the labels provided in the public Dark solitons in BECs dataset 2.0 [16], while the QE metric achieved an \(R^2\) score of \(95.7~\%\). Representative results are shown in Fig. 3 (d) and Fig. 3 (e). Comparison with the legacy SolDet outputs reveals modest shifts in both PIE subclass assignments and QE scores. In the PIE classifier, small changes are observed between longitudinal and partial-soliton PIE subclasses. Larger changes occur in the QE predictions, where some excitations shift toward higher or lower estimated soliton quality. These differences arise primarily from improvements introduced in the modified fitting procedures, particularly with changes designed to better support images containing multiple excitations. Because the physically informed analysis methods depend directly on the distributions of fitted soliton parameters, modifications to the fitting workflow naturally alter the resulting statistical classifications.

4.2 Vortex detection in ring BECs↩︎

The vortex detector sub-package illustrates how Q-GAIN can be extended to support new scientific-analysis applications. Whereas the SolDet sub-package demonstrated the modernization and integration of an existing workflow, the vortex detector was developed directly within Q-GAIN as a specialized object-detection pipeline for identifying quantized vortices in time-of-flight images of ring BECs. Only the application-specific components—including the preprocessing routine, ML model, and inference logic—required implementation, while data management, model training, deployment, and export are handled through the shared Q-GAIN infrastructure. The detector’s object-detection workflow and neural-network architecture are inspired by the vortex-localization approach introduced by Metz et al. [9].

4.2.1 VorDet implementation and workflow↩︎

The [1.0]VortexDetector subclasses the base [1.0]Detector interface and integrates custom preprocessing routines, dataset wrappers, ML models, and loss functions through the controller architecture provided by Q-GAIN. A simplified overview of the detector configuration is shown in Listing [lst:vortex]. As with the other applications presented in this work, development is largely limited to implementing application-specific analysis modules, while the underlying infrastructure for data management, model training, deployment, and export is provided by the framework.

class VortexDetector(Detector):
    def __init__(self, **kwargs: dict) -> None:
        super().__init__(process_fn=vortex_process_fn,
                         od_model=ObjectDetector2D,
                         od_dataset_fn=VortexODDataset,
                         od_loss_fn=MetzLoss2D,
                         od_aug=True,
                         **kwargs)
        idx = self.controllers["ML Controller"].get_id(name="OD")
        self.controllers["ML Controller"].tools[idx]["tool"].metrics += [{"name": "Accuracy", "metric": accu_metric}]

    def use_models(self, model_paths: list, data: list | dict | None = None) -> None:
        super().use_models(model_list=("object detector"),
                           model_paths=model_paths, data=data)
        target_data = self.data if data is None else data
        for item in target_data:
            if "OD_pred" in item:
                item["OD_pred"] = vortex_labels_to_data(item["OD_pred"][0], threshold=(0.5, 8))
        if data is None:
            self.data = target_data

The organization of the vortex datasets differs from the default structure expected by Q-GAIN, requiring a custom preprocessing function, [1.0]vortex_process_fn, to import and transform the raw data. During preprocessing, time-of-flight images are cropped, masked, and optionally rescaled, while available metadata and vortex-position labels are extracted and associated with each measurement. The imported samples are then organized into training and testing subsets using the detector’s tagging infrastructure.

The detector’s ML workflow consists of a convolutional object-detection network, a PyTorch-compatible dataset wrapper, and a custom localization loss function. The dataset wrapper converts detector data into tensors suitable for training and optionally applies data augmentation through Gaussian noise, random rotations, horizontal and vertical reflections, and random translations. The object detector is trained using the localization loss introduced in Ref. [9], enabling simultaneous prediction of vortex presence and spatial location.

Like the [1.0]SolitonDetector, the [1.0]VortexDetector represents detections internally using a compressed cell-space representation. Consequently, the inherited [1.0]use_models() method is overridden to convert the network output into physical image coordinates after inference. The converted vortex positions are then stored directly within the detector data entries, allowing them to be processed by the remaining Q-GAIN workflow without additional application-specific code.

4.2.2 Performance↩︎

The vortex detector is trained and evaluated using synthetic time-of-flight images of ring BECs, with representative examples from the test dataset shown in Fig. 4. Prior to training, the dataset is partitioned into 80:20 training and testing subsets. Model development employs 10-fold cross-validation on the training subset using a 90:10 training-validation split, with the final model selected based on the highest weighted \(F_1\) score on the held-out test set.

The current VorDet implementation achieves a mean weighted \(F_1\) score of \(99.1(5)\,\%\) on the test dataset, with the best-performing fold reaching \(100\,\%\). Representative detector outputs are shown in Fig. 4 (a). Examination of the lower-performing models provides insight into the remaining failure modes. False-positive detections occasionally occur when the network identifies ambiguous depletion features as vortices, as illustrated in Fig. 4 (b). Conversely, false-negative predictions arise when valid vortex signatures are missed, as shown in Fig. 4 (c). These examples suggest that increasing the diversity of the training data, particularly by including more ambiguous depletion features, could further improve the detector’s robustness.

Compared with the SolDet object detector, the current VorDet architecture is substantially larger, with 440340 trainable parameters versus 80306 for SolDet. Consequently, a single 500-epoch fold requires about one hour of training on an NVIDIA A100 GPU.

Although additional optimization is desirable to reduce model size and computational cost, the present implementation demonstrates that complex, application-specific analysis workflows can be developed in Q-GAIN using only a small number of detector-specific components. Future work will focus on improving computational efficiency and extending the detector to more diverse experimental datasets and alternative trapping geometries.

5 Conclusion↩︎

We have presented Q-GAIN, a modular Python framework for constructing reusable ML and statistical-analysis workflows for scientific imaging applications. The framework provides shared infrastructure for data management, preprocessing, ML training and inference, downstream statistical analysis, visualization, and export, while allowing application-specific analysis logic to be incorporated through extensible [1.0]Detector and [1.0]Controller interfaces. By separating workflow infrastructure from domain-specific analysis modules, Q-GAIN enables scientific-analysis pipelines to be adapted, extended, and reused across a broad range of applications.

The examples presented in this work demonstrate complementary aspects of the framework architecture. The MNIST workflow illustrates how existing ML models and statistical analysis tools can be integrated into Q-GAIN via lightweight application-specific wrappers. The SolDet sub-package demonstrates how a mature scientific-analysis workflow can be incorporated into the framework while preserving domain-specific functionality and reusing shared infrastructure. This updated revision of the SolDet implementation provides a PyTorch-based ML pipeline, greater flexibility in handling heterogeneous experimental data formats, and enhanced support for multiple excitations across the physically informed statistical analysis stages. Finally, the vortex detector demonstrates how new scientific-analysis workflows can be developed within the framework using only a small number of application-specific modules layered on top of the reusable Q-GAIN infrastructure.

Together, these examples illustrate how Q-GAIN supports layered scientific-analysis workflows in which ML-based feature identification can be combined with downstream physics-informed statistical analysis within a unified execution environment. The public release of the framework provides a foundation for continued development of new detectors, ML models, statistical-analysis tools, and scientific workflows through the Q-GAIN GitHub organization [10]. As additional applications mature, they can be incorporated into the framework as reusable sub-packages while continuing to share common infrastructure for reproducible scientific analysis.

Acknowledgments↩︎

5.0.0.1 Author contributions

M. Doris: Software, Conceptualization, Validation, Methodology, Writing - original draft. S. Guo: Software, Conceptualization, Investigation. S.M. Koh: Software, Formal Analysis, Methodology. L. Ritter: Software, Methodology, Conceptualization. A.R. Fritsch: Software, Data curation, Investigation. S. Mukherjee: Data curation, Investigation. I.B. Spielman: Project administration, Resources, Writing - review & editing. J.P. Zwolak: Project administration, Resources, Software, Writing - review & editing.

5.0.0.2 Funding information

This work was partially supported by the National Institute of Standards and Technology; the Quantum Leap Challenge Institute for Robust Quantum Simulation (OMA-2120757); and the Air Force Office of Scientific Research Multidisciplinary University Research Initiative “RAPSYDY in Q” (FA9550-22-1-0339).

5.0.0.3 Note

The views, conclusions, and recommendations contained in this paper are those of the authors and are not necessarily endorsed nor should they be interpreted as representing the official policies, either expressed or implied, of the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright noted herein. Any mention of commercial products is for information only; it does not imply recommendation or endorsement by the National Institute of Standards and Technology.

References↩︎