11952572376

Committed 21 Nov 2024 11:46AM UTC coverage: 74.188% (-0.04%) from 74.226%

Build # 11952572376

Build Type

push

github

Committed by

web-flow

Commit Message

2375 convergence plots (#2399)

Run Details

2651 of 3909 branches covered (67.82%)

Branch coverage included in aggregate %.

13 of 24 new or added lines in 4 files covered. (54.17%)

3 existing lines in 2 files now uncovered.

9567 of 12560 relevant lines covered (76.17%)

0.76 hits per line

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

19.75

/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, Algorithm
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.optimisation.utilities.callbacks import Callback
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 MIProgressCallback(Callback):

    def __init__(self, verbose=1, progress: Progress | None = None) -> None:
        super().__init__(verbose)
        self.progress = progress
        self.iteration_count = 1

    def __call__(self, algo: Algorithm) -> None:
        if self.progress:
            extra_info = {'iterations': algo.iterations, 'losses': algo.loss}
            self.progress.update(
                steps=1,
                msg=f'CIL: Iteration {self.iteration_count } of {algo.max_iteration}'
                f': Objective {algo.get_last_objective():.2f}',
                force_continue=False,
                extra_info=extra_info,
            )
            self.iteration_count += 1


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) -> BlockDataContainer | 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}'")

            # 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, 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,
                            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
                algo.max_iteration = num_iter
                algo.run(num_iter, callbacks=[MIProgressCallback(progress=progress)])

            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}'")

            # 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, 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,
                            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
                algo.max_iteration = num_iter
                algo.run(num_iter, callbacks=[MIProgressCallback(progress=progress)])

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

mantidproject / mantidimaging / 11952572376

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