nabu.resources.dataset_tomoscan

[docs] module nabu.resources.dataset_tomoscan

"""
Datasets parsing classes that use the tomoscan library .
Used for Nexus-NXTomo and EDF scans.
"""

from packaging.version import parse as parse_version
from enum import Enum
from warnings import warn

import numpy as np
from h5py import File
from silx.io import get_data
from silx.io.url import DataUrl

from tomoscan.esrf.scan.edfscan import EDFTomoScan
from tomoscan.esrf.scan.nxtomoscan import NXtomoScan
from tomoscan import __version__ as __tomoscan_version__

from ..utils import (
    check_supported,
    indices_to_slices,
    search_sorted,
    deprecated_class,
    deprecation_warning,
    docstring,
)
from ..io.reader import EDFStackReader, NXDarksFlats, NXTomoReader
from ..io.utils import get_compacted_dataslices, get_first_hdf5_entry
from ..pipeline.utils import nabu_env_settings
from .utils import get_values_from_file
from .dataset_base import Dataset, get_angle_at_index


# We could import the 1000+ LoC  nxtomo.nxobject.nxdetector.ImageKey... or we can do this
class ImageKey(Enum):
[docs]
    ALIGNMENT = -1
    PROJECTION = 0
    FLAT_FIELD = 1
    DARK_FIELD = 2
    INVALID = 3


# ---

_image_type = {
    "projections": ImageKey.PROJECTION.value,
    "projection": ImageKey.PROJECTION.value,
    "radios": ImageKey.PROJECTION.value,
    "radio": ImageKey.PROJECTION.value,
    "flats": ImageKey.FLAT_FIELD.value,
    "flat": ImageKey.FLAT_FIELD.value,
    "darks": ImageKey.DARK_FIELD.value,
    "dark": ImageKey.DARK_FIELD.value,
    "static": ImageKey.ALIGNMENT.value,
    "alignment": ImageKey.ALIGNMENT.value,
    "return": ImageKey.ALIGNMENT.value,
    "invalid": ImageKey.INVALID.value,
}


class TomoscanDataset(Dataset):
[docs]
    """
    A base class for scans that use tomoscan behind the scenes.
    """

    _default_extra_options = {
        "exclude_projections": None,
        "hdf5_entry": None,
    }
    _scanner = None

    @docstring(Dataset)
    def __init__(self, location, extra_options=None):
        self.location = location
        self._set_extra_options(extra_options)
        self._get_excluded_projections()
        self._init_dataset_scan()
        self._finish_init()

    # pylint: disable=E1136
    def _get_excluded_projections(self):
        excluded_projs = self.extra_options["exclude_projections"]
        self._ignore_projections = None
        if excluded_projs is None:
            return

        if excluded_projs["type"] == "angular_range":
            excluded_projs["type"] = "range"  # compat with tomoscan #pylint: disable=E1137
            values = excluded_projs["range"]
        for ignore_kind, dtype in {"indices": np.int32, "angles": np.float32}.items():
            if excluded_projs["type"] == ignore_kind:
                values = get_values_from_file(excluded_projs["file"], any_size=True).astype(dtype).tolist()
        self._ignore_projections = {"kind": excluded_projs["type"], "values": values}  # pylint: disable=E0606

    def _init_dataset_scan(self, **kwargs):
        if self._scanner is None:
            raise ValueError("Base class")
        if self._scanner is NXtomoScan:
            if self.extra_options.get("hdf5_entry", None) is not None:
                kwargs["entry"] = self.extra_options["hdf5_entry"]
            if self.extra_options.get("nx_version", None) is not None:
                kwargs["nx_version"] = self.extra_options["nx_version"]
        if self._scanner is EDFTomoScan:
            # Assume 1 frame per file (otherwise too long to open each file)
            kwargs["n_frames"] = 1

        self.dataset_scanner = self._scanner(  # pylint: disable=E1102
            self.location, ignore_projections=self._ignore_projections, **kwargs
        )

        if self._ignore_projections is not None:
            warn("Excluding projections: %s" % str(self._ignore_projections))

        if nabu_env_settings.skip_tomoscan_checks:
            warn(
                " WARNING: according to nabu_env_settings.skip_tomoscan_checks, skipping virtual layout integrity check of tomoscan which is time consuming"
            )
            self.dataset_scanner.set_check_behavior(run_check=False, raise_error=False)

        self._rotation_angles = None
        self.raw_flats = self.dataset_scanner.flats
        self.raw_darks = self.dataset_scanner.darks
        self.n_angles = len(self.dataset_scanner.projections)
        self.radio_dims = (self.dataset_scanner.dim_1, self.dataset_scanner.dim_2)
        self._radio_dims_notbinned = self.radio_dims  # COMPAT

    def _finish_init(self):
        pass

    def _get_energy(self):
        return self.dataset_scanner.energy

    def _get_distance(self):
        if parse_version(__tomoscan_version__) < parse_version("2.2"):
            return abs(self.dataset_scanner.distance)
        else:
            return abs(self.dataset_scanner.sample_detector_distance)

    def _get_pixel_size(self):
        return self.dataset_scanner.pixel_size * 1e6

    def _get_rotation_angles(self):
        return self._rotation_angles  # None by default

    def _get_projections(self):
        return self.dataset_scanner.projections

    def _get_projections_srcurrent(self):
        return self._get_srcurrent("radios")  # pylint: disable=E1128

    def _get_flats_srcurrent(self):
        return self._get_srcurrent("flats")  # pylint: disable=E1128

    def check_defined_attribute(self, name, error_msg=None):
[docs]
        """
        Utility function to check that a given attribute is defined.
        """
        if getattr(self, name, None) is None:
            raise ValueError(error_msg or str("No information on %s was found in the dataset" % name))

    def _get_flats(self):
        return self.get_reduced_flats()

    def _get_darks(self):
        return self.get_reduced_darks()

    def _get_flip_frame_lr(self):
        return getattr(self.dataset_scanner, "detector_is_lr_flip", False)

    def _get_flip_frame_ud(self):
        return getattr(self.dataset_scanner, "detector_is_ud_flip", False)

    def _get_rot_center_first_guess(self):
        cor_rel = self.dataset_scanner.x_rotation_axis_pixel_position
        if cor_rel is None:
            return None
        return (self.radio_dims[0] - 1) / 2 + cor_rel

    def _get_is_helical(self):
        # TODO is there a way, with tomoscan, to know whether a scan is helical ? Bliss tells this information!
        return np.abs(np.ptp(self.translations_z)) > 10


class EDFDataset(TomoscanDataset):
[docs]
    """
    EDF Dataset analyzer for legacy ESRF acquisitions
    """

    _scanner = EDFTomoScan
    kind = "edf"

    def _get_flats_darks(self):
        return

    @property
    def hdf5_entry(self):
[docs]
        """
        Return the HDF5 entry of the current dataset.
        Not applicable for EDF (return None)
        """
        return None

    def _get_is_halftomo(self):
        return None

    def _get_rotation_angles(self):
        return np.deg2rad(self.dataset_scanner.rotation_angle())

    def get_reduced_flats(self, **reader_kwargs):
[docs]
        if self.raw_flats in [None, {}]:
            raise FileNotFoundError("No reduced flat ('refHST') found in %s" % self.location)
        # A few notes:
        # (1) In principle we could do the reduction (mean/median) from raw frames (ref_xxxx_yyyy)
        #   but for legacy datasets it's always already done (by fasttomo3), and EDF support is supposed to be dropped on our side
        # (2) We use EDFStackReader class to handle the possible additional data modifications
        #  (eg. subsampling, binning, distortion correction...)
        # (3) The following spawns one reader instance per file, which is not elegant,
        #   but in principle there are typically 1-2 reduced flats in a scan
        readers = {k: EDFStackReader([self.raw_flats[k].file_path()], **reader_kwargs) for k in self.raw_flats}
        return {k: readers[k].load_data()[0] for k in self.raw_flats}

    def get_reduced_darks(self, **reader_kwargs):
[docs]
        # See notes in get_reduced_flats() above
        if self.raw_darks in [None, {}]:
            raise FileNotFoundError("No reduced dark ('darkend.edf' or 'dark.edf') found in %s" % self.location)
        readers = {k: EDFStackReader([self.raw_darks[k].file_path()], **reader_kwargs) for k in self.raw_darks}
        return {k: readers[k].load_data()[0] for k in self.raw_darks}

    @property
    def files(self):
[docs]
        return sorted([u.file_path() for u in self.dataset_scanner.projections.values()])

    def get_reader(self, **kwargs):
[docs]
        return EDFStackReader(self.files, **kwargs)

    def _get_scan_basename(self):
        # os.path.basename(self.dataset_scanner.path)
        return self.dataset_scanner.get_dataset_basename()

    def _get_scan_dirname(self):
        return self.dataset_scanner.path
[docs]
    def get_excluded_projections_indices(self, including_other_frames_types=True):
        if not (including_other_frames_types):
            raise NotImplementedError
        return self.dataset_scanner.get_ignored_projection_indices()

    @docstring(Dataset.get_frames_indices)
    def get_frames_indices(self, frame_type):
[docs]
        frame_type = _image_type[frame_type]
        if frame_type == ImageKey.PROJECTION.value:
            return sorted(self.projections.keys())
        if frame_type == ImageKey.FLAT_FIELD.value:
            return sorted(self.flats.keys())
        if frame_type == ImageKey.DARK_FIELD.value:
            return sorted(self.darks.keys())
        if frame_type == ImageKey.ALIGNMENT.value:
            return sorted(self.dataset_scanner.alignment_projections.keys())
        return []

    def _get_srcurrent(self, frame_type):
        return None
        # return self.dataset_scanner.machine_current


class NXTomoDataset(TomoscanDataset):
[docs]
    """
    NXTomo dataset analyzer
    """

    _scanner = NXtomoScan
    kind = "nx"

    def _finish_init(self):
        self.data_path = self.dataset_scanner.get_detector_data_path()

    @property
    def z_translation(self):
[docs]
        raw_data = np.array(self.dataset_scanner.z_translation)
        projs_idx = np.array(list(self.projections.keys()))
        filtered_data = raw_data[projs_idx]
        return 1.0e6 * filtered_data / self.pixel_size

    @property
    def x_translation(self):
[docs]
        raw_data = np.array(self.dataset_scanner.x_translation)
        projs_idx = np.array(list(self.projections.keys()))
        filtered_data = raw_data[projs_idx]
        return 1.0e6 * filtered_data / self.pixel_size

    def _get_rotation_angles(self):
        if self._rotation_angles is None:
            angles = np.array(self.dataset_scanner.rotation_angle)
            projs_idx = np.array(list(self.projections.keys()))
            angles = angles[projs_idx]
            self._rotation_angles = np.deg2rad(angles)
        return self._rotation_angles

    def _get_dataset_hdf5_url(self):
        if len(self.projections) > 0:
            frames_to_take = self.projections
        elif len(self.raw_flats) > 0:
            frames_to_take = self.raw_flats
        elif len(self.raw_darks) > 0:
            frames_to_take = self.raw_darks
        else:
            raise ValueError("No projections, no flats and no darks ?!")
        first_proj_idx = sorted(frames_to_take.keys())[0]
        first_proj_url = frames_to_take[first_proj_idx]
        return DataUrl(
            file_path=first_proj_url.file_path(), data_path=first_proj_url.data_path(), data_slice=None, scheme="silx"
        )

    @property
    def dataset_hdf5_url(self):
[docs]
        return self._get_dataset_hdf5_url()

    @property
    def hdf5_entry(self):
[docs]
        """
        Return the HDF5 entry of the current dataset
        """
        return self.dataset_scanner.entry

    def _get_is_halftomo(self):
        ha = getattr(self.dataset_scanner, "half_acquisition", None)
        if ha is not None:
            return ha
        # COMPAT.
        try:
            is_halftomo = self.dataset_scanner.field_of_view.value.lower() == "half"
        except:
            is_halftomo = None
        return is_halftomo
        # ---

    def get_data_slices(self, what):
[docs]
        """
        Return indices in the data volume where images correspond to a given kind.

        Parameters
        ----------
        what: str
            Which keys to get. Can be "projections", "flats", "darks"

        Returns
        --------
        slices: list of slice
            A list where each item is a slice.
        """
        name_to_attr = {
            "projections": self.projections,
            "flats": self.raw_flats,
            "darks": self.raw_darks,
        }
        check_supported(what, name_to_attr.keys(), "image type")
        images = name_to_attr[what]  # dict
        # we can't directly use set() on slice() object (unhashable). Use tuples
        slices = set()
        for du in get_compacted_dataslices(images).values():
            if du.data_slice() is not None:
                # note: du.data_slice is a uint in recent tomoscan version
                s = (int(du.data_slice().start), int(du.data_slice().stop))
            else:
                s = None
            slices.add(s)
        slices_list = [slice(item[0], item[1]) if item is not None else None for item in list(slices)]
        return slices_list

    def _select_according_to_frame_type(self, data, frame_type):
        if data is None:
            return None
        return data[self.dataset_scanner.image_key_control == _image_type[frame_type]]
[docs]
    def get_reduced_flats(self, method="median", force_reload=False, **reader_kwargs):
        if not (np.any(self.dataset_scanner.image_key_control == ImageKey.FLAT_FIELD.value)):
            return None
        dkrf_reader = NXDarksFlats(
            self.dataset_hdf5_url.file_path(), data_path=self.dataset_hdf5_url.data_path(), **reader_kwargs
        )
        return dkrf_reader.get_reduced_flats(method=method, force_reload=force_reload, as_dict=True)
[docs]
    def get_reduced_darks(self, method="mean", force_reload=False, **reader_kwargs):
        if not (np.any(self.dataset_scanner.image_key_control == ImageKey.DARK_FIELD.value)):
            return None
        dkrf_reader = NXDarksFlats(
            self.dataset_hdf5_url.file_path(), data_path=self.dataset_hdf5_url.data_path(), **reader_kwargs
        )
        return dkrf_reader.get_reduced_darks(method=method, force_reload=force_reload, as_dict=True)

    def _get_srcurrent(self, frame_type):
        return self._select_according_to_frame_type(self.dataset_scanner.machine_current, frame_type)

    def frames_slices(self, frame_type):
[docs]
        """
        Return a list of slice objects corresponding to the data corresponding to "frame_type".
        For example, if the dataset flats are located at indices [1, 2, ..., 99], then
        frame_slices("flats") will return [slice(0, 100)].
        """
        return indices_to_slices(np.where(self.dataset_scanner.image_key_control == _image_type[frame_type])[0])

    def get_reader(self, **kwargs):
[docs]
        return NXTomoReader(self.dataset_hdf5_url.file_path(), data_path=self.dataset_hdf5_url.data_path(), **kwargs)

    @property
    def scan_basename(self):
[docs]
        # os.path.splitext(os.path.basename(self.dataset_hdf5_url.file_path()))[0]
        return self.dataset_scanner.get_dataset_basename()

    @property
    def scan_dirname(self):
[docs]
        # os.path.dirname(di.dataset_hdf5_url.file_path())
        return self.dataset_scanner.path

    @docstring(Dataset.get_alignment_projections)
[docs]
    def get_alignment_projections(self, image_sub_region=None, return_in_radians=None):
        sub_region = None
        if image_sub_region is not None:
            sub_region = (None,) + image_sub_region
        reader = self.get_reader(image_key=ImageKey.ALIGNMENT.value, sub_region=sub_region)
        projs = reader.load_data()
        indices = reader.get_frames_indices()
        angles = get_angle_at_index(self.all_angles, indices)
        # COMPAT - to be deleted in the future
        if return_in_radians is None:
            deprecation_warning(
                "Returning angles in degrees by default. This will change in the future, angles will be returned in radians. You will have to use return_in_radians=False",
                func_name="get_alignment_projections",
            )
            return_in_radians = False
        # ---
        if not (return_in_radians):
            angles = np.rad2deg(angles)
        return projs, angles, indices

    def _get_all_angles(self):
        return np.array(np.deg2rad(self.dataset_scanner.rotation_angle))

    @docstring(Dataset.get_index_from_angle)
[docs]
    def get_index_from_angle(self, angle, image_key=0, return_found_angle=False):
        all_angles = self.all_angles
        all_indices = np.arange(len(all_angles))
        all_image_key = self.dataset_scanner.image_key_control

        idx2 = np.where(all_image_key == image_key)[0]
        angles = np.unwrap(all_angles[idx2])  # unwrap !
        idx_angles_sorted = np.argsort(angles)
        angles_sorted = angles[idx_angles_sorted]

        pos = search_sorted(angles_sorted, angle)
        # this gives a position in "idx2", but we need the position in "all_indices"
        idx = all_indices[idx2[idx_angles_sorted[pos]]]
        if return_found_angle:
            return idx, angles_sorted[pos]
        return idx

    @docstring(Dataset.get_image_at_angle)
[docs]
    def get_image_at_angle(self, angle_rad, frame_type="projection", sub_region=None, return_angle_and_index=False):
        image_key = _image_type[frame_type]
        idx, angle_found = self.get_index_from_angle(angle_rad, image_key=image_key, return_found_angle=True)

        # Option 1:
        if sub_region is None:
            sub_region = (None, None)
        # Convert absolute index to index of image_key
        idx2 = np.searchsorted(np.where(self.dataset_scanner.image_key_control == image_key)[0], idx)
        sub_region = (slice(idx2, idx2 + 1),) + sub_region
        reader = self.get_reader(image_key=image_key, sub_region=sub_region)
        img = reader.load_data()[0]
        if return_angle_and_index:
            return img, angle_found, idx
        return img

        # Option 2:
        # return self.get_frame(idx)
        # something like:
        # [fr for fr in self.dataset_scanner.frames if fr.image_key.value == 0 and fr.rotation_angle == 180 and fr._is_control_frame is False]

    @docstring(Dataset.get_frame)
    def get_frame(self, idx):
[docs]
        return get_data(self.dataset_scanner.frames[idx].url)

    @docstring(Dataset.get_frames_indices)
    def get_frames_indices(self, frame_type):
[docs]
        return self._select_according_to_frame_type(np.arange(self.dataset_scanner.image_key_control.size), frame_type)

    @docstring(Dataset.index_to_proj_number)
    def index_to_proj_number(self, proj_index):
[docs]
        all_projs_indices = self.get_frames_indices("projection")
        return search_sorted(all_projs_indices, proj_index)

    def _get_transl(self, name):
        transl = getattr(self.dataset_scanner, name)
        if transl is None:
            transl = np.zeros(self.n_angles, dtype="f")
        else:
            transl = np.array(transl)[self.get_frames_indices(frame_type="projections")]
        transl /= self.pixel_size * 1e-6
        return transl

    def _get_translations_x(self):
        return self._get_transl("x_translation")

    def _get_translations_y(self):
        return self._get_transl("y_translation")

    def _get_translations_z(self):
        return self._get_transl("z_translation")

    def _get_n_frames(self):
        return len(self.dataset_scanner.frames)

    def _get_end_time(self):
        return self.dataset_scanner.end_time

    def _get_sample_name(self):
        return self.dataset_scanner.sample_name
[docs]
    def get_excluded_projections_indices(self, including_other_frames_types=True):
        # the index accounts for flats/darks !
        # Get indices of excluded projs (again, accounting for flats/darks)
        ignored_projs_indices = self.dataset_scanner.get_ignored_projection_indices()
        ignored_projs_indices = [
            idx for idx in ignored_projs_indices if self.dataset_scanner.frames[idx].is_control is False
        ]
        if including_other_frames_types:
            return ignored_projs_indices
        # Get indices of excluded projs, now relative to the pure projections stack
        ignored_projs_indices_rel = [
            self.index_to_proj_number(ignored_proj_idx_abs) for ignored_proj_idx_abs in ignored_projs_indices
        ]
        return ignored_projs_indices_rel


# COMPAT. !

DatasetAnalyzer = deprecated_class("'DatasetAnalyzer' is deprecated. Please use 'Dataset' instead.", do_print=True)(
    Dataset
)
HDF5DatasetAnalyzer = deprecated_class(
    "'HDF5DatasetAnalyzer' is deprecated. Please use 'NXTomoDataset' instead.", do_print=True
)(NXTomoDataset)
EDFDatasetAnalyzer = deprecated_class(
    "'EDFDatasetAnalyzer' is deprecated. Please use 'EDFDataset' instead.", do_print=True
)(EDFDataset)

# ----


def is_nxtomo(fname, entry=None):
[docs]
    if fname.endswith(".nx"):
        return True
    entry = entry or get_first_hdf5_entry(fname)
    with File(fname, "r") as f:
        is_nx = f[entry].attrs.get("definition", "").lower() == "nxtomo"
    return is_nx