mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
344 lines
12 KiB
Python
344 lines
12 KiB
Python
from __future__ import division, print_function, absolute_import
|
|
|
|
import torch as t
|
|
from torch.utils import data as torchdata
|
|
from matplotlib import pyplot as plt
|
|
from matplotlib.widgets import Slider
|
|
from matplotlib import ticker
|
|
import numpy as np
|
|
|
|
#
|
|
# 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 += 1
|
|
|
|
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, schedule=False):
|
|
|
|
# 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)
|
|
|
|
|
|
# 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)
|
|
|
|
|
|
def LBFGS_optimize(self, iterations, dataset, batch_size=None,
|
|
lr=0.1,history_size=2):
|
|
|
|
# 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.LBFGS(self.parameters(),
|
|
lr = lr, history_size=history_size)
|
|
|
|
return self.AD_optimize(iterations, data_loader, optimizer)
|
|
|
|
|
|
# By default, the plot_list is empty
|
|
plot_list = []
|
|
|
|
|
|
def inspect(self, dataset=None, update=True):
|
|
"""Plots all the plots defined in the model's plot_list attribute
|
|
|
|
If update is set to True, it will update any previously plotted set
|
|
of plots, if one exists, and then redraw them. Otherwise, it will
|
|
plot a new set, and any subsequent updates will update the new set
|
|
|
|
Optionally, a dataset can be passed, which then will plot any
|
|
registered plots which need to incorporate some information from
|
|
the dataset (such as geometry or a comparison with measured data).
|
|
|
|
Args:
|
|
dataset (CDataset): Optional, a dataset matched to the model type
|
|
update (bool) : Whether to update existing plots or plot new ones
|
|
|
|
"""
|
|
first_update = False
|
|
if update and hasattr(self, 'figs') and self.figs:
|
|
figs = self.figs
|
|
elif update:
|
|
figs = None
|
|
self.figs = []
|
|
first_update = True
|
|
else:
|
|
figs = None
|
|
self.figs = []
|
|
|
|
idx = 0
|
|
for plots in self.plot_list:
|
|
# If a conditional is included in the plot
|
|
try:
|
|
if len(plots) >=3 and not plots[2](self):
|
|
continue
|
|
except TypeError as e:
|
|
if len(plots) >= 3 and not plots[2](self, dataset):
|
|
continue
|
|
|
|
name = plots[0]
|
|
plotter = plots[1]
|
|
|
|
if figs is None:
|
|
fig = plt.figure()
|
|
self.figs.append(fig)
|
|
else:
|
|
fig = figs[idx]
|
|
|
|
try:
|
|
plotter(self,fig)
|
|
plt.title(name)
|
|
except TypeError as e:
|
|
if dataset is not None:
|
|
try:
|
|
plotter(self, fig, dataset)
|
|
plt.title(name)
|
|
except (IndexError, KeyError, AttributeError) as e:
|
|
pass
|
|
except (IndexError, KeyError, AttributeError) as e:
|
|
pass
|
|
|
|
idx += 1
|
|
|
|
if update:
|
|
plt.draw()
|
|
fig.canvas.start_event_loop(0.001)
|
|
|
|
if first_update:
|
|
plt.pause(0.05 * len(self.figs))
|
|
|
|
|
|
def compare(self, dataset):
|
|
"""Opens a tool for comparing simulated and measured diffraction patterns
|
|
|
|
Args:
|
|
dataset (CDataset) : A dataset containing the simulated diffraction patterns to compare agains
|
|
"""
|
|
|
|
fig, axes = plt.subplots(1,3,figsize=(12,5.3))
|
|
fig.tight_layout(rect=[0.02, 0.09, 0.98, 0.96])
|
|
axslider = plt.axes([0.15,0.06,0.75,0.03])
|
|
|
|
|
|
def update_colorbar(im):
|
|
# If the update brought the colorbar out of whack
|
|
# (say, from clicking back in the navbar)
|
|
# Holy fuck this was annoying. Sorry future for how
|
|
# crappy this solution is.
|
|
#if not np.allclose(im.colorbar.ax.get_xlim(),
|
|
# (np.min(im.get_array()),
|
|
# np.max(im.get_array()))):
|
|
if hasattr(im, 'norecurse') and im.norecurse:
|
|
im.norecurse=False
|
|
return
|
|
|
|
im.norecurse=True
|
|
im.colorbar.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array()))
|
|
im.colorbar.ax.set_ylim(0,1)
|
|
im.colorbar.set_ticks(ticker.LinearLocator(numticks=5))
|
|
im.colorbar.draw_all()
|
|
|
|
|
|
def update(idx):
|
|
idx = int(idx) % len(dataset)
|
|
fig.pattern_idx = idx
|
|
updating = True if len(axes[0].images) >= 1 else False
|
|
|
|
inputs, output = dataset[idx]
|
|
sim_data = self.forward(*inputs).detach().cpu().numpy()
|
|
sim_data = sim_data
|
|
meas_data = output.detach().cpu().numpy()
|
|
if hasattr(self, 'mask') and self.mask is not None:
|
|
mask = self.mask.detach().cpu().numpy()
|
|
else:
|
|
mask = 1
|
|
|
|
if not updating:
|
|
axes[0].set_title('Simulated')
|
|
axes[1].set_title('Measured')
|
|
axes[2].set_title('Difference')
|
|
|
|
sim = axes[0].imshow(sim_data)
|
|
meas = axes[1].imshow(meas_data * mask)
|
|
diff = axes[2].imshow((sim_data-meas_data) * mask)
|
|
|
|
cb1 = plt.colorbar(sim, ax=axes[0], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1)
|
|
cb1.ax.tick_params(labelrotation=20)
|
|
cb1.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(sim))
|
|
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1)
|
|
cb2.ax.tick_params(labelrotation=20)
|
|
cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
|
|
cb3 = plt.colorbar(diff, ax=axes[2], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1)
|
|
cb3.ax.tick_params(labelrotation=20)
|
|
cb3.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(diff))
|
|
|
|
else:
|
|
sim = axes[0].images[-1]
|
|
sim.set_data(sim_data)
|
|
update_colorbar(sim)
|
|
|
|
meas = axes[1].images[-1]
|
|
meas.set_data(meas_data * mask)
|
|
update_colorbar(meas)
|
|
|
|
diff = axes[2].images[-1]
|
|
diff.set_data((sim_data-meas_data) * mask)
|
|
update_colorbar(diff)
|
|
|
|
|
|
# This is dumb but the slider doesn't work unless a reference to it is
|
|
# kept somewhere...
|
|
self.slider = Slider(axslider, 'Pattern #', 0, len(dataset)-1, valstep=1, valfmt="%d")
|
|
self.slider.on_changed(update)
|
|
|
|
def on_action(event):
|
|
if not hasattr(event, 'button'):
|
|
event.button = None
|
|
if not hasattr(event, 'key'):
|
|
event.key = None
|
|
|
|
if event.key == 'up' or event.button == 'up':
|
|
update(fig.pattern_idx - 1)
|
|
elif event.key == 'down' or event.button == 'down':
|
|
update(fig.pattern_idx + 1)
|
|
self.slider.set_val(fig.pattern_idx)
|
|
plt.draw()
|
|
|
|
fig.canvas.mpl_connect('key_press_event',on_action)
|
|
fig.canvas.mpl_connect('scroll_event',on_action)
|
|
update(0)
|
|
|
|
|
|
|
|
|
|
from CDTools.models.simple_ptycho import SimplePtycho
|
|
from CDTools.models.fancy_ptycho import FancyPtycho
|