Add multi-energy feature to tike

[YudongYao]

Purpose

Add multi-wavelength (MW) method for broadband illumination ptychography reconstruction.

Approach

Multi-wavelength method is developed to deal with the deteriorated coherence property caused by the broadband illumination. The MW approach assumes that the diffraction pattern with a broadband source is the superposition of diffracted intensities from a series of wavelengths within the illumination bandwidth.

Therefore, the probe is decomposed into a set of spectral probe modes corresponding to different wavelengths. These spectral probe modes are generated using the Fresnel propagation from the focusing optics, which is implemented in 'src/tike/ptycho/probe_MW.py'.

To keep both the MW feature and the multiple probe modes to jointly solve the partial spatial and temporal coherence problem, another dimension (represented by energy) is added to the probe array. The probe is composed of a set of spectral modes, and each spectral mode includes multiple spatial probe modes. Now the probe shape is (ntheta, nscan // fly, fly, energy, nmodes, probe_shape, probe_shape). Correspondingly, changes are made to 'tike/operators', where the intensity, grad and cost are calculated.

In 'ptycho/solver', for the probe update, another loop over the energy is added to update each probe mode, each energy sequentially. For the position update,only the probe modes related to central wavelength is used to update the postion.

Pre-Merge Checklists

Submitter

• Write a helpfully descriptive pull request title.
• Organize changes into logically grouped commits with descriptive commit messages.
• Document all new functions.
• Write tests for new functions or explain why they are not needed.
• Build the documentation successfully
• Use yapf to format python code.

Reviewer

• Actually read all of the code.
• Run the new code yourself.
• Write a summary of the changes as you understand them.
• Thank the submitter.

[pep8speaks]

Hello @YudongYao! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

• In the file src/tike/operators/numpy/ptycho.py:

Line 140:30: E251 unexpected spaces around keyword / parameter equals
Line 140:32: E251 unexpected spaces around keyword / parameter equals
Line 140:81: E501 line too long (98 > 80 characters)

• In the file src/tike/ptycho/position.py:

Line 107:5: E265 block comment should start with '# '

• In the file src/tike/ptycho/probe_MW.py:

Line 1:1: E265 block comment should start with '# '
Line 1:81: E501 line too long (83 > 80 characters)
Line 13:81: E501 line too long (85 > 80 characters)
Line 15:81: E501 line too long (88 > 80 characters)
Line 18:81: E501 line too long (96 > 80 characters)
Line 20:81: E501 line too long (115 > 80 characters)
Line 123:60: W291 trailing whitespace

• In the file src/tike/ptycho/solvers/combined.py:

Line 78:21: W292 no newline at end of file

• In the file tests/test_MW.py:

Line 14:10: E225 missing whitespace around operator
Line 14:43: E231 missing whitespace after ','
Line 14:76: E261 at least two spaces before inline comment
Line 14:77: E262 inline comment should start with '# '
Line 14:81: E501 line too long (96 > 80 characters)
Line 15:10: E225 missing whitespace around operator
Line 15:43: E231 missing whitespace after ','
Line 15:76: E261 at least two spaces before inline comment
Line 15:77: E262 inline comment should start with '# '
Line 15:81: E501 line too long (87 > 80 characters)
Line 17:33: E231 missing whitespace after ','
Line 18:49: E231 missing whitespace after ','
Line 19:49: E231 missing whitespace after ','
Line 20:1: W293 blank line contains whitespace
Line 24:33: W291 trailing whitespace
Line 27:1: W293 blank line contains whitespace
Line 28:1: E302 expected 2 blank lines, found 1
Line 34:10: E231 missing whitespace after ','
Line 34:38: E231 missing whitespace after ','
Line 34:45: E231 missing whitespace after ','
Line 35:9: E128 continuation line under-indented for visual indent
Line 35:20: E231 missing whitespace after ','
Line 45:36: E225 missing whitespace around operator
Line 48:5: E265 block comment should start with '# '
Line 49:16: E261 at least two spaces before inline comment
Line 49:17: E262 inline comment should start with '# '
Line 50:22: E261 at least two spaces before inline comment
Line 52:52: W291 trailing whitespace
Line 53:21: E261 at least two spaces before inline comment
Line 53:22: E262 inline comment should start with '# '
Line 54:51: E231 missing whitespace after ','
Line 63:16: E231 missing whitespace after ','
Line 63:20: E231 missing whitespace after ','
Line 63:26: E231 missing whitespace after ','
Line 67:1: E303 too many blank lines (3)
Line 73:5: E265 block comment should start with '# '
Line 74:9: E231 missing whitespace after ','
Line 74:13: E231 missing whitespace after ','
Line 74:19: E231 missing whitespace after ','
Line 74:49: E231 missing whitespace after ','
Line 74:56: E231 missing whitespace after ','
Line 79:22: E251 unexpected spaces around keyword / parameter equals
Line 79:24: E251 unexpected spaces around keyword / parameter equals
Line 80:21: E251 unexpected spaces around keyword / parameter equals
Line 80:23: E251 unexpected spaces around keyword / parameter equals
Line 84:23: E251 unexpected spaces around keyword / parameter equals
Line 84:25: E251 unexpected spaces around keyword / parameter equals
Line 90:14: W291 trailing whitespace

Comment last updated at 2021-08-12 21:39:20 UTC

[carterbox]

@YudongYao, I'm a little confused. Since the refractive indices of an object depend on the energy, shouldn't there also be a new dimension for the object (psi)? Otherwise, this implementation is the same as increasing the number of modes for a single energy.

[YudongYao]

@YudongYao, I'm a little confused. Since the refractive indices of an object depend on the energy, shouldn't there also be a new dimension for the object (psi)? Otherwise, this implementation is the same as increasing the number of modes for a single energy.

@carterbox As the X-ray energy is away from the elementary absorption edge of the sample in our case, so the object function is assumed to be constant with broadband illumination.
The specral probes are different from the existing probe modes in following ways:

1. The Fresnel zone plate used to focus the beam is chromatic optics that produces variant focal length for different wavelength, therefore, the probe profile and size on the object plane are wavelength-dependent. In the code, these spectral probe modes are generated with the numerical propagation from the zone plate.
2. The diffraction propagation from the sample plane to the detector plane is also wavelength-dependent, which means there is a scaling factor among different wavelengths. For this part, I haven't add it to the current code, and will do it later.

[carterbox]

To clarify, the main reason this PR was not merged is that the proposed changes to existing code are not necessary to support the multi-wavelength feature (which can be accomplished only by adding additional probe modes of a specific structure since the special scaling factor for different wavelengths has not been included).