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

Fixed merging conflict - log scale in ptycho_2d_dataset.py

Browse Source
This commit is contained in:
David Rower
2019-10-03 15:25:27 -04:00
parent 61f8bf733e 69be77cc09
commit 8197eed50e
19 changed files with 144 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/datasets/__init__.py
+9 -3
View File
@@ -32,8 +32,11 @@ import numpy as np
import torch as t
from copy import copy
import h5py
import pathlib
try:
import pathlib
except ImportError:
import pathlib2 as pathlib
from CDTools.tools import data as cdtdata
from CDTools.tools import plotting
from torch.utils import data as torchdata
@@ -98,7 +101,10 @@ class CDataset(torchdata.Dataset):
self.wavelength = wavelength
self.detector_geometry = copy(detector_geometry)
if mask is not None:
self.mask = t.tensor(mask)
if isinstance(mask, t.Tensor):
self.mask = mask.detach().to(dtype=t.bool)
else:
self.mask = t.BoolTensor(mask)
else:
self.mask = None
if background is not None:
CDTools/datasets/ptycho_2d_dataset.py
+15 -5
View File
@@ -3,7 +3,10 @@ import numpy as np
import torch as t
from copy import copy
import h5py
import pathlib
try:
import pathlib
except ImportError:
import pathlib2 as pathlib
from CDTools.datasets import CDataset
from CDTools.tools import data as cdtdata
@@ -62,7 +65,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 +184,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
@@ -252,7 +255,10 @@ class Ptycho2DDataset(CDataset):
cb1.ax.set_title('Integrated Intensity', size="medium", pad=5)
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.17,fraction=0.1)
cb2.ax.tick_params(labelrotation=20)
@@ -280,7 +286,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)
CDTools/models/__init__.py
+50 -6
View File
@@ -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 = []
CDTools/models/fancy_ptycho.py
+2 -2
View File
@@ -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
CDTools/models/simple_ptycho.py
+1 -1
View File
@@ -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
CDTools/tools/analysis.py
+5 -8
View File
@@ -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)
CDTools/tools/data.py
+1 -1
View File
@@ -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
CDTools/tools/initializers.py
+1
View File
@@ -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)
CDTools/tools/losses.py
+7 -3
View File
@@ -21,7 +21,11 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
This function calculates the mean squared error between their
associated amplitudes. Because this is not well defined for negative
numbers, make sure that all the intensities are >0 before using this
loss.
loss. Note that this is actually a sum-squared error, because this
formulation makes it vastly simpler to compare error calculations
between reconstructions with different minibatch size. I hope to
find a better way to do this that is more honest with this
cost function, though.
It can accept intensity and simulated intensity tensors of any shape
as long as their shapes match, and the provided mask array can be
@@ -50,11 +54,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)