Source code for mantidimaging.core.rotation.data_model

# Copyright (C) 2023 ISIS Rutherford Appleton Laboratory UKRI
# SPDX - License - Identifier: GPL-3.0-or-later
from __future__ import annotations

from logging import getLogger
from typing import Optional, List, Iterator, NamedTuple

import numpy as np
import scipy as sp

from mantidimaging.core.operation_history import const
from ..utility.data_containers import ScalarCoR, Degrees, Slope

LOG = getLogger(__name__)
Point = NamedTuple("Point", [("slice_index", int), ("cor", float)])

[docs]class CorTiltDataModel: """ Model for finding COR/Tilt from (slice index, centre of rotation) data points """ _cached_gradient: Optional[float] _cached_cor: Optional[float] _points: List[Point] def __init__(self): self._points = [] self._cached_gradient = None self._cached_cor = None
[docs] def populate_slice_indices(self, begin, end, count, cor=0.0): self.clear_results() self._points = [Point(int(idx), cor) for idx in np.linspace(begin, end, count, dtype=int)] LOG.debug('Populated slice indices: {}'.format(self.slices))
[docs] def linear_regression(self): LOG.debug('Running linear regression with {} points'.format(self.num_points)) self._cached_gradient, self._cached_cor, *_ = sp.stats.linregress(self.slices, self.cors)
[docs] def add_point(self, idx=None, slice_idx=0, cor=0.0): self.clear_results() if idx is None: self._points.append(Point(slice_idx, cor)) else: self._points.insert(idx, Point(slice_idx, cor))
[docs] def set_point(self, idx, slice_idx: int = None, cor: float = None, reset_results=True): if reset_results: self.clear_results() if len(self._points) > idx: if slice_idx is not None: self._points[idx] = Point(int(slice_idx), self._points[idx].cor) if cor is not None: self._points[idx] = Point(self._points[idx].slice_index, float(cor))
def _get_data_idx_from_slice_idx(self, slice_idx) -> int: for i, p in enumerate(self._points): if p.slice_index == slice_idx: return i raise ValueError(f"Slice {slice_idx} that is not in COR table")
[docs] def set_cor_at_slice(self, slice_idx: int, cor: float): data_idx = self._get_data_idx_from_slice_idx(slice_idx) self.set_point(data_idx, cor=cor)
[docs] def remove_point(self, idx): self.clear_results() del self._points[idx]
[docs] def clear_points(self): self._points = [] self.clear_results()
[docs] def clear_results(self): self._cached_gradient = None self._cached_cor = None
[docs] def point(self, idx): return self._points[idx] if idx < self.num_points else None
[docs] def sort_points(self): self._points.sort(key=lambda p: p.slice_index)
[docs] def get_cor_from_regression(self, slice_idx) -> float: cor = (self.gradient.value * slice_idx) + self.cor.value return cor
[docs] def get_all_cors_from_regression(self, image_height) -> List[ScalarCoR]: """ :param image_height: How many cors will be generated, this should be equal to the image height (i.e. number of sinograms that will be reconstructed) :return: List of cors for every slice of the image height """ cors = [] for i in range(image_height): cors.append(ScalarCoR(self.get_cor_from_regression(i))) return cors
@property def slices(self): return [p.slice_index for p in self._points] @property def cors(self): return [float(p.cor) for p in self._points] @property def gradient(self) -> Slope: if self._cached_gradient is not None: return Slope(self._cached_gradient) else: return Slope(0.0) @property def cor(self) -> ScalarCoR: if self._cached_cor is not None: return ScalarCoR(self._cached_cor) else: return ScalarCoR(0.0) @property def angle_in_degrees(self) -> Degrees: return Degrees(-np.rad2deg(np.arctan(self.gradient.value))) @property def has_results(self): return self._cached_gradient is not None and self._cached_cor is not None @property def empty(self): return not self._points @property def num_points(self): return len(self._points) @property def stack_properties(self): return { # TODO remove float casts const.COR_TILT_ROTATION_CENTRE: float(self.cor.value), const.COR_TILT_FITTED_GRADIENT: float(self.gradient.value), const.COR_TILT_TILT_ANGLE_DEG: float(self.angle_in_degrees.value), const.COR_TILT_SLICE_INDICES: self.slices, const.COR_TILT_ROTATION_CENTRES: self.cors }
[docs] def set_precalculated(self, cor: ScalarCoR, tilt: Degrees): self._cached_cor = cor.value # reverse the tilt calculation to get the slope of the regression back self._cached_gradient = -np.tan(np.deg2rad(tilt.value)) for idx in range(len(self._points)): current = self._points[idx] self._points[idx] = Point(current.slice_index, self.get_cor_from_regression(current.slice_index))
[docs] def iter_points(self) -> Iterator[Point]: return iter(self._points)