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

Merge branch 'master' of github.mit.edu:Scattering/CDTools

Browse Source
This commit is contained in:
Abe Levitan committed 2019-04-24 15:29:15 -04:00
2 parents d4576d1939 84ac2cc058
commit 6035841874
4 files changed
+103 -62
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

CDTools/models/__init__.py
+26 -25
View File
@@ -13,8 +13,8 @@ from torch.utils import data as torchdata
# I gain in it being easy to reload the non-learned aspects of the model,
# like the wavelength and sample geometry. Remember that it's important
# that the final outputs of the reconstructions are transferrable to other
# places
#
# places
#
#
@@ -26,14 +26,14 @@ from torch.utils import data as torchdata
class CDIModel(t.nn.Module):
"""This base model defines all the functions that must be exposed for a valid CDIModel subclass
Most of the functions only raise a NotImplementedError at this level and
must be explicitly defined by any subclass. The functions required can be
split into several subsections:
Creation:
from_dataset : a function to create a CDIModel from an appropriate CDataset
Simulation:
interaction : a function to simulate exit waves from experimental parameters
forward_propagator : the propagator from the experiment plane to the detector plane
@@ -47,33 +47,33 @@ class CDIModel(t.nn.Module):
Reconstruction:
AD_optimize : predefined, a generic automatic differentiation reconstruction
Adam_optimize : predefined, sensible automatic differentiation reconstruction using ADAM
The work of defining the various subclasses boils down to creating an
appropriate implementation for this set of functions.
"""
def from_dataset(self, dataset):
raise NotImplementedError()
def interaction(self, *args):
raise NotImplementedError()
def forward_propagator(self, exit_wave):
raise NotImplementedError()
def backward_propagator(self, detector_wave):
raise NotImplementedError()
def measurement(self, detector_wave):
raise NotImplementedError()
def forward(self, *args):
return self.measurement(self.forward_propagator(self.interaction(*args)))
@@ -86,15 +86,15 @@ class CDIModel(t.nn.Module):
def to(self, *args, **kwargs):
super(CDIModel,self).to(*args,**kwargs)
def simulate(self, args_list):
return t.Tensor([self.forward(*args) for args in args_list])
def simulate_to_dataset(self, args_list):
raise NotImplementedError()
def AD_optimize(self, iterations, data_loader, optimizer, scheduler=None):
for it in range(iterations):
@@ -102,7 +102,7 @@ class CDIModel(t.nn.Module):
N = 0
for inputs, patterns in data_loader:
N += 1
def closure():
optimizer.zero_grad()
sim_patterns = self.forward(*inputs)
@@ -113,7 +113,7 @@ class CDIModel(t.nn.Module):
loss.backward()
return loss
loss += optimizer.step(closure).detach().cpu().numpy()
loss /= N
@@ -121,7 +121,7 @@ class CDIModel(t.nn.Module):
scheduler.step(loss)
yield loss
def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False):
@@ -132,13 +132,13 @@ class CDIModel(t.nn.Module):
# Define the optimizer
optimizer = t.optim.Adam(self.parameters(), lr = lr)
# Define the scheduler
if schedule:
scheduler = t.optim.ReduceLROnPlateau(optimizer, factor=0.2)
else:
scheduler = None
return self.AD_optimize(iterations, data_loader, optimizer, scheduler=scheduler)
@@ -152,15 +152,16 @@ class CDIModel(t.nn.Module):
else:
data_loader = torchdata.DataLoader(dataset)
# Define the optimizer
optimizer = t.optim.LBFGS(self.parameters(),
lr = lr, history_size=history_size)
return self.AD_optimize(iterations, data_loader, optimizer)
from CDTools.models.simple_ptycho import SimplePtycho
from CDTools.models.fancy_ptycho import FancyPtycho
from CDTools.models.incoherent_ptycho import IncoherentPtycho
CDTools/models/simple_ptycho.py
+56 -14
View File
@@ -4,13 +4,14 @@ 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):
def __init__(self, wavelength, detector_geometry,
probe_basis, detector_slice,
probe_guess, obj_guess, min_translation = t.Tensor([0,0]),
probe_guess, obj_guess, min_translation = [0,0],
mask=None):
super(SimplePtycho,self).__init__()
@@ -36,7 +37,7 @@ class SimplePtycho(CDIModel):
# We rescale the probe here so it learns at the same rate as the
# object
self.probe_norm = t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
@@ -56,7 +57,7 @@ class SimplePtycho(CDIModel):
dataset.get_as(*get_as_args[0],**get_as_args[1])
center = tools.image_processing.centroid(t.sum(patterns,dim=0))
# Then, generate the probe geometry from the dataset
ewg = tools.initializers.exit_wave_geometry
probe_basis, probe_shape, det_slice = ewg(det_basis,
@@ -64,7 +65,7 @@ class SimplePtycho(CDIModel):
wavelength,
distance,
center=center)
# Next generate the object geometry from the probe geometry and
# the translations
pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations)
@@ -72,7 +73,7 @@ class SimplePtycho(CDIModel):
# Finally, initialize the probe and object using this information
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice)
obj = t.ones(obj_size+(2,))
det_geo = dataset.detector_geometry
@@ -80,10 +81,10 @@ class SimplePtycho(CDIModel):
mask = dataset.mask.to(t.uint8)
else:
mask = None
return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, mask=mask)
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations)
@@ -104,12 +105,12 @@ class SimplePtycho(CDIModel):
def measurement(self, wavefields):
return tools.measurements.intensity(wavefields,
detector_slice=self.detector_slice)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data,mask=mask)
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
def to(self, *args, **kwargs):
super(SimplePtycho, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args,**kwargs)
@@ -124,13 +125,54 @@ class SimplePtycho(CDIModel):
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.probe_norm = self.probe_norm.to(*args,**kwargs)
def sim_to_dataset(self, args_list):
pass
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:
:arg ``iterations``: Controls the number of iterations run, defaults to 1.
:arg ``beta``: Algorithmic parameter described in Maiden's implementation of rPIE. Defaults to 0.15.
:arg ``probe``: Initial probe wavefunction.
: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.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[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):
loss = []
for (i, [translations]), [patterns] in data_loader:
probe = self.probe.clone()
object = self.obj.clone()
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()
CDTools/tools/image_processing.py
+12 -12
View File
@@ -23,7 +23,8 @@ def centroid(im, dims=2):
Returns:
t.Tensor : An (i,j) index or stack of indices
"""
indices = (t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims))
# For some reason this needs to be a list
indices = [t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims)]
indices = t.meshgrid(*indices)
use_dims = [-dims+i for i in range(dims)]
@@ -61,7 +62,7 @@ def centroid_sq(im, dims=2, comp=False):
def sinc_subpixel_shift(im, shift):
"""Performs a subpixel shift with sinc interpolation on the given tensor
The subpixel shift is done circularly via a multiplication with a linear
phase mask in Fourier space.
@@ -120,15 +121,15 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)
# Now, I need to shift the array to pull out a contiguous window
# around the correlation maximum
try:
search_around = search_around.cpu()
except:
search_around = t.tensor(search_around)
window_size = 15
shift_zero = tuple(-search_around + t.tensor([window_size,window_size]))
cor_window = t.roll(cor, shift_zero, dims=(0,1))[:2*window_size,:2*window_size]
@@ -148,11 +149,11 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
cormax = t.tensor([t.argmax(upsampled) // sh[1],
t.argmax(upsampled) % sh[1]]).to(device=upsampled.device)
subpixel_shift = ((cormax + sh // 2) % sh - sh//2).to(dtype=upsampled.dtype)
return search_around.to(device=upsampled.device, dtype=upsampled.dtype) + \
subpixel_shift / resolution
subpixel_shift / resolution
def find_pixel_shift(im1, im2):
"""Calculates the integer pixel shift between two images by maximizing the autocorrelation
@@ -179,8 +180,8 @@ def find_pixel_shift(im1, im2):
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
cor = cmath.cabs(t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2))
sh = t.tensor(cor.shape).to(device=im1.device)
cormax = t.tensor([t.argmax(cor) // sh[1],
t.argmax(cor) % sh[1]]).to(device=im1.device)
@@ -205,4 +206,3 @@ def find_shift(im1, im2, resolution=10):
resolution=resolution)
return subpixel_shift
examples/simple_ptycho.py
+9 -11
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,17 +23,14 @@ 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')
#model.probe.requires_grad = False
for loss in model.Adam_optimize(10, dataset):
for loss in model.ePIE(1, dataset):
print(loss)
#for loss in model.Adam_optimize(20, dataset, lr=0.0005):
# print(loss)
from matplotlib import pyplot as plt
plot_amplitude(model.probe)
@@ -39,4 +38,3 @@ plot_phase(model.probe)
plot_amplitude(model.obj)
plot_phase(model.obj)
plt.show()