Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add JAX 3D X-ray CT projector #529

Merged
merged 40 commits into from
Sep 16, 2024
Merged
Changes from 1 commit
Commits
Show all changes
40 commits
Select commit Hold shift + click to select a range
4627595
Rough in 3D X-ray projector
Michael-T-McCann May 20, 2024
453a53e
Start on example
Michael-T-McCann May 21, 2024
c51370c
Rough in scico to astra conversion
Michael-T-McCann May 21, 2024
9493ff2
Work on example, still unfinished
Michael-T-McCann Jun 4, 2024
3625324
Work on example
Michael-T-McCann Jun 7, 2024
9a4ea59
Prepare to build example
Michael-T-McCann Jun 7, 2024
caef6eb
Add example to list of GPU-required
Michael-T-McCann Jun 7, 2024
3d0dd60
Modify example script titles
bwohlberg Jun 11, 2024
b3ce7e0
Update example script titles in examples README
bwohlberg Jun 11, 2024
ac7d34c
Rename example script
bwohlberg Jun 11, 2024
10be36e
Switch to using ContextTimer for more compact code
bwohlberg Jun 11, 2024
99d2131
Update change summary
bwohlberg Jun 11, 2024
1869259
Add missing __all__ entry
bwohlberg Jun 11, 2024
7bdcce5
Update copyright date
bwohlberg Jun 11, 2024
579df96
Refactor, improve docs
Michael-T-McCann Jun 14, 2024
f4fac55
Start to address memory problem
Michael-T-McCann Jun 21, 2024
46cde7e
Split up input to address memory bottleneck
Michael-T-McCann Jun 21, 2024
8661aa9
Rough in new conversion functions
Michael-T-McCann Jul 9, 2024
0a3aaa9
Point to correct data branch
Michael-T-McCann Jul 9, 2024
00d785c
Fix detector center transposition
Michael-T-McCann Jul 10, 2024
45daca1
Update submodule
bwohlberg Jul 19, 2024
28fe495
Merge branch 'main' into mike/3d_xray
bwohlberg Jul 19, 2024
b95e398
Start on back projector
Michael-T-McCann Jul 15, 2024
8a59ea4
Add manual back projector
Michael-T-McCann Jul 22, 2024
3b6a911
Update submodule
bwohlberg Jul 24, 2024
dc5dd5b
Merge branch 'main' into mike/3d_xray
bwohlberg Jul 24, 2024
beba882
Again correct unhashable argument (ndarray -> tuple for shapes)
Michael-T-McCann Jul 24, 2024
5df1c75
Merge branch 'main' into mike/3d_xray
bwohlberg Jul 30, 2024
4add502
Merge branch 'main' into mike/3d_xray
bwohlberg Aug 15, 2024
4351941
Work on comments
Michael-T-McCann Aug 26, 2024
854d3ad
Merge branch 'main' into mike/3d_xray
bwohlberg Sep 6, 2024
eee399f
Rename new projectors to XRayTransform2D and XRayTransform3D
Michael-T-McCann Sep 9, 2024
7acbaa4
Merge branch 'main' into mike/3d_xray
bwohlberg Sep 10, 2024
0b63130
Fix import errors in examples
bwohlberg Sep 10, 2024
6c01c4d
Changes to `linop.xray.astra` module (#551)
bwohlberg Sep 10, 2024
7afd4b6
Loosen tolerances to make tests pass
Michael-T-McCann Sep 11, 2024
9bb7b68
Remove unfinished example
Michael-T-McCann Sep 11, 2024
9da091d
Rerun examples that use the new projectors
Michael-T-McCann Sep 11, 2024
1b215e9
Update submodule
Michael-T-McCann Sep 11, 2024
bcba849
Add docstrings, typing
Michael-T-McCann Sep 11, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Fix detector center transposition
Michael-T-McCann committed Jul 10, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
commit 00d785cc4fca3abf6eaa7a2085205fd05cff88fb
8 changes: 4 additions & 4 deletions scico/linop/xray/astra.py
Original file line number Diff line number Diff line change
@@ -153,14 +153,14 @@ def project_world_coordinates(x, ray, d, u, v, det_shape):

"""
Phi = np.stack((ray, u, v), axis=1)
x = x - d # express with respect to detector center
alpha = np.linalg.pinv(Phi) @ x[..., :, np.newaxis] # (3,3) times <stack of> (3,1)
alpha = alpha[..., 0] # squash to (..., 3)
offset_from_center = alpha - d
projected_offset = offset_from_center[..., 1:] # throw away ray coordinate
alpha = alpha[..., 0] # squash from (..., 3, 1) to (..., 3)
Palpha = alpha[..., 1:] # throw away ray coordinate
det_center_idx = (
np.array(det_shape)[::-1] / 2 - 0.5
) # center of length-2 is index 0.5, length-3 -> index 1
ind_xy = projected_offset + det_center_idx
ind_xy = Palpha + det_center_idx
ind_ij = ind_xy[..., ::-1]
return ind_ij

36 changes: 27 additions & 9 deletions scico/test/linop/xray/test_astra.py
Original file line number Diff line number Diff line change
@@ -172,6 +172,15 @@ def test_angle_to_vector():
## conversion functions
@pytest.fixture(scope="module")
def test_geometry():
"""
In this geometry, if vol[i, j, k]==1, we expect proj[j-2, k-1]==1.

Because:
- We project along z, i.e. `ray=(0,0,1)`, i.e., we remove axis=0.
- We set `v=(0, 1, 0)`, so detector rows go with y axis, axis=1.
- We set `u=(1, 0, 0)`, so detector columns go with x axis, axis=2.
- We shift the detector by (x=1, y=2, z=3) <-> i-3, j-2, k-1
"""
in_shape = (30, 31, 32)
# in ASTRA terminology:
n_rows = in_shape[1] # y
@@ -183,11 +192,11 @@ def test_geometry():
assert vol_geom["option"]["WindowMinY"] == -n_rows / 2
assert vol_geom["option"]["WindowMinZ"] == -n_slices / 2

# project along z (slices)
# project along z, axis=0
det_row_count = n_rows
det_col_count = n_cols
ray = (0, 0, 1)
d = (0, 0, 0)
d = (1, 2, 3) # axis=2 offset by 1, axis=1 offset by 2, axis=0 offset by 3
u = (1, 0, 0) # increments columns, goes with X
v = (0, 1, 0) # increments rows, goes with Y
vectors = np.array(ray + d + u + v)[np.newaxis, :]
@@ -200,6 +209,11 @@ def test_geometry():

@pytest.mark.skipif(jax.devices()[0].platform != "gpu", reason="GPU required for test")
def test_projection_convention(test_geometry):
"""
If vol[i, j, k]==1, test that astra puts proj[j-2, k-1]==1.

See `test_geometry` for the setup.
"""
vol_geom, proj_geom = test_geometry
in_shape = scico.linop.xray.astra.astra.functions.geom_size(vol_geom)
vol = np.zeros(in_shape)
@@ -213,18 +227,22 @@ def test_projection_convention(test_geometry):
assert len(np.unique(proj) == 2)

idx_proj_i, idx_proj_j = np.nonzero(proj)
np.testing.assert_array_equal(idx_proj_i, j)
np.testing.assert_array_equal(idx_proj_j, k)
np.testing.assert_array_equal(idx_proj_i, j - 2)
np.testing.assert_array_equal(idx_proj_j, k - 1)


def test_project_coords(test_geometry):
"""See `test_geometry` for the setup and
`test_projection_convention` for proof ASTRA works this way."""
"""
If vol[i, j, k]==1, test that we predict proj[j-2, k-1]==1.

See `test_geometry` for the setup and `test_projection_convention`
for proof ASTRA works this way.
"""
vol_geom, proj_geom = test_geometry
in_shape = scico.linop.xray.astra.astra.functions.geom_size(vol_geom)
x_vol = np.array([np.random.randint(0, s) for s in in_shape])
x_proj_gt = np.array([[x_vol[1], x_vol[2]]]) # projection along slices removes first index
x_proj_gt = np.array(
[[x_vol[1] - 2, x_vol[2] - 1]]
) # projection along slices removes first index
x_proj = scico.linop.xray.astra._project_coords(x_vol, vol_geom, proj_geom)
np.testing.assert_array_equal(x_proj_gt, x_proj)

scico.linop.xray.astra.convert_to_scico_geometry(vol_geom, proj_geom)