Merge pull request #8 from gschramm/main

add simulation of two different time frames

python/00_simulate_lm_data.py

@@ -6,6 +6,7 @@
 import prd
 import parallelproj
 import parallelproj_utils
+import utils
 import array_api_compat.numpy as np
 import matplotlib.pyplot as plt
 from array_api_compat import to_device
@@ -60,15 +61,20 @@
 img_shape = (num_trans, num_trans, num_ax)
 n0, n1, n2 = img_shape
 
-# setup an image containing a box
-
-img = xp.asarray(
+# setup the image for the 1st time frame
+img_f1 = xp.asarray(
     np.tile(np.load("../data/SL.npy")[..., None], (1, 1, num_ax)),
     device=dev,
     dtype=xp.float32,
 )
-img[:, :, :2] = 0
-img[:, :, -2:] = 0
+img_f1[:, :, :2] = 0
+img_f1[:, :, -2:] = 0
+
+img_f1 /= xp.max(img_f1)
+
+# setup the image for the 2nd time frame
+img_f2 = xp.sqrt(img_f1)
+
 
 # ----------------------------------------------------------------------------
 # ----------------------------------------------------------------------------
@@ -84,68 +90,45 @@
 )
 projector.tof = True  # set this to True to get a time of flight projector
 
+# repeat number of TOF bin times here
+num_tof_bins = projector.tof_parameters.num_tofbins
+
 # forward project the image
-noise_free_sinogram = projector(img)
+noise_free_sinogram_f1 = projector(img_f1)
+noise_free_sinogram_f2 = projector(img_f2)
 
 # rescale the forward projection and image such that we get the expected number of trues
-scale = expected_num_trues / float(xp.sum(noise_free_sinogram))
-noise_free_sinogram *= scale
-img *= scale
+scale = expected_num_trues / float(xp.sum(noise_free_sinogram_f1))
+noise_free_sinogram_f1 *= scale
+noise_free_sinogram_f2 *= scale
+img_f1 *= scale
+img_f2 *= scale
 
 # calculate the sensitivity image
 sens_img = projector.adjoint(
-    xp.ones(noise_free_sinogram.shape, device=dev, dtype=xp.float32)
+    xp.ones(noise_free_sinogram_f1.shape, device=dev, dtype=xp.float32)
 )
 
 # add poisson noise to the noise free sinogram
-noisy_sinogram = xp.asarray(
-    np.random.poisson(np.asarray(to_device(noise_free_sinogram, "cpu"))), device=dev
+noisy_sinogram_f1 = xp.asarray(
+    np.random.poisson(np.asarray(to_device(noise_free_sinogram_f1, "cpu"))), device=dev
 )
-# ravel the noisy sinogram and the detector start and end "index" sinograms
-noisy_sinogram = xp.reshape(noisy_sinogram, (noisy_sinogram.size,))
-
-# get the two dimensional indices of all sinogram bins
-start_mods, end_mods, start_inds, end_inds = lor_descriptor.get_lor_indices()
-
-# generate two sinograms that contain the linearized detector start and end indices
-sino_det_start_index = (
-    lor_descriptor.scanner.num_lor_endpoints_per_module[0] * start_mods + start_inds
-)
-sino_det_end_index = (
-    lor_descriptor.scanner.num_lor_endpoints_per_module[0] * end_mods + end_inds
+noisy_sinogram_f2 = xp.asarray(
+    np.random.poisson(np.asarray(to_device(noise_free_sinogram_f2, "cpu"))), device=dev
 )
 
-# repeat number of TOF bin times here
-num_tof_bins = projector.tof_parameters.num_tofbins
+# -------------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------------
+# -------------------------------------------------------------------------------------
 
-sino_det_start_index = xp.reshape(
-    xp.stack(num_tof_bins * [sino_det_start_index], axis=-1),
-    sino_det_start_index.size * num_tof_bins,
+event_det_id_1_f1, event_det_id_2_f1, event_tof_bin_f1 = utils.noisy_tof_sinogram_to_lm(
+    noisy_sinogram_f1, lor_descriptor, xp, dev
 )
-
-sino_det_end_index = xp.reshape(
-    xp.stack(num_tof_bins * [sino_det_end_index], axis=-1),
-    sino_det_end_index.size * num_tof_bins,
+event_det_id_1_f2, event_det_id_2_f2, event_tof_bin_f2 = utils.noisy_tof_sinogram_to_lm(
+    noisy_sinogram_f2, lor_descriptor, xp, dev
 )
-
-# ----------------------------------------------------------------------------
-# ----------------------------------------------------------------------------
-# --- convert the index sinograms in to listmode data ------------------------
-# ----------------------------------------------------------------------------
-# ----------------------------------------------------------------------------
-
-# generate listmode data from the noisy sinogram
-event_sino_inds = np.repeat(np.arange(noisy_sinogram.shape[0]), noisy_sinogram)
-# shuffle the event sinogram indices
-np.random.shuffle(event_sino_inds)
-# convert event sino indices to xp array
-event_sino_inds = xp.asarray(event_sino_inds, device=dev)
-
-event_det_id_1 = xp.take(sino_det_start_index, event_sino_inds)
-event_det_id_2 = xp.take(sino_det_end_index, event_sino_inds)
-event_tof_bin = event_sino_inds % num_tof_bins
-
-print(f"number of simulated events: {event_det_id_1.shape[0]}")
+print(f"number of simulated events f1: {event_det_id_1_f1.shape[0]}")
+print(f"number of simulated events f2: {event_det_id_1_f2.shape[0]}")
 
 # ----------------------------------------------------------------------------
 # ----------------------------------------------------------------------------
@@ -185,10 +168,10 @@
 
 # write the data to PETSIRD
 write_prd_from_numpy_arrays(
-    event_det_id_1,
-    event_det_id_2,
+    [event_det_id_1_f1, event_det_id_1_f2],
+    [event_det_id_2_f1, event_det_id_2_f2],
     scanner_information,
-    tof_idx_array=event_tof_bin,
+    tof_idx_array_blocks=[event_tof_bin_f1, event_tof_bin_f2],
     output_file=str(Path(output_dir) / output_prd_file),
 )
 print(f"wrote PETSIRD LM file to {str(Path(output_dir) / output_prd_file)}")
@@ -211,13 +194,22 @@
 fig_dir = Path("../figs")
 fig_dir.mkdir(exist_ok=True)
 
-vmax = 1.2 * xp.max(img)
-fig, ax = plt.subplots(1, img.shape[2], figsize=(img.shape[2] * 2, 2))
-for i in range(img.shape[2]):
-    ax[i].imshow(
-        xp.asarray(to_device(img[:, :, i], "cpu")), vmin=0, vmax=vmax, cmap="Greys"
+vmax = 1.2 * xp.max(img_f1)
+fig, ax = plt.subplots(
+    2, img_shape[2], figsize=(img_shape[2] * 2, 2 * 2), sharex=True, sharey=True
+)
+for i in range(img_shape[2]):
+    ax[0, i].imshow(
+        xp.asarray(to_device(img_f1[:, :, i], "cpu")), vmin=0, vmax=vmax, cmap="Greys"
+    )
+    ax[0, i].set_title(f"sl {i+1}", fontsize="small")
+
+    ax[1, i].imshow(
+        xp.asarray(to_device(img_f2[:, :, i], "cpu")), vmin=0, vmax=vmax, cmap="Greys"
     )
-    ax[i].set_title(f"ground truth sl {i+1}", fontsize="small")
+
+ax[0, 0].set_ylabel("ground truth frame 1", fontsize="small")
+ax[1, 0].set_ylabel("ground truth frame 2", fontsize="small")
 
 fig.tight_layout()
 fig.savefig(fig_dir / "simulated_phantom.png")

python/01_reconstruct_lm_data.py

@@ -31,7 +31,12 @@
 
 # read the LM file header and all event attributes
 header, event_attributes = read_prd_to_numpy_arrays(
-    str(Path(lm_data_dir) / prd_file), xp, dev, read_tof=None, read_energy=False
+    str(Path(lm_data_dir) / prd_file),
+    xp,
+    dev,
+    read_tof=None,
+    read_energy=False,
+    time_block_ids=range(1, 2),
 )
 
 # read the detector coordinates into a 2D lookup table