Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add object regularization for multi-slice reconstruction #525

Draft
wants to merge 15 commits into
base: dev
Choose a base branch
from

Conversation

yiranlus
Copy link
Contributor

I add object regularization which is mentioned in

Zhen Chen et al. ,Electron ptychography achieves atomic-resolution limits set by lattice vibrations.Science372,826-831(2021).DOI:10.1126/science.abg2533

I tested with the two random images and have the following results on phase (regularization = 0.05),

image

and before it was like this

image

For now, it only supports single mode and slice distances should be uniform.

However, the current code is very slow. I hope someone can give some hints how to make it faster.

@yiranlus yiranlus requested a review from daurer January 30, 2024 09:49
Copy link
Contributor

@daurer daurer left a comment

Choose a reason for hiding this comment

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

@yiranlus thanks for your contribution. It's seems to make to add such an optional regularizer to the 3PIE engine. Here are a some suggestions to improve the code

  1. Avoid hard-coding storage IDs like "Sscan_00G00" as they can be changed by the user script. Instead, you can get things like shape and psize from the first object view:
shape = list(self.ob.views.values())[0].shape
psize = list(self.ob.views.values())[0].psize
  1. Instead of np.fft.fftn and np.fft.ifftn you can use the internal propagater which you can get from the first pod
fft  = list(P.pods.values())[0].fw
ifft = list(P.pods.values())[0].bw

and you should probably iterate over the pods when you apply the regularizer update (similar to what is done in the multislice update) - this way your code should also work for multiple probe modes.}

  1. See other in-line comments below

Comment on lines +236 to +237
self.apply_object_regularization()
Copy link
Contributor

Choose a reason for hiding this comment

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

It would probably make sense to move this a bit higher up to make sure that self.ob is calculated for plotting after the regulariser has been applied...

Comment on lines +247 to +254
psize = self._object[0].S["Sscan_00G00"].psize[0]
kz = np.fft.fftfreq(self.p.number_of_slices, self.p.slice_thickness)[..., np.newaxis, np.newaxis]
ky = np.fft.fftfreq(shape[0], psize)[..., np.newaxis]
kx = np.fft.fftfreq(shape[1], psize)

# calculate the weight array
w = 1 - 2*np.arctan2(self.p.object_regularization_rate**2 * kz**2, kx**2+ky**2+np.spacing(1))/np.pi
Copy link
Contributor

Choose a reason for hiding this comment

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

this part is basically just calculating some weights w and does not depend on any current update so can be moved into a separate function, e.g. initialize_regularizer and called once in the constructor. In this current implementation the weights are re-calculated for every iteration which seems unnecessary.

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.

None yet

3 participants