MBIRJAX v0.4.2 (#37)

* 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. * add NSI preprocess module and demo script * Add helper functions for the NSI demo script * Update NSI preprocess docstring * Update NSI preprocess docstring * Add preprocess module into documentation * Reorganization of nsi demo script * Add preprocess module * rename function: NSI_load_scans_and_param_jax() -> NSI_load_scans_and_params() * remove redundant functions in demo_utils * Add deppendency striprtf to environment.yml * update documentation of preprocess module * preprocess.correct_det_rotation(): reduce interpolation order for faster preprocess * preprocess.transmission_CT_compute_sino(): catch runtime warning in sinogram calculation and print out customized message. * Fix bug in slider. * Modify download_and_extract_tar() function to automatically parse the top level filename from the tarball * Add more comments to precisely describe the behavior of download_and_extract_tar() * NSI preprocess function: automatically parse file names from a NSI scan directory * Docstring and more print messages for NSI preprocesss function * Minor updates in the print messages * fix docstring for NSI preprocess function * seperate general preprocess functions and NSI specific functions to different files * Include grid partitions. * Improve granularity sequence. * Print and return alpha values. * Improve experiments for spectral and Fourier response. * Improve save and load to file and add tests for save and load. * Remove unneeded file after test. * Add blur model. * Update to preprocess documentation * preprocess.transmission_CT_compute_sino(): automatically correct the defective sino entries. * remove nsi demo scripts from mbirjax. These are moved to mbirjax_applications * syntax correction for warning handling * documentation update * Improved documentation on preprocessing * Refactor preprocess into a submodule. Remove dev docs. * Refactor preprocess into a submodule. Remove dev docs. * Import dependencies for general preprocess utilities. Make read_scan_dir() and read_scan_img() accessible by NSI functions * syntax update corresponding to the preprocess module reorganization * NSI preprocessing: Add wrapper function to compute sinogram data and geometry parameters with one function call * Add NSI wrapper function to usr_api.rst * minor update to NSI preprocess documentation * Remove unneeded files. * Update docstrings for preprocessing utilities * Improve proprocessing doc strings * Include tests for blur model. * Update parameter handling and docs. * Remove det_rotation. * Designate Blur as a development/demo model. * Ignore known warning in test load. * Update version to 0.4.0 * Bug fix for prox map and add prox map experiment. * Update prox experiment to demonstrate fixed point behavior. * Add proximal map test * Implement memory efficient cone-beam projection. * Add large scale recon for testing. * Working version with significantly reduced memory and good speed. * Simplify horizontal projector. * Specify the main device correctly. * Reuse existing variable rather than making a new one. * Tune batch sizes for 512**3. * Minor updates * Simplify the demo options. * Update batch size. * Remove unused import. * Improve the memory estimation. * Reduce default number of iterations. * Change parallel projectors to match conebeam horizontal projectors. * Update to match prerelease. * Update version number. * Improve preprocssing docs. --------- Co-authored-by: Greg Buzzard <buzzard@purdue.edu> Co-authored-by: gbuzzard <54102356+gbuzzard@users.noreply.github.com> Co-authored-by: Diyu Yang <yang1467@gilbreth-fe01.rcac.purdue.edu> Co-authored-by: Diyu Yang <yang1467@gilbreth-g005.rcac.purdue.edu> Co-authored-by: Diyu Yang <yang1467@gilbreth-g004.rcac.purdue.edu> Co-authored-by: Diyu Yang <yang1467@gilbreth-g010.rcac.purdue.edu> Co-authored-by: Diyu Yang <yang1467@purdue.edu> Co-authored-by: Diyu Yang <yang1467@gilbreth-g000.rcac.purdue.edu>
cabouman · Jul 22, 2024 · 02f9c77 · 02f9c77
1 parent 77b778c
commit 02f9c77
Show file tree

Hide file tree

Showing 10 changed files with 497 additions and 138 deletions.
diff --git a/demo/demo_shepp_logan.py b/demo/demo_shepp_logan.py
@@ -60,15 +60,15 @@
     # weights = ct_model.gen_weights(sinogram / sinogram.max(), weight_type='transmission_root')
 
     # Set reconstruction parameter values
-    ct_model.set_params(sharpness=sharpness, verbose=1)
+    ct_model.set_params(sharpness=sharpness)
 
     # Print out model parameters
     ct_model.print_params()
 
     # ##########################
     # Perform VCD reconstruction
     time0 = time.time()
-    recon, recon_params = ct_model.recon(sinogram, weights=weights, compute_prior_loss=True)
+    recon, recon_params = ct_model.recon(sinogram, weights=weights)
 
     recon.block_until_ready()
     elapsed = time.time() - time0

diff --git a/demo/demo_shepp_logan_large.py b/demo/demo_shepp_logan_large.py
@@ -0,0 +1,95 @@
+import jax
+import numpy as np
+import time
+import pprint
+import jax.numpy as jnp
+import mbirjax.parallel_beam
+
+if __name__ == "__main__":
+    """
+    This is a script to develop, debug, and tune the parallel beam mbirjax code.
+    """
+    # Choose the geometry type
+    geometry_type = 'cone'  # 'cone' or 'parallel'
+
+    print('Using {} geometry.'.format(geometry_type))
+
+    # Set parameters
+    num_views = 512
+    num_det_rows = 512
+    num_det_channels = 512
+    sharpness = 0.0
+
+    # These can be adjusted to describe the geometry in the cone beam case.
+    # np.Inf is an allowable value, in which case this is essentially parallel beam
+    source_detector_dist = 4 * num_det_channels
+    source_iso_dist = source_detector_dist
+
+    if geometry_type == 'cone':
+        detector_cone_angle = 2 * np.arctan2(num_det_channels / 2, source_detector_dist)
+    else:
+        detector_cone_angle = 0
+    start_angle = -(np.pi + detector_cone_angle) * (1 / 2)
+    end_angle = (np.pi + detector_cone_angle) * (1 / 2)
+
+    # Initialize sinogram
+    sinogram_shape = (num_views, num_det_rows, num_det_channels)
+    angles = jnp.linspace(start_angle, end_angle, num_views, endpoint=False)
+
+    # Set up the model
+    if geometry_type == 'cone':
+        ct_model = mbirjax.ConeBeamModel(sinogram_shape, angles, source_detector_dist=source_detector_dist,
+                                         source_iso_dist=source_iso_dist)
+    elif geometry_type == 'parallel':
+        ct_model = mbirjax.ParallelBeamModel(sinogram_shape, angles)
+    else:
+        raise ValueError('Invalid geometry type.  Expected cone or parallel, got {}'.format(geometry_type))
+
+    # Generate 3D Shepp Logan phantom
+    print('Creating phantom')
+    phantom = ct_model.gen_modified_3d_sl_phantom()
+
+    # Generate synthetic sinogram data
+    print('Creating sinogram')
+    sinogram = ct_model.forward_project(phantom)
+    phantom = jax.device_put(phantom, jax.devices('cpu')[0])
+    # View sinogram
+    # mbirjax.slice_viewer(sinogram, title='Original sinogram', slice_axis=0, slice_label='View')
+
+    # Generate weights array - for an initial reconstruction, use weights = None, then modify as desired.
+    weights = None
+    # weights = ct_model.gen_weights(sinogram / sinogram.max(), weight_type='transmission_root')
+
+    # Set reconstruction parameter values
+    ct_model.set_params(sharpness=sharpness, verbose=1)
+
+    # Print out model parameters
+    ct_model.print_params()
+
+    # ##########################
+    # Perform VCD reconstruction
+
+    time0 = time.time()
+    recon, recon_params = ct_model.recon(sinogram, weights=weights, compute_prior_loss=False, num_iterations=10)
+
+    recon.block_until_ready()
+    elapsed = time.time() - time0
+    print('Elapsed time for recon is {:.3f} seconds'.format(elapsed))
+    # ##########################
+    mbirjax.get_memory_stats()
+
+    # Print out parameters used in recon
+    # pprint.pprint(recon_params._asdict())
+    #
+    phantom = np.array(phantom)
+    recon = np.array(recon)
+    max_diff = np.amax(np.abs(recon - phantom))
+    print('Geometry = {}'.format(geometry_type))
+    nrmse = np.linalg.norm(recon - phantom) / np.linalg.norm(phantom)
+    pct_95 = np.percentile(np.abs(recon - phantom), 95)
+    print('NRMSE between recon and phantom = {}'.format(nrmse))
+    print('Maximum pixel difference between phantom and recon = {}'.format(max_diff))
+    print('95% of recon pixels are within {} of phantom'.format(pct_95))
+
+    # Display results
+    # mbirjax.slice_viewer(phantom, recon, title='Phantom (left) vs VCD Recon (right)')
diff --git a/docs/source/usr_api.rst b/docs/source/usr_api.rst
@@ -75,13 +75,21 @@ Preprocessing
 ------------------
 
 Preprocessing functions are implemented in :ref:`PreprocessDocs`. This includes various methods to compute and correct the sinogram data as needed.
+The following are functions specific to NSI scanners.
 
 .. autosummary::
 
-   mbirjax.preprocess.compute_sino_transmission
-   mbirjax.preprocess.interpolate_defective_pixels
-   mbirjax.preprocess.correct_det_rotation
-   mbirjax.preprocess.estimate_background_offset
-   mbirjax.preprocess.NSI.compute_sino_and_params
-   mbirjax.preprocess.NSI.load_scans_and_params
+   preprocess.NSI.compute_sino_and_params
+   preprocess.NSI.load_scans_and_params
+
+The remaining functions can be used for multiple types of scan data.
+
+.. autosummary::
+
+   preprocess.compute_sino_transmission
+   preprocess.interpolate_defective_pixels
+   preprocess.correct_det_rotation
+   preprocess.estimate_background_offset
+
+
 
diff --git a/experiments/prox_map_experiment.py b/experiments/prox_map_experiment.py
@@ -0,0 +1,74 @@
+import numpy as np
+import time
+import jax.numpy as jnp
+import mbirjax.parallel_beam
+
+if __name__ == "__main__":
+    """
+    This is a script to develop, debug, and tune the parallel beam mbirjax code.
+    """
+    # Choose the geometry type
+    geometry_type = 'cone'
+
+    print('Using {} geometry.'.format(geometry_type))
+
+    # Set parameters
+    num_views = 1
+    view_angle = np.pi / 2
+    num_det_rows = 40
+    num_det_channels = 64
+    sharpness = 4.0
+
+    # These can be adjusted to describe the geometry in the cone beam case.
+    # np.Inf is an allowable value, in which case this is essentially parallel beam
+    source_detector_dist = 4 * num_det_channels
+    source_iso_dist = source_detector_dist
+
+    detector_cone_angle = 2 * np.arctan2(num_det_channels / 2, source_detector_dist)
+
+    start_angle = view_angle
+    end_angle = view_angle + (np.pi + detector_cone_angle)
+
+    # Initialize sinogram
+    sinogram_shape = (num_views, num_det_rows, num_det_channels)
+    angles = jnp.linspace(start_angle, end_angle, num_views, endpoint=False)
+
+    # Set up the model
+    ct_model = mbirjax.ConeBeamModel(sinogram_shape, angles, source_detector_dist=source_detector_dist, source_iso_dist=source_iso_dist)
+
+    # Generate 3D Shepp Logan phantom
+    print('Creating phantom')
+    phantom = ct_model.gen_modified_3d_sl_phantom()
+
+    # Generate synthetic sinogram data
+    print('Creating sinogram')
+    sinogram = ct_model.forward_project(phantom)
+
+    # View sinogram
+    mbirjax.slice_viewer(sinogram, title='Original sinogram', slice_axis=0, slice_label='View')
+
+    # Generate weights array - for an initial reconstruction, use weights = None, then modify as desired.
+    weights = None
+    # weights = ct_model.gen_weights(sinogram / sinogram.max(), weight_type='transmission_root')
+
+    # Set reconstruction parameter values
+    ct_model.set_params(sharpness=sharpness, verbose=1)
+
+    # Print out model parameters
+    ct_model.print_params()
+
+    # ##########################
+    # Perform VCD reconstruction
+    prox_input = phantom  # jnp.zeros_like(phantom)
+    time0 = time.time()
+    recon, recon_params = ct_model.prox_map(prox_input, sinogram, weights=weights, num_iterations=10, init_recon=prox_input)
+
+    recon.block_until_ready()
+    elapsed = time.time() - time0
+    print('Elapsed time for recon is {:.3f} seconds'.format(elapsed))
+    # ##########################
+
+    # Display results
+    mbirjax.slice_viewer(phantom-recon, title='Phantom (left) vs VCD Recon (right)', vmin=-0.1, vmax=0.5)
+
+
diff --git a/mbirjax/_utils.py b/mbirjax/_utils.py
@@ -36,11 +36,11 @@
     'positivity_flag': {'val': False, 'recompile_flag': False},
     'snr_db': {'val': 30.0, 'recompile_flag': False},
     'sharpness': {'val': 0.0, 'recompile_flag': False},
-    'granularity': {'val': [1, 4, 64, 512, 2048], 'recompile_flag': False},
-    'partition_sequence': {'val': [0, 1, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4], 'recompile_flag': False},
+    'granularity': {'val': [1, 4, 64, 128], 'recompile_flag': False},
+    'partition_sequence': {'val': [0, 1, 2, 2, 3], 'recompile_flag': False},
     'verbose': {'val': 1, 'recompile_flag': False},
     'pixel_batch_size': {'val': 2048, 'recompile_flag': True},  # TODO: Determine batch sizes dynamically.
-    'view_batch_size': {'val': 4, 'recompile_flag': True}
+    'view_batch_size': {'val': 32, 'recompile_flag': True}
 }
 
 # These headings should match the dictionaries