# -*- coding: utf-8 -*-
"""demo_2_large_object.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1-kk_HeR8Y8f6pZ2zjTza8NTEpAgwgVRB

**MBIRJAX: Large Object Demo**

See the [MBIRJAX documentation](https://mbirjax.readthedocs.io/en/latest/) for an overview and details.  

This script demonstrates how to improve the reconstruction when the object does not project completely inside the detector.  For simplicity, we show this only for parallel beam, but the same steps apply for cone beam.  

See [demo_1_shepp_logan.py](https://colab.research.google.com/drive/1zG_H6CDjuQxeMRQHan3XEyX2YVKcSSNC) for the basic steps of synthetic sinogram generation and reconstruction.

Select a GPU as runtime type for best performance.
"""

# Commented out IPython magic to ensure Python compatibility.
# %pip install mbirjax

import numpy as np
import time
import pprint
import jax.numpy as jnp
import mbirjax as mj

"""**Set the geometry parameters**"""

# Set parameters for the problem size - you can vary these, but if you make num_det_rows very small relative to
# channels, then the generated phantom may not have an interior.
num_views = 120
num_det_rows = 80
num_det_channels = 100

start_angle = - np.pi / 2
end_angle = np.pi / 2

"""**Data generation:** For demo purposes, we create a phantom and then project it to create a sinogram.

The default recon shape for parallel beam is

(rows, columns, slices) = (num_det_channels, num_det_channels, num_det_rows),

where we assume that the recon voxels are cubes and have the same size as the detector elements.
Here we generate a phantom that is bigger than the detector to show how to deal with this case.


Note:  the sliders on the viewer won't work in notebook form.  For that you'll need to run the python code with an interactive matplotlib backend, typcially using the command line or a development environment like Spyder or Pycharm to invoke python.  

"""

# Initialize sinogram
sinogram_shape = (num_views, num_det_rows, num_det_channels)
angles = jnp.linspace(start_angle, end_angle, num_views, endpoint=False)

ct_model_for_generation = mj.ParallelBeamModel(sinogram_shape, angles)

# Generate large 3D Shepp Logan phantom
print('Creating phantom that projects partially outside the detector')
phantom_row_scale = 1.0
phantom_col_scale = 1.75
phantom_rows = int(num_det_channels * phantom_row_scale)
phantom_cols = int(num_det_channels * phantom_col_scale)
phantom_slices = num_det_rows
phantom_shape = (phantom_rows, phantom_cols, phantom_slices)
phantom = mj.generate_3d_shepp_logan_low_dynamic_range(phantom_shape)

# Generate synthetic sinogram data
print('Creating sinogram')
ct_model_for_generation.set_params(recon_shape=phantom_shape)
sinogram = ct_model_for_generation.forward_project(phantom)
sinogram = np.asarray(sinogram)

# View sinogram
mj.slice_viewer(sinogram, title='Original sinogram\nChange view to see projections in and outside detector',
                 slice_axis=0, slice_label='View')

"""**Do a baseline reconstruction**

First we do a reconstruction with the default settings.  This will have significant artifacts because some of the information in the sinogram comes from voxels that are projected to the detector on only some of the views.  In the default reconstruction, all of the voxels in the reconstruction project to the detector in all of the views, so there is only way to account for these partially projected voxels leads to a corrupted reconstruciton.  
"""

# Initialize model for default reconstruction.
weights = None
ct_model_for_recon = mj.ParallelBeamModel(sinogram_shape, angles)

# Print model parameters
ct_model_for_recon.print_params()

# Default VCD reconstruction
print('Starting default recon - will have significant artifacts because of the missing projections.\n')
time0 = time.time()
recon, recon_dict = ct_model_for_recon.recon(sinogram, weights=weights)
recon_params = recon_dict['recon_params']
recon.block_until_ready()
elapsed = time.time() - time0

# Print out parameters used in recon
pprint.pprint(recon_params, compact=True)
print('Elapsed time for recon is {:.3f} seconds'.format(elapsed))

"""**Display the default reconstruction.**"""

title = 'Default recon: Phantom (left) vs VCD Recon (right)'
title += '\nAdjust intensity range to [0, 1] to see internal artifacts from projection outside detector.'
title += '\nAdjust intensity range to [1.5, 2] to see outer ring from projection outside detector.'
mj.slice_viewer(phantom, recon, data_dicts=[None, recon_dict], title=title, vmin=0.0, vmax=2.0)

"""**Decrease sharpness to reduce artifacts.**

One way to reduce the artifacts seen above is to decrease sharpness, which promotes smoother images.  This does reduce artifacts but also blurs the edges.  

Below we show that we can pad the recon to reduce artifacts without blurring the edges.  
"""

# Increased regularization VCD reconstruction
# We can reduce the artifacts by increasing regularization (decreasing sharpness).
sharpness = -1.5
ct_model_for_recon.set_params(sharpness=sharpness)
print('\nStarting recon with reduced sharpness - will have reduced artifacts but blurred edges.\n')
recon_smooth, recon_dict_smooth = ct_model_for_recon.recon(sinogram, weights=weights)
recon_params_smooth = recon_dict_smooth['recon_params']

# Print out parameters used in recon
pprint.pprint(recon_params_smooth, compact=True)

# Display results
title = 'Recon with sharpness = {:.1f}: Phantom (left) vs VCD Recon (right)'.format(sharpness)
title += '\nAdjust intensity range to [0, 1] to see reduced internal artifacts from projection outside detector.'
title += '\nOuter ring is still evident in intensity range [1, 2], and edges are blurry.'
mj.slice_viewer(phantom, recon_smooth, data_dicts=[None, recon_dict_smooth], title=title, vmin=0.0, vmax=2.0)

"""**Padded recon VCD reconstruction**

Alternatively, we can pad the recon to allow for a partial reconstruction of the pixels that project outside the detector in some views.  This reduces the artifacts without blurring edges and greatly reduces the outer ring seen in the non-padded recon.

Note that the enlarged recon doesn't have to match the phantom size.  Increasing the recon size won't allow us to fully reconstruct the pixels that sometimes project outside the detector.  However, it will provide
room for those partial projections to be absorbed into partial projected pixels, which allows for better reconstruction of the pixels with full projections.
"""

# Increase the recon size.  In this case, we increase just the columns to
# approximate the phantom shape.  Note that it doesn't have to be an exact match.
recon_row_scale = 1.0
recon_col_scale = 1.5
ct_model_for_recon.scale_recon_shape(row_scale=recon_row_scale, col_scale=recon_col_scale)

# Reset the default sharpness
sharpness = 1.0
ct_model_for_recon.set_params(sharpness=sharpness)

print('\nStarting enlarged recon - will have reduced artifacts, sharper edges, some extra pixel estimation.\n')
recon_enlarged, recon_dict_enlarged = ct_model_for_recon.recon(sinogram, weights=weights)
recon_params_enlarged = recon_dict_enlarged['recon_params']

# Print out parameters used in recon
pprint.pprint(recon_params_enlarged, compact=True)

"""**Display the result using the enlarged reconstruction.**"""

title = 'Padded recon with sharpness = {:.1f}: Phantom (left) vs VCD Recon (right)'.format(sharpness)
title += '\nPadding the recon reduces the internal artifacts even with default sharpness.'
title += '\nEdges are sharp, outer ring is mostly gone, and the partially projected pixels are partially recovered.'
mj.slice_viewer(phantom, recon_enlarged, data_dicts=[None, recon_dict_enlarged], title=title, vmin=0.0, vmax=2.0)

"""**Next:** Try changing some of the parameters and re-running or try [some of the other demos](https://mbirjax.readthedocs.io/en/latest/demos_and_faqs.html).  """