Source code for nxtomomill.converter.fluo.blissfluoscan

"""Scan dedicated for bliss fluo-data format - based on h5 files generated by ewoksfluo"""

from __future__ import annotations

import logging
import pint

import numpy

from dataclasses import dataclass, field
import numpy as np
from numpy.typing import NDArray, DTypeLike

import h5py

_ureg = pint.get_application_registry()
_logger = logging.getLogger(__name__)

try:
    import tifffile  # noqa #F401 needed for later possible lazy loading
except ImportError:
    has_tifffile = False
else:
    has_tifffile = True

__all__ = ["BlissFluoTomoScanBase", "BlissFluoTomoScan3D", "BlissFluoTomoScan2D"]


@dataclass
class BlissFluoTomoScanBase:
    """Base class to read XRF data fitted with ewoksfluo."""

    ewoksfluo_filename: str
    detectors: tuple = ()
    dtype: DTypeLike = numpy.float32
    verbose: bool = False
    angles_deg: list[dtype] | None = None
    lines: dict[str, list[str]] = field(default_factory=dict)
    pixel_size: float | None = None
    energy: float | None = None

    @property
    def rot_angles_deg(self) -> NDArray:
        return self.angles_deg

    @property
    def rot_angles_rad(self) -> NDArray:
        return self.rot_angles_deg.to(_ureg.rad)


@dataclass
class BlissFluoTomoScan3D(BlissFluoTomoScanBase):
    """Base class to read XRF data fitted with ewoksfluo."""

    ewoksfluo_filename: str
    detectors: tuple = ()
    dtype: DTypeLike = numpy.float32
    verbose: bool = False
    angles_deg: list[dtype] | None = None
    lines: list[str] = field(default_factory=list)
    pixel_size: float | None = None
    energy: float | None = None

    def _check_entry(self, entry, f):
        data_shape = f[
            f"{entry}/grid/{self.detectors[0]}/results/massfractions/{self.lines[0]}"
        ][()].shape
        flag = self.slow_npoints * self.fast_npoints == data_shape[0] * data_shape[1]
        if not flag:
            msg = f"Entry {entry} has unexpected data shape {data_shape}, expected {(self.slow_npoints, self.fast_npoints)}. It has been ignored (i.e. not converted)."
            _logger.warning(msg)
        return flag

    def __post_init__(self):
        try:
            with h5py.File(self.ewoksfluo_filename, "r") as f:
                self.entries = sorted(list(f.keys()), key=lambda x: float(x))
                entry0 = self.entries[0]

                # Get detector names
                if len(self.detectors) == 0:
                    self.detectors = list(f[f"{entry0}/grid"].keys())

                # Get fitted line names
                detector0 = self.detectors[0]
                self.lines = []
                for name, ds in f[
                    f"{entry0}/grid/{detector0}/results/massfractions"
                ].items():
                    if ds.ndim == 2:
                        self.lines.append(name)

                # Get projections dimensions
                self.slow_npoints = float(
                    f[f"{entry0}/instrument/fscan_parameters/slow_npoints"][()]
                )
                self.fast_npoints = float(
                    f[f"{entry0}/instrument/fscan_parameters/fast_npoints"][()]
                )

                # Filter out interrupted scans:
                self.entries = [
                    entry for entry in self.entries if self._check_entry(entry, f)
                ]

                # Get rotation angles
                angles_tmp = []
                for entry in self.entries:
                    angles_tmp.append(f[f"{entry}/instrument/positioners/somega"][()])
                self.angles_deg = np.array(angles_tmp, dtype=self.dtype) * _ureg.deg

                # Get pixel size
                fast_pixel_size = float(
                    f[f"{entry0}/instrument/fscan_parameters/fast_step_size"][()]
                )
                slow_pixel_size = float(
                    f[f"{entry0}/instrument/fscan_parameters/slow_step_size"][()]
                )
                if slow_pixel_size == fast_pixel_size:
                    self.pixel_size = slow_pixel_size * _ureg.um
                else:
                    msg = f"Found slow_pixel_size ({slow_pixel_size}) different from fast_pixel_size ({fast_pixel_size})"
                    _logger.info(msg)
                    raise ValueError(msg)

                # Get energy
                self.energy = (
                    float(
                        f[
                            f"{entry0}/instrument/metadata/InstrumentMonochromator_energy"
                        ][()]
                    )
                    * _ureg.keV
                )

        except FileNotFoundError:
            self.entries = None
            _logger.info(
                f"Detected {len(self.entries)} projections in {self.ewoksfluo_filename}."
            )

    def load_data(self, det, line):
        if self.entries is None:
            _logger.info(f"No entry was detected in {self.ewoksfluo_filename}.")
        else:
            with h5py.File(self.ewoksfluo_filename, "r") as f:
                projs_tmp = []
                angles_tmp = []
                for entry in self.entries:
                    data = f[f"{entry}/grid/{det}/results/massfractions/{line}"][()]
                    projs_tmp.append(numpy.nan_to_num(data).astype(self.dtype))
                    angles_tmp.append(f[f"{entry}/instrument/positioners/somega"][()])

            projs_tmp = np.array(projs_tmp)
            angles_tmp = np.array(angles_tmp).astype(self.dtype)
            if np.allclose(angles_tmp * _ureg.deg, self.angles_deg):
                return np.ascontiguousarray(projs_tmp)
            else:
                msg = f"Angles found for {det}/{line} do not match the angles found at the class level."
                _logger.info(msg)
                raise ValueError(msg)


@dataclass
class BlissFluoTomoScan2D(BlissFluoTomoScanBase):
    """Base class to read XRF data fitted with ewoksfluo."""

    ewoksfluo_filename: str
    detectors: tuple = ()
    dtype: DTypeLike = numpy.float32
    verbose: bool = False
    angles_deg: list[dtype] | None = None
    lines: list[str] = field(default_factory=list)
    pixel_size: float | None = None
    energy: float | None = None