pmsco-public/pmsco/scan.py

"""
@package pmsco.scan
scan classes and factories.

@author Matthias Muntwiler, matthias.muntwiler@psi.ch

@copyright (c) 2015-21 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

import copy
import logging
import os
from typing import Any, Callable, Dict, Generator, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Union

import numpy as np

import pmsco.config as config
import pmsco.data as md

logger = logging.getLogger(__name__)


class Scan(object):
    """
    class to describe the scanning scheme or store the experimental data set.
    """

    ## @var filename (string)
    # file name from which a scan was loaded

    ## @var raw_data (numpy.ndarray)
    # original scan data (ETPAIS array)

    ## @var dtype (dict)
    # data type of self.raw_data.
    #
    # one of the data.DTYPE_Xxxx constants.

    ## @var modulation (numpy.ndarray)
    # modulation function calculated from original scan (ETPAIS array)

    ## @var mode (list of characters)
    # list of ETPAI column names which are scanned in self.raw_data.
    #
    # example: ['t','p']

    ## @var emitter (string)
    # chemical symbol and, optionally following, further specification (chemical state, environment, ...)
    # of photo-emitting atoms.
    # the interpretation of this string is up to the project and its cluster generator.
    # it should, however, always start with a chemical element symbol.
    #
    # examples: 'Ca' (calcium), 'CA' (carbon A), 'C a' (carbon a), 'C 1' (carbon one), 'N=O', 'FeIII'.

    ## @var initial_state (string)
    # nl term of initial state
    #
    # in the form expected by EDAC, for example: '1s'

    ## @var energies (numpy.ndarray)
    # kinetic energy referenced to Fermi level.
    #
    # one-dimensional array.

    ## @var thetas (numpy.ndarray)
    # polar angle referenced to normal emission
    #
    # one-dimensional array.
    #
    # note: in the case of a hemispherical scan, the values in this array will not be unique.

    ## @var phis (numpy.ndarray)
    # azimuthal angle referenced to arbitrary origin
    #
    # one-dimensional array.
    #
    # note: in the case of a hemispherical scan, the values in this array will not be unique, and not monotonic.

    ## @var alphas (numpy.ndarray)
    # polar angle referenced to normal emission
    #
    # one-dimensional array.

    def __init__(self):
        self.filename = ""
        self.raw_data = None
        self.dtype = None
        self.modulation = None
        self.modulation_func = md.default_modfunc
        self.modulation_args = {"smth": 0.4}
        self.rfactor_func = md.default_rfactor
        self.rfactor_args = {}
        self.mode = []
        self.emitter = ""
        self.initial_state = "1s"
        self.positions = {
            'e': np.empty(0),
            't': np.empty(0),
            'p': np.empty(0),
            'a': np.empty(0),
        }

    def __str__(self):
        return f"({self.__class__.__name__}) {self.filename} ({self.emitter} {self.initial_state})"

    @property
    def energies(self):
        return self.positions['e']

    @energies.setter
    def energies(self, value):
        self.positions['e'] = value

    @property
    def thetas(self):
        return self.positions['t']

    @thetas.setter
    def thetas(self, value):
        self.positions['t'] = value

    @property
    def phis(self):
        return self.positions['p']

    @phis.setter
    def phis(self, value):
        self.positions['p'] = value

    @property
    def alphas(self):
        return self.positions['a']

    @alphas.setter
    def alphas(self, value):
        self.positions['a'] = value

    def copy(self):
        """
        create a copy of the scan.

        @return: new independent scan object with the same attributes as the original one.
        """
        return copy.deepcopy(self)

    def import_scan_file(self, filename, emitter, initial_state):
        """
        import the reference experiment.

        the extension must be one of msc_data.DATATYPES (case insensitive)
        corresponding to the meaning of the columns in the file.

        this method does not calculate the modulation function.

        @attention EDAC can only calculate equidistant, rectangular scans.
        holo scans are transparently mapped to rectangular scans by pmsco.
        this method accepts the following scans:

        * intensity vs energy at fixed theta, phi
        * intensity vs analyser angle vs energy at normal emission (theta = 0, constant phi)
        * intensity vs theta, phi, or alpha
        * holo scan (theta,phi)

        @param filename: (string) file name of the experimental data, possibly including a path.

        @param emitter: (string) chemical symbol of the photo-emitting atom, e.g. "Cu".

        @param initial_state: (string) nl term of the initial state of the atom, e.g. "2p".

        """
        self.filename = filename
        self.emitter = emitter
        self.initial_state = initial_state

        if self.filename:
            self.raw_data = md.load_data(self.filename)
            self.analyse_raw_data()
        else:
            logger.error("empty file name in Scan.import_scan_file.")

    def analyse_raw_data(self):
        """
        analyse raw data and update dependant properties

        must be called after loading or modifying self.raw_data.
        part of import_scan_file.

        @return: None
        """
        self.dtype = self.raw_data.dtype
        self.mode, self.positions = md.detect_scan_mode(self.raw_data)
        if 'e' not in self.mode:
            try:
                self.energies = np.asarray((self.raw_data['e'][0],))
            except ValueError:
                logger.error("missing energy in scan file %s", self.filename)
                raise
        if 't' not in self.mode:
            try:
                self.thetas = np.asarray((self.raw_data['t'][0],))
            except ValueError:
                logger.info("missing theta in scan file %s, defaulting to 0.0", self.filename)
                self.thetas = np.zeros(1)
        if 'p' not in self.mode:
            try:
                self.phis = np.asarray((self.raw_data['p'][0],))
            except ValueError:
                logger.info("missing phi in scan file %s, defaulting to 0.0", self.filename)
                self.phis = np.zeros(1)
        if 'a' not in self.mode:
            try:
                self.alphas = np.asarray((self.raw_data['a'][0],))
            except ValueError:
                logger.info("missing alpha in scan file %s, defaulting to 0.0", self.filename)
                self.alphas = np.zeros(1)

    def define_scan(self, positions, emitter, initial_state):
        """
        define a cartesian (rectangular/grid) scan.

        this method initializes the scan with a one- or two-dimensional cartesian scan
        of the four possible scan dimensions.
        the scan range is given as arguments, the intensity values are initialized as 1.
        the file name and modulation functions are reset to empty and None, respectively.

        the method can create the following scan schemes:

        * intensity vs energy at fixed theta, phi
        * intensity vs analyser angle vs energy at normal emission (theta = 0, constant phi)
        * intensity vs theta, phi, or alpha
        * intensity vs theta and phi (rectangular holo scan)

        @param positions: (dictionary of numpy arrays)
            the dictionary must contain a one-dimensional array for each scan dimension 'e', 't', 'p' and 'a'.
            these array must contain unique, equidistant positions.
            constant dimensions must contain exactly one value.
            missing angle dimensions default to 0,
            a missing energy dimension results in a KeyError.

        @param emitter: (string) chemical symbol of the photo-emitting atom, e.g. "Cu".

        @param initial_state: (string) nl term of the initial state of the atom, e.g. "2p".

        """
        self.filename = ""
        self.emitter = emitter
        self.initial_state = initial_state
        self.mode = []
        shape = 1

        try:
            self.energies = np.copy(positions['e'])
        except KeyError:
            logger.error("missing energy in define_scan arguments")
            raise
        else:
            if self.energies.shape[0] > 1:
                self.mode.append('e')
                shape *= self.energies.shape[0]

        try:
            self.thetas = np.copy(positions['t'])
        except KeyError:
            logger.info("missing theta in define_scan arguments, defaulting to 0.0")
            self.thetas = np.zeros(1)
        else:
            if self.thetas.shape[0] > 1:
                self.mode.append('t')
                shape *= self.thetas.shape[0]

        try:
            self.phis = np.copy(positions['p'])
        except KeyError:
            logger.info("missing phi in define_scan arguments, defaulting to 0.0")
            self.phis = np.zeros(1)
        else:
            if self.phis.shape[0] > 1:
                self.mode.append('p')
                shape *= self.phis.shape[0]

        try:
            self.alphas = np.copy(positions['a'])
        except KeyError:
            logger.info("missing alpha in define_scan arguments, defaulting to 0.0")
            self.alphas = np.zeros(1)
        else:
            if self.alphas.shape[0] > 1:
                self.mode.append('a')
                shape *= self.alphas.shape[0]

        assert 0 < len(self.mode) <= 2, "unacceptable number of dimensions in define_scan"
        assert not ('t' in self.mode and 'a' in self.mode), "unacceptable combination of dimensions in define_scan"

        self.dtype = md.DTYPE_ETPAI
        self.raw_data = np.zeros(shape, self.dtype)
        dimensions = [self.positions[dim] for dim in ['e', 't', 'p', 'a']]
        grid = np.meshgrid(*dimensions)
        for i, dim in enumerate(['e', 't', 'p', 'a']):
            self.raw_data[dim] = grid[i].reshape(-1)
        self.raw_data['i'] = 1

    def generate_holo_scan(self,
                           generator: Callable[..., Iterable[Tuple[float, float]]],
                           generator_args: Dict,
                           other_positions: Dict,
                           emitter: Optional[str] = None,
                           initial_state: Optional[str] = None):
        """
        Generate a hologram (theta-phi) scan.

        The grid algorithm must be specified as the generator function, e.g. pmsco.data.holo_grid.

        @param generator Generator function that yields tuples (theta, phi) for each grid point,
            e.g., pmsco.data.holo_grid.
        @param generator_args Arguments to be passed to the generator.
            See the documentation of the respective generator function.
        @param other_positions:
            Coordinates of the other dimensions ('e' and 'a') in dictionary format.
            As of this version, they must be constant and contain exactly one value.
        @param emitter:
            Chemical symbol of the photo-emitting atom, e.g. "Cu".
            Optional if the corresponding Scan attribute is set elsewhere.
        @param initial_state:
            nlj term of the initial state of the atom, e.g. "2p1/2".
            Optional if the corresponding Scan attribute is set elsewhere.
        """

        if emitter:
            self.emitter = emitter
        if initial_state:
            self.initial_state = initial_state

        self.mode = ["t", "p"]
        self.positions = {'t': None, 'p': None}

        for dim, val in other_positions.items():
            self.positions[dim] = np.atleast_1d(np.asarray(val))
            assert self.positions[dim].shape[0] == 1, f"other_positions[{dim}] must have scalar value"

        tp = np.fromiter(generator(**generator_args), dtype=md.DTYPES['TP'])
        self.positions['t'] = tp['t']
        self.positions['p'] = tp['p']

        fields = set(self.positions.keys())
        fields.add('i')
        self.dtype = [dt for dt in md.DTYPE_ETPAIS if dt[0] in fields]
        self.raw_data = np.zeros(tp.shape, self.dtype)
        self.raw_data['i'] = np.random.default_rng().normal(1, 0.1, self.raw_data.shape)
        for dim in self.positions.keys():
            self.raw_data[dim] = self.positions[dim]

    def load(self):
        return self


class ScanKey(config.ConfigurableObject):
    """
    create a Scan object based on a project-supplied dictionary

    this class can be used in a run file to create a scan object based on the scan_dict attribute of the project.
    this may be convenient if you're project should selectively use scans out of a long list of data files
    and you don't want to clutter up the run file with parameters that don't change.

    to do so, set the key property to match an item of scan_dict.
    the load method will look up the corresponding scan_dict item and construct the final Scan object.
    """
    def __init__(self, project=None):
        super().__init__()
        self.key = ""
        self.project = project

    def __str__(self):
        return f"({self.__class__.__name__}) {self.key}"

    def load(self, dirs: Optional[Mapping] = None) -> Scan:
        """
        load the selected scan as specified in the project's scan dictionary

        the method uses ScanLoader or ScanCreator as an intermediate.

        @return a new Scan object which contains the loaded data.
        """
        scan_spec = self.project.scan_dict[self.key]
        if hasattr(scan_spec, 'positions'):
            loader = ScanCreator()
        elif hasattr(scan_spec, "generator"):
            loader = HoloScanCreator()
        else:
            loader = ScanLoader()
        for k, v in scan_spec.items():
            setattr(loader, k, v)
        scan = loader.load(dirs=dirs)
        return scan


class BaseScanSpec(config.ConfigurableObject):
    """
    Basic scan specification

    Declares and manages common scan attributes of the scan creators and loaders.
    """

    ## @var filename (string)
    # Name of the file from which to load or to which to save the scan data.
    # The file name can contain a format specifier like `${project}`.

    ## @var modulation_func (string or callable)
    # Function to calculate the modulation function.
    # Must have the same signature as pmsco.data.default_modfunc.
    # The function name can be given as a string relative to the global namespace of the project module,
    # e.g. `pmsco.data.calc_modfunc_loess` if `pmsco.data` has been imported.

    ## @var emitter (string)
    # Chemical symbol and, optionally following, further specification (chemical state, environment, ...)
    # of photo-emitting atoms.
    # The interpretation of this string is up to the project and its cluster generator.
    # It should, however, always start with a chemical element symbol.
    #
    # Examples: 'Ca' (calcium), 'CA' (carbon A), 'C a' (carbon a), 'C 1' (carbon one), 'N=O', 'FeIII'.

    ## @var initial_state (string)
    # Term symbol of the initial state
    #
    # nlj term of the initial state in the form expected by EDAC, for example: '2p1/2'

    def __init__(self):
        super().__init__()
        self.filename: Union[str, os.PathLike] = ""
        self.emitter: str = ""
        self.initial_state: str = "1s"
        self.modulation_func: Union[str, Callable] = md.default_modfunc
        self.modulation_args: Dict = {"smth": 0.4}
        self.rfactor_func: Union[str, Callable] = md.default_rfactor
        self.rfactor_args: Dict = {}

    def __str__(self):
        return f"({self.__class__.__name__}) {self.filename} ({self.emitter} {self.initial_state})"

    def load(self, dirs: Optional[Mapping] = None) -> Scan:
        """
        Load the scan according to specification

        Create a new Scan object and initialize the attributes defined in this class.
        This method is extended by derived classes.
        In this class, the method returns a partially initialized object.

        @return a new Scan object which contains the loaded data file.
        """

        scan = Scan()
        scan.filename = self.filename
        scan.emitter = self.emitter
        scan.initial_state = self.initial_state

        if callable(self.modulation_func):
            scan.modulation_func = self.modulation_func
        else:
            scan.modulation_func = eval(self.modulation_func, self.project_symbols)
        scan.modulation_args = self.modulation_args

        if callable(self.rfactor_func):
            scan.rfactor_func = self.rfactor_func
        else:
            scan.rfactor_func = eval(self.rfactor_func, self.project_symbols)
        scan.rfactor_args = self.rfactor_args

        return scan


class ScanLoader(BaseScanSpec):
    """
    create a Scan object from a data file reference

    this class can be used in a run file to create a scan object from an experimental data file.
    to do so, fill the properties with values as documented.
    the load() method is called when the project is run.
    """

    ## @var filename (string)
    # file name from which the scan should be loaded.
    # the file name can contain a format specifier like ${project} to include the base path.

    ## @var patch (dict)
    # patching instructions for raw_data array
    #
    # this attribute lets you modify parts of the raw_data array after loading the scan file.
    # this is useful, for instance, to change the energy in an angle scan.
    #
    # in the most common case, a fixed axis is changed to another value.
    # more complex patches can be declared using arbitrary numpy functions.
    # the resulting array must be compatible with broadcasting into the raw_data array,
    # i.e., its length must be 1 or equal to the one loaded from the file.
    #
    # the dictionary can contain any of the four keys: 'e', 't', 'p', 'a', 'i'
    # representing the four position and one data columns.
    # each key holds a string that contains a python expression.
    # the string is evaluated using python's built-in eval() function.
    # the expression must evaluate to an iterable object or numpy ndarray.
    # the namespaces of the project module can be used.
    # the `raw_data` object contains the loaded file.
    #
    # examples:
    # override the energy of a loaded scan with a constant value:
    # ~~~~~~{.py}
    # self.patch = {'e': '100.'}
    # ~~~~~~
    # subtract a constant offset from the theta angle:
    # ~~~~~~{.py}
    # self.patch = {'t': "raw_data['t'] - 3.4"}
    # ~~~~~~

    ## @var is_modf (bool)
    # declares whether the data file contains the modulation function rather than intensity values
    #
    # if false, the project will calculate a modulation function from the raw data

    def __init__(self):
        super().__init__()
        self.patch = {}
        self.is_modf = False

    def __str__(self):
        return f"({self.__class__.__name__}) {self.filename} ({self.emitter} {self.initial_state})"

    def load(self, dirs: Optional[Mapping] = None) -> Scan:
        """
        load the scan according to specification

        create a new Scan object and load the file by calling Scan.import_scan_file().

        @return a new Scan object which contains the loaded data file.
        """

        scan = super().load()
        filename = config.resolve_path(self.filename, dirs)
        scan.import_scan_file(filename, self.emitter, self.initial_state)

        if self.patch:
            patch_arrays = {}
            locals = {'raw_data': scan.raw_data}
            for axis in self.patch.keys():
                patch_arrays[axis] = np.atleast_1d(np.asarray(eval(self.patch[axis], self.project_symbols, locals)))
            for axis in patch_arrays.keys():
                scan.raw_data[axis][:] = patch_arrays[axis]
                logger.warning(f"scan file {filename.name}, overriding axis '{axis}'")
            scan.analyse_raw_data()

        if self.is_modf:
            scan.modulation = scan.raw_data

        return scan


class ScanCreator(BaseScanSpec):
    """
    Create a linear or rectangular scan object from string expressions

    This class can be used in a run file to create a scan object from python expressions,
    such as lists, ranges or numpy functions.
    To do so, fill the properties with values as documented.
    The load() method is called when the project is run.

    @note The raw_data property of the scan cannot be filled this way.
    Thus, the class is useful in `single` calculation mode only.
    """

    ## @var positions (dict)
    # Dictionary specifying the scan positions.
    #
    # The dictionary must contain four keys: 'e', 't', 'p', 'a' representing the four scan axes.
    # Each value must contain a Python expression that evaluates to a one-dimensional coordinate array.
    # The coordinates are meshed into a one or multidimensional rectangular scan grid.
    # Non-scanning dimensions must be set to a constant scalar value.
    #
    # The values are evaluated using python's built-in eval() function.
    # The expression must evaluate to an iterable object, numpy ndarray or scalar of the scan positions.
    # Functions from the project namespace can be used to build the scan.
    #
    # Example
    # -------
    #
    # Generate a rectangular energy-alpha scan at fixed angle:
    #
    # ~~~~~~{.py}
    # self.positions = {"e": "numpy.linspace(100, 200, 5)", "t": "0", "p": "60", "a": "numpy.linspace(-10, 10, 51)"}
    # ~~~~~~
    #
    # In the run-file:
    #
    # ~~~~~~{.json}
    # "positions": {"e": "numpy.linspace(100, 200, 5)", "t": "0", "p": "60", "a": "numpy.linspace(-10, 10, 51)"}
    # ~~~~~~
    #

    def __init__(self):
        super().__init__()
        self.positions = {'e': None, 't': None, 'p': None, 'a': None}

    def __str__(self):
        return f"({self.__class__.__name__}) {self.filename} ({self.emitter} {self.initial_state})"

    def load(self, dirs: Optional[Mapping] = None) -> Scan:
        """
        Create the scan according to specification

        @return a new Scan object which contains the created scan array.
        """

        scan = super().load()
        positions = {}
        for axis in self.positions.keys():
            positions[axis] = np.atleast_1d(np.asarray(eval(str(self.positions[axis]), self.project_symbols)))
        scan.define_scan(positions, self.emitter, self.initial_state)
        scan.filename = config.resolve_path(self.filename, dirs)
        return scan


class HoloScanCreator(BaseScanSpec):
    """
    Generate a hologram scan

    This scan generator can be used to generate a holo scan.
    The (theta, phi) coordinates are obtained from a generator function, e.g. pmsco.data.holo_grid.
    The remaining coordinates are specified in the other_positions dictionary.
    As of this version, they must be scalar values.

    Example
    -------

    ~~~~~~{.py}
    sg = HoloScanCreator()
    sg.generator = pmsco.data.holo_grid
    sg.generator_args = {"theta_start": 85, "theta_step": 1)
    sg.other_positions = {"e": "250", "a": "0"}
    sg.load()
    ~~~~~~
    """

    def __init__(self):
        super().__init__()
        self.generator: Union[str, Callable[..., Iterable[Tuple[float, float]]]] = md.holo_grid
        self.generator_args: Dict[str, Union[str, int, float]] = {}
        self.other_positions = {}

    def __str__(self):
        return f"({self.__class__.__name__}) {self.filename} ({self.emitter} {self.initial_state})"

    def load(self, dirs: Optional[Mapping] = None) -> Scan:
        """
        Generate the scan according to specification

        @return a new Scan object which contains the created scan array.
        """

        scan = super().load()
        positions = {}

        for axis in self.other_positions.keys():
            positions[axis] = np.atleast_1d(np.asarray(eval(str(self.other_positions[axis]), self.project_symbols)))

        generator = self.generator if callable(self.generator) else eval(self.generator, self.project_symbols)

        scan.generate_holo_scan(generator=generator, generator_args=self.generator_args, other_positions=positions)
        scan.filename = config.resolve_path(self.filename, dirs)

        return scan