Implementing a data loader plugin

Create a loader plugin for an ARPES setup by subclassing LoaderBase. Each subclass is registered automatically.

At minimum, define the name attribute and the load_single method. Add other attributes and methods when the file format requires them. The registry below shows the loaders that are available before the example plugin is defined.

import erlab

erlab.io.loaders

Data loading flow

First, examine how ERLabPy resolves files and processes the loaded data.

The core loader method is load_single. It receives a path to one file and returns an xarray object. In most cases, it returns an xarray.DataArray. For complex data, such as multiple region scans with different axes, it can return an xarray.Dataset or xarray.DataTree. Do not include dimension renaming or reordering in this method. Class attributes can apply these post-processing steps automatically.

ARPES data files from one experiment usually follow a fixed naming scheme, such as file_0001.h5 and file_0002.h5. A defined naming scheme lets the loader infer a file path from a sequence number. Implement identify to support this behavior. The method receives an integer sequence number (identifier) and a directory path (data_dir). It returns a list of matching file paths and a mapping of scan coordinates. For a single-file scan, it returns a one-item file list and an empty coordinate mapping.

The following flowchart shows the process of loading data from a single scan, given the path to the directory (data_dir) and the sequence number or file name (identifier):

Flowchart for loading data from a single file

The default infer_index implementation recovers a sequence number from trailing decimal digits in a bare file name. Override it whenever the naming scheme needs different parsing. Besides supporting file-name-based loading, this provides the numeric fallback for spreadsheet metadata when a logged name does not match directly.

Some setups save the data from one scan over multiple files, such as file_0001_0001.h5, file_0001_0002.h5, and so on. In this case, the final number identifies a file within the scan rather than the scan itself, so the loader must override infer_index to recover the scan number. The following flowchart shows the process of loading data from multiple files:

Flowchart for loading data from multiple files

In this case, the method identify should resolve all files that belong to the given sequence number, and return a list of file paths along with a dictionary of coordinates that are varied across the files. For example, if there are three files for a scan taken at three different beta angles, the method should return a list of three file paths and a dictionary with 'beta' as the sole key and an array of length 3 containing the angle as the value. An empty dictionary should be returned if there are no varying coordinates.

The method infer_index receives a bare file name without its extension or directory. For example, a multi-file loader should infer 3 from file_0003_0123 rather than using the default trailing value 123.

A minimal example

Consider a setup that saves data into a .csv file named data_0001.csv, data_0002.csv, and so on. A simple implementation of a loader for the setup will look something like this:

import os

import pandas as pd

from erlab.io.dataloader import LoaderBase


class MyLoader(LoaderBase):
    name = "my_loader"
    description = "Barebones loader for CSV files"
    extensions = {".csv"}
    skip_validate = False
    always_single = True

    def identify(self, num, data_dir):
        file = os.path.join(data_dir, f"data_{str(num).zfill(4)}.csv")
        return [file], {}

    def load_single(self, file_path, without_values=False):
        return pd.read_csv(file_path).to_xarray()

Some class attributes and methods have been implemented. For a detailed explanation of each attribute and method, see the LoaderBase documentation.

We can see that the loader has been properly registered:

erlab.io.loaders

erlab.io.loaders["my_loader"]

The loader can be used just like the built-in loaders:

data = erlab.io.loaders.my_loader.load(1, data_dir="/path/to/data)

Handling metadata

Unlike the previous example, real ARPES data is more than just a simple array of numbers. It contains metadata such as the experimental geometry, sample temperature, and so on. It is important to store this metadata in the xarray object in a consistent manner as defined by the ARPES data conventions.

To obtain a consistent representation of the data, data loaded by load_single must be post-processed to adhere to the conventions. Typically, this involves manipulating coordinate and attribute names, which is automatically performed based on the following class attributes:

Any post-processing steps that reach beyond renaming and reordering dimensions can be implemented in the post_process method:

def post_process(self, data: xr.DataArray) -> xr.DataArray:
    data = super().post_process(data)
    # Perform additional post-processing steps here
    return data

The loaders perform a basic check for some of the data conventions using validate for every data file loaded. A warning is issued if some are missing. This behavior can be controlled with loader class attributes skip_validate and strict_validation.

Data spanning multiple files

Next, let’s create a more realistic loader for a hypothetical HDF5 setup. The setup saves single scans as data_001.h5 and data_002.h5. It saves multiple-file scans as data_001_S001.h5 and data_001_S002.h5. A separate file named data_001_axis.csv stores the scan-axis information.

Let us first generate a data directory and place some synthetic data in it. Before saving, we rename and set some attributes that resemble real ARPES data.

import csv
import datetime
import tempfile

import numpy as np

import erlab
from erlab.io.exampledata import generate_data_angles


def make_data(beta=5.0, temp=20.0, hv=50.0, bandshift=0.0):
    data = generate_data_angles(
        shape=(250, 1, 300),
        angrange={"alpha": (-15, 15), "beta": (beta, beta)},
        hv=hv,
        configuration=1,
        temp=temp,
        bandshift=bandshift,
        assign_attributes=False,
        seed=1,
    ).T

    # Rename coordinates. The loader must rename them back to the original names.
    data = data.rename(
        {
            "alpha": "ThetaX",
            "beta": "Polar",
            "eV": "BindingEnergy",
            "hv": "PhotonEnergy",
            "xi": "Tilt",
            "delta": "Azimuth",
        }
    )
    dt = datetime.datetime.now()

    # Assign some attributes that real data would have
    data = data.assign_attrs(
        {
            "LensMode": "Angular30",  # Lens mode of the analyzer
            "SpectrumType": "Fixed",  # Acquisition mode of the analyzer
            "PassEnergy": 10,  # Pass energy of the analyzer
            "UndPol": 0,  # Undulator polarization
            "Date": dt.strftime(r"%d/%m/%Y"),  # Date of the measurement
            "Time": dt.strftime("%I:%M:%S %p"),  # Time of the measurement
            "TB": temp,
            "X": 0.0,
            "Y": 0.0,
            "Z": 0.0,
        }
    )
    return data


# Create a temporary directory
tmp_dir = tempfile.TemporaryDirectory()

# Define coordinates for the scan
beta_coords = np.linspace(2, 7, 10)

# Generate and save cuts with different beta values
for i, beta in enumerate(beta_coords):
    data = make_data(beta=beta, temp=20.0, hv=50.0)
    filename = f"{tmp_dir.name}/data_001_S{str(i + 1).zfill(3)}.h5"
    data.to_netcdf(filename, engine="h5netcdf")

# Write scan coordinates to a csv file
with open(f"{tmp_dir.name}/data_001_axis.csv", "w", newline="") as file:
    writer = csv.writer(file)
    writer.writerow(["Index", "Polar"])

    for i, beta in enumerate(beta_coords):
        writer.writerow([i + 1, beta])

# Generate some cuts with different band shifts
for i in range(4):
    data = make_data(beta=5.0, temp=20.0, hv=50.0, bandshift=-i * 0.05)
    filename = f"{tmp_dir.name}/data_{str(i + 2).zfill(3)}.h5"
    data.to_netcdf(filename, engine="h5netcdf")

The generated folder resembles typical data from an ARPES experiment. List the contents of the folder:

sorted(os.listdir(tmp_dir.name))

Each HDF5 file represents a single ARPES cut. data_001_S001.h5 to data_001_S010.h5 represents an ARPES map with 10 cuts, with the scan axis recorded in data_001_axis.csv. Inspect the raw data.

xr.load_dataarray(f"{tmp_dir.name}/data_002.h5")

The data has been properly loaded, but the coordinates and attributes have names that are specific to the beamline.

Our loader should do three things: rename the coordinates and attributes to standard names, add metadata to the dataset, and combine related cuts into a single DataArray that contains the ARPES mapping.

Note

Here, we easily loaded the data into an xarray object directly, but that is not the case for most experimental setups. Properly loading raw data into an xarray object is a complex process that requires knowledge of the data format and the experimental setup, and this is what must be implemented in the load_single.

ERLabPy provides convenient functions to ease this process. See implementations of existing data loaders for examples.

Now, we will implement the loader. Make identify handle multiple files for one scan. Also implement infer_index to extract the scan number from the file name.

import pathlib
import re

import erlab


class ExampleLoader(erlab.io.dataloader.LoaderBase):
    name = "example"
    description = "Example loader for multiple files"
    extensions = {".h5"}

    name_map = {
        "eV": "BindingEnergy",
        "alpha": "ThetaX",
        "beta": ["Polar", "Polar Compens"],
        # Can have multiple names assigned to the same name
        # If both are present in the data, a ValueError will be raised
        "delta": "Azimuth",
        "xi": "Tilt",
        "hv": "PhotonEnergy",
        "polarization": "UndPol",
        "sample_temp": "TB",
    }
    # Map the names of the coordinates or attributes in the resulting data to the names
    # present in the data returned by `load_single`. Note that the order of
    # non-dimension coordinates in the output data will follow the order of the keys in
    # this dictionary.

    coordinate_attrs: tuple[str, ...] = (
        "beta",
        "delta",
        "xi",
        "hv",
        "X",
        "Y",
        "Z",
        "polarization",
        "photon_flux",
        "sample_temp",
    )
    # Attributes to be used as coordinates. Place all attributes that we do not want to
    # lose when merging multiple file scans here.

    additional_attrs = {
        "configuration": 1,  # Experimental geometry. Required for momentum conversion
        "sample_workfunction": 4.3,
    }
    # Any additional metadata you want to add to the data. Note that attributes defined
    # here will not be transformed into coordinates. If you wish to promote some fixed
    # attributes to coordinates, add them to additional_coords.

    additional_coords = {}
    # Additional non-dimension coordinates to be added to the data, for instance the
    # photon energy for lab-based ARPES.

    always_single = False

    def identify(self, num, data_dir):
        data_dir = pathlib.Path(data_dir)

        coord_dict = {}

        # Look for scans with data_###_S###.h5, and sort them
        files = sorted(data_dir.glob(f"data_{str(num).zfill(3)}_S*.h5"))

        if len(files) == 0:
            # If no files found, look for data_###.h5
            files = sorted(data_dir.glob(f"data_{str(num).zfill(3)}.h5"))
            if len(files) > 1:
                # More than one file found with the same scan number, show warning
                erlab.utils.misc.emit_user_level_warning(
                    f"Multiple files found for scan {num}, using {files[0]}"
                )
                files = files[:1]
        else:
            # If files found, extract coordinate values from the filenames
            axis_file = data_dir / f"data_{str(num).zfill(3)}_axis.csv"
            with axis_file.open("r") as f:
                header = f.readline().strip().split(",")

            # Load the coordinates from the csv file
            coord_arr = np.loadtxt(axis_file, delimiter=",", skiprows=1)

            # Each header entry will contain a dimension name
            for i, hdr in enumerate(header[1:]):
                coord_dict[hdr] = coord_arr[: len(files), i + 1].astype(np.float64)

        if len(files) == 0:
            # If no files found up to this point, return None
            return None

        return files, coord_dict

    def load_single(self, file_path, without_values=False):
        return xr.open_dataarray(file_path, engine="h5netcdf")

    def infer_index(self, name):
        # Get the scan number from file name
        try:
            scan_num: str = re.match(r".*?(\d{3})(?:_S\d{3})?", name).group(1)
        except (AttributeError, IndexError):
            return None, None

        if scan_num.isdigit():
            # The second return value, a dictionary, is reserved for more complex
            # setups. See tips below for a brief explanation.
            return int(scan_num), {}
        return None, None

erlab.io.loaders

The example loader is registered. Test the loader by loading and plotting some data.

erlab.io.set_loader("example")
erlab.io.set_data_dir(tmp_dir.name)
erlab.io.load(1)

erlab.io.load(5).qplot()

Brilliant! We now have a working loader for our hypothetical setup.

Note

  • There are more class attributes and methods that can be inherited or overridden to customize the loader’s behavior.

  • For single-file loaders which save data in well-known formats such as outputs from Scienta Omicron DA30 analyzers, SES, or NeXus, the implementation can be much more straightforward. See the implementations of existing data loaders for examples.

Summary generation

Note

Summary generation is discouraged. For new data loader plugins, we recommend using Data Explorer to browse, preview, and load data. The interfaces below remain documented for existing loaders and code that uses erlab.io.summarize().

An existing loader requires two attributes and one method to support summary generation:

  • formatters: A dictionary that maps attribute or coordinate names in the data to functions that convert the coordinate or attribute value into a human-readable form.

  • summary_attrs: A dictionary that maps summary column names to attribute or coordinate names in the data. A callable can also be used to generate entries for attributes that are not directly present in the data.

  • files_for_summary: A method that takes a path to a directory and returns a list of file paths in the directory that are associated with the loader.

You can also choose to implement the following attribute to further customize the summary:

  • summary_sort: A string that determines the column name to sort the summary table with.

    If not provided, the table will respect the order of the files returned by files_for_summary.

To improve the performance of summary generation, you can optionally implement load_single to utilize the without_values argument. If it is True, it means that the values in the returned data of load_single will not be accessed, so you can return the data with its values set to arbitrary numbers. This is useful when only the metadata is needed for the summary. An example of this will be shown below.

def _format_polarization(val) -> str:
    val = round(float(val))
    return {0: "LH", 2: "LV", -1: "RC", 1: "LC"}.get(val, str(val))


def _parse_time(darr: xr.DataArray) -> datetime.datetime:
    return datetime.datetime.strptime(
        f"{darr.attrs['Date']} {darr.attrs['Time']}", "%d/%m/%Y %I:%M:%S %p"
    )


def _determine_kind(darr: xr.DataArray) -> str:
    data_type = "xps"
    if "alpha" in darr.dims:
        data_type = "cut"
    if "beta" in darr.dims:
        data_type = "map"
    if "hv" in darr.dims:
        data_type = "hvdep"
    return data_type


class ExampleLoaderComplete(ExampleLoader):
    name = "example_complete"
    description = "Example loader that supports summary generation"

    formatters = {
        "polarization": _format_polarization,
        "LensMode": lambda x: x.replace("Angular", "A"),
    }

    summary_attrs = {
        "Time": _parse_time,
        "Type": _determine_kind,
        "Lens Mode": "LensMode",
        "Scan Type": "SpectrumType",
        "T(K)": "sample_temp",
        "Pass E": "PassEnergy",
        "Polarization": "polarization",
        "hv": "hv",
        "x": "X",
        "y": "Y",
        "z": "Z",
        "polar": "beta",
        "tilt": "xi",
        "azi": "delta",
    }

    summary_sort = "Time"

    def load_single(self, file_path, without_values=False):
        darr = xr.open_dataarray(file_path, engine="h5netcdf")

        if without_values:
            # Prevent loading values into memory
            return xr.DataArray(
                np.zeros(darr.shape, darr.dtype),
                coords=darr.coords,
                dims=darr.dims,
                attrs=darr.attrs,
                name=darr.name,
            )

        return darr

    def files_for_summary(self, data_dir):
        return erlab.io.utils.get_files(data_dir, extensions=[".h5"])


erlab.io.loaders

View the resulting summary.

Note

  • If ipywidgets is not installed, only the DataFrame will be displayed.

  • If you are viewing this documentation online, the summary will not be interactive. Run the code locally to try it out.

erlab.io.set_loader("example_complete")
erlab.io.summarize()

Each cell in the summary table is formatted with formatter after applying the formatters.

Testing a contributed loader

Add regression tests when you contribute a new loader or change an existing one. ERLabPy stores large ARPES test files in the separate erlabpy-data repository.

  1. Fork and clone the erlabpy-data repository.

  2. Create a directory in the repository root. Use the loader name for the directory name.

  3. Add the minimum raw files required by the tests. When practical, also add a processed result that the test can use as an expected value.

  4. Set ERLAB_TEST_DATA_DIR to the path of the cloned erlabpy-data repository.

  5. Add the loader tests to tests/io/plugins/test_<plugin_name>.py. Use the test_data_dir fixture to locate the test files.

  6. Run the loader tests:

    uv run pytest tests/io/plugins/test_<plugin_name>.py

  7. Push the test files to your erlabpy-data fork and open a pull request.

  8. After the data pull request is merged, update these values in tests/test-data.env:

    • Set ERLAB_TEST_DATA_COMMIT to the full commit hash that contains the test files.

    • Set ERLAB_TEST_DATA_ARCHIVE_SHA256 to the SHA-256 hash reported by the checksum workflow for that commit.

  9. Include the loader, its tests, and the updated tests/test-data.env in the ERLabPy pull request.

Tips

  • The data loading framework is designed to be simple and flexible, but it may not cover all possible setups. If you encounter a setup that cannot be loaded with the existing API, please let us know by opening an issue!

  • Before implementing a loader, see erlab.io.dataloader for descriptions about each attribute, and the values and types of the expected outputs. The implementation of existing loaders in the erlab.io.plugins module is a good starting point; see the source code on github.

  • If you wish to add general post-processing steps such as fixing the sign of the binding energy coordinates, you can reimplement post_process which by default handles coordinate and attribute renaming.

  • For complex data structures, constructing a full path from just the sequence number and the data directory can be difficult. In this case, identify can be implemented to take additional keyword arguments. All additional keyword arguments passed to load are passed to identify.

    For example, consider A_001.h5, A_002.h5, and B_001.h5 in the same directory. The sequence number alone does not identify the prefix. Add a prefix argument to identify to remove this ambiguity. Then load A_001.h5 with erlab.io.load(1, prefix="A").

    For multiple-file scans, file names can include A_001_S001.h5 and A_001_S002.h5. Use the second return value of infer_index to pass prefix to load. Return a dictionary that contains the argument.

    For an example of this, see the implementation of erlab.io.plugins.erpes.ERPESLoader.

  • If you have implemented a new loader or have improved an existing one, consider contributing it to the ERLabPy project by opening a pull request. We are always looking for new loaders to support more experimental setups! See Development setup and workflow for the pull-request workflow.

Don’t forget to cleanup the temporary directory!

tmp_dir.cleanup()