Source code for arpes.endstations.plugin.HERS

"""Preliminary implementation of data loading at the ALS HERS beamline."""
import copy
import itertools
import os.path
import warnings

import numpy as np
from astropy.io import fits

import arpes.config
import xarray as xr
from arpes.endstations import HemisphericalEndstation, SynchrotronEndstation, find_clean_coords
from arpes.provenance import provenance_from_file
from arpes.utilities import rename_keys

__all__ = ("HERSEndstation",)


class HERSEndstation(SynchrotronEndstation, HemisphericalEndstation):
    """Implements data loading at the ALS HERS beamline.

    This should be unified with the FITs endstation code, but I don't have any projects at BL10
    at the moment so I will defer the complexity of unifying them for now
    """

    PRINCIPAL_NAME = "ALS-BL1001"
    ALIASES = ["ALS-BL1001", "HERS", "ALS-HERS", "BL1001"]

    def load(self, scan_desc: dict = None, **kwargs):
        """Loads HERS data from FITS files. Shares a lot in common with the Lanzara group formats."""
        if scan_desc is None:
            warnings.warn("Attempting to make due without user associated scan_desc for the file")
            raise TypeError("Expected a dictionary of scan_desc with the location of the file")

        scan_desc = dict(copy.deepcopy(scan_desc))

        data_loc = scan_desc.get("path", scan_desc.get("file"))
        data_loc = (
            data_loc if data_loc.startswith("/") else os.path.join(arpes.config.DATA_PATH, data_loc)
        )

        hdulist = fits.open(data_loc)

        hdulist[0].verify("fix+warn")
        _header_hdu, hdu = hdulist[0], hdulist[1]

        coords, dimensions, spectrum_shape = find_clean_coords(hdu, scan_desc)
        columns = hdu.columns  # pylint: disable=no-member

        column_renamings = {}
        take_columns = columns

        spectra_names = [name for name in take_columns if name in columns.names]

        skip_frags = {}
        skip_predicates = {lambda k: any(s in k for s in skip_frags)}
        scan_desc = {
            k: v for k, v in scan_desc.items() if not any(pred(k) for pred in skip_predicates)
        }

        data_vars = {
            k: (
                dimensions[k],
                hdu.data[k].reshape(spectrum_shape[k]),
                scan_desc,
            )  # pylint: disable=no-member
            for k in spectra_names
        }
        data_vars = rename_keys(data_vars, column_renamings)

        hdulist.close()

        relevant_dimensions = {
            k
            for k in coords.keys()
            if k in set(itertools.chain(*[l[0] for l in data_vars.values()]))
        }
        relevant_coords = {k: v for k, v in coords.items() if k in relevant_dimensions}

        deg_to_rad_coords = {"beta", "psi", "chi", "theta"}
        relevant_coords = {
            k: c * (np.pi / 180) if k in deg_to_rad_coords else c
            for k, c in relevant_coords.items()
        }

        dataset = xr.Dataset(
            data_vars,
            relevant_coords,
            scan_desc,
        )

        provenance_from_file(
            dataset,
            data_loc,
            {
                "what": "Loaded BL10 dataset",
                "by": "load_DLD",
            },
        )

        return dataset