mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
First pass at models, including the model for classic Ptychography
This commit is contained in:
Abe Levitan
@@ -7,3 +7,4 @@
|
||||
docs/build
|
||||
build/*
|
||||
dist
|
||||
example_data/*
|
||||
@@ -2,3 +2,4 @@ from __future__ import division, print_function, absolute_import
|
||||
|
||||
from CDTools import tools
|
||||
from CDTools import datasets
|
||||
from CDTools import models
|
||||
|
||||
+2
-80
@@ -3,86 +3,11 @@ import numpy as np
|
||||
import torch as t
|
||||
from copy import copy
|
||||
|
||||
# The naming overlap here is definitely going to get confusing.
|
||||
from CDTools.tools import data as cdtdata
|
||||
from torch.utils import data as torchdata
|
||||
|
||||
__all__ = ['CDataset', 'Ptycho_2D_Dataset']
|
||||
|
||||
#
|
||||
# I think that here should live a variety of datasets that all
|
||||
# subclass the basic dataset class. They should be able to contain
|
||||
# all the information that would be available in a CDI-type experiment.
|
||||
#
|
||||
# It's important for each kind of dataset to have tools to:
|
||||
#
|
||||
# * Pass too and from the GPU (where to store the data)
|
||||
# * Be able to load itself intelligently from a cxi file
|
||||
# * Be able to save itself intelligently to a cxi file
|
||||
#
|
||||
# And in addition, it should be relatively simple to initialize from
|
||||
# data held as python structures
|
||||
#
|
||||
# The datasets need to be able to work just like pytorch datasets,
|
||||
# in fact, they should subclass them, where they return the relevant
|
||||
# information for each diffraction pattern when sliced.
|
||||
#
|
||||
|
||||
|
||||
#
|
||||
# One question, though, is how to deal with information that models need
|
||||
# to know as well as the datasets? The basic issue is that it would be
|
||||
# nice to be able to write down a forward model and have it populate
|
||||
# a dataset object with the simulated data. But you also want to be able
|
||||
# to write a forward model that pulls the relevant information from a
|
||||
# dataset that has been read from a real data's cxi file.
|
||||
#
|
||||
# One workaround would be to just have the things defined twice, but to
|
||||
# make sure that every model can load/initialize itself from a dataset.
|
||||
# If it also knows how to define itself from an initialization function
|
||||
# or a python dataset, then it can easily also have a generic function to
|
||||
# simulate the action of the model and create a dataset from that simulation.
|
||||
#
|
||||
# This is unrelated, but it will then be important to be able to save and load
|
||||
# models easily from a predefined format. To be honest, this could just
|
||||
# literally be a snippet of python code that redefines the specific model
|
||||
# using only the base CDTools packages. That way a saved model could be a
|
||||
# .py file that just creates the model on the spot. Not sure I love this
|
||||
# though
|
||||
#
|
||||
# Is the trio of data, model, and reconstruction the correct thing though?
|
||||
# Perhaps the reconstruction doesn't need to exist. The model can have
|
||||
# a few more things involved in it, and the reconstructions could be a
|
||||
# single function, for example, that only uses what's in the model.
|
||||
#
|
||||
# The issue with the automatic differentiation stuff is that it's history
|
||||
# dependent in a sense. So you generate an optimization object and you take
|
||||
# a step with it.
|
||||
#
|
||||
# Can all reconstruction algorithms be formulated in terms of an explicit
|
||||
# forward step, backward step, and gradient with respect to the forward
|
||||
# step? No, for example the position annealing step needs to do it's own
|
||||
# thing. Okay, well then it could be possible to write an automatic
|
||||
# differentiation sequence that can generate n steps of a given automatic
|
||||
# differentiation solved. So given a model, you could just generically
|
||||
# create n steps of automatic differentiation. But then also write an
|
||||
# explicit ePIE step, etc. It could be interesting if the reconstruction
|
||||
# algorithms defined generators which would run one step and then yield
|
||||
# the loss at the end of that step. This would let you easily save out
|
||||
# the loss but also would make it simple to include live plotting
|
||||
#
|
||||
#
|
||||
# Then instead of a plan, you would just write a python script to follow
|
||||
# the steps you wanted to follow.
|
||||
#
|
||||
#
|
||||
|
||||
|
||||
# I think there should be a base CDataSet that has entry_info, sample_info, detector_info, and mask attributes. It can know how to load this data from a cxi file and write it out to a cxi file. It also can know how to pass this info to and from the GPU (not all of it needs to be passed in that way).
|
||||
|
||||
# Then, I will write a few common datasets as a demonstration. First is
|
||||
# A 2D CDI dataset, second is a 2D Ptycho dataset.
|
||||
|
||||
|
||||
#
|
||||
# This loads and stores all the kinds of metadata that are common to
|
||||
@@ -106,11 +31,8 @@ class CDataset(torchdata.Dataset):
|
||||
self.mask = t.tensor(mask)
|
||||
else:
|
||||
self.mask = None
|
||||
|
||||
if t.cuda.is_available():
|
||||
self.get_as(device='cuda:0')
|
||||
else:
|
||||
self.get_as(device='cpu')
|
||||
|
||||
self.get_as(device='cpu')
|
||||
|
||||
|
||||
def to(self,*args,**kwargs):
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import torch as t
|
||||
from torch.utils import data as torchdata
|
||||
|
||||
|
||||
#
|
||||
# This is unrelated, but it will then be important to be able to save and load
|
||||
# models easily from a predefined format. To be honest, this could just
|
||||
# literally be by pickling the model. They could also be saved out as
|
||||
# state_dicts or via torch.save. I think it's best to just save the whole
|
||||
# model - I lose out on the modularity of just saving the state_dict, but
|
||||
# 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
|
||||
#
|
||||
|
||||
|
||||
#
|
||||
# For now, just save/load model via the built-in t.save() and t.load()
|
||||
# functions
|
||||
#
|
||||
|
||||
|
||||
|
||||
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
|
||||
backward_propagator : the propagator from the detector plane to the experiment plane
|
||||
measurement : a function to simulate the detector readout from a detector plane wavefront
|
||||
forward : predefined, the entire stacked forward model
|
||||
loss : the loss function to report and use for automatic differentiation
|
||||
simulation : predefined, simulates a stack of detector images from the forward model
|
||||
simulate_to_dataset : a function to create a CDataset from the simulation defined in the model
|
||||
|
||||
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)))
|
||||
|
||||
def loss(self, sim_data, real_data):
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
# I know this is silly but it makes it clear this should be explicitly
|
||||
# overwritten
|
||||
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):
|
||||
loss = 0
|
||||
N = 0
|
||||
for inputs, patterns in data_loader:
|
||||
N += patterns.shape[0]
|
||||
|
||||
def closure():
|
||||
optimizer.zero_grad()
|
||||
sim_patterns = self.forward(*inputs)
|
||||
if hasattr(self, 'mask'):
|
||||
loss = self.loss(patterns,sim_patterns, mask=self.mask)
|
||||
else:
|
||||
loss = self.loss(patterns,sim_patterns)
|
||||
|
||||
loss.backward()
|
||||
return loss
|
||||
|
||||
loss += optimizer.step(closure).detach().cpu().numpy()
|
||||
|
||||
loss /= N
|
||||
if scheduler is not None:
|
||||
scheduler.step(loss)
|
||||
|
||||
yield loss
|
||||
|
||||
|
||||
def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005):
|
||||
|
||||
# Make a dataloader
|
||||
data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
|
||||
shuffle=True)
|
||||
|
||||
|
||||
# Define the optimizer
|
||||
optimizer = t.optim.Adam(self.parameters(), lr = lr)
|
||||
|
||||
return self.AD_optimize(iterations, data_loader, optimizer)
|
||||
|
||||
|
||||
|
||||
from CDTools.models.simple_ptycho import SimplePtycho
|
||||
@@ -0,0 +1,122 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import torch as t
|
||||
from CDTools.models import CDIModel
|
||||
from CDTools import tools
|
||||
from copy import copy
|
||||
|
||||
|
||||
class SimplePtycho(CDIModel):
|
||||
|
||||
def __init__(self, wavelength, detector_geometry,
|
||||
probe_basis, detector_slice,
|
||||
probe_guess, obj_guess, min_translation = t.Tensor([0,0]),
|
||||
background = None):
|
||||
|
||||
super(SimplePtycho,self).__init__()
|
||||
self.wavelength = t.Tensor([wavelength])
|
||||
self.detector_geometry = copy(detector_geometry)
|
||||
det_geo = self.detector_geometry
|
||||
if hasattr(det_geo, 'distance'):
|
||||
det_geo['distance'] = t.Tensor(det_geo['distance'])
|
||||
if hasattr(det_geo, 'basis'):
|
||||
det_geo['basis'] = t.Tensor(det_geo['basis'])
|
||||
if hasattr(det_geo, 'corner'):
|
||||
det_geo['corner'] = t.Tensor(det_geo['corner'])
|
||||
|
||||
self.min_translation = t.Tensor(min_translation)
|
||||
|
||||
self.probe_basis = t.Tensor(probe_basis)
|
||||
self.detector_slice = detector_slice
|
||||
|
||||
# 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))
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset):
|
||||
wavelength = dataset.wavelength
|
||||
det_basis = dataset.detector_geometry['basis']
|
||||
det_shape = dataset[0][1].shape
|
||||
distance = dataset.detector_geometry['distance']
|
||||
|
||||
# always do this on the cpu
|
||||
get_as_args = dataset.get_as_args
|
||||
dataset.get_as(device='cpu')
|
||||
(indices, translations), patterns = dataset[:]
|
||||
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,
|
||||
det_shape,
|
||||
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)
|
||||
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations)
|
||||
|
||||
# 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
|
||||
|
||||
return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation)
|
||||
|
||||
|
||||
def interaction(self, index, translations):
|
||||
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
|
||||
translations)
|
||||
pix_trans -= self.min_translation
|
||||
return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe,
|
||||
self.obj,
|
||||
pix_trans)
|
||||
|
||||
def forward_propagator(self, wavefields):
|
||||
return tools.propagators.far_field(wavefields)
|
||||
|
||||
|
||||
def backward_propagator(self, wavefields):
|
||||
return tools.propagators.inverse_far_field(wavefields)
|
||||
|
||||
|
||||
def measurement(self, wavefields):
|
||||
return tools.measurements.intensity(wavefields,
|
||||
detector_slice=self.detector_slice)
|
||||
|
||||
|
||||
def loss(self, sim_data, real_data):
|
||||
return tools.losses.amplitude_mse(real_data, sim_data)
|
||||
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
super(SimplePtycho, self).to(*args, **kwargs)
|
||||
self.wavelength = self.wavelength.to(*args,**kwargs)
|
||||
# move the detector geometry too
|
||||
det_geo = self.detector_geometry
|
||||
if hasattr(det_geo, 'distance'):
|
||||
det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
|
||||
if hasattr(det_geo, 'basis'):
|
||||
det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
|
||||
if hasattr(det_geo, 'corner'):
|
||||
det_geo['corner'] = det_geo['corner'].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
|
||||
|
||||
|
||||
@@ -7,4 +7,6 @@ from CDTools.tools import image_processing
|
||||
from CDTools.tools import initializers
|
||||
from CDTools.tools import losses
|
||||
from CDTools.tools import projectors
|
||||
from CDTools.tools import interactions
|
||||
from CDTools.tools import propagators
|
||||
from CDTools.tools import measurements
|
||||
|
||||
+15
-3
@@ -123,8 +123,16 @@ def get_sample_info(cxi_file):
|
||||
|
||||
s1 = cxi_file['entry_1/sample_1']
|
||||
metadata_attrs = ['name','description','unit_cell_group']
|
||||
metadata = {attr: str(s1[attr][()].decode()) for attr in metadata_attrs
|
||||
if attr in s1}
|
||||
|
||||
metadata = {}
|
||||
for attr in metadata_attrs:
|
||||
# Somehow different ways of saving can lead to different ways to
|
||||
# decode it here, so we try both
|
||||
if attr in s1:
|
||||
try:
|
||||
metadata[attr] = str(s1[attr][()].decode())
|
||||
except AttributeError as e:
|
||||
metadata[attr] = str(np.array(s1[attr][:])[0].decode())
|
||||
|
||||
float_attrs = ['concentration',
|
||||
'mass',
|
||||
@@ -276,7 +284,7 @@ def get_mask(cxi_file):
|
||||
return None
|
||||
|
||||
|
||||
def get_data(cxi_file):
|
||||
def get_data(cxi_file, cut_zeroes = True):
|
||||
"""Returns an array with the full stack of detector data defined in the cxi file object
|
||||
|
||||
This function will make sure to check all the various places that it's
|
||||
@@ -306,6 +314,10 @@ def get_data(cxi_file):
|
||||
else:
|
||||
raise KeyError('Data is not defined within cxi file')
|
||||
data = np.array(cxi_file[pull_from]).astype(np.float32)
|
||||
|
||||
if cut_zeroes:
|
||||
data[data < 0] = 0
|
||||
|
||||
if 'axes' in cxi_file[pull_from].attrs:
|
||||
axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':')
|
||||
axes = [axis.strip().lower() for axis in axes]
|
||||
|
||||
@@ -35,17 +35,25 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
|
||||
torch.Tensor : The exit wave's shape
|
||||
tuple(slice) : The slice corresponding to the physical detector
|
||||
"""
|
||||
det_shape = t.Tensor(tuple(det_shape))
|
||||
|
||||
det_shape = t.Tensor(tuple(det_shape)).to(t.int32)
|
||||
det_basis = t.Tensor(det_basis)
|
||||
# First, set the center if it's not already specified
|
||||
# This definition matches the center pixel of an fftshifted array
|
||||
if center is None:
|
||||
center = det_shape // 2
|
||||
|
||||
else:
|
||||
center = t.Tensor(center).to(t.int32)
|
||||
|
||||
# Then, calculate the required detector size from the centering
|
||||
# This is a bit opaque but was worth doing accurately
|
||||
min_left = center * 2
|
||||
min_right = (det_shape - center) * 2 - 1
|
||||
full_shape = t.max(min_left,min_right).to(t.int32) + 2 * padding
|
||||
|
||||
# In some edge cases this shape can be smaller than the detector shape
|
||||
full_shape = t.max(full_shape, det_shape)
|
||||
|
||||
if opt_for_fft:
|
||||
full_shape = t.Tensor([next_fast_len(dim) for dim in full_shape]).to(t.int32)
|
||||
# Then, generate a slice that pops the actual detector from the full
|
||||
@@ -64,7 +72,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
|
||||
basis_dirs = det_basis / t.norm(det_basis, dim=0)
|
||||
real_space_basis = basis_dirs * wavelength * distance / \
|
||||
(full_shape.to(t.float32) * t.norm(det_basis,dim=0))
|
||||
|
||||
|
||||
# Finally, convert the shape back to a torch.Size
|
||||
full_shape = t.Size([dim for dim in full_shape])
|
||||
|
||||
@@ -83,23 +91,27 @@ def calc_object_setup(probe_shape, translations, padding=0):
|
||||
correspond to (padding,padding)
|
||||
|
||||
Args:
|
||||
probe_shape (t.Size) : The size of the probe array
|
||||
translations (t.Tensor) : Jx2 stack of pixel-valued (i,j) translations
|
||||
probe_shape (torch.Size) : The size of the probe array
|
||||
translations (torch.Tensor) : Jx2 stack of pixel-valued (i,j) translations
|
||||
padding (int) : Optional, the size of an extra border to include
|
||||
Returns:
|
||||
torch.Size : required size of object array
|
||||
torch.Tensor : minimum pixel-valued translation
|
||||
"""
|
||||
|
||||
# First we look at the translations to find the minimum translation
|
||||
# and the range of translations
|
||||
min_translation = t.min(translations, dim=0)[0]
|
||||
translation_range = t.max(translations, dim=0)[0] - min_translation
|
||||
|
||||
|
||||
# Calculate the required shape
|
||||
translation_range = t.ceil(translation_range).numpy().astype(np.int32)
|
||||
shape = translation_range + np.array(probe_shape) + 2 * padding
|
||||
shape = t.Size(shape)
|
||||
|
||||
|
||||
# And the minimum translation
|
||||
min_translation = min_translation - padding
|
||||
|
||||
|
||||
return shape, min_translation
|
||||
|
||||
|
||||
@@ -213,19 +225,23 @@ def SHARP_style_probe(dataset, shape, det_slice):
|
||||
We make a small tweak to this procedure to lower the central pixel of
|
||||
the probe generated this way, which can often overwhelm the rest of the
|
||||
probe if there is significant noise on the detector
|
||||
|
||||
|
||||
|
||||
|
||||
Args:
|
||||
dataset (Ptycho_2D_Dataset) : The dataset to work from
|
||||
shape (torch.Size) : The size of the probe array to simulate
|
||||
det_slice (slice) : A slice or tuple of slices corresponding to the detector region in Fourier space
|
||||
"""
|
||||
|
||||
|
||||
intensities = np.zeros(shape)
|
||||
for params, im in dataset:
|
||||
intensities[det_slice] += im.cpu().numpy()
|
||||
intensities /= len(dataset)
|
||||
|
||||
|
||||
probe_fft = cmath.complex_to_torch(np.sqrt(intensities))
|
||||
|
||||
|
||||
probe_guess = cmath.torch_to_complex(inverse_far_field(probe_fft))
|
||||
|
||||
|
||||
# Now we remove the central pixel
|
||||
|
||||
center = np.array(probe_guess.shape) // 2
|
||||
@@ -235,7 +251,7 @@ def SHARP_style_probe(dataset, shape, det_slice):
|
||||
probe_guess[center[0]+1, center[1]],
|
||||
probe_guess[center[0], center[1]-1],
|
||||
probe_guess[center[0], center[1]+1]])
|
||||
|
||||
|
||||
return cmath.complex_to_torch(probe_guess)
|
||||
|
||||
|
||||
|
||||
@@ -3,6 +3,8 @@ from __future__ import division, print_function, absolute_import
|
||||
from CDTools.tools.cmath import *
|
||||
import torch as t
|
||||
|
||||
|
||||
|
||||
#
|
||||
# This file will host tools to turn various kinds of model information
|
||||
# (probe, 2D object, 3D object, etc) into exit waves leaving the sample
|
||||
@@ -31,11 +33,11 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0,0,1]))
|
||||
|
||||
projection_1 = t.Tensor([[1,0,0],
|
||||
[0,1,0],
|
||||
[0,0,0]])
|
||||
[0,0,0]]).to(device=basis.device,dtype=basis.dtype)
|
||||
projection_2 = t.inverse(t.Tensor([[1,0,0],
|
||||
[0,1,0],
|
||||
-surface_normal/
|
||||
surface_normal[2]]))
|
||||
surface_normal[2]])).to(device=basis.device,dtype=basis.dtype)
|
||||
basis_vectors_inv = t.pinverse(basis)
|
||||
projection = t.mm(basis_vectors_inv,
|
||||
t.mm(projection_2,projection_1))
|
||||
|
||||
@@ -86,5 +86,5 @@ def quadratic_background(wavefield, background, detector_slice=None, measurement
|
||||
return measurement(wavefield) + background**2
|
||||
else:
|
||||
return measurement(wavefield, detector_slice) \
|
||||
+ background[detector_slice]**2
|
||||
+ background**2
|
||||
|
||||
|
||||
@@ -33,7 +33,7 @@ def far_field(wavefront):
|
||||
torch.Tensor : The JxNxMx2 propagated wavefield
|
||||
"""
|
||||
|
||||
return fftshift(t.fft(wavefront, 2, normalized=True))
|
||||
return fftshift(t.fft(ifftshift(wavefront), 2, normalized=True))
|
||||
|
||||
|
||||
def inverse_far_field(wavefront):
|
||||
@@ -54,7 +54,7 @@ def inverse_far_field(wavefront):
|
||||
Returns:
|
||||
torch.Tensor : The JxNxMx2 exit wavefield
|
||||
"""
|
||||
return t.ifft(ifftshift(wavefront), 2, normalized=True)
|
||||
return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True))
|
||||
|
||||
|
||||
def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, **kwargs):
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from CDTools.tools import cmath
|
||||
import h5py
|
||||
import torch as t
|
||||
import numpy as np
|
||||
|
||||
with h5py.File('../example_data/NiCr.cxi','r') as f:
|
||||
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f)
|
||||
|
||||
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')
|
||||
dataset.to(device='cuda')
|
||||
dataset.get_as(device='cuda')
|
||||
|
||||
for loss in model.Adam_optimize(500, dataset):
|
||||
print(loss)
|
||||
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
probe = cmath.torch_to_complex(model.probe.detach().cpu())
|
||||
obj = cmath.torch_to_complex(model.obj.detach().cpu())
|
||||
|
||||
plt.imshow(np.abs(probe))
|
||||
plt.colorbar()
|
||||
plt.figure()
|
||||
plt.imshow(np.abs(obj))
|
||||
plt.colorbar()
|
||||
plt.show()
|
||||
|
||||
@@ -52,7 +52,7 @@ def test_quadratic_background():
|
||||
np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + background.numpy()**2
|
||||
det_slice = np.s_[3:,5:8]
|
||||
|
||||
result = measurements.quadratic_background(wavefields,background,
|
||||
result = measurements.quadratic_background(wavefields,background[det_slice],
|
||||
detector_slice=det_slice,
|
||||
measurement=measurements.intensity)
|
||||
assert t.allclose(result, t.tensor(np_result[(np.s_[:],)+det_slice]))
|
||||
|
||||
@@ -28,7 +28,7 @@ def exit_waves_1():
|
||||
|
||||
def test_far_field(exit_waves_1):
|
||||
# Far field diffraction patterns calculated by numpy with zero frequency in center
|
||||
np_result = np.fft.fftshift(np.fft.fft2(cmath.torch_to_complex(exit_waves_1),norm='ortho'))
|
||||
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cmath.torch_to_complex(exit_waves_1)),norm='ortho'))
|
||||
|
||||
assert(np.allclose(np_result, cmath.torch_to_complex(propagators.far_field(exit_waves_1))))
|
||||
|
||||
@@ -37,7 +37,7 @@ def test_far_field(exit_waves_1):
|
||||
def test_inverse_far_field(exit_waves_1):
|
||||
# We want the inverse far field to map back to the exit waves with no intensity corrections
|
||||
# Far field result for exit waves calculated with numpy
|
||||
far_field_np_result = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(cmath.torch_to_complex(exit_waves_1),norm='ortho')))
|
||||
far_field_np_result = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cmath.torch_to_complex(exit_waves_1)),norm='ortho')))
|
||||
|
||||
assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user