mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
misc small changes
This commit is contained in:
Abe Levitan
@@ -62,7 +62,7 @@ class Ptycho2DDataset(CDataset):
|
||||
self.translations = t.tensor(translations)
|
||||
self.patterns = t.tensor(patterns)
|
||||
if self.mask is None:
|
||||
self.mask = t.ones(self.patterns.shape[-2:]).to(dtype=t.uint8)
|
||||
self.mask = t.ones(self.patterns.shape[-2:]).to(dtype=t.bool)
|
||||
self.mask.masked_fill_(t.isnan(t.sum(self.patterns,dim=(0,))),0)
|
||||
self.patterns.masked_fill_(t.isnan(self.patterns),0)
|
||||
|
||||
@@ -181,7 +181,7 @@ class Ptycho2DDataset(CDataset):
|
||||
cdtdata.add_ptycho_translations(cxi_file, self.translations)
|
||||
|
||||
|
||||
def inspect(self):
|
||||
def inspect(self, logarithmic=True):
|
||||
"""Launches an interactive plot for perusing the data
|
||||
|
||||
This launches an interactive plotting tool in matplotlib that
|
||||
@@ -249,9 +249,12 @@ class Ptycho2DDataset(CDataset):
|
||||
axes[0].set_ylabel('Translation y (um)')
|
||||
cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.15,fraction=0.1)
|
||||
cb1.ax.tick_params(labelrotation=20)
|
||||
|
||||
meas = axes[1].imshow(np.log(meas_data) / np.log(10) * mask)
|
||||
|
||||
if logarithmic:
|
||||
meas = axes[1].imshow(np.log(meas_data) / np.log(10) * mask)
|
||||
else:
|
||||
meas = axes[1].imshow(meas_data * mask)
|
||||
|
||||
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.15,fraction=0.1)
|
||||
cb2.ax.tick_params(labelrotation=20)
|
||||
cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
|
||||
@@ -277,7 +280,11 @@ class Ptycho2DDataset(CDataset):
|
||||
|
||||
|
||||
meas = axes[1].images[-1]
|
||||
meas.set_data(np.log(meas_data) / np.log(10) * mask)
|
||||
if logarithmic:
|
||||
meas.set_data(np.log(meas_data) / np.log(10) * mask)
|
||||
else:
|
||||
meas.set_data(meas_data * mask)
|
||||
|
||||
update_colorbar(meas)
|
||||
|
||||
|
||||
|
||||
@@ -131,13 +131,17 @@ class CDIModel(t.nn.Module):
|
||||
scheduler : torch.optim.lr_scheduler._LRScheduler
|
||||
Optional, a learning rate scheduler to use
|
||||
"""
|
||||
|
||||
# First, calculate the normalization
|
||||
normalization = 0
|
||||
for inputs, patterns in data_loader:
|
||||
normalization += t.sum(patterns).cpu().numpy()
|
||||
|
||||
|
||||
for it in range(iterations):
|
||||
loss = 0
|
||||
N = 0
|
||||
for inputs, patterns in data_loader:
|
||||
N += 1
|
||||
|
||||
def closure():
|
||||
optimizer.zero_grad()
|
||||
sim_patterns = self.forward(*inputs)
|
||||
@@ -151,14 +155,14 @@ class CDIModel(t.nn.Module):
|
||||
|
||||
loss += optimizer.step(closure).detach().cpu().numpy()
|
||||
|
||||
loss /= N
|
||||
loss /= normalization
|
||||
if scheduler is not None:
|
||||
scheduler.step(loss)
|
||||
|
||||
|
||||
yield loss
|
||||
|
||||
|
||||
def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False):
|
||||
def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005, schedule=False, amsgrad=False):
|
||||
"""Runs a round of reconstruction using the Adam optimizer
|
||||
|
||||
This is generally accepted to be the most robust algorithm for use
|
||||
@@ -184,7 +188,7 @@ class CDIModel(t.nn.Module):
|
||||
shuffle=True)
|
||||
|
||||
# Define the optimizer
|
||||
optimizer = t.optim.Adam(self.parameters(), lr = lr)
|
||||
optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad)
|
||||
|
||||
|
||||
# Define the scheduler
|
||||
@@ -234,6 +238,46 @@ class CDIModel(t.nn.Module):
|
||||
return self.AD_optimize(iterations, data_loader, optimizer)
|
||||
|
||||
|
||||
def SGD_optimize(self, iterations, dataset, batch_size=None,
|
||||
lr=0.01, momentum=0, dampening=0, weight_decay=0,
|
||||
nesterov=False):
|
||||
"""Runs a round of reconstruction using the SGDoptimizer
|
||||
|
||||
This algorithm is often less stable that Adam, but it is simpler
|
||||
and is the basic workhorse of gradience descent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
iterations : int
|
||||
How many epochs of the algorithm to run
|
||||
dataset : CDataset
|
||||
The dataset to reconstruct against
|
||||
batch_size : int
|
||||
Optional, the size of the minibatches to use
|
||||
lr : float
|
||||
Optional, the learning rate to use
|
||||
momentum : float
|
||||
Optional, the length of the history to use.
|
||||
"""
|
||||
|
||||
# Make a dataloader
|
||||
if batch_size is not None:
|
||||
data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
|
||||
shuffle=True)
|
||||
else:
|
||||
data_loader = torchdata.DataLoader(dataset)
|
||||
|
||||
|
||||
# Define the optimizer
|
||||
optimizer = t.optim.SGD(self.parameters(),
|
||||
lr = lr, momentum=momentum,
|
||||
dampening=dampening,
|
||||
weight_decay=weight_decay,
|
||||
nesterov=nesterov)
|
||||
|
||||
return self.AD_optimize(iterations, data_loader, optimizer)
|
||||
|
||||
|
||||
# By default, the plot_list is empty
|
||||
plot_list = []
|
||||
|
||||
|
||||
@@ -46,7 +46,7 @@ class FancyPtycho(CDIModel):
|
||||
if mask is None:
|
||||
self.mask = mask
|
||||
else:
|
||||
self.mask = t.ByteTensor(mask)
|
||||
self.mask = t.BoolTensor(mask)
|
||||
|
||||
# We rescale the probe here so it learns at the same rate as the
|
||||
# object
|
||||
@@ -175,7 +175,7 @@ class FancyPtycho(CDIModel):
|
||||
weights = t.ones(len(dataset))
|
||||
|
||||
if hasattr(dataset, 'mask') and dataset.mask is not None:
|
||||
mask = dataset.mask.to(t.uint8)
|
||||
mask = dataset.mask.to(t.bool)
|
||||
else:
|
||||
mask = None
|
||||
|
||||
|
||||
@@ -98,7 +98,7 @@ class SimplePtycho(CDIModel):
|
||||
|
||||
|
||||
if hasattr(dataset, 'mask') and dataset.mask is not None:
|
||||
mask = dataset.mask.to(t.uint8)
|
||||
mask = dataset.mask.to(t.bool)
|
||||
else:
|
||||
mask = None
|
||||
|
||||
|
||||
@@ -148,15 +148,13 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
|
||||
|
||||
|
||||
if correct_ramp:
|
||||
# Need to check if this is actually working and, if noy, why not
|
||||
center_freq = ip.centroid_sq(cmath.fftshift(t.fft(probe[0],2)),comp=True)
|
||||
# Need to check if this is actually working and, if not, why not
|
||||
center_freq = ip.centroid(cmath.cabssq(cmath.fftshift(t.fft(probe[0],2))))
|
||||
center_freq -= (t.tensor(probe[0].shape[:-1]) // 2).to(t.float32)
|
||||
center_freq /= t.tensor(probe[0].shape[:-1]).to(t.float32)
|
||||
|
||||
|
||||
|
||||
Is, Js = np.mgrid[:probe[0].shape[0],:probe[0].shape[1]]
|
||||
probe_phase_ramp = cmath.expi(2*np.pi *
|
||||
probe_phase_ramp = cmath.expi(2 * np.pi *
|
||||
(center_freq[0] * t.tensor(Is).to(t.float32) +
|
||||
center_freq[1] * t.tensor(Js).to(t.float32)))
|
||||
probe = cmath.cmult(probe, cmath.cconj(probe_phase_ramp))
|
||||
@@ -165,7 +163,6 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
|
||||
(center_freq[0] * t.tensor(Is).to(t.float32) +
|
||||
center_freq[1] * t.tensor(Js).to(t.float32)))
|
||||
obj = cmath.cmult(obj, obj_phase_ramp)
|
||||
|
||||
|
||||
# Then, we set them to consistent absolute phases
|
||||
|
||||
@@ -475,7 +472,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
|
||||
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
|
||||
|
||||
if nbins is None:
|
||||
nbins = np.max(synth_obj[im_slice].shape) // 4
|
||||
nbins = np.max(im1[im_slice].shape) // 4
|
||||
|
||||
|
||||
cor_fft = cmath.cmult(cmath.fftshift(t.fft(im1[im_slice],2)),
|
||||
@@ -490,7 +487,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
|
||||
|
||||
i_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[0],d=di))
|
||||
j_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[1],d=dj))
|
||||
|
||||
|
||||
Js,Is = np.meshgrid(j_freqs,i_freqs)
|
||||
Rs = np.sqrt(Is**2+Js**2)
|
||||
|
||||
|
||||
@@ -278,7 +278,7 @@ def get_mask(cxi_file):
|
||||
mask = np.array(i1['detector_1/mask']).astype(np.uint32)
|
||||
mask_on = np.equal(mask,np.uint32(0))
|
||||
mask_has_signal = np.equal(mask,np.uint32(0x00001000))
|
||||
return np.logical_or(mask_on,mask_has_signal).astype(np.uint8)
|
||||
return np.logical_or(mask_on,mask_has_signal).astype(np.bool)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
@@ -73,6 +73,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
|
||||
|
||||
# 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)
|
||||
|
||||
@@ -50,11 +50,11 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
|
||||
|
||||
if mask is None:
|
||||
return t.sum((t.sqrt(sim_intensities) -
|
||||
t.sqrt(intensities))**2) / intensities.view(-1).shape[0]
|
||||
t.sqrt(intensities))**2)
|
||||
else:
|
||||
masked_intensities = intensities.masked_select(mask)
|
||||
return t.sum((t.sqrt(sim_intensities.masked_select(mask)) -
|
||||
t.sqrt(masked_intensities))**2) / masked_intensities.shape[0]
|
||||
t.sqrt(masked_intensities))**2)
|
||||
|
||||
|
||||
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -2,6 +2,7 @@ from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from matplotlib import pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
|
||||
# This file is too large to be distributed via Github.
|
||||
|
||||
Reference in New Issue
Block a user