Nanosurveyor: a framework for real-time data processing

When new drugs are synthesized [1], dust particles are brought back from space [2], or new superconductors are discovered [3], a variety of sophisticated X-ray microscopes, spectrometers, and scattering instruments are often summoned to characterize their structure and properties. High-resolution and hyperspectral X-ray imaging, scattering and tomography instruments at modern synchrotrons are among the workhorses of modern discovery to study nano-materials and characterize chemical interactions or electronic properties at their interfaces.

A new generation of microscopes are being pioneered, commissioned, and planned at several US Department of Energy (DOE) user facilities [4–6] and elsewhere to achieve superior resolution and contrast in three dimensions, encompassing a macroscopic field of view and chemical or magnetic sensitivity, by coupling together the brightest sources of tunable X-rays, nanometer positioning, nanofocusing lenses, and faster detectors. Existing soft X-ray detector technology in use at the Advanced Light Source (ALS) for example generates 350 MBytes/s per instrument [7]; commercial detectors for hard X-rays can record 6 GB/s or raw data per detector  [8, 9], and a synchrotron light source can support 40 or more experiments simultaneously 24 hours a day. Accelerator technologies such as multi-bend achromat [10] are poised to increase brightness by two orders of magnitude around the globe  [11, 12]. Next generation microscopes may exploit multi-color sources, increased detector parallelism, increased frame rate, or stroboscopic structured illumination to extract higher-dimensional, higher resolution, higher frame-rate characterization of a specimen. There is a need for reducing data into meaningful images as rapidly as it is acquired, using low-cost algorithms and computational resources.

Modern synchrotron experiments often have quite complex processing pipelines, iterating through many different steps until reaching the final output. One example for such an experiment is ptychography [13–15], which enables one to build up very large images by combining the large field of view of a high-precision scanning microscope system with the resolution provided by diffraction measurements.

Ptychography uses a small step size relative to the size of the illuminating beam when scanning the sample, continuously generating large redundant datasets that can be reduced into a high-resolution image. Resolution of a ptychography image does not depend directly on the size or shape of the illumination. X-ray wavelengths can probe atomic and subatomic scales, although resolution in scattering experiments is limitated by other factors such as radiation damage, exposure, and brightness of the source to a few nanometers except in special cases (such as periodic crystals). To reconstruct an image of the object from a series of X-ray scattering experiments, one needs to solve a difficult phase retrieval problem, because at short wavelengths it is only possible to measure the intensity of the photons on a detector. The phase retrieval problem is made tractable in ptychography by recording multiple diffraction patterns from overlapping regions of the object, providing redundant datasets to compensate for the lack of the phase information. The problem is made even more challenging in the presence of noise, experimental uncertainties, optical aberrations, and perturbations of the experimental geometry which require specialized solvers and software [16–18].

In addition to its reconstruction pipeline, a ptychography experiment involves additional I/O operations such as calibrating the detector, filtering raw data, and communicating parameters (such as X-ray wavelength, scan positions, detector distance, and flux or exposure times) to the analysis infrastructure.

Large community driven projects have developed frameworks optimized for distributed data stream processing. Map-reduce-based solutions such as Hadoop [19, 20] and Spark [21] provide distributed I/O, a unified environment, and hooks for running map and reduce operations over a cloud-based network. Other frameworks such as Flink [22], Samza [23], and Storm [24] are more tailored for real-time stream processing of tasks executing a directed acyclic graph (DAG) [25] of operations as fast as possible. Workflow graphs such a Luigi [26] and Dask distributed [27, 28] provide an iterative component, but are either optimized for batch processing and workers are treated as a singular entity able to execute the DAG in its entirety.

Such frameworks target operations as a unit of tasks and generalize the notion of resources; however, the ecosystem is harder to decentralize. These paradigms are not easily mappable to a production beamline environment, where processing algorithms working on data from a detector might be running on a field-programmable gate array (FPGA), the motion control system on a real-time MCU, the acquisition control on a windows operating system, and the scientist a macOS laptop. The rest of the pipeline tasks might hop to several different architectures including CPUs for latency bound tasks, and GPUs for high-throughput image processing and visualization. While frameworks such as Flink along with Kafka [29] (high-throughput distributed message system) and ZooKeeper [30] (distributed coordination and management) can be adopted to fit the described processing environment, our solution at a lower level accomplishes the same task with less computational and human resources.

Nanosurveyor is a modular framework to support distributed real-time analysis and visualization of data. The framework makes use of a modular infrastructure similar to Hummingbird [31] developed to monitor flash X-ray imaging experiments at free electron lasers (FELs) with high data rates in real time over multiple cores and nodes. Similar frameworks and pipelines have been also implemented in other research fields such as serial crystallography [32], cryo-electron microscopy [33], or functional magnetic resonance imaging [34].

Within this framework, we developed a streamlined processing pipeline for ptychography which unifies all components involved and allows scientists to monitor and quickly act upon changes along the experimental and computational pipeline.

Nanosurveyor was developed to provide real-time feedback through analysis and visualization for experiments performed at synchrotron facilities, and execute a complex set of operations within a production environment. Its design is such that it can be effectively adapted to different beamline environments. It is built around a client–server infrastructure allowing scientists to use facility resources while located at a beamline or remotely, operating on live data streamed from the beamline. Additionally, one can use the Nanosurveyor user interface for off-line processing of experimental data saved on disk. In this section, we describe the resources and capabilities provided by the modular streaming infrastructure.

Modular framework

As described above, Nanosurveyor is designed to be adaptable and modular. Therefore, we designed it with a client–server infrastructure (Fig. 1) enabling scientists to run their experiment while at the beamline or remotely from their institution. This strategy also allows the client to be very light and flexible while the server can be scaled according to the resources needed.

The Nanosurveyor infrastructure equips each module with two fundamental capabilities. First, a description format language of key-value pairs allows every module to describe its input and output. Second, it provides the ability to describe the connection between the modules, including the front-end.

The capability to connect the communication path between modules allows the end-to-end pipeline to be constructed and described seamlessly. This is done through a proxy communication layer allowing the modules to run either closely together or on completely separate machines. This strategy is transparent to the beamline user and accommodates both environments with centralized resources as well as those where resources are spread across a network.

Additionally, as each module in the pipeline can be executed in its own environment, Nanosurveyor provides dynamic parallelism by allowing the user to scale the number of resources available to each step: this is done by treating each stage as a worker process that can be scaled up or down to address bottleneck or performance issues.

We adapted the outlined streaming framework described above for the specific needs of ptychography and are currently implementing this ptychography streaming pipeline at the beamline for scanning transmission X-ray microscopy (STXM) at the ALS. The main motivation for this project is to make high-resolution ptychographic reconstructions available to the scientist in real-time. To achieve this goal, we streamlined all relevant processing components of ptychography into a single unit. A detailed outline of our pipeline is sketched in Fig. 2.

As described in the previous sections, we follow the idea of a modular streaming network using a client–server architecture, with a back-end for ptychographic processing pipeline and a front-end for configuration, control and visualization purposes.

On the back-end side, the streaming infrastructure is composed of a communication handler and four different kinds of workers addressing dark frames, diffraction frames, reduced and downsampled images and the ptychographic reconstruction using a software package for scalable heterogeneous adaptive real-time ptychography (SHARP [18]). The handler bridges the back-end with the front-end and controls the communication and data flow among the different back-end workers. The dark worker accumulates dark frames and provides statistical maps (mean and variance) of the noise structure on the detector. The frame workers transform raw into clean (pre-processed) diffraction frames. This involves a subtraction of the average dark, filtering, photon counting, and downsampling. Depending on the computing capacities of the back-end, it is possible to run as many frame workers simultaneously as needed. The image worker reduces a collection of clean diffraction frames, producing low-resolution image reconstructions and an initial estimate of the illumination function which, together with the clean diffraction frames, is then feeded as an input for the high-resolution ptychographic reconstruction worker (SHARP).

The front-end consists of a worker that reads raw data frames from a fast charge-coupled device (FCCD) [42], coordinating with a separately developed interface for controlling the experiment (such as motors and shutters) and a graphical user interface (GUI) which is used both for visualizing and controlling the ongoing reconstruction. An example view of the GUI for streaming ptychography is shown in Fig. 3.

Following the data flow along the streaming pipeline, the starting trigger comes from the control interface which initiates a new ptychographic scan providing information about the scan (step size, scan pattern, number of scan points) and other relevant information (e.g., wavelength) to the back-end handler. Simultaneously, the control sends triggers to the scanning motors and the FCCD. A typical ptychographic scan combines the accumulation of a given number of dark frames together with scanning the sample in a region of interest. The frame-grabber, already waiting for raw data packets to arrive, assembles the data and sends it frame-by-frame to the back-end handler. When dealing with an acquisition control system that runs independently, the handler can distinguish between dark and data frames using counters. Dark and data frames are distributed to the corresponding workers. Having clean diffraction frames and an initial guess for the illumination ready, the SHARP worker is able to start the iterative reconstruction process. SHARP initializes and allocates space to hold all frames in a scan, computes a decomposition scheme, initializes the image and starts the reconstruction process. Unmeasured frames are either set to a bright-field frame (measured by removing the sample) or their weight is set to 0 until real data are received.

Depending on the configuration, data at different states within the streaming flow can be displayed in the GUI and/or saved to a CXI file via the nscxiwrite worker module.

All components of the streaming interface run independent event loops and use asynchronous (non-blocking) socket communication. To maximize performance, the front-end operates very close to the actual experiment, while the back-end runs remotely on a powerful GPU/CPU cluster.

This work introduced Nanosurveyor—a framework for real-time processing at synchrotron facilities. The infrastructure provides a modular framework, support for load-balancing operations, the ability to run in a distributed client–server mode, and gives feedback on each stage of a complex pipeline.

The framework was adapted to support streamlined pipelines for ptychography. In this case, expensive stages such as pre-processing are load-balanced with multiple workers, and image reconstruction are parallelized over MPI to compute efficiently in a distributed manner. Results from every stage of the pipeline are then transmitted to the front-end, providing users at the beamline comprehensive knowledge of the experiment and of how the data are transformed from start of acquisition to end output. Although the Nanosurveyor framework provides several core capabilities that are necessary for operating at typical beamlines, there are several key advances that we are currently working on to make the computational pipeline complete. A couple of highlights include the following:

Iterative execution, instrument control Adding support for controlling the beamline itself will complete the current pipeline and provide an iterative execution loop enabling future pipelines to adaptively acquire and analyze data from the operating beamline, and automatically request more data when necessary. For example, if the reconstruction detects bad frames, or that the sample has drifted, then more frames can be automatically requested on the fly without interrupting the overall experiment. If the reconstruction determines that part of the image being acquired is empty or uninteresting it could request fewer frames and focus on the relevant part of the sample.

Optimizing pipeline execution Currently, communication occurs over ZeroMQ providing many benefits, including dealing with backlog, automated load-balancing, and the ability to interleave work running different stages of the execution pipeline. We are also investigating ways to fuse modules to optimize execution times. Making communication agnostic by using handles enables efficient use of memory optimization strategies, socket communication, or saving on data movement costs, e.g., transferring data between GPU-based modules by moving a pointer rather than copying data.

In conclusion, we have presented a framework that is built to run at modern beamlines, can handle the geographic considerations between users and experiments running at synchrotron facilities, and supports real-time feedback. These features, along with the modular design, provide a foundation that can be extended and readily deployed on many of the beamlines in use today.