cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-09 21:12:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

added ePIE

Browse Source
This commit is contained in:
Maddie Cain
2019-04-24 12:26:32 -04:00
parent fcf17d9355
commit 84ac2cc058
2 changed files with 32 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/simple_ptycho.py
+24 -10
View File
@@ -4,6 +4,7 @@ import torch as t
from CDTools.models import CDIModel
from CDTools import tools
from copy import copy
from torch.utils import data as torchdata
class SimplePtycho(CDIModel):
@@ -134,7 +135,7 @@ class SimplePtycho(CDIModel):
pass
def ePIE(iterations, dataset, beta = 1.0):
def ePIE(self, iterations, dataset, beta = 1.0):
"""Runs an ePIE reconstruction as described in `Maiden et al. (2017) <https://www.osapublishing.org/optica/abstract.cfm?uri=optica-4-7-736>`_.
Optional parameters are:
@@ -144,21 +145,34 @@ class SimplePtycho(CDIModel):
:arg ``object``: Initial object wavefunction.
"""
probe_shape = self.probe.shape
if self.mask is not None:
mask = self.mask[...,None]
def probe_update(exit_wave, exit_wave_corrected, probe, object, translation):
return probe+tools.cmult(beta*(tools.cconj(object)/ \
t.max(tools.cabssq(object)))[translations[0]:translations[0]+probe_shape[0],translations[1]:translations[1]+probe_shape[1]], \
return probe+tools.cmath.cmult(beta*(tools.cmath.cconj(object)/t.max(tools.cmath.cabssq(object)))[translation[0]:translation[0]+probe_shape[0],translation[1]:translation[1]+probe_shape[1]], \
exit_wave_corrected-exit_wave)
def object_update(exit_wave, exit_wave_corrected, probe, object, translation):
object[translations[0]:translations[0]+probe_shape[0],translations[1]:translations[1]+probe_shape[1]]\
+=tools.cmult(tools.cconj(probe)/t.max(tools.cabssq(probe)),exit_wave_corrected-exit_wave)
object[translation[0]:translation[0]+probe_shape[0],translation[1]:translation[1]+probe_shape[1]]\
+=tools.cmath.cmult(tools.cmath.cconj(probe)/t.max(tools.cmath.cabssq(probe)),exit_wave_corrected-exit_wave)
return object
with t.no_grad():
data_loader = torchdata.DataLoader(dataset, shuffle=True)
for it in range(iterations):
for translations, patterns in data_loader:
loss = []
for (i, [translations]), [patterns] in data_loader:
probe = self.probe.clone()
object = self.obj.clone()
exit_wave = tools.cmult(probe, object[translations:translations[0]+probe_shape[0],translations[1]:translations[1]+probe_shape[1]]).clone()
exit_wave_corrected = projectors.modulus(self.forward_propagator(exit_wave), patterns, mask = self.mask)
self.probe = probe_update(exit_wave, exit_wave_corrected, probe, object, translations)
self.obj = object_update(exit_wave, exit_wave_corrected, probe, object, translations)
exit_wave = self.interaction(i, translations)
exit_wave_corrected = exit_wave.clone()
exit_wave_corrected[self.detector_slice] = tools.projectors.modulus(self.forward_propagator(exit_wave)[self.detector_slice], patterns, mask = mask)
integer_translations = t.round(translations).to(dtype=t.int32)
self.probe.data = probe_update(exit_wave, exit_wave_corrected, probe, object, integer_translations)
self.obj.data = object_update(exit_wave, exit_wave_corrected, probe, object, integer_translations)
loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns))
yield t.mean(t.Tensor(loss)).cpu().numpy()
examples/simple_ptycho.py
+8 -6
View File
@@ -7,9 +7,11 @@ import h5py
import torch as t
import numpy as np
filename = '../../../Projects/CSX_3_19/cxis/processed/114429_p.cxi'
filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi'
filename = '../../../Downloads/AuBalls_700ms_30nmStep_3_3SS_filter.cxi'
filename = '../../Downloads/114429_p.cxi'
#filename = '../../../Projects/CSX_3_19/cxis/processed/114429_p.cxi'
#filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi'
#filename = '../../../Downloads/AuBalls_700ms_30nmStep_3_3SS_filter.cxi'
with h5py.File(filename,'r') as f:
@@ -21,12 +23,12 @@ model = CDTools.models.SimplePtycho.from_dataset(dataset)
# Uncomment these to use on the CPU
# default is CPU with 32-bit floats
model.to(device='cuda')
#model.to(device='cuda')
#dataset.to(device='cuda')
dataset.get_as(device='cuda')
#dataset.get_as(device='cuda')
for loss in model.ePIE(100, dataset):
for loss in model.ePIE(1, dataset):
print(loss)
from matplotlib import pyplot as plt