MBIRJAX v0.3.3 (#33)

* Convert from cos_sin_angles to angles.

* Working version of new install.

* Update install to work with readthedocs (I hope).

* Update docs and doc installation.

* Improve developer docs.

* Add gpu memory and timing scripts.

* Improve and revise docs

* Minor correction

* buzzard_dev (#10)

* Refactor num_iterations as a method input rather than a class parameter.

* Implement voxel_batch_size for forward projection.

* Implement voxel_batch_size for back projection.

* Add initialization feature to recon

* Prepare for merge of add_prox.

* Improve error message.

---------

Co-authored-by: Charles Bouman <Charles.Bouman@gmail.com>

* Add proximal map (#12)

* Possibly buggy but functioning commit of refactored VCD functions

* Partially tested initial implementation of prox map

* Prox bug causes nan when initial condition is perfect

* Update gpu evaluatation for new interface.

* Add epsilon to avoid alpha=nan.

* Pull auto_set_regularization_params() out of prox()

* Update demo script

---------

Co-authored-by: Greg Buzzard <buzzard@purdue.edu>
Co-authored-by: gbuzzard <54102356+gbuzzard@users.noreply.github.com>

* Remove unused attributes: proximal_map_flag, max_resolutions, initialization

* Add batch size management to recon_test.py

* Add advanced features to documentation

* Fix typo.

* Add prox_map to docs

* Convert to use public instance variables in Projectors.

* Refactor to use recon_shape.

* Refactor to use recon_shape and projector_params.

* Remove prox_recon parameter

* Improve documentation of Projectors.

* Improve documentation of Projectors.

* Basic Conebeam (#16)

* Refactor parallel_beam.py for clarity

* Remove references to svmbir.

* More renaming in parallel_beam.py

* Add partial implementation of cone_beam.py

* Update incomplete version of cone_beam.py

* Updates to incomplete version cone beam projector

* Updates to incomplete version cone beam projector

* Progress towards working version.

* More progress towards working version.

* More progress towards working version.

* More progress towards working version.

* More progress towards working version.

* More progress towards working version.

* Working version - still needs to be tested in multiple cases.

* Add cone beam demo and experiment scripts

* Refactor per-pixel magnification to allow for infinite source-detector distance.

* Include more options to exercise conebeam model.

* Remove reshape from test and demo scripts

* Bug fix to get the correct system matrix row entries.

* Add link to powerpoint slides.

* Rename variables and functions for clarity

* Rename various variables and methods.

* Rename voxel to pixel as needed.

* Fix naming problems and restore vcd_figs_for_abst.py to working condition.

* Minor improvement to vcd_figs_for_abst.py script

* Add variable intensity window to slice_viewer and add 'modified' to the function that generates the modified Shepp-Logan phantom.

* Partial progress towards improving efficiency in conebeam.

* Update cvpr scripts

* Partial progress towards improving efficiency in conebeam.

* Rename delta_voxel_xy, delta_voxel_z, and delta_pixel_recon to delta_voxel.

* First implementation of forward vertical fan beam projection.

* Fix bugs in sphinx docs

* Add TomographyModel parameter definitions

* Refactor sphinx docs names

* Add documentation for cone beam class

* Fix bug in conebeam backprojection and simplify parameter handling.

---------

Co-authored-by: gbuzzard <54102356+gbuzzard@users.noreply.github.com>
Co-authored-by: Greg Buzzard <buzzard@purdue.edu>

* Efficient Cone Beam and Parameter Handling (#18)

* Partial implementation of horizontal fan beam projection.

* Add cone beam developer sphinx doc

* Working refactored forward projection for cone beam.

* Remove unused code.

* Change from magnification to source_iso_dist.

* Refactor magnification from a parameter to a function.

* Additional comment.

* Working version of vertical fan beam back projection.

* Working version of full refactored back projection.

* Working version of full refactored back projection.

* Fix Hessian calculation for cone beam.

* Add cone beam figure

* Remove logo

* Add logo

* Add pdf, jpg, pptx to gitignore

* Fix bug in docs.

* Change interface for basic geometry class to operate on a batch of pixels and one view for forward and back projections.
Improve docstrings.

* Refactor to vmap over pixels in vertical back projection.

* Move p to the list of geometry parameters.

* Include default batch sizes.  Are you happy?

* Refactor top level data flow in forward projector.

* Add projectors. Yes :-)

* Add cone beam tests and pytest support

* Refactor top level data flow in back projector.

* Add vmap over slices in horizontal fan beam.

* Add warnings when tests fail.

* Update pixel and voxel batch sizes and handling.
Make source_detector_distance a multiple of num_det_channels.

* Change auto_set_recon_size to be geometry specific.

* Refactored auto_set_sigma functions

* Refactor TomographModel to put parameter handling in a superclass.
Perform some code cleanup.

* Fix docs to reflect ParameterHandler.

* Change p to psf_radius and include sinogram viewer.

* Add psf_radius support

* Add psf_radius for both parallel and cone geometry

* Add ability to save and load to file.

* Check that loaded parameters match the type of the loading class.

* Update documentation.

* Edit documentation

---------

Co-authored-by: Greg Buzzard <buzzard@purdue.edu>

* Add bug-fixes and PyPI support (#20)

* Refactor to separate out the calculation of quantities needed for projection.
Improve documentation.

* Provide a display default when vmin=vmax.

* Provide a display default when vmin=vmax.

* Allow the slice label to be set.

* Fix bug to prevent zooming in sliders.

* Add cpu memory stats.

* Refactor for joint pan and zoom with two images.

* Minor modification to pyproject.toml

* Add PyPI dependencies

---------

Co-authored-by: Greg Buzzard <buzzard@purdue.edu>

* Update step size in vcd and other various improvements.   (#26)

* Update PyPI instructions.

* Correct docstring, tweak auto_regularization, improve demos, and update version number

* Update gitignore for jupyter notebooks and add LLNL experiments directory

* Minor bug fixes.

* Add files to reconstruct and view nersc data.

* Fix bug in backproject and improve parameter return for recon.

* Fix bug in default weights.

* Fix step size in vcd updates.
Clean up docs.
Change default partitions.
Fix bug in row and channel offets in cone and parallel.
Combine 2 demos into one.
Include more tests.
Fix bug in default weights.
Use reflected boundary conditions for qggmrf top and bottom slices.

* Fix minor bug.

* Update code for cvpr abstract to use the original partition sequences.

* Remove lbnl files for PR.

* Update version number

---------

Co-authored-by: Charles Bouman <Charles.Bouman@gmail.com>

* Improve package maintenance instructions.

* Refactor get_delta as a separate function and include a test for it.

* Refactor to use apply_map_in_batches.

* Refactor to use sum_function_in_batches and rename apply_map_in batches to concatenate_function_in_batches.

* Label regularization parameters as static for more efficient compiling and execution.

* Reorganize qggmrf names and functions in preparation for refactoring to use vmap.

* Fix error in doc related to recon_shape

* Add parameter checking to raise an exception if the user tries to set an invalid parameter.

* Add parameter checking to raise an exception if the user tries to set an invalid parameter.

* Relax the tolerances a little to account for randomness in vcd partition selection.

* Add cube phantom.

* Refactor qggmrf gradient and hessian to use vmap.

* Update qggmrf test.

* Remove randomness to ensure reproducibility.

* Partially working version - not fully debugged:
Refactor sum and concatenate in batches to work on tuples of inputs and outputs and include tests of these functions.
Implement infrastructure to return prior cost and to choose optimal alpha using prior gradient and hessian.

* Add QR code generation script

* Update tests.

* Include cost function for full recon and test for qggmrf gradient correctness.

* Start to remove cost function.

* Finish refactoring prior cost.

* Change default number of iterations.

* Update cvpr figures.

* Include test for surrogate hessian.

* Update displayed total loss.

* Rename cost to loss, sigma_p to sigma_prox.

* Move qggmrf functions to separate file.

* Move qggmrf functions to separate file.

* Working version of alpha using approximate prior quadratic.

* Remove deprecated keywords.

* Change version number.

* Buzzard dev (#31)

* Fix transparancy of corners of QR code

* Put names on variables.

* Add spectral plots for cvpr.

* Update docs to remove **kwargs from constructor.

* Include blue noise sampling and updated partition sequence for demo.

* Change default granularity sequence and handle case of more subsets than pixels.

* Allow display of 2D images.

* Minor changes to experiments.

* Update and correct documentation

* Update docs and output messages.

* Update docs and output messages.

* Make random samples the default and blue noise the alternate.

---------

Co-authored-by: Charles Bouman <Charles.Bouman@gmail.com>

* Include nrms update values.

* Include nrms update values.

* Update version number to v0.3.2

* Fix warning messages.

* Update documentation with VCD information

* Update version to v0.3.3

* Fix bug in tests.

---------

Co-authored-by: Greg Buzzard <buzzard@purdue.edu>
Co-authored-by: gbuzzard <54102356+gbuzzard@users.noreply.github.com>

docs/source/index.rst

@@ -11,7 +11,7 @@ MBIRJAX: High-performance tomographic reconstruction
 **Key features:**
 
 
-* Vectorized Coordinate Descent algorithm for fast, robust convergence.
+* Vectorized Coordinate Descent algorithm for fast, robust convergence :cite:`2024CV4SciencePoster`.
 * Automatic parameter selection, with fine-tuning using intuitive meta-parameters.
 * Support for Plug-and-Play prior models that can dramatically improve image quality :cite:`venkatakrishnan2013plug` :cite:`sreehari2016plug`.
 * Modular, extensible, easy-to-use, object-oriented Python interface.

docs/source/refs.bib

@@ -21,4 +21,14 @@ @article{sreehari2016plug
   doi={10.1109/TCI.2016.2599778}
 }
 
+@inproceedings{2024CV4SciencePoster,
+  title={Vectorized Coordinate Descent for Fast {CT} Reconstruction},
+  author={Bouman, Charles A and Bouman, Gregery T},
+  booktitle={2024 CVPR Workshop: Computer Vision for Science},
+  year={2024},
+  address = {Seattle, WA},
+  month = {June 17},
+  url = {https://engineering.purdue.edu/~bouman/publications/pdf/2024-CVPR-poster.pdf}
+}
+
 

docs/source/theory.rst

@@ -99,6 +99,13 @@ where the quantities correspond to the following python variables:
 * :math:`\sigma_p` corresponds to ``sigma_p``
 
 
+Vectorized Coordinate Descent (VCD)
+-----------------------------------
+
+At its core, MBIRJAX is based on multi-granular VCD (MG-VCD) optimization as described in :cite:`2024CV4SciencePoster`.
+For the user, this is all "under the hood", but it is critically important because it results in fast robust convergence that can be efficiently implemented in JAX and on modern GPU architectures.
+Moreover, MG-VCD does not require the selection of a geometry specific pre-conditioner, as would be typically required with gradient-based methods, so it enables MBIRJAX's unique support for multiple geometries.
+
 
 Arbitrary Length Units (ALU) Conversion
 ---------------------------------------

mbirjax/qggmrf.py

@@ -84,8 +84,6 @@ def qggmrf_gradient_and_hessian_at_indices(flat_recon, recon_shape, pixel_indice
     # Initialize the neighborhood weights for averaging surrounding pixel values.
     # Order is [row+1, row-1, col+1, col-1, slice+1, slice-1] - see definition in _utils.py
     b, sigma_x, p, q, T = qggmrf_params
-    b = jnp.array(b)
-    b = tuple(b / jnp.sum(b))
 
     # First work on cylinders - determine the contributions from neighbors in the voxel cylinder
     qggmrf_params = (b, sigma_x, p, q, T)

mbirjax/tomography_model.py

@@ -649,10 +649,10 @@ def create_vcd_subset_iterator(self, fm_hessian, weights=None, prox_input=None):
         positivity_flag = self.get_params('positivity_flag')
         fm_constant = 1.0 / (self.get_params('sigma_y') ** 2.0)
         b, sigma_x, p, q, T = self.get_params(['b', 'sigma_x', 'p', 'q', 'T'])
-        sharpness = self.get_params('sharpness')
-        alpha_clip_value = jnp.clip(1.3 - 0.2 * sharpness, 0.8, 1.6)
+        b = np.array(b)
+        b = b / np.sum(b)
         b = tuple(b)
-        qggmrf_params = (b, sigma_x, p, q, T)
+        qggmrf_params = tuple((b, sigma_x, p, q, T))
         sigma_prox = self.get_params('sigma_prox')
         pixel_batch_size = self.get_params('pixel_batch_size')
         recon_shape = self.get_params('recon_shape')

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mbirjax"
-version = "0.3.2"
+version = "0.3.3"
 description = "High-performance tomographic reconstruction"
 keywords = ["tomography", "tomographic reconstruction", "computed tomography"]
 readme = "README.rst"

tests/test_qggmrf.py

@@ -109,6 +109,8 @@ def test_gradient_and_hessian(self):
         T = 1.102763376071644  # 0.5 + np.random.rand(1)[0]
         sigma_x = 0.6448831829968968  # 0.1 + np.random.rand(1)[0]
         b = (1, 1, 1, 1, 1, 1)
+        b = normalize_b(b)
+
         qggmrf_params = (b, sigma_x, p, q, T)
 
         recon_shape = (3, 3, 3)
@@ -139,6 +141,7 @@ def test_loss_and_gradient(self):
         T = 1.46
         sigma_x = 0.789
         b = (1, 1, 1, 1, 1, 1)
+        b = normalize_b(b)
         qggmrf_params = (b, sigma_x, p, q, T)
 
         # Get a random recon, x
@@ -169,5 +172,10 @@ def test_loss_and_gradient(self):
         assert (jnp.allclose(grad_direct, grad0.reshape(recon_shape)))
 
 
+def normalize_b(b):
+    b_sum = np.sum(np.array(b))
+    b = tuple([b_entry / b_sum for b_entry in b])
+    return b
+
 if __name__ == '__main__':
     unittest.main()