Skip to content

Commit

Permalink
added semi-functioning switching on of slices at arbitrary iterations
Browse files Browse the repository at this point in the history
  • Loading branch information
kahntm committed Nov 5, 2023
1 parent 55a7b26 commit 7a65a77
Showing 1 changed file with 50 additions and 23 deletions.
73 changes: 50 additions & 23 deletions ptypy/custom/ThreePIE.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,12 @@ class ThreePIE(stochastic.EPIE):
help = Thickness of a single slice in meters
doc =
[slice_start_iteration]
default = 0
type = int, list, tuple
help = iteration number to start using a specific slice
doc =
[fslices]
default = slices.h5
type = str
Expand Down Expand Up @@ -74,7 +80,13 @@ def engine_initialize(self):
# ToDo:
# - allow for non equal slice spacing
# - allow for start_slice_update at a freely chosen iteration
# for each slice separately
# for each slice separately - works, but not if the
# most downstream slice is switched off

if isinstance(self.p.slice_start_iteration, int):
self.p.slice_start_iteration = np.ones(self.p.number_of_slices) * self.p.slice_start_iteration
#if ĺen(self.p.slice_start_iteration) != self.p.number_of_slices:
# logger.info(f'dimension of given slice_start_iteration ({ĺen(self.p.slice_start_iteration)}) does not match number of slices ({self.p.number_of_slices})')

scan = list(self.ptycho.model.scans.values())[0]
geom = scan.geometries[0]
Expand Down Expand Up @@ -124,6 +136,7 @@ def engine_finalize(self):
slices_info.slice_thickness = self.p.slice_thickness
slices_info.objects = {ob.ID: {ID: S._to_dict() for ID, S in ob.storages.items()}
for ob in self._object}
slices_info.slice_start_iteration = self.p.slice_start_iteration

header = {'description': 'multi-slices result details.'}

Expand All @@ -142,46 +155,60 @@ def multislice_update(self, view):
"""

for i in range(self.p.number_of_slices-1):
# Forward multislice, calculate exit waves
for name, pod in view.pods.items():
# exit wave for this slice
self._exits[i][pod.pr_view] = self._probe[i][pod.pr_view] * self._object[i][pod.ob_view]
if self.curiter >= self.p.slice_start_iteration[i]:
self._exits[i][pod.pr_view] = self._probe[i][pod.pr_view] * self._object[i][pod.ob_view]
else:
self._exits[i][pod.pr_view] = self._probe[i][pod.pr_view] * 1.
# incident wave for next slice
self._probe[i+1][pod.pr_view] = self.fw(self._exits[i][pod.pr_view])

for name, pod in view.pods.items():
# Exit wave for last slice
self._exits[-1][pod.pr_view] = self._probe[-1][pod.pr_view] * self._object[-1][pod.ob_view]
if self.curiter >= self.p.slice_start_iteration[-1]:
self._exits[-1][pod.pr_view] = self._probe[-1][pod.pr_view] * self._object[-1][pod.ob_view]
else:
self._exits[-1][pod.pr_view] = self._probe[-1][pod.pr_view] * 1.
# Save final state into pod (need for ptypy fourier update)
pod.probe = self._probe[-1][pod.pr_view]
pod.object = self._object[-1][pod.ob_view]
pod.exit = self._exits[-1][pod.pr_view]

# Fourier update
error = self.fourier_update(view)

# Object/probe update for the last slice
self.object_update(view, {pod.ID:self._exits[-1][pod.pr_view] for name, pod in view.pods.items()})
self.probe_update(view, {pod.ID:self._exits[-1][pod.pr_view] for name, pod in view.pods.items()})
for name, pod in view.pods.items():
self._object[-1][pod.ob_view] = pod.object
self._probe[-1][pod.pr_view] = pod.probe
if self.curiter >= self.p.slice_start_iteration[-1]:
self.object_update(view, {pod.ID:self._exits[-1][pod.pr_view] for name, pod in view.pods.items()})
self.probe_update(view, {pod.ID:self._exits[-1][pod.pr_view] for name, pod in view.pods.items()})
for name, pod in view.pods.items():
self._object[-1][pod.ob_view] = pod.object
self._probe[-1][pod.pr_view] = pod.probe
else:
for name, pod in view.pods.items():
self._probe[-1][pod.pr_view] = pod.exit * 1.

# Object/probe update for other slices (backwards)
for i in range(self.p.number_of_slices-2, -1, -1):

for name, pod in view.pods.items():
# Backwards propagation of the probe
pod.exit = self.bw(self._probe[i+1][pod.pr_view])
# Save state into pods
pod.probe = self._probe[i][pod.pr_view]
pod.object = self._object[i][pod.ob_view]

# Actual object/probe update
self.object_update(view, {pod.ID:self._exits[i][pod.pr_view] for name, pod in view.pods.items()})
self.probe_update(view, {pod.ID:self._exits[i][pod.pr_view] for name, pod in view.pods.items()})
for name, pod in view.pods.items():
self._object[i][pod.ob_view] = pod.object
self._probe[i][pod.pr_view] = pod.probe
if self.curiter >= self.p.slice_start_iteration[i]:

for name, pod in view.pods.items():
# Backwards propagation of the probe
pod.exit = self.bw(self._probe[i+1][pod.pr_view])
# Save state into pods
pod.probe = self._probe[i][pod.pr_view]
pod.object = self._object[i][pod.ob_view]

# Actual object/probe update
self.object_update(view, {pod.ID:self._exits[i][pod.pr_view] for name, pod in view.pods.items()})
self.probe_update(view, {pod.ID:self._exits[i][pod.pr_view] for name, pod in view.pods.items()})
for name, pod in view.pods.items():
self._object[i][pod.ob_view] = pod.object
self._probe[i][pod.pr_view] = pod.probe
else:
for name, pod in view.pods.items():
self._probe[i][pod.pr_view] = self.bw(self._probe[i+1][pod.pr_view])

# set the object as the product of all slices for better live plotting
self.ob.fill(self._object[0])
Expand Down

1 comment on commit 7a65a77

@kahntm
Copy link
Contributor Author

@kahntm kahntm commented on 7a65a77 Nov 5, 2023

Choose a reason for hiding this comment

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

The switching on object slices at arbitrary individual iteration numbers does work, except for when the most downstream slice is not initialized from the start.

Example1: 3 slices. p.engines.engine00.slice_start_iteration = [10,20,0] slice 1 starts to be updated from iteration 10, slice 2 starts to be updated from iteration 20 and the most downstream slice 3 is used from the start. This works just fine.
Example2: 3 slices. p.engines.engine00.slice_start_iteration = [0,0,20] slice 1 and slice 2 are updated from the start and the most downstream slice 3 should update starting at iteration 20. This runs through, but the reconstruction fails.

Please sign in to comment.