8829910821

Committed 25 Apr 2024 09:06AM UTC coverage: 73.005% (-0.4%) from 73.358%

Build # 8829910821

Build Type

Pull #1997

github

Committed by

web-flow

Commit Message

Merge 0d07c184f into 6c64af36f

Pull Request Pull Request #1997: Explicit tgv backend

Run Details

1908 of 3167 branches covered (60.25%)

Branch coverage included in aggregate %.

12 of 52 new or added lines in 3 files covered. (23.08%)

188 existing lines in 5 files now uncovered.

9407 of 12332 relevant lines covered (76.28%)

0.76 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

13.04

/mantidimaging/core/reconstruct/cil_recon.py

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

import time
from logging import getLogger, DEBUG
from math import sqrt, ceil
from threading import Lock
from typing import TYPE_CHECKING

import numpy as np

from cil.framework import (AcquisitionData, AcquisitionGeometry, DataOrder, ImageGeometry, BlockGeometry,
                           BlockDataContainer)
from cil.optimisation.algorithms import PDHG, SPDHG
from cil.optimisation.operators import GradientOperator, BlockOperator
from cil.optimisation.operators import SymmetrisedGradientOperator, ZeroOperator, IdentityOperator

from cil.optimisation.functions import MixedL21Norm, L2NormSquared, BlockFunction, ZeroFunction, IndicatorBox, Function
from cil.plugins.astra.operators import ProjectionOperator

from mantidimaging.core.data import ImageStack
from mantidimaging.core.reconstruct.base_recon import BaseRecon
from mantidimaging.core.utility.optional_imports import safe_import
from mantidimaging.core.utility.progress_reporting import Progress
from mantidimaging.core.utility.size_calculator import full_size_KB
from mantidimaging.core.utility.memory_usage import system_free_memory

if TYPE_CHECKING:
    from mantidimaging.core.utility.data_containers import ProjectionAngles, ReconstructionParameters, ScalarCoR

LOG = getLogger(__name__)
tomopy = safe_import('tomopy')
cil_mutex = Lock()


class CILRecon(BaseRecon):

    @staticmethod
    def set_up_TV_regularisation(
            image_geometry: ImageGeometry, acquisition_data: AcquisitionData,
            recon_params: ReconstructionParameters) -> tuple[BlockOperator, BlockFunction, Function]:

        # Forward operator
        A2d = ProjectionOperator(image_geometry, acquisition_data.geometry, 'gpu')

        if recon_params.stochastic:
            for partition_geometry, partition_operator in zip(acquisition_data.geometry, A2d, strict=True):
                CILRecon.set_approx_norm(partition_operator, partition_geometry, image_geometry)
        else:
            CILRecon.set_approx_norm(A2d, acquisition_data.geometry, image_geometry)

        # Define Gradient Operator and BlockOperator
        alpha = recon_params.alpha
        Grad = GradientOperator(image_geometry)

        if recon_params.stochastic:
            # now, A2d is a BlockOperator as acquisition_data is a BlockDataContainer
            fs = []
            for i, _ in enumerate(acquisition_data.geometry):
                fs.append(L2NormSquared(b=acquisition_data.get_item(i)))
            fs.append(MixedL21Norm())

            F = BlockFunction(*fs)

            # needs to unrol the A2d BlockOperator and put it in another, followed by
            # the gradient operator, as in the deterministic case
            K = BlockOperator(*A2d.get_as_list(), alpha * Grad)

        else:
            # Define BlockFunction F using the MixedL21Norm() and the L2NormSquared()
            f1 = MixedL21Norm()
            f2 = L2NormSquared(b=acquisition_data)

            F = BlockFunction(f1, f2)

            # define the BlockOperator
            K = BlockOperator(alpha * Grad, A2d)

        if recon_params.non_negative:
            G = IndicatorBox(lower=0)
        else:
            # Define Function G simply as zero
            G = ZeroFunction()

        return (K, F, G)

    @staticmethod
    def set_up_TGV_regularisation(
            image_geometry: ImageGeometry, acquisition_data: AcquisitionData,
            recon_params: ReconstructionParameters) -> tuple[BlockOperator, BlockFunction, Function]:

        # Forward operator
        A2d = ProjectionOperator(image_geometry, acquisition_data.geometry, 'gpu')

        if recon_params.stochastic:
            for partition_geometry, partition_operator in zip(acquisition_data.geometry, A2d, strict=True):
                CILRecon.set_approx_norm(partition_operator, partition_geometry, image_geometry)
        else:
            CILRecon.set_approx_norm(A2d, acquisition_data.geometry, image_geometry)

        # Define Gradient Operator and BlockOperator
        alpha = recon_params.alpha
        gamma = recon_params.gamma
        beta = alpha * gamma

        f2 = MixedL21Norm()
        f3 = MixedL21Norm()

        if recon_params.stochastic:
            raise ValueError("TGV reconstruction does not yet support stochastic mode")
            # now, A2d is a BlockOperator as acquisition_data is a BlockDataContainer
            fs = []
            for i, _ in enumerate(acquisition_data.geometry):
                fs.append(L2NormSquared(b=acquisition_data.get_item(i)))

            F = BlockFunction(*fs, f2, f3)

        else:
            # Define BlockFunction F using the MixedL21Norm() and the L2NormSquared()
            # mathematicians like to multiply 1/2 in front of L2NormSquared. This is not necessary
            # it will mean that the regularisation parameter alpha is doubled
            f1 = L2NormSquared(b=acquisition_data)

            F = BlockFunction(f1, f2, f3)

        # Define BlockOperator K

        # Set up the 3 operator A, Grad and Symmetrised Gradient
        K11 = A2d
        K21 = alpha * GradientOperator(K11.domain)
        K32 = beta * SymmetrisedGradientOperator(K21.range)
        # these define the domain and range of the other operators
        K12 = ZeroOperator(K32.domain, K11.range)
        K22 = -alpha * IdentityOperator(domain_geometry=K21.range, range_geometry=K32.range)
        K31 = ZeroOperator(K11.domain, K32.range)

        K = BlockOperator(K11, K12, K21, K22, K31, K32, shape=(3, 2))

        if recon_params.non_negative:
            G = BlockFunction(IndicatorBox(lower=0, upper=None), ZeroFunction())

        else:
            # Define Function G simply as zero
            G = ZeroFunction()

        return (K, F, G)

    @staticmethod
    def set_approx_norm(A2d: BlockOperator, acquisition_data: AcquisitionGeometry,
                        image_geometry: ImageGeometry) -> None:
        """
        Use an analytic approximation of the norm of the operator
        Otherwise this is slow to calculate, this approximation is good to with in 5%
        When running in debug mode, check the approximation and raise an error if it is bad
        """
        assert all(s == 1.0 for s in image_geometry.spacing), "Norm approximations assume voxel size == 1"
        approx_a2d_norm = sqrt(image_geometry.voxel_num_x * acquisition_data.num_projections)
        if LOG.isEnabledFor(DEBUG):
            num_a2d_norm = A2d.PowerMethod(A2d, max_iteration=100)
            diff = abs(approx_a2d_norm - num_a2d_norm) / max(approx_a2d_norm, num_a2d_norm)
            LOG.debug(f"ProjectionOperator approx norm: {diff=} {approx_a2d_norm=} {num_a2d_norm=}")
            if diff > 0.05:
                raise RuntimeError(f"Bad ProjectionOperator norm: {diff=} {approx_a2d_norm=} {num_a2d_norm=}\n")
        A2d.set_norm(approx_a2d_norm)

    @staticmethod
    def get_data(sino: np.ndarray, ag: AcquisitionGeometry, recon_params: ReconstructionParameters,
                 num_subsets: int) -> AcquisitionData:
        data = AcquisitionData(sino, deep_copy=False, geometry=ag, suppress_warning=True)

        if recon_params.stochastic:
            # split the data and put it in a BlockDataContainer
            # unfortunately now the data will be duplicated in memory
            data = data.partition(num_subsets, 'staggered')
            geo = []
            for i in range(len(data)):
                data.get_item(i).reorder('astra')
                geo.append(data.get_item(i).geometry)
            # COMPAT - workaround for https://github.com/TomographicImaging/CIL/issues/1445
            data.geometry = BlockGeometry(*geo)
        else:
            data.reorder('astra')

        return data

    @staticmethod
    def find_cor(images: ImageStack, slice_idx: int, start_cor: float, recon_params: ReconstructionParameters) -> float:
        return tomopy.find_center(images.sinograms,
                                  images.projection_angles(recon_params.max_projection_angle).value,
                                  ind=slice_idx,
                                  init=start_cor,
                                  sinogram_order=True)

    @staticmethod
    def single_sino(sino: np.ndarray,
                    cor: ScalarCoR,
                    proj_angles: ProjectionAngles,
                    recon_params: ReconstructionParameters,
                    progress: Progress | None = None) -> np.ndarray:
        """
        Reconstruct a single slice from a single sinogram. Used for the preview and the single slice button.
        Should return a numpy array,
        """

        num_iter = recon_params.num_iter
        num_subsets = ceil(sino.shape[0] / recon_params.projections_per_subset)
        if recon_params.stochastic:
            # The UI will pass the number of epochs in this case
            num_iter *= num_subsets

        if progress:
            progress.add_estimated_steps(num_iter + 1)
            progress.update(steps=1, msg='CIL: Setting up reconstruction', force_continue=False)

        if cil_mutex.locked():
            LOG.warning("CIL recon already in progress")

        with cil_mutex:
            t0 = time.perf_counter()
            sino = BaseRecon.prepare_sinogram(sino, recon_params)
            pixel_num_h = sino.shape[1]
            pixel_size = 1.
            rot_pos_x = (cor.value - pixel_num_h / 2) * pixel_size
            ag = AcquisitionGeometry.create_Parallel2D(rotation_axis_position=[rot_pos_x, 0])

            ag.set_panel(pixel_num_h, pixel_size=pixel_size)
            ag.set_labels(DataOrder.ASTRA_AG_LABELS)
            ag.set_angles(angles=proj_angles.value, angle_unit='radian')

            # let's create a CIL AcquisitionData or BlockDataContainer
            data = CILRecon.get_data(sino, ag, recon_params, num_subsets)

            ig = ag.get_ImageGeometry()

            if recon_params.regulariser == 'TV':
                K, F, G = CILRecon.set_up_TV_regularisation(ig, data, recon_params)
            elif recon_params.regulariser == 'TGV':
                K, F, G = CILRecon.set_up_TGV_regularisation(ig, data, recon_params)
            else:
                raise ValueError(f"Regulariser must be one of 'TV', 'TGV'. Received '{recon_params.regulariser}'")

            max_iteration = 100000
            # this should set to a sensible number as evaluating the objective is costly
            update_objective_interval = 10
            if recon_params.stochastic:
                reg_percent = recon_params.regularisation_percent
                probs = [(1 - reg_percent / 100) / num_subsets] * num_subsets + [reg_percent / 100]
                algo = SPDHG(f=F,
                             g=G,
                             operator=K,
                             prob=probs,
                             max_iteration=max_iteration,
                             update_objective_interval=update_objective_interval)
            else:
                normK = K.norm()
                sigma = 1
                tau = 1 / (sigma * normK**2)
                algo = PDHG(f=F,
                            g=G,
                            operator=K,
                            tau=tau,
                            sigma=sigma,
                            max_iteration=max_iteration,
                            update_objective_interval=update_objective_interval)

            try:
                # this may be confusing for the user in case of SPDHG, because they will
                # input num_iter and they will run num_iter * num_subsets
                for iter in range(num_iter):
                    if progress:
                        progress.update(steps=1,
                                        msg=f'CIL: Iteration {iter + 1} of {num_iter}'
                                        f': Objective {algo.get_last_objective():.2f}',
                                        force_continue=False)
                    algo.next()
            finally:
                if progress:
                    progress.mark_complete()
            t1 = time.perf_counter()
            LOG.info(f"single_sino time: {t1-t0}s for shape {sino.shape}")

            if isinstance(algo.solution, BlockDataContainer):
                # TGV case
                return algo.solution[0].as_array()

            return algo.solution.as_array()

    @staticmethod
    def full(images: ImageStack,
             cors: list[ScalarCoR],
             recon_params: ReconstructionParameters,
             progress: Progress | None = None) -> ImageStack:
        """
        Performs a volume reconstruction using sample data provided as sinograms.

        :param images: Array of sinogram images
        :param cors: Array of centre of rotation values
        :param proj_angles: Array of projection angles in radians
        :param recon_params: Reconstruction Parameters
        :param progress: Optional progress reporter
        :return: 3D image data for reconstructed volume
        """

        num_iter = recon_params.num_iter
        num_subsets = ceil(images.num_projections / recon_params.projections_per_subset)
        if recon_params.stochastic:
            # The UI will pass the number of epochs in this case
            num_iter *= num_subsets

        progress = Progress.ensure_instance(progress, task_name='CIL reconstruction', num_steps=num_iter + 1)
        shape = images.data.shape
        if images.is_sinograms:
            data_order = DataOrder.ASTRA_AG_LABELS
            pixel_num_h, pixel_num_v = shape[2], shape[0]
        else:
            data_order = DataOrder.TIGRE_AG_LABELS
            pixel_num_h, pixel_num_v = shape[2], shape[1]

        projection_size = full_size_KB(images.data.shape, images.dtype)
        recon_volume_shape = pixel_num_h, pixel_num_h, pixel_num_v
        recon_volume_size = full_size_KB(recon_volume_shape, images.dtype)
        if recon_params.stochastic:
            estimated_mem_required = 3 * projection_size + 14 * recon_volume_size
        else:
            estimated_mem_required = 5 * projection_size + 13 * recon_volume_size
        free_mem = system_free_memory().kb()

        if (estimated_mem_required > free_mem):
            estimate_gb = estimated_mem_required / 1024 / 1024
            raise RuntimeError(
                "The machine does not have enough physical memory available to allocate space for this data."
                f" Estimated RAM needed is {estimate_gb:.2f} GB")

        if cil_mutex.locked():
            LOG.warning("CIL recon already in progress")

        with cil_mutex:
            t0 = time.perf_counter()
            LOG.info(f"Starting 3D PDHG-TV reconstruction: input shape {images.data.shape}"
                     f"output shape {recon_volume_shape}\n"
                     f"Num iter {recon_params.num_iter}, alpha {recon_params.alpha}, "
                     f"Non-negative {recon_params.non_negative},"
                     f"Stochastic {recon_params.stochastic}, subsets {num_subsets}")
            progress.update(steps=1, msg='CIL: Setting up reconstruction', force_continue=False)
            angles = images.projection_angles(recon_params.max_projection_angle).value

            pixel_size = 1.
            if recon_params.tilt is None:
                raise ValueError("recon_params.tilt is not set")
            rot_pos = [(cors[pixel_num_v // 2].value - pixel_num_h / 2) * pixel_size, 0, 0]
            slope = -np.tan(np.deg2rad(recon_params.tilt.value))
            rot_angle = [slope, 0, 1]

            ag = AcquisitionGeometry.create_Parallel3D(rotation_axis_position=rot_pos,
                                                       rotation_axis_direction=rot_angle)
            ag.set_panel([pixel_num_h, pixel_num_v], pixel_size=(pixel_size, pixel_size))
            ag.set_angles(angles=angles, angle_unit='radian')
            ag.set_labels(data_order)

            data = CILRecon.get_data(BaseRecon.prepare_sinogram(images.data, recon_params), ag, recon_params,
                                     num_subsets)

            ig = ag.get_ImageGeometry()

            if recon_params.regulariser == 'TV':
                K, F, G = CILRecon.set_up_TV_regularisation(ig, data, recon_params)
            elif recon_params.regulariser == 'TGV':
                K, F, G = CILRecon.set_up_TGV_regularisation(ig, data, recon_params)
            else:
                raise ValueError(f"Regulariser must be one of 'TV', 'TGV'. Received '{recon_params.regulariser}'")

            max_iteration = 100000
            # this should set to a sensible number as evaluating the objective is costly
            update_objective_interval = 10
            if recon_params.stochastic:
                reg_percent = recon_params.regularisation_percent
                probs = [(1 - reg_percent / 100) / num_subsets] * num_subsets + [reg_percent / 100]
                algo = SPDHG(f=F,
                             g=G,
                             operator=K,
                             prob=probs,
                             max_iteration=max_iteration,
                             update_objective_interval=update_objective_interval)
            else:
                normK = K.norm()
                sigma = 1
                tau = 1 / (sigma * normK**2)
                algo = PDHG(f=F,
                            g=G,
                            operator=K,
                            tau=tau,
                            sigma=sigma,
                            max_iteration=max_iteration,
                            update_objective_interval=update_objective_interval)

            with progress:
                # this may be confusing for the user in case of SPDHG, because they will
                # input num_iter and they will run num_iter * num_subsets
                for iter in range(num_iter):
                    if progress:
                        progress.update(steps=1,
                                        msg=f'CIL: Iteration {iter + 1} of {num_iter}'
                                        f': Objective {algo.get_last_objective():.2f}',
                                        force_continue=False)
                    algo.next()

                if isinstance(algo.solution, BlockDataContainer):
                    # TGV case
                    volume = algo.solution[0].as_array()
                else:
                    volume = algo.solution.as_array()
                LOG.info(f'Reconstructed 3D volume with shape: {volume.shape}')
            t1 = time.perf_counter()
            LOG.info(f"full reconstruction time: {t1-t0}s for shape {images.data.shape}")
            return ImageStack(volume)


def allowed_recon_kwargs() -> dict[str, list[str]]:
    return {
        'CIL: PDHG-TV':
        ['alpha', 'num_iter', 'non_negative', 'stochastic', 'projections_per_subset', 'regularisation_percent']
    }

1	# Copyright (C) 2024 ISIS Rutherford Appleton Laboratory UKRI
2	# SPDX - License - Identifier: GPL-3.0-or-later
3	from __future__ import annotations	1✔
4
5	import time	1✔
6	from logging import getLogger, DEBUG	1✔
7	from math import sqrt, ceil	1✔
8	from threading import Lock	1✔
9	from typing import TYPE_CHECKING	1✔
10
11	import numpy as np	1✔
12
13	from cil.framework import (AcquisitionData, AcquisitionGeometry, DataOrder, ImageGeometry, BlockGeometry,	1✔
14	BlockDataContainer)
15	from cil.optimisation.algorithms import PDHG, SPDHG	1✔
16	from cil.optimisation.operators import GradientOperator, BlockOperator	1✔
17	from cil.optimisation.operators import SymmetrisedGradientOperator, ZeroOperator, IdentityOperator	1✔
18
19	from cil.optimisation.functions import MixedL21Norm, L2NormSquared, BlockFunction, ZeroFunction, IndicatorBox, Function	1✔
20	from cil.plugins.astra.operators import ProjectionOperator	1✔
21
22	from mantidimaging.core.data import ImageStack	1✔
23	from mantidimaging.core.reconstruct.base_recon import BaseRecon	1✔
24	from mantidimaging.core.utility.optional_imports import safe_import	1✔
25	from mantidimaging.core.utility.progress_reporting import Progress	1✔
26	from mantidimaging.core.utility.size_calculator import full_size_KB	1✔
27	from mantidimaging.core.utility.memory_usage import system_free_memory	1✔
28
29	if TYPE_CHECKING:
30	from mantidimaging.core.utility.data_containers import ProjectionAngles, ReconstructionParameters, ScalarCoR
31
32	LOG = getLogger(__name__)	1✔
33	tomopy = safe_import('tomopy')	1✔
34	cil_mutex = Lock()	1✔
35
36
37	class CILRecon(BaseRecon):	1✔
38
39	@staticmethod	1✔
40	def set_up_TV_regularisation(	1✔
41	image_geometry: ImageGeometry, acquisition_data: AcquisitionData,
42	recon_params: ReconstructionParameters) -> tuple[BlockOperator, BlockFunction, Function]:
43
44	# Forward operator
45	A2d = ProjectionOperator(image_geometry, acquisition_data.geometry, 'gpu')	×
46
47	if recon_params.stochastic:	×
48	for partition_geometry, partition_operator in zip(acquisition_data.geometry, A2d, strict=True):	×
49	CILRecon.set_approx_norm(partition_operator, partition_geometry, image_geometry)	×
50	else:
51	CILRecon.set_approx_norm(A2d, acquisition_data.geometry, image_geometry)	×
52
53	# Define Gradient Operator and BlockOperator
54	alpha = recon_params.alpha	×
55	Grad = GradientOperator(image_geometry)	×
56
57	if recon_params.stochastic:	×
58	# now, A2d is a BlockOperator as acquisition_data is a BlockDataContainer
59	fs = []	×
60	for i, _ in enumerate(acquisition_data.geometry):	×
61	fs.append(L2NormSquared(b=acquisition_data.get_item(i)))	×
62	fs.append(MixedL21Norm())	×
63
64	F = BlockFunction(*fs)	×
65
66	# needs to unrol the A2d BlockOperator and put it in another, followed by
67	# the gradient operator, as in the deterministic case
68	K = BlockOperator(A2d.get_as_list(), alpha Grad)	×
69
70	else:
71	# Define BlockFunction F using the MixedL21Norm() and the L2NormSquared()
72	f1 = MixedL21Norm()	×
73	f2 = L2NormSquared(b=acquisition_data)	×
74
75	F = BlockFunction(f1, f2)	×
76
77	# define the BlockOperator
78	K = BlockOperator(alpha * Grad, A2d)	×
79
80	if recon_params.non_negative:	×
81	G = IndicatorBox(lower=0)	×
82	else:
83	# Define Function G simply as zero
84	G = ZeroFunction()	×
85
86	return (K, F, G)	×
87
88	@staticmethod	1✔
89	def set_up_TGV_regularisation(	1✔
90	image_geometry: ImageGeometry, acquisition_data: AcquisitionData,
91	recon_params: ReconstructionParameters) -> tuple[BlockOperator, BlockFunction, Function]:
92
93	# Forward operator
NEW 94	A2d = ProjectionOperator(image_geometry, acquisition_data.geometry, 'gpu')	×
95
NEW 96	if recon_params.stochastic:	×
NEW 97	for partition_geometry, partition_operator in zip(acquisition_data.geometry, A2d, strict=True):	×
NEW 98	CILRecon.set_approx_norm(partition_operator, partition_geometry, image_geometry)	×
99	else:
NEW 100	CILRecon.set_approx_norm(A2d, acquisition_data.geometry, image_geometry)	×
101
102	# Define Gradient Operator and BlockOperator
NEW 103	alpha = recon_params.alpha	×
NEW 104	gamma = recon_params.gamma	×
NEW 105	beta = alpha * gamma	×
106
NEW 107	f2 = MixedL21Norm()	×
NEW 108	f3 = MixedL21Norm()	×
109
NEW 110	if recon_params.stochastic:	×
NEW 111	raise ValueError("TGV reconstruction does not yet support stochastic mode")	×
112	# now, A2d is a BlockOperator as acquisition_data is a BlockDataContainer
113	fs = []
114	for i, _ in enumerate(acquisition_data.geometry):
115	fs.append(L2NormSquared(b=acquisition_data.get_item(i)))
116
117	F = BlockFunction(*fs, f2, f3)
118
119	else:
120	# Define BlockFunction F using the MixedL21Norm() and the L2NormSquared()
121	# mathematicians like to multiply 1/2 in front of L2NormSquared. This is not necessary
122	# it will mean that the regularisation parameter alpha is doubled
NEW 123	f1 = L2NormSquared(b=acquisition_data)	×
124
NEW 125	F = BlockFunction(f1, f2, f3)	×
126
127	# Define BlockOperator K
128
129	# Set up the 3 operator A, Grad and Symmetrised Gradient
NEW 130	K11 = A2d	×
NEW 131	K21 = alpha * GradientOperator(K11.domain)	×
NEW 132	K32 = beta * SymmetrisedGradientOperator(K21.range)	×
133	# these define the domain and range of the other operators
NEW 134	K12 = ZeroOperator(K32.domain, K11.range)	×
NEW 135	K22 = -alpha * IdentityOperator(domain_geometry=K21.range, range_geometry=K32.range)	×
NEW 136	K31 = ZeroOperator(K11.domain, K32.range)	×
137
NEW 138	K = BlockOperator(K11, K12, K21, K22, K31, K32, shape=(3, 2))	×
139
NEW 140	if recon_params.non_negative:	×
NEW 141	G = BlockFunction(IndicatorBox(lower=0, upper=None), ZeroFunction())	×
142
143	else:
144	# Define Function G simply as zero
NEW 145	G = ZeroFunction()	×
146
NEW 147	return (K, F, G)	×
148
149	@staticmethod	1✔
150	def set_approx_norm(A2d: BlockOperator, acquisition_data: AcquisitionGeometry,	1✔
151	image_geometry: ImageGeometry) -> None:
152	"""
153	Use an analytic approximation of the norm of the operator
154	Otherwise this is slow to calculate, this approximation is good to with in 5%
155	When running in debug mode, check the approximation and raise an error if it is bad
156	"""
157	assert all(s == 1.0 for s in image_geometry.spacing), "Norm approximations assume voxel size == 1"	×
158	approx_a2d_norm = sqrt(image_geometry.voxel_num_x * acquisition_data.num_projections)	×
159	if LOG.isEnabledFor(DEBUG):	×
160	num_a2d_norm = A2d.PowerMethod(A2d, max_iteration=100)	×
161	diff = abs(approx_a2d_norm - num_a2d_norm) / max(approx_a2d_norm, num_a2d_norm)	×
162	LOG.debug(f"ProjectionOperator approx norm: {diff=} {approx_a2d_norm=} {num_a2d_norm=}")	×
163	if diff > 0.05:	×
164	raise RuntimeError(f"Bad ProjectionOperator norm: {diff=} {approx_a2d_norm=} {num_a2d_norm=}\n")	×
165	A2d.set_norm(approx_a2d_norm)	×
166
167	@staticmethod	1✔
168	def get_data(sino: np.ndarray, ag: AcquisitionGeometry, recon_params: ReconstructionParameters,	1✔
169	num_subsets: int) -> AcquisitionData:
170	data = AcquisitionData(sino, deep_copy=False, geometry=ag, suppress_warning=True)	×
171
172	if recon_params.stochastic:	×
173	# split the data and put it in a BlockDataContainer
174	# unfortunately now the data will be duplicated in memory
175	data = data.partition(num_subsets, 'staggered')	×
176	geo = []	×
177	for i in range(len(data)):	×
178	data.get_item(i).reorder('astra')	×
179	geo.append(data.get_item(i).geometry)	×
180	# COMPAT - workaround for https://github.com/TomographicImaging/CIL/issues/1445
181	data.geometry = BlockGeometry(*geo)	×
182	else:
183	data.reorder('astra')	×
184
185	return data	×
186
187	@staticmethod	1✔
188	def find_cor(images: ImageStack, slice_idx: int, start_cor: float, recon_params: ReconstructionParameters) -> float:	1✔
189	return tomopy.find_center(images.sinograms,	×
190	images.projection_angles(recon_params.max_projection_angle).value,
191	ind=slice_idx,
192	init=start_cor,
193	sinogram_order=True)
194
195	@staticmethod	1✔
196	def single_sino(sino: np.ndarray,	1✔
197	cor: ScalarCoR,
198	proj_angles: ProjectionAngles,
199	recon_params: ReconstructionParameters,
200	progress: Progress \| None = None) -> np.ndarray:
201	"""
202	Reconstruct a single slice from a single sinogram. Used for the preview and the single slice button.
203	Should return a numpy array,
204	"""
205
206	num_iter = recon_params.num_iter	×
207	num_subsets = ceil(sino.shape[0] / recon_params.projections_per_subset)	×
208	if recon_params.stochastic:	×
209	# The UI will pass the number of epochs in this case
210	num_iter *= num_subsets	×
211
212	if progress:	×
213	progress.add_estimated_steps(num_iter + 1)	×
214	progress.update(steps=1, msg='CIL: Setting up reconstruction', force_continue=False)	×
215
216	if cil_mutex.locked():	×
217	LOG.warning("CIL recon already in progress")	×
218
219	with cil_mutex:	×
220	t0 = time.perf_counter()	×
221	sino = BaseRecon.prepare_sinogram(sino, recon_params)	×
222	pixel_num_h = sino.shape[1]	×
223	pixel_size = 1.	×
224	rot_pos_x = (cor.value - pixel_num_h / 2) * pixel_size	×
225	ag = AcquisitionGeometry.create_Parallel2D(rotation_axis_position=[rot_pos_x, 0])	×
226
227	ag.set_panel(pixel_num_h, pixel_size=pixel_size)	×
228	ag.set_labels(DataOrder.ASTRA_AG_LABELS)	×
229	ag.set_angles(angles=proj_angles.value, angle_unit='radian')	×
230
231	# let's create a CIL AcquisitionData or BlockDataContainer
232	data = CILRecon.get_data(sino, ag, recon_params, num_subsets)	×
233
234	ig = ag.get_ImageGeometry()	×
235
NEW 236	if recon_params.regulariser == 'TV':	×
NEW 237	K, F, G = CILRecon.set_up_TV_regularisation(ig, data, recon_params)	×
NEW 238	elif recon_params.regulariser == 'TGV':	×
NEW 239	K, F, G = CILRecon.set_up_TGV_regularisation(ig, data, recon_params)	×
240	else:
NEW 241	raise ValueError(f"Regulariser must be one of 'TV', 'TGV'. Received '{recon_params.regulariser}'")	×
242
243	max_iteration = 100000	×
244	# this should set to a sensible number as evaluating the objective is costly
245	update_objective_interval = 10	×
246	if recon_params.stochastic:	×
247	reg_percent = recon_params.regularisation_percent	×
248	probs = [(1 - reg_percent / 100) / num_subsets] * num_subsets + [reg_percent / 100]	×
249	algo = SPDHG(f=F,	×
250	g=G,
251	operator=K,
252	prob=probs,
253	max_iteration=max_iteration,
254	update_objective_interval=update_objective_interval)
255	else:
256	normK = K.norm()	×
257	sigma = 1	×
258	tau = 1 / (sigma * normK**2)	×
259	algo = PDHG(f=F,	×
260	g=G,
261	operator=K,
262	tau=tau,
263	sigma=sigma,
264	max_iteration=max_iteration,
265	update_objective_interval=update_objective_interval)
266
267	try:	×
268	# this may be confusing for the user in case of SPDHG, because they will
269	# input num_iter and they will run num_iter * num_subsets
270	for iter in range(num_iter):	×
271	if progress:	×
272	progress.update(steps=1,	×
273	msg=f'CIL: Iteration {iter + 1} of {num_iter}'
274	f': Objective {algo.get_last_objective():.2f}',
275	force_continue=False)
276	algo.next()	×
277	finally:
278	if progress:	×
279	progress.mark_complete()	×
280	t1 = time.perf_counter()	×
281	LOG.info(f"single_sino time: {t1-t0}s for shape {sino.shape}")	×
282
NEW 283	if isinstance(algo.solution, BlockDataContainer):	×
284	# TGV case
NEW 285	return algo.solution[0].as_array()	×
286
UNCOV 287	return algo.solution.as_array()	×
288
289	@staticmethod	1✔
290	def full(images: ImageStack,	1✔
291	cors: list[ScalarCoR],
292	recon_params: ReconstructionParameters,
293	progress: Progress \| None = None) -> ImageStack:
294	"""
295	Performs a volume reconstruction using sample data provided as sinograms.
296
297	:param images: Array of sinogram images
298	:param cors: Array of centre of rotation values
299	:param proj_angles: Array of projection angles in radians
300	:param recon_params: Reconstruction Parameters
301	:param progress: Optional progress reporter
302	:return: 3D image data for reconstructed volume
303	"""
304
305	num_iter = recon_params.num_iter	×
306	num_subsets = ceil(images.num_projections / recon_params.projections_per_subset)	×
307	if recon_params.stochastic:	×
308	# The UI will pass the number of epochs in this case
309	num_iter *= num_subsets	×
310
311	progress = Progress.ensure_instance(progress, task_name='CIL reconstruction', num_steps=num_iter + 1)	×
312	shape = images.data.shape	×
313	if images.is_sinograms:	×
314	data_order = DataOrder.ASTRA_AG_LABELS	×
315	pixel_num_h, pixel_num_v = shape[2], shape[0]	×
316	else:
317	data_order = DataOrder.TIGRE_AG_LABELS	×
318	pixel_num_h, pixel_num_v = shape[2], shape[1]	×
319
320	projection_size = full_size_KB(images.data.shape, images.dtype)	×
321	recon_volume_shape = pixel_num_h, pixel_num_h, pixel_num_v	×
322	recon_volume_size = full_size_KB(recon_volume_shape, images.dtype)	×
323	if recon_params.stochastic:	×
324	estimated_mem_required = 3 * projection_size + 14 * recon_volume_size	×
325	else:
326	estimated_mem_required = 5 * projection_size + 13 * recon_volume_size	×
327	free_mem = system_free_memory().kb()	×
328
329	if (estimated_mem_required > free_mem):	×
330	estimate_gb = estimated_mem_required / 1024 / 1024	×
331	raise RuntimeError(	×
332	"The machine does not have enough physical memory available to allocate space for this data."
333	f" Estimated RAM needed is {estimate_gb:.2f} GB")
334
335	if cil_mutex.locked():	×
336	LOG.warning("CIL recon already in progress")	×
337
338	with cil_mutex:	×
339	t0 = time.perf_counter()	×
340	LOG.info(f"Starting 3D PDHG-TV reconstruction: input shape {images.data.shape}"	×
341	f"output shape {recon_volume_shape}\n"
342	f"Num iter {recon_params.num_iter}, alpha {recon_params.alpha}, "
343	f"Non-negative {recon_params.non_negative},"
344	f"Stochastic {recon_params.stochastic}, subsets {num_subsets}")
345	progress.update(steps=1, msg='CIL: Setting up reconstruction', force_continue=False)	×
346	angles = images.projection_angles(recon_params.max_projection_angle).value	×
347
348	pixel_size = 1.	×
349	if recon_params.tilt is None:	×
350	raise ValueError("recon_params.tilt is not set")	×
351	rot_pos = [(cors[pixel_num_v // 2].value - pixel_num_h / 2) * pixel_size, 0, 0]	×
352	slope = -np.tan(np.deg2rad(recon_params.tilt.value))	×
353	rot_angle = [slope, 0, 1]	×
354
355	ag = AcquisitionGeometry.create_Parallel3D(rotation_axis_position=rot_pos,	×
356	rotation_axis_direction=rot_angle)
357	ag.set_panel([pixel_num_h, pixel_num_v], pixel_size=(pixel_size, pixel_size))	×
358	ag.set_angles(angles=angles, angle_unit='radian')	×
359	ag.set_labels(data_order)	×
360
361	data = CILRecon.get_data(BaseRecon.prepare_sinogram(images.data, recon_params), ag, recon_params,	×
362	num_subsets)
363
364	ig = ag.get_ImageGeometry()	×
365
NEW 366	if recon_params.regulariser == 'TV':	×
NEW 367	K, F, G = CILRecon.set_up_TV_regularisation(ig, data, recon_params)	×
NEW 368	elif recon_params.regulariser == 'TGV':	×
NEW 369	K, F, G = CILRecon.set_up_TGV_regularisation(ig, data, recon_params)	×
370	else:
NEW 371	raise ValueError(f"Regulariser must be one of 'TV', 'TGV'. Received '{recon_params.regulariser}'")	×
372
373	max_iteration = 100000	×
374	# this should set to a sensible number as evaluating the objective is costly
375	update_objective_interval = 10	×
376	if recon_params.stochastic:	×
377	reg_percent = recon_params.regularisation_percent	×
378	probs = [(1 - reg_percent / 100) / num_subsets] * num_subsets + [reg_percent / 100]	×
379	algo = SPDHG(f=F,	×
380	g=G,
381	operator=K,
382	prob=probs,
383	max_iteration=max_iteration,
384	update_objective_interval=update_objective_interval)
385	else:
386	normK = K.norm()	×
387	sigma = 1	×
388	tau = 1 / (sigma * normK**2)	×
389	algo = PDHG(f=F,	×
390	g=G,
391	operator=K,
392	tau=tau,
393	sigma=sigma,
394	max_iteration=max_iteration,
395	update_objective_interval=update_objective_interval)
396
397	with progress:	×
398	# this may be confusing for the user in case of SPDHG, because they will
399	# input num_iter and they will run num_iter * num_subsets
400	for iter in range(num_iter):	×
401	if progress:	×
402	progress.update(steps=1,	×
403	msg=f'CIL: Iteration {iter + 1} of {num_iter}'
404	f': Objective {algo.get_last_objective():.2f}',
405	force_continue=False)
406	algo.next()	×
407
NEW 408	if isinstance(algo.solution, BlockDataContainer):	×
409	# TGV case
NEW 410	volume = algo.solution[0].as_array()	×
411	else:
NEW 412	volume = algo.solution.as_array()	×
413	LOG.info(f'Reconstructed 3D volume with shape: {volume.shape}')	×
414	t1 = time.perf_counter()	×
415	LOG.info(f"full reconstruction time: {t1-t0}s for shape {images.data.shape}")	×
416	return ImageStack(volume)	×
417
418
419	def allowed_recon_kwargs() -> dict[str, list[str]]:	1✔
420	return {	1✔
421	'CIL: PDHG-TV':
422	['alpha', 'num_iter', 'non_negative', 'stochastic', 'projections_per_subset', 'regularisation_percent']
423	}

mantidproject / mantidimaging / 8829910821

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous