mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-19 09:02:09 +02:00
Consistently get Bragg/specular geometry implemented across both models/datasets and dataloaders
This commit is contained in:
Abe Levitan
@@ -90,7 +90,7 @@ class FancyPtycho(CDIModel):
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1):
|
||||
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None):
|
||||
|
||||
wavelength = dataset.wavelength
|
||||
det_basis = dataset.detector_geometry['basis']
|
||||
@@ -123,7 +123,19 @@ class FancyPtycho(CDIModel):
|
||||
surface_normal = dataset.sample_info['orientation'][2]
|
||||
else:
|
||||
surface_normal = np.array([0.,0.,1.])
|
||||
|
||||
|
||||
|
||||
# If this information is supplied when the function is called,
|
||||
# then we override the information in the .cxi file
|
||||
if scattering_mode in {'t', 'transmission'}:
|
||||
surface_normal = np.array([0.,0.,1.])
|
||||
elif scattering_mode in {'r', 'reflection'}:
|
||||
outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
|
||||
outgoing_dir /= np.linalg.norm(outgoing_dir)
|
||||
surface_normal = outgoing_dir + np.array([0.,0.,1.])
|
||||
surface_normal /= np.linalg.norm(outgoing_dir)
|
||||
|
||||
|
||||
# Next generate the object geometry from the probe geometry and
|
||||
# the translations
|
||||
pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
|
||||
@@ -202,9 +214,6 @@ class FancyPtycho(CDIModel):
|
||||
all_exit_waves = []
|
||||
for i in range(self.probe.shape[0]):
|
||||
pr = self.probe[i] * self.probe_support
|
||||
#exit_waves = self.probe_norm * tools.interactions.ptycho_2D_round(self.probe[i],
|
||||
# self.obj,
|
||||
# pix_trans)
|
||||
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr,
|
||||
self.obj_support * self.obj,
|
||||
pix_trans,
|
||||
@@ -263,8 +272,9 @@ class FancyPtycho(CDIModel):
|
||||
self.probe_norm = self.probe_norm.to(*args,**kwargs)
|
||||
self.probe_support = self.probe_support.to(*args,**kwargs)
|
||||
self.obj_support = self.obj_support.to(*args,**kwargs)
|
||||
|
||||
self.surface_normal = self.surface_normal.to(*args, **kwargs)
|
||||
|
||||
|
||||
def sim_to_dataset(self, args_list):
|
||||
# In the future, potentially add more control
|
||||
# over what metadata is saved (names, etc.)
|
||||
@@ -275,7 +285,16 @@ class FancyPtycho(CDIModel):
|
||||
'instrument_n': 'Simulated Data',
|
||||
'start_time': datetime.now()}
|
||||
|
||||
sample_info = {'description': 'A simulated sample'}
|
||||
surface_normal = self.surface_normal.detach().cpu().numpy()
|
||||
xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
|
||||
xsurfacevec /= np.linalg.norm(xsurfacevec)
|
||||
ysurfacevec = np.cross(surface_normal, xsurfacevec)
|
||||
ysurfacevec /= np.linalg.norm(ysurfacevec)
|
||||
orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
|
||||
|
||||
sample_info = {'description': 'A simulated sample',
|
||||
'orientation': orientation}
|
||||
|
||||
|
||||
detector_geometry = self.detector_geometry
|
||||
mask = self.mask
|
||||
|
||||
@@ -1,188 +0,0 @@
|
||||
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
|
||||
import numpy as np
|
||||
|
||||
|
||||
class IncoherentPtycho(CDIModel):
|
||||
|
||||
def __init__(self, wavelength, detector_geometry,
|
||||
probe_basis, detector_slice,
|
||||
probe_guess, obj_guess, min_translation = t.Tensor([0,0]),
|
||||
translation_offsets=None,
|
||||
background = None, mask=None, weights = None):
|
||||
|
||||
super(IncoherentPtycho,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
|
||||
|
||||
|
||||
if mask is None:
|
||||
self.mask = mask
|
||||
else:
|
||||
self.mask = t.ByteTensor(mask)
|
||||
|
||||
# We rescale the probe here so it learns at the same rate as the
|
||||
# object
|
||||
probe_norm = t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
|
||||
|
||||
self.probe = t.nn.Parameter(probe_guess.to(t.float32)/probe_norm)
|
||||
|
||||
self.probe_norm = float(probe_norm.numpy())
|
||||
|
||||
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
|
||||
|
||||
if background is None:
|
||||
background = 1e-6 * t.ones(self.probe[(np.s_[0],)+self.detector_slice].shape[:-1])
|
||||
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
|
||||
|
||||
if weights is None:
|
||||
self.weights = None
|
||||
else:
|
||||
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
|
||||
|
||||
if translation_offsets is None:
|
||||
self.translation_offsets = None
|
||||
else:
|
||||
self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32))
|
||||
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0):
|
||||
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])
|
||||
|
||||
# Set to none to avoid issues with things outside the detector
|
||||
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,
|
||||
padding=padding,
|
||||
opt_for_fft=False)
|
||||
|
||||
# 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, padding=20)
|
||||
|
||||
# Finally, initialize the probe and object using this information
|
||||
if probe_size is None:
|
||||
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice)
|
||||
else:
|
||||
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size)
|
||||
|
||||
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
|
||||
|
||||
# For incoherent probe mixing
|
||||
probe = t.stack((probe,0.05*t.rand(probe.shape).to(probe.dtype)))
|
||||
|
||||
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
|
||||
|
||||
det_geo = dataset.detector_geometry
|
||||
|
||||
weights = t.ones(len(dataset))
|
||||
|
||||
if hasattr(dataset, 'mask') and dataset.mask is not None:
|
||||
mask = dataset.mask.to(t.uint8)
|
||||
else:
|
||||
mask = None
|
||||
|
||||
return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, translation_offsets=translation_offsets, weights=weights, mask=mask)
|
||||
|
||||
|
||||
def interaction(self, index, translations):
|
||||
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
|
||||
translations)
|
||||
# The 10x term is to condition the translation offsets
|
||||
pix_trans -= self.min_translation
|
||||
pix_trans = pix_trans + self.translation_offsets[index]
|
||||
|
||||
all_exit_waves = []
|
||||
for i in range(self.probe.shape[0]):
|
||||
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(self.probe[i],
|
||||
self.obj,
|
||||
pix_trans,
|
||||
shift_probe=True)
|
||||
|
||||
if exit_waves.dim() == 4:
|
||||
exit_waves = self.weights[index][:,None,None,None] * exit_waves
|
||||
else:
|
||||
exit_waves = self.weights[index] * exit_waves
|
||||
|
||||
all_exit_waves.append(exit_waves)
|
||||
|
||||
return t.stack(all_exit_waves)
|
||||
|
||||
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.quadratic_background(wavefields,
|
||||
self.background,
|
||||
detector_slice=self.detector_slice,
|
||||
measurement=tools.measurements.incoherent_sum)
|
||||
|
||||
|
||||
def loss(self, sim_data, real_data, mask=None):
|
||||
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
|
||||
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
super(IncoherentPtycho, 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)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def sim_to_dataset(self, args_list):
|
||||
pass
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ class SimplePtycho(CDIModel):
|
||||
def __init__(self, wavelength, detector_geometry,
|
||||
probe_basis, detector_slice,
|
||||
probe_guess, obj_guess, min_translation = [0,0],
|
||||
mask=None):
|
||||
surface_normal=np.array([0.,0.,1.]), mask=None):
|
||||
|
||||
super(SimplePtycho,self).__init__()
|
||||
self.wavelength = t.Tensor([wavelength])
|
||||
@@ -33,6 +33,9 @@ class SimplePtycho(CDIModel):
|
||||
|
||||
self.probe_basis = t.Tensor(probe_basis)
|
||||
self.detector_slice = detector_slice
|
||||
|
||||
self.surface_normal = t.Tensor(surface_normal)
|
||||
|
||||
if mask is None:
|
||||
self.mask = None
|
||||
else:
|
||||
@@ -71,9 +74,16 @@ class SimplePtycho(CDIModel):
|
||||
distance,
|
||||
center=center)
|
||||
|
||||
if hasattr(dataset, 'sample_info') and \
|
||||
dataset.sample_info is not None and \
|
||||
'orientation' in dataset.sample_info:
|
||||
surface_normal = dataset.sample_info.orientation[2]
|
||||
else:
|
||||
surface_normal = np.array([0.,0.,1.])
|
||||
|
||||
# Next generate the object geometry from the probe geometry and
|
||||
# the translations
|
||||
pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations)
|
||||
pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
|
||||
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations)
|
||||
|
||||
# Finally, initialize the probe and object using this information
|
||||
@@ -82,17 +92,19 @@ class SimplePtycho(CDIModel):
|
||||
obj = t.ones(obj_size+(2,))
|
||||
det_geo = dataset.detector_geometry
|
||||
|
||||
|
||||
if hasattr(dataset, 'mask') and dataset.mask is not None:
|
||||
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)
|
||||
return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, mask=mask, surface_normal=surface_normal)
|
||||
|
||||
|
||||
def interaction(self, index, translations):
|
||||
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
|
||||
translations)
|
||||
translations,
|
||||
surface_normal=self.surface_normal)
|
||||
pix_trans -= self.min_translation
|
||||
return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe,
|
||||
self.obj,
|
||||
@@ -130,11 +142,11 @@ 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)
|
||||
|
||||
self.surface_normal = self.surface_normal.to(*args, **kwargs)
|
||||
|
||||
def sim_to_dataset(self, args_list):
|
||||
# In the future, potentially add more control
|
||||
@@ -146,7 +158,15 @@ class SimplePtycho(CDIModel):
|
||||
'instrument_n': 'Simulated Data',
|
||||
'start_time': datetime.now()}
|
||||
|
||||
sample_info = {'description': 'A simulated sample'}
|
||||
surface_normal = self.surface_normal.detach().cpu().numpy()
|
||||
xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
|
||||
xsurfacevec /= np.linalg.norm(xsurfacevec)
|
||||
ysurfacevec = np.cross(surface_normal, xsurfacevec)
|
||||
ysurfacevec /= np.linalg.norm(ysurfacevec)
|
||||
orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
|
||||
|
||||
sample_info = {'description': 'A simulated sample',
|
||||
'orientation': orientation}
|
||||
|
||||
detector_geometry = self.detector_geometry
|
||||
mask = self.mask
|
||||
|
||||
+10
-2
@@ -154,8 +154,16 @@ def get_sample_info(cxi_file):
|
||||
orient = np.array(s1['geometry_1/orientation']).astype(np.float32)
|
||||
xvec = orient[:3] / np.linalg.norm(orient[:3])
|
||||
yvec = orient[3:] / np.linalg.norm(orient[3:])
|
||||
metadata['orientation'] = np.array([xvec,yvec,
|
||||
np.cross(xvec,yvec)])
|
||||
metadata['orientation'] = np.array([xvec,yvec,
|
||||
np.cross(xvec,yvec)])
|
||||
|
||||
if 'geometry_1/surface_normal' in s1:
|
||||
snorm = np.array(s1['geometry_1/surface_normal']).astype(np.float32)
|
||||
xvec = np.cross(np.array([0.,1.,0.]), snorm)
|
||||
xvec /= np.linalg.norm(xvec)
|
||||
yvec = np.cross(snorm, xvec)
|
||||
yvec /= np.linalg.norm(yvec)
|
||||
metadata['orientation'] = np.array([xvec, yvec, snorm])
|
||||
|
||||
# Check if the metadata is empty
|
||||
if metadata == {}:
|
||||
|
||||
@@ -15,7 +15,7 @@ __all__ = ['translations_to_pixel', 'pixel_to_translations',
|
||||
#
|
||||
|
||||
|
||||
def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0,0,1])):
|
||||
def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])):
|
||||
"""Takes real space translations and outputs them in pixel space
|
||||
|
||||
This works for any 2D ptychography geometry. It takes in
|
||||
|
||||
@@ -22,11 +22,11 @@ with h5py.File(filename,'r') as f:
|
||||
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f)
|
||||
|
||||
|
||||
# A hack for the specular geometry
|
||||
dataset.detector_geometry['basis'] = np.array([[0,-30e-6],[30e-6,0],[0,0]])
|
||||
|
||||
# Figure out why the padding doesn't work here
|
||||
model = CDTools.models.FancyPtycho.from_dataset(dataset,randomize_ang = np.pi/4, padding=0, translation_scale=10)#, n_modes=2, propagation_distance=-5e-4)
|
||||
model = CDTools.models.FancyPtycho.from_dataset(dataset,
|
||||
randomize_ang = np.pi/4,
|
||||
padding=0,
|
||||
translation_scale=10,
|
||||
scattering_mode='reflection')
|
||||
|
||||
|
||||
# Uncomment these to use on the CPU
|
||||
|
||||
Reference in New Issue
Block a user