Note
Go to the end to download the full example code.
TOF-MLEM with projection data
This example demonstrates the use of the MLEM algorithm to minimize the negative Poisson log-likelihood function.
\[f(x) = \sum_{i=1}^m \bar{y}_i (x) - y_i \log(\bar{y}_i (x))\]
subject to
\[x \geq 0\]
using the linear forward model
\[\bar{y}(x) = A x + s\]
Tip
parallelproj is python array API compatible meaning it supports different
array backends (e.g. numpy, cupy, torch, …) and devices (CPU or GPU).
Choose your preferred array API xp and device dev below.
30 from __future__ import annotations
31 from parallelproj import Array
32
33 import array_api_compat.numpy as xp
34
35 # import array_api_compat.cupy as xp
36 # import array_api_compat.torch as xp
37
38 import parallelproj
39 from array_api_compat import to_device
40 import array_api_compat.numpy as np
41 import matplotlib.pyplot as plt
42 import matplotlib.animation as animation
43
44 # choose a device (CPU or CUDA GPU)
45 if "numpy" in xp.__name__:
46 # using numpy, device must be cpu
47 dev = "cpu"
48 elif "cupy" in xp.__name__:
49 # using cupy, only cuda devices are possible
50 dev = xp.cuda.Device(0)
51 elif "torch" in xp.__name__:
52 # using torch valid choices are 'cpu' or 'cuda'
53 if parallelproj.cuda_present:
54 dev = "cuda"
55 else:
56 dev = "cpu"
Setup of the forward model \(\bar{y}(x) = A x + s\)
We setup a linear forward operator \(A\) consisting of an image-based resolution model, a non-TOF PET projector and an attenuation model
Note
The MLEM implementation below works with all linear operators that
subclass LinearOperator (e.g. the high-level projectors).
setup the LOR descriptor that defines the sinogram
85 img_shape = (40, 40, 8)
86 voxel_size = (2.0, 2.0, 2.0)
87
88 lor_desc = parallelproj.RegularPolygonPETLORDescriptor(
89 scanner,
90 radial_trim=10,
91 max_ring_difference=2,
92 sinogram_order=parallelproj.SinogramSpatialAxisOrder.RVP,
93 )
94
95 proj = parallelproj.RegularPolygonPETProjector(
96 lor_desc, img_shape=img_shape, voxel_size=voxel_size
97 )
98
99 # setup a simple test image containing a few "hot rods"
100 x_true = xp.ones(proj.in_shape, device=dev, dtype=xp.float32)
101 c0 = proj.in_shape[0] // 2
102 c1 = proj.in_shape[1] // 2
103 x_true[(c0 - 2) : (c0 + 2), (c1 - 2) : (c1 + 2), :] = 5.0
104 x_true[4, c1, 2:] = 5.0
105 x_true[c0, 4, :-2] = 5.0
106
107 x_true[:2, :, :] = 0
108 x_true[-2:, :, :] = 0
109 x_true[:, :2, :] = 0
110 x_true[:, -2:, :] = 0
Attenuation image and sinogram setup
116 # setup an attenuation image
117 x_att = 0.01 * xp.astype(x_true > 0, xp.float32)
118 # calculate the attenuation sinogram
119 att_sino = xp.exp(-proj(x_att))
Complete PET forward model setup
We combine an image-based resolution model, a non-TOF or TOF PET projector and an attenuation model into a single linear operator.
130 # enable TOF - comment if you want to run non-TOF
131 proj.tof_parameters = parallelproj.TOFParameters(
132 num_tofbins=13, tofbin_width=12.0, sigma_tof=12.0
133 )
134
135 # setup the attenuation multiplication operator which is different
136 # for TOF and non-TOF since the attenuation sinogram is always non-TOF
137 if proj.tof:
138 att_op = parallelproj.TOFNonTOFElementwiseMultiplicationOperator(
139 proj.out_shape, att_sino
140 )
141 else:
142 att_op = parallelproj.ElementwiseMultiplicationOperator(att_sino)
143
144 res_model = parallelproj.GaussianFilterOperator(
145 proj.in_shape, sigma=4.5 / (2.35 * proj.voxel_size)
146 )
147
148 # compose all 3 operators into a single linear operator
149 pet_lin_op = parallelproj.CompositeLinearOperator((att_op, proj, res_model))
Simulation of projection data
We setup an arbitrary ground truth \(x_{true}\) and simulate noise-free and noisy data \(y\) by adding Poisson noise.
159 # simulated noise-free data
160 noise_free_data = pet_lin_op(x_true)
161
162 # generate a contant contamination sinogram
163 contamination = xp.full(
164 noise_free_data.shape,
165 0.5 * float(xp.mean(noise_free_data)),
166 device=dev,
167 dtype=xp.float32,
168 )
169
170 noise_free_data += contamination
171
172 # add Poisson noise
173 np.random.seed(1)
174 y = xp.asarray(
175 np.random.poisson(parallelproj.to_numpy_array(noise_free_data)),
176 device=dev,
177 dtype=xp.float64,
178 )
EM update to minimize \(f(x)\)
The EM update that can be used in MLEM or OSEM is given by cite:p:Dempster1977 [SV82] [LC84] [HL94]
\[x^+ = \frac{x}{(A^k)^H 1} (A^k)^H \frac{y^k}{A^k x + s^k}\]
to calculate the minimizer of \(f(x)\) iteratively.
To monitor the convergence we calculate the relative cost
\[\frac{f(x) - f(x^*)}{|f(x^*)|}\]
and the distance to the optimal point
\[\frac{\|x - x^*\|}{\|x^*\|}.\]
We setup a function that calculates a single MLEM/OSEM update given the current solution, a linear forward operator, data, contamination and the adjoint of ones.
207 def em_update(
208 x_cur: Array,
209 data: Array,
210 op: parallelproj.LinearOperator,
211 s: Array,
212 adjoint_ones: Array,
213 ) -> Array:
214 """EM update
215
216 Parameters
217 ----------
218 x_cur : Array
219 current solution
220 data : Array
221 data
222 op : parallelproj.LinearOperator
223 linear forward operator
224 s : Array
225 contamination
226 adjoint_ones : Array
227 adjoint of ones
228
229 Returns
230 -------
231 Array
232 _description_
233 """
234 ybar = op(x_cur) + s
235 return x_cur * op.adjoint(data / ybar) / adjoint_ones
Run the MLEM iterations
242 # number of MLEM iterations
243 num_iter = 20
244
245 # initialize x
246 x = xp.ones(pet_lin_op.in_shape, dtype=xp.float32, device=dev)
247 # calculate A^H 1
248 adjoint_ones = pet_lin_op.adjoint(
249 xp.ones(pet_lin_op.out_shape, dtype=xp.float32, device=dev)
250 )
251
252 for i in range(num_iter):
253 print(f"MLEM iteration {(i + 1):03} / {num_iter:03}", end="\r")
254 x = em_update(x, y, pet_lin_op, contamination, adjoint_ones)
MLEM iteration 001 / 020
MLEM iteration 002 / 020
MLEM iteration 003 / 020
MLEM iteration 004 / 020
MLEM iteration 005 / 020
MLEM iteration 006 / 020
MLEM iteration 007 / 020
MLEM iteration 008 / 020
MLEM iteration 009 / 020
MLEM iteration 010 / 020
MLEM iteration 011 / 020
MLEM iteration 012 / 020
MLEM iteration 013 / 020
MLEM iteration 014 / 020
MLEM iteration 015 / 020
MLEM iteration 016 / 020
MLEM iteration 017 / 020
MLEM iteration 018 / 020
MLEM iteration 019 / 020
MLEM iteration 020 / 020
Calculation of the negative Poisson log-likelihood function of the reconstruction
261 # calculate the negative Poisson log-likelihood function of the reconstruction
262 exp = pet_lin_op(x) + contamination
263 # calculate the relative cost and distance to the optimal point
264 cost = float(xp.sum(exp - y * xp.log(exp)))
265 print(f"\nMLEM cost {cost:.6E} after {num_iter:03} iterations")
MLEM cost -2.667688E+05 after 020 iterations
Visualize the results
272 def _update_img(i):
273 img0.set_data(x_true_np[:, :, i])
274 img1.set_data(x_np[:, :, i])
275 ax[0].set_title(f"true image - plane {i:02}")
276 ax[1].set_title(f"MLEM iteration {num_iter} - plane {i:02}")
277 return (img0, img1)
278
279
280 x_true_np = parallelproj.to_numpy_array(x_true)
281 x_np = parallelproj.to_numpy_array(x)
282
283 fig, ax = plt.subplots(1, 2, figsize=(10, 5))
284 vmax = x_np.max()
285 img0 = ax[0].imshow(x_true_np[:, :, 0], cmap="Greys", vmin=0, vmax=vmax)
286 img1 = ax[1].imshow(x_np[:, :, 0], cmap="Greys", vmin=0, vmax=vmax)
287 ax[0].set_title(f"true image - plane {0:02}")
288 ax[1].set_title(f"MLEM iteration {num_iter} - plane {0:02}")
289 fig.tight_layout()
290 ani = animation.FuncAnimation(fig, _update_img, x_np.shape[2], interval=200, blit=False)
291 fig.show()
Total running time of the script: (0 minutes 30.603 seconds)
Related examples
