Skip to content

Conversation

@romainvo
Copy link
Contributor

@romainvo romainvo commented May 3, 2025

I have implemented a TomographyWithAstra operator using the https://astra-toolbox.com/. It uses the GPULink functionality of astra to share the underlying CUDA memory of torch.Tensor and arrays used by astra. By wrapping astra, the operator supports 2d geometries (parallel and fanbeam) as well as 3d (parallel and conebeam).

Astra can be installed easily with conda install -c astra-toolbox -c nvidia astra-toolbox.

The contribution consists in two main classes:

  • a deepinv.physics.functional.astra.XrayTransform object, which wraps the astra-computation, both forward and adjoint computation, with torch.Tensor as inputs.
  • a deepinv.physics.TomographyWithAstra LinearPhysics operator, which instanciates an XrayTransform as attribute (similar to how the Radon module is used). The calls to XrayTransform are wrapped in a custom torch.autograd.Function (AutogradTransform in astra.py) to make the computation of A() and A_adjoint() differentiable with torch.

The operator is more memory efficient, though slower for batched computations in parallel 2d. This is because ASTRA does not handle batched computations, therefore the custom autograd.Function processes batch elements sequentially. It is a priori, not easily vmappable (https://pytorch.org/docs/stable/notes/extending.func.html) because the XrayTransform needs to access pointers of batch elements, which is not possible for Tensors produced during vmap.

I have also adapted the following examples, to check for an astra-toolbox install and adapt the CT operator used automatically.

  • examples/plug-and-play/demo_vanilla_PnP.py
  • examples/patch-priors/demo_patch_priors_CT.py

I haven't touched the learned primal-dual example which runs significantly slower while using the astra operator.

Other changes:

  • I have simplified and adapted the AbstractFiler module in deepinv.physics.functional.radon.py to handle sinogram data for both the Tomography operator and the TomographyWithAstra operator.

Test:

  • I have implemented an independant test for the TomographyWithAstra operator, which is definitely not going to pass the adjointness and pseudoinverse tests of linear OPERATORS. Not sure if this what is intended ? They seem to pass, and skipped by pytest if astra is not installed.

Checks to be done before submitting your PR

  • python3 -m pytest tests/ runs successfully.
  • black . runs successfully.
  • make html runs successfully (in the docs/ directory).
  • Updated docstrings related to the changes (as applicable).
  • Added an entry to the CHANGELOG.rst.

r = physics.A_adjoint(physics.A(x))
y = physics.A(r)
error = torch.linalg.norm(physics.A_dagger(y) - r) / torch.linalg.norm(r)
assert error < 0.2
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems like a quite high threshold given that you (correctly) test on the range of A... Do you think it's possible to reduce this?

Copy link
Contributor Author

@romainvo romainvo May 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Although we test on the range of A, I guess because the adjoint is not matched we still get a somewhat high error.

I have fixed the detector and angular parameters of the fan/cone setup to set a nicer geometry with a proper field of view. At such low-resolution, the error does not seem to go below 0.05 for parallel, though it does decrease with higher resolution.. For fan/cone beam the error is a little higher (0.1/0.15), which is expected.

Copy link
Collaborator

@Andrewwango Andrewwango left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR and it's really cool to provide this functionality! Some intermediate minor comments:

@codecov
Copy link

codecov bot commented May 12, 2025

Codecov Report

Attention: Patch coverage is 79.71831% with 72 lines in your changes missing coverage. Please review.

Project coverage is 80.77%. Comparing base (b0fd882) to head (872bc60).
Report is 3 commits behind head on main.

✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
deepinv/physics/functional/astra.py 74.61% 27 Missing and 23 partials ⚠️
deepinv/tests/test_external_libraries.py 65.38% 17 Missing and 1 partial ⚠️
deepinv/physics/functional/radon.py 92.30% 0 Missing and 2 partials ⚠️
deepinv/physics/tomography.py 97.43% 0 Missing and 2 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #474      +/-   ##
==========================================
+ Coverage   79.85%   80.77%   +0.91%     
==========================================
  Files         187      189       +2     
  Lines       16822    17226     +404     
  Branches     2213     2282      +69     
==========================================
+ Hits        13434    13914     +480     
+ Misses       2661     2532     -129     
- Partials      727      780      +53     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@romainvo
Copy link
Contributor Author

romainvo commented May 14, 2025

Hi everyone. I've tried to resolve the last checks failing. The doc won't build for some files because TomographyWithAstra only supports CUDA operations, this concerns:

  • the code snippets in the docstring
  • and the example in external-libraries/demo_astra_tomography.py

Is there a quick way to skip those checks if a CUDA device is not available ?

  • in the doctest_global_setup I've added an import for astra and a check for cuda device, then the snippets in the docstring check for it and execute the code dependently
  • I've renamed demo_astra_tomography.py -> _demo_astra_tomography.py, this way it is parsed but not executed.

Overall, I've implemented the tests and docs to be skipped if there is no GPU, I guess this is why the code coverage marks most of the implemented functions as not tested ?

Another point concerns the supported platform, at the moment, ASTRA only provides a pip package for linux. To get support for windows, you need to install it via their conda channel. I've restricted the dependency install to Linux platform for the moment.

@Andrewwango
Copy link
Collaborator

One solution to the lack of CPU tests that we discussed with @tachella is to mock the Astra library calls, like @jscanvic does to create dataset tests when we can't actually download the datasets in #490. Do you think this could be possible?

The philosophy here is that as Astra is an external library, our tests should not have to cover its functionality, only our own wrapper. Then it is also ok that we don't have tests for adjointness etc, only that it runs ok and outputs correct shapes etc. Would you agree with this idea?

@romainvo
Copy link
Contributor Author

One solution to the lack of CPU tests that we discussed with @tachella is to mock the Astra library calls, like @jscanvic does to create dataset tests when we can't actually download the datasets in #490. Do you think this could be possible?

The philosophy here is that as Astra is an external library, our tests should not have to cover its functionality, only our own wrapper. Then it is also ok that we don't have tests for adjointness etc, only that it runs ok and outputs correct shapes etc. Would you agree with this idea?

It seems like a good alternative while there is no GPU for testing ! I'll dig into it and come up with a solution.

Although, I think keeping the already written tests in the codebase might be a good idea, at least it gives the ability to check on local machines that next versions of ASTRA or torch don't change the intented behaviour, in particular for the pseudo-inverse which is half torch, half astra code.

@romainvo
Copy link
Contributor Author

@Andrewwango I've patched the calls to the astra kernel with dummy functions in a new test, it should be good, tell me what you think. I put it in a test_external_libraries.py.

Copy link
Collaborator

@Andrewwango Andrewwango left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good so far!

Copy link
Contributor

@tachella tachella left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PR looks very nice!

some minor comments

Copy link
Collaborator

@Andrewwango Andrewwango left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very minor, but apart from these LGTM and happy to approve once everything resolved

@Andrewwango
Copy link
Collaborator

@romainvo Please can you check out the bugs in the workflows?

@Andrewwango Andrewwango merged commit 0c6f5da into deepinv:main Jun 17, 2025
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants