From 997ebf3a7ff3a418e90a91eab6169d4334434ba8 Mon Sep 17 00:00:00 2001 From: Michael McCann Date: Fri, 7 Jun 2024 07:51:43 -0600 Subject: [PATCH] Build example --- data | 2 +- docs/source/examples.rst | 1 + examples/scripts/README.rst | 2 + .../scripts/ct_projector_comparison_3d.py | 96 +++++++++++++++---- examples/scripts/index.rst | 1 + scico/linop/xray/_xray.py | 8 +- 6 files changed, 86 insertions(+), 24 deletions(-) diff --git a/data b/data index b085cef32..87e7d17db 160000 --- a/data +++ b/data @@ -1 +1 @@ -Subproject commit b085cef32b1866cb87f5d10734355c665491e400 +Subproject commit 87e7d17db717758d999270ac102ee35276cf90d1 diff --git a/docs/source/examples.rst b/docs/source/examples.rst index da89c8289..583603760 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -38,6 +38,7 @@ Computed Tomography examples/ct_astra_odp_train_foam2 examples/ct_astra_unet_train_foam2 examples/ct_projector_comparison + examples/ct_projector_comparison_3d examples/ct_multi_cs_tv_admm examples/ct_multi_tv_admm diff --git a/examples/scripts/README.rst b/examples/scripts/README.rst index 95d6aed8c..0e0b6aa0e 100644 --- a/examples/scripts/README.rst +++ b/examples/scripts/README.rst @@ -41,6 +41,8 @@ Computed Tomography CT Training and Reconstructions with UNet `ct_projector_comparison.py `_ X-ray Transform Comparison + `ct_projector_comparison_3d.py `_ + X-ray Transform Comparison in 3D `ct_multi_cs_tv_admm.py `_ TV-Regularized Sparse-View CT Reconstruction (Multiple Projectors, Common Sinogram) `ct_multi_tv_admm.py `_ diff --git a/examples/scripts/ct_projector_comparison_3d.py b/examples/scripts/ct_projector_comparison_3d.py index 5ea4062d0..a8808b2b0 100644 --- a/examples/scripts/ct_projector_comparison_3d.py +++ b/examples/scripts/ct_projector_comparison_3d.py @@ -9,8 +9,8 @@ X-ray Transform Comparison in 3D ================================ -This example compares SCICO's native 3D X-ray transform algorithm -to that of the ASTRA toolbox. +This example shows how to define a SCICO native 3D X-ray transform using +ASTRA toolbox conventions and vice versa. """ import numpy as np @@ -35,13 +35,12 @@ x = create_block_phantom(in_shape) x = jnp.array(x) -diagonal_length = int(jnp.ceil(jnp.sqrt(3) * N)) # use rectangular detector to check whether it is handled correctly -out_shape = (diagonal_length, diagonal_length + 1) +out_shape = (N, N + 1) """ -Set up SCICO projection +Set up SCICO projection. """ num_angles = 7 @@ -58,7 +57,7 @@ """ -Specify geometry using SCICO conventions and project +Specify geometry using SCICO conventions and project. """ num_repeats = 3 @@ -83,19 +82,30 @@ timer_scico.stop("avg_fwd") timer_scico.td["avg_fwd"] /= num_repeats +timer_scico.start("first_back") +HTy_scico = H_scico.T @ y_scico +timer_scico.stop("first_back") + +timer_scico.start("avg_back") +for _ in range(num_repeats): + HTy_scico = H_scico.T @ y_scico + jax.block_until_ready(HTy_scico) +timer_scico.stop("avg_back") +timer_scico.td["avg_back"] /= num_repeats + """ -Convert SCICO geometry to ASTRA and project +Convert SCICO geometry to ASTRA and project. """ P_to_astra_vectors = scico.linop.xray.P_to_vectors(in_shape, P, out_shape) timer_astra = Timer() -timer_astra.start("astra_init") +timer_astra.start("init") H_astra_from_scico = astra.XRayTransform3D( input_shape=in_shape, det_count=out_shape, vectors=P_to_astra_vectors ) -timer_astra.stop("astra_init") +timer_astra.stop("init") timer_astra.start("first_fwd") y_astra_from_scico = H_astra_from_scico @ x @@ -103,19 +113,30 @@ timer_astra.stop("first_fwd") timer_astra.start("first_fwd") -y_astra_from_scico = H_scico @ x +y_astra_from_scico = H_astra_from_scico @ x timer_astra.stop("first_fwd") timer_astra.start("avg_fwd") for _ in range(num_repeats): - y_astra_from_scico = H_scico @ x + y_astra_from_scico = H_astra_from_scico @ x jax.block_until_ready(y_astra_from_scico) timer_astra.stop("avg_fwd") timer_astra.td["avg_fwd"] /= num_repeats +timer_astra.start("first_back") +HTy_astra_from_scico = H_astra_from_scico.T @ y_astra_from_scico +timer_astra.stop("first_back") + +timer_astra.start("avg_back") +for _ in range(num_repeats): + HTy_astra_from_scico = H_astra_from_scico.T @ y_astra_from_scico + jax.block_until_ready(HTy_astra_from_scico) +timer_astra.stop("avg_back") +timer_astra.td["avg_back"] /= num_repeats + """ -Specify geometry with ASTRA conventions and project +Specify geometry with ASTRA conventions and project. """ angles = np.linspace(0, np.pi, num_angles) # evenly spaced projection angles @@ -126,8 +147,10 @@ y_astra = H_astra @ x +HTy_astra = H_astra.T @ y_astra + """ -Convert ASTRA geometry to SCICO and project +Convert ASTRA geometry to SCICO and project. """ P_from_astra = scico.linop.xray.astra_to_scico(H_astra.vol_geom, H_astra.proj_geom) @@ -135,6 +158,21 @@ y_scico_from_astra = H_scico_from_astra @ x +HTy_scico_from_astra = H_scico_from_astra.T @ y_scico_from_astra + +""" +Print timing results. +""" +print(f"init astra {timer_astra.td['init']:.2e} s") +print(f"init scico {timer_scico.td['init']:.2e} s") +print("") +for tstr in ("first", "avg"): + for dstr in ("fwd", "back"): + for timer, pstr in zip((timer_astra, timer_scico), ("astra", "scico")): + print(f"{tstr:5s} {dstr:4s} {pstr} {timer.td[tstr + '_' + dstr]:.2e} s") + print() + + """ Show projections. """ @@ -142,9 +180,9 @@ plot.imview(y_scico[0], title="SCICO projections", cbar=None, fig=fig, ax=ax[0, 0]) plot.imview(y_scico[2], cbar=None, fig=fig, ax=ax[1, 0]) plot.imview(y_scico[4], cbar=None, fig=fig, ax=ax[2, 0]) -plot.imview(y_astra_from_scico[0], title="ASTRA projections", cbar=None, fig=fig, ax=ax[0, 1]) -plot.imview(y_astra_from_scico[2], cbar=None, fig=fig, ax=ax[1, 1]) -plot.imview(y_astra_from_scico[4], cbar=None, fig=fig, ax=ax[2, 1]) +plot.imview(y_astra_from_scico[:, 0], title="ASTRA projections", cbar=None, fig=fig, ax=ax[0, 1]) +plot.imview(y_astra_from_scico[:, 2], cbar=None, fig=fig, ax=ax[1, 1]) +plot.imview(y_astra_from_scico[:, 4], cbar=None, fig=fig, ax=ax[2, 1]) fig.suptitle("Using SCICO conventions") fig.tight_layout() fig.show() @@ -153,12 +191,32 @@ plot.imview(y_scico_from_astra[0], title="SCICO projections", cbar=None, fig=fig, ax=ax[0, 0]) plot.imview(y_scico_from_astra[2], cbar=None, fig=fig, ax=ax[1, 0]) plot.imview(y_scico_from_astra[4], cbar=None, fig=fig, ax=ax[2, 0]) -plot.imview(y_astra[0], title="ASTRA projections", cbar=None, fig=fig, ax=ax[0, 1]) -plot.imview(y_astra[2], cbar=None, fig=fig, ax=ax[1, 1]) -plot.imview(y_astra[4], cbar=None, fig=fig, ax=ax[2, 1]) +plot.imview(y_astra[:, 0], title="ASTRA projections", cbar=None, fig=fig, ax=ax[0, 1]) +plot.imview(y_astra[:, 2], cbar=None, fig=fig, ax=ax[1, 1]) +plot.imview(y_astra[:, 4], cbar=None, fig=fig, ax=ax[2, 1]) fig.suptitle("Using ASTRA conventions") fig.tight_layout() fig.show() +""" +Show back projections. +""" +fig, ax = plot.subplots(nrows=1, ncols=2, figsize=(8, 6)) +plot.imview(HTy_scico[N // 2], title="SCICO back projection", cbar=None, fig=fig, ax=ax[0]) +plot.imview( + HTy_astra_from_scico[N // 2], title="ASTRA back projection", cbar=None, fig=fig, ax=ax[1] +) +fig.suptitle("Using SCICO conventions") +fig.tight_layout() +fig.show() + +fig, ax = plot.subplots(nrows=1, ncols=2, figsize=(8, 6)) +plot.imview( + HTy_scico_from_astra[N // 2], title="SCICO back projection", cbar=None, fig=fig, ax=ax[0] +) +plot.imview(HTy_astra[N // 2], title="ASTRA back projection", cbar=None, fig=fig, ax=ax[1]) +fig.suptitle("Using ASTRA conventions") +fig.tight_layout() +fig.show() input("\nWaiting for input to close figures and exit") diff --git a/examples/scripts/index.rst b/examples/scripts/index.rst index bbc697a95..0f036b5d7 100644 --- a/examples/scripts/index.rst +++ b/examples/scripts/index.rst @@ -25,6 +25,7 @@ Computed Tomography - ct_astra_odp_train_foam2.py - ct_astra_unet_train_foam2.py - ct_projector_comparison.py + - ct_projector_comparison_3d.py - ct_multi_cs_tv_admm.py - ct_multi_tv_admm.py diff --git a/scico/linop/xray/_xray.py b/scico/linop/xray/_xray.py index 00a87e3b7..6ba94d275 100644 --- a/scico/linop/xray/_xray.py +++ b/scico/linop/xray/_xray.py @@ -285,8 +285,8 @@ def P_to_vectors(in_shape: Shape, P: ArrayLike, det_shape: Shape) -> ArrayLike: y_center = np.array(det_shape) / 2 x_center = np.einsum("...mn,n->...m", P[..., :3], np.array(in_shape) / 2) + P[..., 3] d = np.einsum("...mn,...m->...n", P[..., :3], y_center - x_center) # (V, 2, 3) x (V, 2) - u = P[:, 1, :3] - v = P[:, 0, :3] + u = -P[:, 1, :3] + v = -P[:, 0, :3] vectors = np.concatenate((ray, d, u, v), axis=1) # (v, 12) return vectors @@ -295,11 +295,11 @@ def astra_to_scico(vol_geom, proj_geom): """ Convert ASTRA volume and projection geometry into a SCICO X-ray projection matrix. """ - in_shape = (vol_geom["GridColCount"], vol_geom["GridRowCount"], vol_geom["GridSliceCount"]) + in_shape = (vol_geom["GridSliceCount"], vol_geom["GridRowCount"], vol_geom["GridColCount"]) det_shape = (proj_geom["DetectorRowCount"], proj_geom["DetectorColCount"]) vectors = proj_geom["Vectors"] _, d, u, v = vectors[:, 0:3], vectors[:, 3:6], vectors[:, 6:9], vectors[:, 9:12] - P = np.stack((v, u), axis=1) + P = -np.stack((v, u), axis=1) center_diff = np.einsum("...mn,...n->...m", P, d) # y_center - x_center y_center = np.array(det_shape) / 2 Px_center_t = -(center_diff - y_center)