cSAXS/cdtools
0
0 Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-11-11 16:53:57 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1 from allevitan/organization

Browse Source 
Organization
This commit is contained in:
allevitan authored and GitHub committed 2024-09-02 11:20:37 +02:00
2 parents 35f140d348 13a2ba86df
commit 3b6b0534af
47 files changed
+1015 -6691
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/source/conf.py
+3 -1
View File
@@ -49,7 +49,9 @@ extensions = [
# One only works in 1.8+, the other is depricated in >1.8
autodoc_default_options = {
'show-inheritance': True
'show-inheritance': True,
'special-members': '__init__',
'member-order': 'bysource',
}
autodoc_default_flags = [ 'show-inheritance']
docs/source/intro.rst
+14 -2
View File
@@ -17,6 +17,9 @@ CDTools is a python library for ptychography and CDI reconstructions, using an a
# Load a data file
dataset = Ptycho2DDataset.from_cxi('ptycho_data.cxi')
# Look at the data inside
dataset.inspect()
# Initialize a model from the data
model = FancyPtycho.from_dataset(dataset)
@@ -36,7 +39,16 @@ CDTools is a python library for ptychography and CDI reconstructions, using an a
CDTools makes it simple to load and inspect data stored in .cxi files. Reconstruction models for common geometries are included "out of the box". For more advanced users, it includes a collection of differentiable functions which can beused to construct new forward models.
The high-level interface to CDTools - datasets and models - is built on a set oflower-level tools. These include functions for accessing stored data in .cxi files, tools to visualize data and reconstructions, and tools that implement basic operations - such as light propagation - relevant to coherent diffraction. These functions can be used alongside the high-level interface, when needed.
The high-level interface to CDTools - datasets and models - is built on a set of lower-level tools. These lower level tools include:
Enough blabber. If you're interested, read the docs!
- functions for accessing stored data in .cxi files
- plotting tools to visualize data and reconstructions
- basic operations, like light propagators, needed for coherent diffraction
- tools that implement basic operations - such as light propagation - relevant to coherent diffraction.
- analysis functions for assessing the quality of reconstructions
These functions can be used alongside the high-level interface to make fun and fancy reconstruction scripts for challenging data.
If you're interested, read the docs!
docs/source/models.rst
+1 -1
View File
@@ -3,5 +3,5 @@ Models
.. automodule:: cdtools.models
:members:
:private-members:
docs/source/tools/analysis.rst
+1
View File
@@ -3,3 +3,4 @@ Analysis
.. automodule:: cdtools.tools.analysis
:members:
docs/source/tutorial.rst
+151 -200
View File
@@ -1,111 +1,15 @@
Tutorial
========
WARNING: THE FOLLOWING TUTORIAL IS OUT OF DATE! It has some helpful information, but it is not trustworthy. Please rely on the examples, which are current as of February 2024.
This tutorial builds on the examples, leading to a more complete understanding of how to use CDTools and - importantly - how to extend it. First, we will cover the details of writing a useful reconstruction script for a particular experiment. Next, we will discuss how to implement a new dataset type for different kinds of coherent diffraction. Finally, we will go over how to make new models to cover specific types of ptychography which aren't described by any of the built-in models.
Reconstruction Scripts
----------------------
In this section, we will write a script to run a reconstruction on a dataset collected from our benchtop optical ptychography playground. This mirrors very closely the reconstruction examples, however I encourage everyone to follow along, writing this script out line-by-line, to help you learn more permanently the process of writing a custom reconstruction script.
Our first step will be creating the file and filling out the boilerplate: All the imports we'll need.
.. code-block:: python
import cdtools
from matplotlib import pyplot as plt
from scipy import io
You can always import more libraries, like numpy, or pytorch, or pandas, or what have you, as needed. Next, we load the dataset and give it a look-over
.. code-block:: python
filename = 'example_data/lab_ptycho_data.cxi'
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)
dataset.inspect()
plt.show()
Now, run this script! You should see a window pop up, showing a nanomap of the integrated intensities at each scan point on one side, and an individual diffraction pattern on the other. You can then click around to make sure that everything is in order.
Now that we know we have the data loaded and it looks good, we can go ahead and comment out the dataset inspecting code, and move on to creating a model. It's usually a good idea to start by loading a standard :code:`FancyPtycho` model without any special changes, sending it to the GPU.
.. code-block:: python
model = cdtools.models.FancyPtycho.from_dataset(dataset)
model.to(device='cuda')
dataset.get_as(device='cuda')
We then try a basic Adam reconstruction with this model, with no changes to the defaults, to see how it works.
.. code-block:: python
for loss in model.Adam_optimize(50, dataset):
model.inspect(dataset)
print(model.report())
model.compare(dataset)
plt.show()
It is worth noting here exactly how this code is working. The reconstruction methods are actually returning generators. Generators in pythons are objects that work like lists, or tuples, but have to be read out one item at a time, from left to right. The catch is that, instead of just reading out objects from a list, they can run arbitrary code each time they are asked for the "next" item.
In CDTools, every reconstruction method will return a generator. Whenever the generator is asked for the next item, it runs a single epoch of the reconstructionalgorithm, and then returns the average loss over that epoch as that next item. This allows the execution of the reconstruction algorithm to pause once every epoch, allowing some time for the user to run a small snippet of code to inspect how the reconstruction is coming along.
From the end user perspective, all this means is: follow the format above, or more generally put the :code:`model.Adam_optimize(n, dataset)` call anywhere that you would feel comfortable putting a call to :code:`range(n)` - list comprehensions, for loops, etc. In this case, we have called a function to plot out the current state of the reconstruction, and a function to print out the current loss and iteration time.
Once we run this, we can take a look at the result. What we see is pretty good, but we can see that there are some issues with the reconstruction near the edge, and the probe itself seems to be larger than the "stage" on which we're reconstructing it. So, we can make two tweaks to this code in response. First, we increase the oversampling ratio, which doubles the size of the stage (this often can cause other issues as well, but generally works well in situations like this where the probe is honestly too large.
.. code-block:: python
model = cdtools.models.FancyPtycho.from_dataset(dataset, oversampling=2)
And secondly, we note that there don't seem to be any errors with the positioning. So we can just not reconstruct the probe positions, knowing that the initial guesses are already accurate enough. We can do this by writing the following line, just before we run the reconstruction for loop.
.. code-block:: python
model.translation_offsets.requires_grad = False
What is going on here is that, when running the optimization algorithm, pytorch will automatically calculate gradients for and then optimize over a number of parameters defined in the model - this includes parameters like :code:`model.probe`, :code:`model.obj`, :code:`model.background`, etc. We can tell pytorch to stop calculating gradients for (and stop updating) any of these parameters by setting their :code:`requires_grad` property to :code:`False`.
After running this reconstruction, we can see that we're getting a little improvement (and a larger field of view) by using oversampling, but out in the corners we're nucleating extra probes! We can fix this by adding a probe support - that is, declating that the probe has to be defined only within a certain box. This can be done most easily with an argument to the model constructor:
.. code-block:: python
model = cdtools.models.FancyPtycho.from_dataset(dataset, oversampling=2,
probe_support_radius=90)
It also seems like we need a few more iterations to finish converging, so we up the iteration count to 100.
.. code-block:: python
for i, loss in enumerate(model.Adam_optimize(100, dataset)):
Now we expect to get a nice reconstruction, so we can save the data. You can save the data in any form you like, once the relevant information is extracted from the model and put into a dictionary. The standard method for saving out this information is as follows:
.. code-block:: python
io.savemat('example_reconstructions/lab_ptycho.pickle',
model.save_results(dataset))
This is usually placed before the call to :code:`plt.show()`, to make sure that if the user manually exits the program once all the plot windows are opened, the data will still have been saved.
Now, your file should match the example file in examples/lab_ptycho_data.py.
The following tutorial gives a peek under the hood, and is intended for someone who might want to write their own variant of a ptychography model or modify an existing model to meet a specific need. If you just need to use CDTools for a reconstruction, or are just starting to work with the code, the examples section is a great first introduction.
In the first section of the tutorial, we will discuss how the datasets are defined and go through the process of defining a new dataset type. Following that, we will go through the process of defining a simplified model for standard ptychography.
Datasets
--------
In this section, we will write a bare-bones dataset class for 2D ptychography data to demonstrate the process of writing a new dataset class. At the end of the tutorial, we will have written the file examples/basic_ptycho_dataset.py, which can be consulted for reference.
In this section, we will write a bare-bones dataset class for 2D ptychography data to demonstrate the process of writing a new dataset class. At the end of the tutorial, we will have written the file examples/tutorial_basic_ptycho_dataset.py, which can be consulted for reference.
Basic Idea
++++++++++
@@ -149,7 +53,7 @@ The next thing to implement is the initialization code. Here we can leverage som
.. code-block:: python
def __init__(self, *args, **kwargs):
super(BasicPtychoDataset,self).__init__(*args, **kwargs)
super().__init__(*args, **kwargs)
Of course, there is also some data that are unique to this kind of dataset. In this case, those data are the probe translations and the measured diffraction patterns. Therefore, we extend this definition to the following:
@@ -159,7 +63,7 @@ Of course, there is also some data that are unique to this kind of dataset. In t
def __init__(self, translations, patterns, *args, **kwargs):
"""Initialize the dataset from python objects"""
super(BasicPtychoDataset,self).__init__(*args, **kwargs)
super().__init__(*args, **kwargs)
self.translations = t.Tensor(translations).clone()
self.patterns = t.Tensor(patterns).clone()
@@ -187,7 +91,7 @@ Remember that it's not needed to worry about what device or datatype the data is
def to(self, *args, **kwargs):
"""Sends the relevant data to the given device and dtype"""
super(BasicPtychoDataset,self).to(*args,**kwargs)
super().to(*args,**kwargs)
self.translations = self.translations.to(*args, **kwargs)
self.patterns = self.patterns.to(*args, **kwargs)
@@ -228,7 +132,7 @@ Now to save the data, we override :code:`to_cxi()`, in a fairly self-explanatory
def to_cxi(self, cxi_file):
"""Saves out a BasicPtychoDataset as a .cxi file"""
super(BasicPtychoDataset,self).to_cxi(cxi_file)
super().to_cxi(cxi_file)
cdtdata.add_data(cxi_file, self.patterns, axes=self.axes)
cdtdata.add_ptycho_translations(cxi_file, self.translations)
@@ -259,17 +163,21 @@ This is a bare-bones class, set up to demonstrate the minimum neccessary to deve
Models
------
In this section, we will write a basic model for 2D ptychography reconstructions. At the end of this tutorial, we will have written the class defined in examples/simple_ptycho_model.py
In this section, we will write a basic model for 2D ptychography reconstructions. At the end of this tutorial, we will have written the class defined in examples/tutorial_simple_ptycho_model.py
Basic Idea
++++++++++
Just like CDTools Datasets subclass pytorch Datasets, CDTools models subclass pytorch modules (yes, I know they are different words - we use the word "model" in CDTools to conform to usage in the world of ptychography/CDI). The major difference is that the base CDTools models also contains a few standard methods to run automatic differentiation reconstructions on itself. This isn't necessarily the cleanest or most portable approach, but we've found that it feels very natural from the perspective of an end user interacting with the toolbox only through some basic reconstruction scripts.
Just like CDTools Datasets subclass pytorch Datasets, CDTools models subclass pytorch modules. However, the concept of a CDTools model does differ slightly from that of a pytorch module, because the CDTools models also contain a few standard methods to run automatic differentiation reconstructons on themselves.
The models themselves have a :code:`model.forward()` function which contains the real meat. In any CDTools model, this forward function takes in a set of parameters describing the specific diffraction pattern to simulate, and outputs a simulated diffraction pattern. The inputs could be as simple as a diffraction pattern index, or could explicitly include other information like the probe position.
This isn't necessarily the cleanest or most portable approach, but we've found that it feels very natural from the perspective of an end user interacting with the toolbox only through the reconstruction scripts.
In practice, the forward model is defined in the top level :code:`CDIModel` class from which all other models are derived. The definition is quite simple:
The heart of each model is a :code:`model.forward()` function. In any CDTools model, this forward function maps a set of parameters describing the specific diffraction pattern to simulate to the simulated result. When it's paired with an appropriate dataset for a reconstruction, it maps from the "inputs" defined by the dataset to the "outputs".
For ptychography, this information is usually index of the exposure within the dataset (which is used to retrieve exposure-to-exposure information, like probe intensity factors) and the object translation.
A simple forward model is defined in the top level :code:`CDIModel` class from which all other models are derived, and rarely needs to be overridden. The definition is quite simple:
.. code-block:: python
@@ -278,7 +186,7 @@ In practice, the forward model is defined in the top level :code:`CDIModel` clas
So we can see that to fully implement this forward model, we have to define the three functions :code:`model.interaction()`, :code:`model.forward_propagator()`, and :code:`model.measurement()`, which simulate conceptual stages in the diffraction process.
Beyond the basic model definition, a few other tools need to be defined. The model has to be able to create itself from a dataset, has to have a loss function defined for use with automtic differentiation, has to know how to plot out it's progress, and has to be able to save out the results of a reconstruction. The details of how to implement all of this in a model are shown below.
In addition to the core model definition, a few other functions need to be defined to make the model useful. The model needs an *initializer* to create itself from a dataset, it must have an appropriate *loss function* defined for use with automtic differentiation, a way of plotting the progress of a reconstruction, and must know how to save the results of a reconstruction in a useful format.
Writing the Skeleton
@@ -288,16 +196,16 @@ Once again, we start with the basic skeleton
.. code-block:: python
import numpy as np
import torch as t
from cdtools.models import CDIModel
from cdtools import tools
import torch as t
from cdtools.models import CDIModel
from cdtools import tools
from cdtools.tools import plotting as p
__all__ = ['SimplePtycho']
class SimplePtycho(CDIModel):
"""A simple ptychography model for exploring ideas and extensions"""
pass
"""A simple ptychography model to demonstrate the structure of a model
"""
Note that we imported the full tools package, as we will find ourselves using many low-level functions defined there to implement the model.
@@ -305,101 +213,137 @@ Note that we imported the full tools package, as we will find ourselves using ma
Initialization from Python
++++++++++++++++++++++++++
Two initialization functions need to be written. First, we write the :code:`__init__()` function, which initializes the model from a collection of python objects describing the system. We then write an initializer that creates a model using a dataset to define the various parameters.
Two initialization functions need to be written. First, we write the :code:`__init__()` function, which initializes the model from a collection of python objects describing the system. We then write an initializer that creates a model using a dataset to initialize the various parameters.
It's important to note that there's not requirement for what the arguments to the initialization function of any particular model should be, only that they contain enough information to run the simulations! It should be chosen in a model-by-model way to allow for the most transparent code.
There is no requirement for what the arguments to the initialization function of any particular model should be, only that they contain enough information to run the simulations! It should be chosen in a model-by-model basis to allow for the most sensible code.
.. code-block:: python
def __init__(self, probe_basis, probe_guess, obj_guess,
min_translation = [0,0]):
def __init__(
self,
wavelength,
probe_basis,
probe_guess,
obj_guess,
min_translation = [0,0],
):
# We initialize the superclass
super().__init__()
# We register all the constants, like wavelength, as buffers. This
# lets the model hook into some nice pytorch features, like using
# model.to, and broadcasting the model state across multiple GPUs
self.register_buffer('wavelength', t.as_tensor(wavelength))
self.register_buffer('min_translation', t.as_tensor(min_translation))
self.register_buffer('probe_basis', t.as_tensor(probe_basis))
# We cast the probe and object to 64-bit complex tensors
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
# We rescale the probe here so it learns at the same rate as the
# object when using optimizers, like Adam, which set the stepsize
# to a fixed maximum
self.register_buffer('probe_norm', t.max(t.abs(probe_guess)))
# And we store the probe and object guesses as parameters, so
# they can get optimized by pytorch
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
# We have to initialize the Module
super(SimplePtycho,self).__init__()
# We first save the relevant information
self.min_translation = t.Tensor(min_translation)
self.probe_basis = t.Tensor(probe_basis)
The first thing to notice about this model is that all the fixed, geometric information is stored with the :code:`module.register_buffer()` function. This is what makes it possible to move all the relevant tensors between devices using a single call to :code:`module.to()`, for example. It stores thetensor as an object attribute, but it also registers it so that pytorch is aware that this attribute helps to encode the state of the model.
# We rescale the probe so it learns at the same rate as the object
self.probe_norm = t.max(t.abs(probe_guess).to(t.float32))
self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
The supporting information we need is the wavelength of the illumination, the basis of the probe array in real space, and an offset to define the zero point of the translation.
Here, we chose to define the model based on a basis matrix describing the probe array, an initial guess at the probe, and an initial object. In addition, an optional offset for the translations is included.
The final two pieces of information that we need to save are the probe and object, and both of these get defined as :code:`t.nn.Parameter` objects instead of Tensors. As a result, they get registered as parameters in the pytorch module, and will therefore be optimized over in any later reconstructions. In addition, the :code:`requires_grad` flag is set to :code:`True`, which means that the information needed for gradient calculations will be stored on every Tensor that results from a calculation including a Parameter.
The first part of this initialization is quite straightforward - we create some tensors for the minimum translation and the probe basis. But then, the next two pieces of information that we save are defined as :code:`t.nn.Parameter` objects, not Tensors! Parameters are different from Tensors in two ways.
A list of all parameters associated with the module can be found by calling :code:`module.parameters()`.
The first way that they are different is that, by default, they have the :code:`requires_grad` flag set to :code:`True`, which means that the information needed for gradient calculations will be stored on every Tensor that results from a calculation including a Parameter. The second difference is that, when a Parameter is added to a Module, the Module adds that parameter to a list, which can be accessed by calling :code:`module.parameters()`.
Any additional targets of reconstruction - such as exposure-to-exposure illumination weights, translation offsets, or a detector background - would be added to the model as a parameter in a similar way.
The key here is that the model itself is subclassing a pytorch Module. So, every parameter that we will attempt to reconstruct, we add to the CDTools model as a Parameter. This way, the model automatically knows which variables to update with gradient descent, and which to keep as they are. Here, we only need to reconstruct the probe and object.
One final note is that we actually store a scaled version of the probe. This is a specific case of a general policy designed around making it easy to use the Adam optimizer.
One final note is that we actually store a scaled version of the probe. This is a hack. It is simply because the Adam optimization method, which is the most commonly used, uses learning rates that scale with the amplitude of the parameter, rather than with the amplitude of it's gradients. Unti pytorch implements separate parameter-group learning rates in Adam, rescaling all the parameters to have a typical amplitude near 1 is the best way to get well-behaved reconstructions.
The Adam optimizer is designed so that the learning rate sets the maximum stepsize which will be taken in any single iteration. Therefore, it is important to make sure that *all parameters of the model are of order unity*. To enable this, we scale the probe so that the typical pixel value within the probe array is of order 1.
This is important to remember when adding additional error models. Rescaling all the parameters to have a typical amplitude near 1 is the best way to get well-behaved reconstructions.
Initialization from Dataset
+++++++++++++++++++++++++++
To initialize the object from a dataset, we need to start by extracting the relevant information from the dataset. Then we can simply call the constructor we defined earlier.
To initialize the object from a dataset, we need to start by extracting the relevant information from the dataset, before calling the constructor we defined above:
.. code-block:: python
@classmethod
def from_dataset(cls, dataset):
# First, load the information from the dataset
# We get the key geometry information from the dataset
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
det_shape = dataset[0][1].shape
distance = dataset.detector_geometry['distance']
(indices, translations), patterns = dataset[:]
# Then, generate the probe geometry
# Then, we generate the probe geometry
ewg = tools.initializers.exit_wave_geometry
probe_basis, probe_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance)
probe_basis = ewg(det_basis, det_shape, wavelength, distance)
# Next generate the object geometry
# Next generate the object geometry from the probe geometry and
# the translations
(indices, translations), patterns = dataset[:]
pix_translations = tools.interactions.translations_to_pixel(
probe_basis, translations)
probe_basis,
translations,
)
obj_size, min_translation = tools.initializers.calc_object_setup(
probe_shape, pix_translations)
det_shape,
pix_translations,
)
# Finally, initialize the probe and object using this information
probe = tools.initializers.SHARP_style_probe(dataset,
probe_shape,
det_slice)
# Finally, initialize the probe and object using this information
probe = tools.initializers.SHARP_style_probe(dataset)
obj = t.ones(obj_size).to(dtype=t.complex64)
obj = t.ones(obj_size, dtype=t.complex64)
return cls(probe_basis, probe, obj, min_translation=min_translation)
return cls(
wavelength,
probe_basis,
probe,
obj,
min_translation=min_translation
)
Here, we start by pulling the basic geometric information from the dataset. Then, we use a number of the basic tools to do calculations such as finding the probe basis from the detector geometry, or calculating how big our object array should be.
Once we have the basic setup ready, we then use one of the initialization functions - in this case, :code:`tools.initializers.SHARP_style_probe`, to find a sensible initialization for the probe. This particular initialization is based on the approach used in the SHARP package. Once all the needed information has been collected, we initialize the object.
Once we have the basic setup ready, we then use one of the initialization functions - in this case, :code:`tools.initializers.SHARP_style_probe`, to find a sensible initialization for the probe. This particular initialization is based on the approach used in the SHARP package, where the square-root of the mean diffraction pattern intensity is used to estimate the structure of the illumination at focus.
Once all the needed information has been collected, we initialize the object.
The Forward Model
+++++++++++++++++
First, we have to implement the interaction model. This function should take the inputs defining the specific diffraction pattern or collection of patterns, and return an exit wave or set of exit waves that would be expected from that set of inputs. This models the interaction of the probe with the sample
First, we have to implement the interaction model, as below:
.. code-block:: python
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations)
# We map from real-space to pixel-space units
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)
# This function extracts the appropriate window from the object and
# multiplies the object and probe functions
return tools.interactions.ptycho_2D_round(
self.probe_norm * self.probe,
self.obj,
pix_trans)
Here, we take input in the form of an index and a translation. Note that this has to match the format that is output by the associated datasets that we will use for reconstruction.
Here, we take input in the form of an index and a translation. Note that this input format much match the format that is output by the associated datasets that we will use for reconstruction, in this case BasicPtychoDataset.
We start by mapping the translation, given in real space, into pixel coordinates. Then, we use an "off-the-shelf" interaction model - :code:`ptycho_2d_round`, which models a standard 2D ptychography interaction, but rounds the translations to the nearest whole pixel (does not attempt subpixel translations).
@@ -417,23 +361,7 @@ The next three definitions amount to just choosing an off-the-shelf function to
return tools.losses.amplitude_mse(real_data, sim_data)
The forward propagator maps the exit wave to the wave at the surface of the detector, here using a far-field propagator. The measurement maps that exit wave to a measured pixel value, and the loss defines a loss function to attempt to minimize. The loss function we've chosen - the amplitude mean squared error - is the most broadly applicable one.
Device Management
+++++++++++++++++
Next we need to implement a :code:`to()` function, just like we did in the dataset, to allow the entire model to be moved between the GPU and CPU.
.. code-block:: python
def to(self, *args, **kwargs):
super(SimplePtycho, self).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)
We start by calling the same function for the superclass, which will take care of moving every parameter attached to the model. Then we manually take care of moving over the other variables which, while not being updated in the gradient descent, still need to be moved over to the new device!
The forward propagator maps the exit wave to the wave at the surface of the detector, here using a far-field propagator. The measurement maps that exit wave to a measured pixel value, and the loss defines a loss function to attempt to minimize. The loss function we've chosen - the amplitude mean squared error - is the most reliable one, and can also easily be overridden by an end user.
Plotting
@@ -447,62 +375,85 @@ The base CDIModel class has a function, :code:`model.inspect()`, which looks for
.. code-block:: python
# This lists all the plots to display on a call to model.inspect()
plot_list = [
('Probe Amplitude',
lambda self, fig: tools.plotting.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
('Probe Phase',
lambda self, fig: tools.plotting.plot_phase(self.probe, fig=fig, basis=self.probe_basis)),
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)),
('Object Amplitude',
lambda self, fig: tools.plotting.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
lambda self, fig: tools.plotting.plot_phase(self.obj, fig=fig, basis=self.probe_basis))
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis))
]
In this case, we've made use of the convenience plotting functions defined in :code:`tools.plotting`.
Saving
++++++
At the moment, there is no consistent way to save out the results across the board. However, a function :code:`save_results()` should be defined, which should save out the results of the reconstruction into some reasonably formatted python object. Here we return a dictionary with the probe and object:
By default, a function :code:`model.save_results()` is defined, which returns a python dictionary with an entry, :code:`'state_dict'`, containing all the registered parameters and buffers in the model. It also contains a basic record of the model's training history. This function is used internally by :code:`model.save_to_h5()`, as well as all other convenience functions for saving results.
Sometimes, it is also useful to return a more user-friendly version of the results, such as a properly rescaled version of the probe. To make this possible, :code:`model.save_results()` is often overridden:
.. code-block:: python
def save_results(self):
def save_results(self, dataset):
# This will save out everything needed to recreate the object
# in the same state, but it's not the best formatted.
base_results = super().save_results()
# So we also save out the main results in a more useable format
probe_basis = self.probe_basis.detach().cpu().numpy()
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = self.obj.detach().cpu().numpy()
return {'probe':probe,'obj':obj}
wavelength = self.wavelength.cpu().numpy()
results = {
'probe_basis': probe_basis,
'probe': probe,
'obj': obj,
'wavelength': wavelength,
}
return {**base_results, **results}
However, it is perfectly possible to write a new ptychography model without overriding :code:`model.save_results()`
Testing
+++++++
We can test this model with a simple script, shown below. By filling in the backend here, we've been able to create a ptychography model that can be accessed and used in reconstructions via the same interface as the models we discussed in the examples section.
We can test this model with a simple script, in examples/tutorial_finale.py. By filling in the backend here, we've been able to create a ptychography model that can be accessed and used in reconstructions via the same interface as the models we discussed in the examples section.
.. code-block:: python
from basic_ptycho_dataset import BasicPtychoDataset
from tutorial_basic_ptycho_dataset import BasicPtychoDataset
from tutorial_simple_ptycho import SimplePtycho
from h5py import File
from matplotlib import pyplot as plt
filename = 'example_data/lab_ptycho_data.cxi'
with File(filename, 'r') as f:
dataset = BasicPtychoDataset.from_cxi(f)
dataset = BasicPtychoDataset.from_cxi(f)
dataset.inspect()
model = SimplePtycho.from_dataset(dataset)
model.to(device='cuda')
dataset.get_as(device='cuda')
for loss in model.Adam_optimize(100, dataset):
model.inspect(dataset)
print(model.report())
model.to(device='mps')#cuda')
dataset.get_as(device='mps')#cuda')
for loss in model.Adam_optimize(10, dataset):
model.inspect(dataset)
print(model.report())
model.inspect(dataset)
model.compare(dataset)
plt.show()
Happy modeling!
examples/basic_ptycho_dataset.py → examples/tutorial_basic_ptycho_dataset.py
+4 -5
View File
@@ -1,4 +1,3 @@
import numpy as np
import torch as t
from matplotlib import pyplot as plt
from cdtools.datasets import CDataset
@@ -14,7 +13,7 @@ class BasicPtychoDataset(CDataset):
def __init__(self, translations, patterns, *args, **kwargs):
"""Initialize the dataset from python objects"""
super(BasicPtychoDataset,self).__init__(*args, **kwargs)
super().__init__(*args, **kwargs)
self.translations = t.Tensor(translations).clone()
self.patterns = t.Tensor(patterns).clone()
@@ -27,7 +26,7 @@ class BasicPtychoDataset(CDataset):
def to(self, *args, **kwargs):
"""Sends the relevant data to the given device and dtype"""
super(BasicPtychoDataset,self).to(*args,**kwargs)
super().to(*args,**kwargs)
self.translations = self.translations.to(*args, **kwargs)
self.patterns = self.patterns.to(*args, **kwargs)
@@ -55,7 +54,7 @@ class BasicPtychoDataset(CDataset):
def to_cxi(self, cxi_file):
"""Saves out a BasicPtychoDataset as a .cxi file"""
super(BasicPtychoDataset,self).to_cxi(cxi_file)
super().to_cxi(cxi_file)
cdtdata.add_data(cxi_file, self.patterns, axes=self.axes)
cdtdata.add_ptycho_translations(cxi_file, self.translations)
@@ -63,6 +62,6 @@ class BasicPtychoDataset(CDataset):
def inspect(self):
"""Plots a random diffraction pattern"""
index = np.random.randint(len(self))
index = t.randint(len(self), (1,))[0]
plt.figure()
plt.imshow(self.patterns[index,:,:].cpu().numpy())
examples/tutorial_finale.py
+23
View File
@@ -0,0 +1,23 @@
from tutorial_basic_ptycho_dataset import BasicPtychoDataset
from tutorial_simple_ptycho import SimplePtycho
from h5py import File
from matplotlib import pyplot as plt
filename = 'example_data/lab_ptycho_data.cxi'
with File(filename, 'r') as f:
dataset = BasicPtychoDataset.from_cxi(f)
dataset.inspect()
model = SimplePtycho.from_dataset(dataset)
model.to(device='cuda')
dataset.get_as(device='cuda')
for loss in model.Adam_optimize(10, dataset):
model.inspect(dataset)
print(model.report())
model.inspect(dataset)
model.compare(dataset)
plt.show()
examples/tutorial_simple_ptycho.py
+140
View File
@@ -0,0 +1,140 @@
import torch as t
from cdtools.models import CDIModel
from cdtools import tools
from cdtools.tools import plotting as p
__all__ = ['SimplePtycho']
class SimplePtycho(CDIModel):
"""A simple ptychography model to demonstrate the structure of a model
"""
def __init__(
self,
wavelength,
probe_basis,
probe_guess,
obj_guess,
min_translation = [0,0],
):
# We initialize the superclass
super().__init__()
# We register all the constants, like wavelength, as buffers. This
# lets the model hook into some nice pytorch features, like using
# model.to, and broadcasting the model state across multiple GPUs
self.register_buffer('wavelength', t.as_tensor(wavelength))
self.register_buffer('min_translation', t.as_tensor(min_translation))
self.register_buffer('probe_basis', t.as_tensor(probe_basis))
# We cast the probe and object to 64-bit complex tensors
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
# We rescale the probe here so it learns at the same rate as the
# object when using optimizers, like Adam, which set the stepsize
# to a fixed maximum
self.register_buffer('probe_norm', t.max(t.abs(probe_guess)))
# And we store the probe and object guesses as parameters, so
# they can get optimized by pytorch
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
@classmethod
def from_dataset(cls, dataset):
# We get the key geometry information from the dataset
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
det_shape = dataset[0][1].shape
distance = dataset.detector_geometry['distance']
# Then, we generate the probe geometry
ewg = tools.initializers.exit_wave_geometry
probe_basis = ewg(det_basis, det_shape, wavelength, distance)
# Next generate the object geometry from the probe geometry and
# the translations
(indices, translations), patterns = dataset[:]
pix_translations = tools.interactions.translations_to_pixel(
probe_basis,
translations,
)
obj_size, min_translation = tools.initializers.calc_object_setup(
det_shape,
pix_translations,
)
# Finally, initialize the probe and object using this information
probe = tools.initializers.SHARP_style_probe(dataset)
obj = t.ones(obj_size).to(dtype=t.complex64)
return cls(
wavelength,
probe_basis,
probe,
obj,
min_translation=min_translation
)
def interaction(self, index, translations):
# We map from real-space to pixel-space units
pix_trans = tools.interactions.translations_to_pixel(
self.probe_basis,
translations)
pix_trans -= self.min_translation
# This function extracts the appropriate window from the object and
# multiplies the object and probe functions
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 measurement(self, wavefields):
return tools.measurements.intensity(wavefields)
def loss(self, real_data, sim_data):
return tools.losses.amplitude_mse(real_data, sim_data)
# This lists all the plots to display on a call to model.inspect()
plot_list = [
('Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
('Probe Phase',
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis)),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis))
]
def save_results(self, dataset):
# This will save out everything needed to recreate the object
# in the same state, but it's not the best formatted.
base_results = super().save_results()
# So we also save out the main results in a more useable format
probe_basis = self.probe_basis.detach().cpu().numpy()
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = self.obj.detach().cpu().numpy()
wavelength = self.wavelength.cpu().numpy()
results = {
'probe_basis': probe_basis,
'probe': probe,
'obj': obj,
'wavelength': wavelength,
}
return {**base_results, **results}
src/cdtools/datasets/__init__.py
+1 -2
View File
@@ -29,8 +29,7 @@ attempting to do so
# I don't believe that __all__ really needed, but it's nice to define it
# to be explicit that import * is safe
__all__ = ['CDataset','Ptycho2DDataset','PolarizationSweptPtycho2DDataset']
__all__ = ['CDataset','Ptycho2DDataset']
from cdtools.datasets.base import CDataset
from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
from cdtools.datasets.polarization_swept_ptycho_2d_dataset import PolarizationSweptPtycho2DDataset
src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py
-194
View File
@@ -1,194 +0,0 @@
import torch as t
import h5py
import pathlib
from cdtools.datasets import Ptycho2DDataset
from cdtools.tools import data as cdtdata
__all__ = ['PolarizationSweptPtycho2DDataset']
class PolarizationSweptPtycho2DDataset(Ptycho2DDataset):
"""The standard dataset for a 2D ptychography scan
Subclasses datasets.CDataset
This class loads and saves 2D ptychography scan data from .cxi files.
It should save and load files compatible with most reconstruction
programs, although it is only tested against SHARP.
"""
def __init__(self, translations, patterns, polarization_indices,
polarization_states,
*args, **kwargs):
"""The __init__ function allows construction from python objects.
The detector_geometry dictionary is defined to have the
entries defined by the outputs of data.get_detector_geometry.
Note that the created dataset object will not copy the data in the
patterns parameter in order to avoid doubling the memory requiement
for large datasets.
Parameters
----------
translations : array
An nx3 array containing the probe translations at each scan point
patterns : array
An nxmxl array containing the full stack of measured diffraction
patterns
polarization_indices : array(int)
A length-n array containing the index of the polarization state
attached to each pattern
polarization_states : array
An mx2 array, where m is the number of polarization indices,
encoding the polarzation state associated with each index
axes : list(str)
A list of names for the axes of the probe translations
entry_info : dict
A dictionary containing the entry_info metadata
sample_info : dict
A dictionary containing the sample_info metadata
wavelength : float
The wavelength of light used in the experiment
detector_geometry : dict
A dictionary containing the various detector geometry
parameters
mask : array
A mask for the detector, defined as 1 for live pixels, 0
for dead
background : array
An initial guess for the not-previously-subtracted
detector background
intensities : array
A list of measured shot-to-shot intensities
"""
super().__init__(translations, patterns,
*args, **kwargs)
self.polarization_indices = t.tensor(polarization_indices, dtype=t.int32)
self.polarization_states = t.tensor(polarization_states, dtype=t.complex64)
def _load(self, index):
""" Internal function to load data
This function is used internally by the global __getitem__ function
defined in the base class, which handles moving data around when
the dataset is (for example) storing the data on the CPU but
getting data as GPU tensors.
It loads data in the format (inputs, output)
The inputs for a 2D ptychogaphy data set are:
1) The indices of the patterns to use
2) The recorded probe positions associated with those points
3) The angles of the polarizers if polarized=True
Parameters
----------
index (polarized=False): int or slice
The index or indices of the scan points to use
index (polarized=True):
tuple ((phi1, phi2), ind)
ind - index or indices of the scan points to use (in a (phi1, phi2) polarization state), int or slice
(phi1, phi2) - angles of the 1st and 2nd polarizers, ints
Returns
-------
inputs : tuple
A tuple of the inputs to the related forward models
if polarized: inputs = ((phi1, phi2), ind, transl[(phi1, phi2)][ind])
outputs : tuple
The output pattern or stack of output patterns
"""
return ((index, self.translations[index],
self.polarization_indices[index]),
self.patterns[index])
def to(self, *args, **kwargs):
"""Sends the relevant data to the given device and dtype
This function sends the stored translations, patterns,
mask and background to the specified device and dtype
Accepts the same parameters as torch.Tensor.to
"""
super().to(*args, **kwargs)
self.polarization_states = self.polarization_states.to(*args, **kwargs)
# It sucks that I can't reuse the base factory method here,
# perhaps there is a way but I couldn't figure it out.
@classmethod
def from_cxi(cls, cxi_file, cut_zeros=True):
"""Generates a new PolarizationSweptPtycho2DDataset from a .cxi file
This generates a new PolarizationSweptPtycho2DDataset from a .cxi file
storing the 2D ptychography scan.
Parameters
----------
file : str, pathlib.Path, or h5py.File
The .cxi file to load from
cut_zeros : bool
Default True, whether to set all negative data to zero
Returns
-------
dataset : PolarizationSweptPtycho2DDataset
The constructed dataset object
"""
# If a bare string is passed
if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
with h5py.File(cxi_file, 'r') as f:
return cls.from_cxi(f, cut_zeros=cut_zeros)
# Generate a base dataset
dataset = Ptycho2DDataset.from_cxi(cxi_file, cut_zeros=cut_zeros)
# Mutate the class to this subclass (PolarizedPtycho2DDataset)
dataset.__class__ = cls
# Now, we save out the polarizer and analyzer states
polarization_indices = cdtdata.get_shot_to_shot_info(
cxi_file, 'polarization_indices')
polarization_states = dataset.entry_info['polarization_states']
dataset.polarization_indices = t.tensor(polarization_indices,
dtype=t.int32)
dataset.polarization_states = t.tensor(polarization_states,
dtype=t.complex64)
return dataset
def to_cxi(self, cxi_file):
"""Saves out a PolarizationSweptPtycho2DDataset as a .cxi file
This function saves all the compatible information in a
PolarizationSweptPtycho2DDataset object into a .cxi file. This saved
.cxi file should be compatible with any standard .cxi file based
reconstruction tool, such as SHARP.
Parameters
----------
cxi_file : str, pathlib.Path, or h5py.File
The .cxi file to write to
"""
# If a bare string is passed
if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
with cdtdata.create_cxi(cxi_file) as f:
return self.to_cxi(f)
# This saves the translations, patterns, etc.
super().to_cxi(cxi_file)
# Now, we save out the polarization states
cdtdata.add_shot_to_shot_info(cxi_file, self.polarization_indices,
'polarization_indices')
cdtdata.add_entry_info(cxi_file, {'polarization_states':
self.polarization_states})
src/cdtools/models/__init__.py
+10 -12
View File
@@ -20,24 +20,22 @@ defining a new ptychography model before attempting to do so.
"""
# I don't believe that __all__ really needed, but it's nice to define it
# to be explicit that import * is safe
__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'PolarizationSweptPtycho', 'PolarizedFancyPtycho', 'Bragg2DPtycho', 'Multislice2DPtycho', 'RPI', 'TimeResolvedPtychoCalibration', 'TimeResolvedRPI']
# We define __all__ to be sure that import * only imports what we want
__all__ = [
'CDIModel',
'SimplePtycho',
'FancyPtycho',
'Bragg2DPtycho',
'Multislice2DPtycho',
'MultislicePtycho',
'RPI',
]
from cdtools.models.base import CDIModel
from cdtools.models.simple_ptycho import SimplePtycho
from cdtools.models.fancy_ptycho import FancyPtycho
from cdtools.models.polarized_fancy_ptycho import PolarizedFancyPtycho
from cdtools.models.polarization_swept_ptycho import PolarizationSweptPtycho
from cdtools.models.bragg_2d_ptycho import Bragg2DPtycho
from cdtools.models.multislice_2d_ptycho import Multislice2DPtycho
from cdtools.models.multislice_ptycho import MultislicePtycho
from cdtools.models.rpi import RPI
from cdtools.models.multimode_rpi import MultimodeRPI
from cdtools.models.time_resolved_ptycho_calibration import TimeResolvedPtychoCalibration
from cdtools.models.time_resolved_rpi import TimeResolvedRPI
from cdtools.models.fastccd_ptycho import FastCCDPtycho
# Still needs to be updated for the new complex numbers
#from cdtools.models.s_matrix_ptycho import SMatrixPtycho
src/cdtools/models/base.py
+3 -6
View File
@@ -39,7 +39,6 @@ import queue
import time
from scipy import io
from contextlib import contextmanager
from .complex_lbfgs import MyLBFGS
from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch
__all__ = ['CDIModel']
@@ -118,18 +117,18 @@ class CDIModel(t.nn.Module):
The datatype to convert the values to before registering
"""
self.register_buffer('det_basis',
t.tensor(detector_geometry['basis'],
t.as_tensor(detector_geometry['basis'],
dtype=dtype))
if 'distance' in detector_geometry \
and detector_geometry['distance'] is not None:
self.register_buffer('det_distance',
t.tensor(detector_geometry['distance'],
t.as_tensor(detector_geometry['distance'],
dtype=dtype))
if 'corner' in detector_geometry \
and detector_geometry['corner'] is not None:
self.register_buffer('det_corner',
t.tensor(detector_geometry['corner'],
t.as_tensor(detector_geometry['corner'],
dtype=dtype))
def get_detector_geometry(self):
@@ -629,8 +628,6 @@ class CDIModel(t.nn.Module):
optimizer = t.optim.LBFGS(self.parameters(),
lr = lr, history_size=history_size,
line_search_fn=line_search_fn)
#optimizer = MyLBFGS(self.parameters(),
# lr = lr, history_size=history_size)
return self.AD_optimize(iterations, data_loader, optimizer,
regularization_factor=regularization_factor,
src/cdtools/models/bragg_2d_ptycho.py
+21 -21
View File
@@ -92,23 +92,23 @@ class Bragg2DPtycho(CDIModel):
super(Bragg2DPtycho, self).__init__()
self.register_buffer('wavelength',
t.tensor(wavelength, dtype=dtype))
t.as_tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
dtype=dtype)
self.register_buffer('min_translation',
t.tensor(min_translation, dtype=dtype))
t.as_tensor(min_translation, dtype=dtype))
self.register_buffer('median_propagation',
t.tensor(median_propagation, dtype=dtype))
t.as_tensor(median_propagation, dtype=dtype))
self.register_buffer('obj_basis',
t.tensor(obj_basis, dtype=dtype))
t.as_tensor(obj_basis, dtype=dtype))
if probe_basis is None:
self.register_buffer('probe_basis',
t.tensor(obj_basis, dtype=dtype))
t.as_tensor(obj_basis, dtype=dtype))
else:
self.register_buffer('probe_basis',
t.tensor(probe_basis, dtype=dtype))
t.as_tensor(probe_basis, dtype=dtype))
self.units = units
@@ -119,22 +119,22 @@ class Bragg2DPtycho(CDIModel):
np.array(obj_basis)[:,0])
surface_normal /= np.linalg.norm(surface_normal)
self.register_buffer('surface_normal',
t.tensor(surface_normal, dtype=dtype))
t.as_tensor(surface_normal, dtype=dtype))
if saturation is None:
self.saturation = None
else:
self.register_buffer('saturation',
t.tensor(saturation, dtype=dtype))
t.as_tensor(saturation, dtype=dtype))
if mask is None:
self.mask = None
else:
self.register_buffer('mask',
t.tensor(mask, dtype=t.bool))
t.as_tensor(mask, dtype=t.bool))
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
# We rescale the probe here so it learns at the same rate as the
# object
@@ -150,7 +150,7 @@ class Bragg2DPtycho(CDIModel):
if probe_support is None:
probe_support = t.ones_like(self.probe[0], dtype=t.bool)
self.register_buffer('probe_support',
t.tensor(probe_support, dtype=t.bool))
t.as_tensor(probe_support, dtype=t.bool))
self.probe.data *= self.probe_support
if background is None:
@@ -164,26 +164,26 @@ class Bragg2DPtycho(CDIModel):
self.weights = None
else:
# No incoherent + unstable here yet
self.weights = t.nn.Parameter(t.tensor(weights,
self.weights = t.nn.Parameter(t.as_tensor(weights,
dtype=t.float32))
if translation_offsets is None:
self.translation_offsets = None
else:
t_o = t.tensor(translation_offsets, dtype=t.float32)
t_o = t.as_tensor(translation_offsets, dtype=t.float32)
t_o = t_o / translation_scale
self.translation_offsets = t.nn.Parameter(t_o)
self.register_buffer('translation_scale',
t.tensor(translation_scale, dtype=dtype))
t.as_tensor(translation_scale, dtype=dtype))
self.register_buffer('oversampling',
t.tensor(oversampling, dtype=int))
t.as_tensor(oversampling, dtype=int))
self.register_buffer('propagate_probe',
t.tensor(propagate_probe, dtype=bool))
t.as_tensor(propagate_probe, dtype=bool))
self.register_buffer('correct_tilt',
t.tensor(correct_tilt, dtype=bool))
t.as_tensor(correct_tilt, dtype=bool))
if correct_tilt:
k_map, intensity_map = \
@@ -195,9 +195,9 @@ class Bragg2DPtycho(CDIModel):
self.wavelength,dtype=t.float32,
lens=lens)
self.register_buffer('k_map',
t.tensor(k_map, dtype=dtype))
t.as_tensor(k_map, dtype=dtype))
self.register_buffer('intensity_map',
t.tensor(intensity_map, dtype=dtype))
t.as_tensor(intensity_map, dtype=dtype))
else:
self.k_map = None
@@ -205,7 +205,7 @@ class Bragg2DPtycho(CDIModel):
# The propagation direction of the probe
self.register_buffer('prop_dir',
t.tensor([0, 0, 1], dtype=dtype))
t.as_tensor([0, 0, 1], dtype=dtype))
# This propagator should be able to be multiplied by the propagation
# distance each time to get a propagator
src/cdtools/models/complex_lbfgs.py
-485
View File
@@ -1,485 +0,0 @@
import torch
from functools import reduce
from torch.optim.optimizer import Optimizer
def _cubic_interpolate(x1, f1, g1, x2, f2, g2, bounds=None):
# ported from https://github.com/torch/optim/blob/master/polyinterp.lua
# Compute bounds of interpolation area
if bounds is not None:
xmin_bound, xmax_bound = bounds
else:
xmin_bound, xmax_bound = (x1, x2) if x1 <= x2 else (x2, x1)
# Code for most common case: cubic interpolation of 2 points
# w/ function and derivative values for both
# Solution in this case (where x2 is the farthest point):
# d1 = g1 + g2 - 3*(f1-f2)/(x1-x2);
# d2 = sqrt(d1^2 - g1*g2);
# min_pos = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2));
# t_new = min(max(min_pos,xmin_bound),xmax_bound);
d1 = g1 + g2 - 3 * (f1 - f2) / (x1 - x2)
d2_square = d1**2 - g1 * g2
if d2_square >= 0:
d2 = d2_square.sqrt()
if x1 <= x2:
min_pos = x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2 * d2))
else:
min_pos = x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2 * d2))
return min(max(min_pos, xmin_bound), xmax_bound)
else:
return (xmin_bound + xmax_bound) / 2.
def _strong_wolfe(obj_func,
x,
t,
d,
f,
g,
gtd,
c1=1e-4,
c2=0.9,
tolerance_change=1e-9,
max_ls=25):
# ported from https://github.com/torch/optim/blob/master/lswolfe.lua
d_norm = d.abs().max()
g = g.clone(memory_format=torch.contiguous_format)
# evaluate objective and gradient using initial step
f_new, g_new = obj_func(x, t, d)
ls_func_evals = 1
gtd_new = g_new.dot(d)
# bracket an interval containing a point satisfying the Wolfe criteria
t_prev, f_prev, g_prev, gtd_prev = 0, f, g, gtd
done = False
ls_iter = 0
while ls_iter < max_ls:
# check conditions
if f_new > (f + c1 * t * gtd) or (ls_iter > 1 and f_new >= f_prev):
bracket = [t_prev, t]
bracket_f = [f_prev, f_new]
bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)]
bracket_gtd = [gtd_prev, gtd_new]
break
if abs(gtd_new) <= -c2 * gtd:
bracket = [t]
bracket_f = [f_new]
bracket_g = [g_new]
done = True
break
if gtd_new >= 0:
bracket = [t_prev, t]
bracket_f = [f_prev, f_new]
bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)]
bracket_gtd = [gtd_prev, gtd_new]
break
# interpolate
min_step = t + 0.01 * (t - t_prev)
max_step = t * 10
tmp = t
t = _cubic_interpolate(
t_prev,
f_prev,
gtd_prev,
t,
f_new,
gtd_new,
bounds=(min_step, max_step))
# next step
t_prev = tmp
f_prev = f_new
g_prev = g_new.clone(memory_format=torch.contiguous_format)
gtd_prev = gtd_new
f_new, g_new = obj_func(x, t, d)
ls_func_evals += 1
gtd_new = g_new.dot(d)
ls_iter += 1
# reached max number of iterations?
if ls_iter == max_ls:
bracket = [0, t]
bracket_f = [f, f_new]
bracket_g = [g, g_new]
# zoom phase: we now have a point satisfying the criteria, or
# a bracket around it. We refine the bracket until we find the
# exact point satisfying the criteria
insuf_progress = False
# find high and low points in bracket
low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[-1] else (1, 0)
while not done and ls_iter < max_ls:
# line-search bracket is so small
if abs(bracket[1] - bracket[0]) * d_norm < tolerance_change:
break
# compute new trial value
t = _cubic_interpolate(bracket[0], bracket_f[0], bracket_gtd[0],
bracket[1], bracket_f[1], bracket_gtd[1])
# test that we are making sufficient progress:
# in case `t` is so close to boundary, we mark that we are making
# insufficient progress, and if
# + we have made insufficient progress in the last step, or
# + `t` is at one of the boundary,
# we will move `t` to a position which is `0.1 * len(bracket)`
# away from the nearest boundary point.
eps = 0.1 * (max(bracket) - min(bracket))
if min(max(bracket) - t, t - min(bracket)) < eps:
# interpolation close to boundary
if insuf_progress or t >= max(bracket) or t <= min(bracket):
# evaluate at 0.1 away from boundary
if abs(t - max(bracket)) < abs(t - min(bracket)):
t = max(bracket) - eps
else:
t = min(bracket) + eps
insuf_progress = False
else:
insuf_progress = True
else:
insuf_progress = False
# Evaluate new point
f_new, g_new = obj_func(x, t, d)
ls_func_evals += 1
gtd_new = g_new.dot(d)
ls_iter += 1
if f_new > (f + c1 * t * gtd) or f_new >= bracket_f[low_pos]:
# Armijo condition not satisfied or not lower than lowest point
bracket[high_pos] = t
bracket_f[high_pos] = f_new
bracket_g[high_pos] = g_new.clone(memory_format=torch.contiguous_format)
bracket_gtd[high_pos] = gtd_new
low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[1] else (1, 0)
else:
if abs(gtd_new) <= -c2 * gtd:
# Wolfe conditions satisfied
done = True
elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0:
# old high becomes new low
bracket[high_pos] = bracket[low_pos]
bracket_f[high_pos] = bracket_f[low_pos]
bracket_g[high_pos] = bracket_g[low_pos]
bracket_gtd[high_pos] = bracket_gtd[low_pos]
# new point becomes new low
bracket[low_pos] = t
bracket_f[low_pos] = f_new
bracket_g[low_pos] = g_new.clone(memory_format=torch.contiguous_format)
bracket_gtd[low_pos] = gtd_new
# return stuff
t = bracket[low_pos]
f_new = bracket_f[low_pos]
g_new = bracket_g[low_pos]
return f_new, g_new, t, ls_func_evals
class MyLBFGS(Optimizer):
"""Implements L-BFGS algorithm, heavily inspired by `minFunc
<https://www.cs.ubc.ca/~schmidtm/Software/minFunc.html>`_.
.. warning::
This optimizer doesn't support per-parameter options and parameter
groups (there can be only one).
.. warning::
Right now all parameters have to be on a single device. This will be
improved in the future.
.. note::
This is a very memory intensive optimizer (it requires additional
``param_bytes * (history_size + 1)`` bytes). If it doesn't fit in memory
try reducing the history size, or use a different algorithm.
Args:
lr (float): learning rate (default: 1)
max_iter (int): maximal number of iterations per optimization step
(default: 20)
max_eval (int): maximal number of function evaluations per optimization
step (default: max_iter * 1.25).
tolerance_grad (float): termination tolerance on first order optimality
(default: 1e-5).
tolerance_change (float): termination tolerance on function
value/parameter changes (default: 1e-9).
history_size (int): update history size (default: 100).
line_search_fn (str): either 'strong_wolfe' or None (default: None).
"""
def __init__(self,
params,
lr=1,
max_iter=20,
max_eval=None,
tolerance_grad=1e-7,
tolerance_change=1e-9,
history_size=100,
line_search_fn=None):
if max_eval is None:
max_eval = max_iter * 5 // 4
defaults = dict(
lr=lr,
max_iter=max_iter,
max_eval=max_eval,
tolerance_grad=tolerance_grad,
tolerance_change=tolerance_change,
history_size=history_size,
line_search_fn=line_search_fn)
super(MyLBFGS, self).__init__(params, defaults)
if len(self.param_groups) != 1:
raise ValueError("LBFGS doesn't support per-parameter options "
"(parameter groups)")
self._params = self.param_groups[0]['params']
self._numel_cache = None
def _numel(self):
if self._numel_cache is None:
self._numel_cache = reduce(lambda total, p: total + p.numel(), self._params, 0)
return self._numel_cache
def _gather_flat_grad(self):
views = []
for p in self._params:
if p.grad is None:
view = p.new(p.numel()).zero_()
elif p.grad.is_sparse:
view = p.grad.to_dense().view(-1)
else:
view = p.grad.view(-1)
views.append(view)
return torch.cat(views, 0)
def _add_grad(self, step_size, update):
offset = 0
for p in self._params:
numel = p.numel()
if ((update.dtype==torch.complex64 or update.dtype==torch.complex128)
and (p.dtype==torch.float32 or p.dtype==torch.float64)):
p.add_(update[offset:offset + numel].real.view_as(p), alpha=step_size)
else:
p.add_(update[offset:offset + numel].view_as(p), alpha=step_size)
# view as to avoid deprecated pointwise semantics
#try:
# print(p.dtype)
# print(update.dtype)
# p.add_(update[offset:offset + numel].view_as(p), alpha=step_size)
# print('Worked fine')
# print(update[offset:offset + numel].view_as(p))
#except:
# print('Failed')
# print(update[offset:offset + numel].view_as(p))
# exit()
offset += numel
assert offset == self._numel()
def _clone_param(self):
return [p.clone(memory_format=torch.contiguous_format) for p in self._params]
def _set_param(self, params_data):
for p, pdata in zip(self._params, params_data):
p.copy_(pdata)
def _directional_evaluate(self, closure, x, t, d):
self._add_grad(t, d)
loss = float(closure())
flat_grad = self._gather_flat_grad()
self._set_param(x)
return loss, flat_grad
@torch.no_grad()
def step(self, closure):
"""Performs a single optimization step.
Args:
closure (callable): A closure that reevaluates the model
and returns the loss.
"""
assert len(self.param_groups) == 1
# Make sure the closure is always called with grad enabled
closure = torch.enable_grad()(closure)
group = self.param_groups[0]
lr = group['lr']
max_iter = group['max_iter']
max_eval = group['max_eval']
tolerance_grad = group['tolerance_grad']
tolerance_change = group['tolerance_change']
line_search_fn = group['line_search_fn']
history_size = group['history_size']
# NOTE: LBFGS has only global state, but we register it as state for
# the first param, because this helps with casting in load_state_dict
state = self.state[self._params[0]]
state.setdefault('func_evals', 0)
state.setdefault('n_iter', 0)
# evaluate initial f(x) and df/dx
orig_loss = closure()
loss = float(orig_loss)
current_evals = 1
state['func_evals'] += 1
flat_grad = self._gather_flat_grad()
opt_cond = flat_grad.abs().max() <= tolerance_grad
# optimal condition
if opt_cond:
return orig_loss
# tensors cached in state (for tracing)
d = state.get('d')
t = state.get('t')
old_dirs = state.get('old_dirs')
old_stps = state.get('old_stps')
ro = state.get('ro')
H_diag = state.get('H_diag')
prev_flat_grad = state.get('prev_flat_grad')
prev_loss = state.get('prev_loss')
n_iter = 0
# optimize for a max of max_iter iterations
while n_iter < max_iter:
# keep track of nb of iterations
n_iter += 1
state['n_iter'] += 1
############################################################
# compute gradient descent direction
############################################################
if state['n_iter'] == 1:
d = flat_grad.neg()
old_dirs = []
old_stps = []
ro = []
H_diag = 1
else:
# do lbfgs update (update memory)
y = flat_grad.sub(prev_flat_grad)
s = d.mul(t)
ys = y.dot(s) # y*s
if ys.abs() > 1e-10:
# updating memory
if len(old_dirs) == history_size:
# shift history by one (limited-memory)
old_dirs.pop(0)
old_stps.pop(0)
ro.pop(0)
# store new direction/step
old_dirs.append(y)
old_stps.append(s)
ro.append(1. / ys)
# update scale of initial Hessian approximation
H_diag = ys / y.dot(y) # (y*y)
# compute the approximate (L-BFGS) inverse Hessian
# multiplied by the gradient
num_old = len(old_dirs)
if 'al' not in state:
state['al'] = [None] * history_size
al = state['al']
# iteration in L-BFGS loop collapsed to use just one buffer
q = flat_grad.neg()
for i in range(num_old - 1, -1, -1):
al[i] = old_stps[i].dot(q) * ro[i]
q.add_(old_dirs[i], alpha=-al[i])
# multiply by initial Hessian
# r/d is the final direction
d = r = torch.mul(q, H_diag)
for i in range(num_old):
be_i = old_dirs[i].dot(r) * ro[i]
r.add_(old_stps[i], alpha=al[i] - be_i)
if prev_flat_grad is None:
prev_flat_grad = flat_grad.clone(memory_format=torch.contiguous_format)
else:
prev_flat_grad.copy_(flat_grad)
prev_loss = loss
############################################################
# compute step length
############################################################
# reset initial guess for step size
if state['n_iter'] == 1:
t = min(1., 1. / flat_grad.abs().sum()) * lr
else:
t = lr
# directional derivative
gtd = flat_grad.dot(d) # g * d
# Why did this used to be gtd > -tol_change?
# directional derivative is below tolerance
if gtd.abs() < tolerance_change:
break
# optional line search: user function
ls_func_evals = 0
if line_search_fn is not None:
# perform line search, using user function
if line_search_fn != "strong_wolfe":
raise RuntimeError("only 'strong_wolfe' is supported")
else:
x_init = self._clone_param()
def obj_func(x, t, d):
return self._directional_evaluate(closure, x, t, d)
loss, flat_grad, t, ls_func_evals = _strong_wolfe(
obj_func, x_init, t, d, loss, flat_grad, gtd)
self._add_grad(t, d)
opt_cond = flat_grad.abs().max() <= tolerance_grad
else:
# no line search, simply move with fixed-step
self._add_grad(t, d)
if n_iter != max_iter:
# re-evaluate function only if not in last iteration
# the reason we do this: in a stochastic setting,
# no use to re-evaluate that function here
with torch.enable_grad():
loss = float(closure())
flat_grad = self._gather_flat_grad()
opt_cond = flat_grad.abs().max() <= tolerance_grad
ls_func_evals = 1
# update func eval
current_evals += ls_func_evals
state['func_evals'] += ls_func_evals
############################################################
# check conditions
############################################################
if n_iter == max_iter:
break
if current_evals >= max_eval:
break
# optimal condition
if opt_cond:
break
# lack of progress
if d.mul(t).abs().max() <= tolerance_change:
break
if abs(loss - prev_loss) < tolerance_change:
break
state['d'] = d
state['t'] = t
state['old_dirs'] = old_dirs
state['old_stps'] = old_stps
state['ro'] = ro
state['H_diag'] = H_diag
state['prev_flat_grad'] = prev_flat_grad
state['prev_loss'] = prev_loss
return orig_loss
src/cdtools/models/fancy_ptycho.py
+81 -31
View File
@@ -36,39 +36,47 @@ class FancyPtycho(CDIModel):
units='um',
simulate_probe_translation=False,
simulate_finite_pixels=False,
exponentiate_obj=False,
phase_only=False,
dtype=t.float32,
obj_view_crop=0
):
super(FancyPtycho, self).__init__()
self.register_buffer('wavelength',
t.tensor(wavelength, dtype=dtype))
t.as_tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
dtype=dtype)
self.register_buffer('min_translation',
t.tensor(min_translation, dtype=dtype))
t.as_tensor(min_translation, dtype=dtype))
self.register_buffer('obj_basis',
t.tensor(obj_basis, dtype=dtype))
t.as_tensor(obj_basis, dtype=dtype))
if probe_basis is None:
self.register_buffer('probe_basis',
t.tensor(obj_basis, dtype=dtype))
t.as_tensor(obj_basis, dtype=dtype))
else:
self.register_buffer('probe_basis',
t.tensor(probe_basis, dtype=dtype))
t.as_tensor(probe_basis, dtype=dtype))
self.register_buffer('surface_normal',
t.tensor(surface_normal, dtype=dtype))
t.as_tensor(surface_normal, dtype=dtype))
if saturation is None:
self.saturation = None
else:
self.register_buffer('saturation',
t.tensor(saturation, dtype=dtype))
t.as_tensor(saturation, dtype=dtype))
self.register_buffer('fourier_probe',
t.tensor(fourier_probe, dtype=bool))
t.as_tensor(fourier_probe, dtype=bool))
self.register_buffer('exponentiate_obj',
t.as_tensor(exponentiate_obj, dtype=bool))
self.register_buffer('phase_only',
t.as_tensor(phase_only, dtype=bool))
# Not sure how to make this a buffer...
self.units = units
@@ -77,11 +85,12 @@ class FancyPtycho(CDIModel):
self.mask = None
else:
self.register_buffer('mask',
t.tensor(mask, dtype=t.bool))
t.as_tensor(mask, dtype=t.bool))
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 2:
@@ -112,34 +121,36 @@ class FancyPtycho(CDIModel):
# weights and complex-valued per-mode weight matrices
if len(weights.shape) == 1:
# This is if it's just a list of numbers
self.weights = t.nn.Parameter(t.tensor(weights,
self.weights = t.nn.Parameter(t.as_tensor(weights,
dtype=t.float32))
else:
# Now this is a matrix of weights, so it needs to be complex
self.weights = t.nn.Parameter(t.tensor(weights,
self.weights = t.nn.Parameter(t.as_tensor(weights,
dtype=t.complex64))
if translation_offsets is None:
self.translation_offsets = None
else:
t_o = t.tensor(translation_offsets, dtype=t.float32)
t_o = t.as_tensor(translation_offsets, dtype=t.float32)
t_o = t_o / translation_scale
self.translation_offsets = t.nn.Parameter(t_o)
self.register_buffer('translation_scale',
t.tensor(translation_scale, dtype=dtype))
t.as_tensor(translation_scale, dtype=dtype))
if probe_support is None:
probe_support = t.ones_like(self.probe[0], dtype=t.bool)
self.register_buffer('probe_support',
t.tensor(probe_support, dtype=t.bool))
t.as_tensor(probe_support, dtype=t.bool))
self.probe.data *= self.probe_support
self.register_buffer('oversampling',
t.tensor(oversampling, dtype=int))
t.as_tensor(oversampling, dtype=int))
self.register_buffer('simulate_probe_translation',
t.tensor(simulate_probe_translation, dtype=bool))
self.register_buffer(
'simulate_probe_translation',
t.as_tensor(simulate_probe_translation, dtype=bool)
)
if simulate_probe_translation:
Is = t.arange(self.probe.shape[-2], dtype=dtype)
@@ -153,7 +164,7 @@ class FancyPtycho(CDIModel):
self.register_buffer('simulate_finite_pixels',
t.tensor(simulate_finite_pixels, dtype=bool))
t.as_tensor(simulate_finite_pixels, dtype=bool))
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
@@ -186,6 +197,8 @@ class FancyPtycho(CDIModel):
units='um',
simulate_probe_translation=False,
simulate_finite_pixels=False,
exponentiate_obj=False,
phase_only=False,
obj_view_crop=None,
obj_padding=200,
):
@@ -292,10 +305,16 @@ class FancyPtycho(CDIModel):
probe = t.stack([probe, ] + probe_stack)
obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
obj = (randomize_ang * (t.rand(obj_size)-0.5)).to(dtype=t.complex64)
if not exponentiate_obj:
obj = t.exp(1j * obj)
if n_obj_modes != 1:
obj = t.stack([obj,] + [0.05*t.ones_like(obj),]*(n_obj_modes-1))
if phase_only:
obj.imag[:] = 0
pfc = (probe_fourier_crop if probe_fourier_crop else 0)
if obj_view_crop is None:
obj_view_crop = min(probe.shape[-2], probe.shape[-1]) // 2 + pfc
@@ -363,6 +382,8 @@ class FancyPtycho(CDIModel):
loss=loss, units=units,
simulate_probe_translation=simulate_probe_translation,
simulate_finite_pixels=simulate_finite_pixels,
phase_only=phase_only,
exponentiate_obj=exponentiate_obj,
obj_view_crop=obj_view_crop)
@@ -437,12 +458,23 @@ class FancyPtycho(CDIModel):
prs = t.nn.functional.pad(prs, padding)
prs = tools.propagators.inverse_far_field(prs)
if self.exponentiate_obj:
if self.phase_only:
obj = t.exp(1j*self.obj.real)
else:
obj = t.exp(1j*self.obj)
else:
obj = self.obj
# Now we actually do the interaction, using the sinc subpixel
# translation model as per usual
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj, pix_trans,
shift_probe=True, multiple_modes=True)
prs, obj, pix_trans,
shift_probe=True,
multiple_modes=True,
probe_support=self.probe_support)
return exit_waves
@@ -707,12 +739,12 @@ class FancyPtycho(CDIModel):
('Basis Probe Fourier Space Amplitudes',
lambda self, fig: p.plot_amplitude(
(self.probe if self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
else tools.propagators.far_field(self.probe)),
fig=fig)),
('Basis Probe Fourier Space Phases',
lambda self, fig: p.plot_phase(
('Basis Probe Fourier Space Colorized',
lambda self, fig: p.plot_colorized(
(self.probe if self.fourier_probe
else tools.propagators.inverse_far_field(self.probe))
else tools.propagators.far_field(self.probe))
, fig=fig)),
('Basis Probe Real Space Amplitudes',
lambda self, fig: p.plot_amplitude(
@@ -721,8 +753,8 @@ class FancyPtycho(CDIModel):
fig=fig,
basis=self.probe_basis,
units=self.units)),
('Basis Probe Real Space Phases',
lambda self, fig: p.plot_phase(
('Basis Probe Real Space Colorized',
lambda self, fig: p.plot_colorized(
(self.probe if not self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig,
@@ -750,13 +782,31 @@ class FancyPtycho(CDIModel):
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units)),
units=self.units),
lambda self: not self.exponentiate_obj),
('Object Phase',
lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units)),
units=self.units),
lambda self: not self.exponentiate_obj),
('Real Part of T',
lambda self, fig: p.plot_real(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units,
cmap='cividis'),
lambda self: self.exponentiate_obj),
('Imaginary Part of T',
lambda self, fig: p.plot_imag(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: self.exponentiate_obj),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
src/cdtools/models/fastccd_ptycho.py
-691
View File
@@ -1,691 +0,0 @@
import torch as t
from cdtools.models import CDIModel
from cdtools.datasets import Ptycho2DDataset
from cdtools import tools
from cdtools.tools import plotting as p
from cdtools.tools import analysis
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from scipy import linalg as sla
from copy import copy
__all__ = ['FastCCDPtycho']
class FastCCDPtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess,
obj_guess,
detector_slice=None,
threshold=None,
surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
min_translation=t.tensor([0, 0], dtype=t.float32),
background=None,
translation_offsets=None,
mask=None,
weights=None,
translation_scale=1,
saturation=None,
probe_support=None,
oversampling=1,
fourier_probe=False,
loss='amplitude mse',
units='um',
simulate_probe_translation=False,
simulate_finite_pixels=False,
):
super(FastCCDPtycho, self).__init__()
self.wavelength = t.tensor(wavelength)
self.detector_geometry = copy(detector_geometry)
det_geo = self.detector_geometry
if 'distance' in det_geo:
det_geo['distance'] = t.tensor(det_geo['distance'], dtype=t.float32)
if 'basis' in det_geo:
det_geo['basis'] = t.tensor(det_geo['basis'], dtype=t.float32)
if 'corner' in det_geo and det_geo['corner'] is not None:
det_geo['corner'] = t.tensor(det_geo['corner'], dtype=t.float32)
self.min_translation = t.tensor(min_translation)
self.probe_basis = t.tensor(probe_basis)
self.detector_slice = copy(detector_slice)
self.surface_normal = t.tensor(surface_normal)
self.saturation = saturation
self.units = units
self.fourier_probe = fourier_probe
if mask is None:
self.mask = mask
else:
self.mask = t.tensor(mask, dtype=t.bool)
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 2:
self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
else:
self.probe_norm = 1 * t.max(t.abs(probe_guess))
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
if background is None:
if detector_slice is not None:
dummy_det = t.empty([s//oversampling
for s in self.probe[0].shape])
shape = dummy_det[self.detector_slice].shape
#shape = self.probe[0][self.detector_slice].shape
else:
shape = [s//oversampling for s in self.probe[0]]
background = 1e-6 * t.ones(shape, dtype=t.float32)
self.background = t.nn.Parameter(background)
if threshold is None:
threshold = t.zeros([960,2], dtype=t.float32)
self.threshold = t.nn.Parameter(threshold)
if weights is None:
self.weights = None
else:
# We now need to distinguish between real-valued per-image
# weights and complex-valued per-mode weight matrices
if len(weights.shape) == 1:
# This is if it's just a list of numbers
self.weights = t.nn.Parameter(t.tensor(weights,
dtype=t.float32))
else:
# Now this is a matrix of weights, so it needs to be complex
self.weights = t.nn.Parameter(t.tensor(weights,
dtype=t.complex64))
if translation_offsets is None:
self.translation_offsets = None
else:
t_o = t.tensor(translation_offsets, dtype=t.float32)
t_o = t_o / translation_scale
self.translation_offsets = t.nn.Parameter(t_o)
self.translation_scale = translation_scale
if probe_support is not None:
self.probe_support = probe_support
else:
self.probe_support = t.ones_like(self.probe[0], dtype=t.bool)
self.oversampling = oversampling
self.simulate_probe_translation = simulate_probe_translation
if simulate_probe_translation:
Is = t.arange(self.probe.shape[-2], dtype=t.float32)
Js = t.arange(self.probe.shape[-1], dtype=t.float32)
Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
self.I_phase = 2 * np.pi* Is
self.J_phase = 2 * np.pi* Js
self.simulate_finite_pixels = simulate_finite_pixels
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
or loss.lower().strip() == 'amplitude_mse'):
self.loss = tools.losses.amplitude_mse
elif (loss.lower().strip() == 'poisson nll'
or loss.lower().strip() == 'poisson_nll'):
self.loss = tools.losses.poisson_nll
else:
raise KeyError('Specified loss function not supported')
@classmethod
def from_dataset(cls,
dataset,
probe_size=None,
randomize_ang=0,
padding=0,
n_modes=1,
dm_rank=None,
translation_scale=1,
saturation=None,
probe_support_radius=None,
propagation_distance=None,
scattering_mode=None,
oversampling=1,
auto_center=False,
fourier_probe=False,
loss='amplitude mse',
units='um',
simulate_probe_translation=False,
simulate_finite_pixels=False,
):
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')
# We include the *extras to make this work even with datasets, like
# polarization dependent datasets, that might toss out extra inputs
(indices, translations, *extras), patterns = dataset[:]
dataset.get_as(*get_as_args[0], **get_as_args[1])
# Set to none to avoid issues with things outside the detector
if auto_center:
center = tools.image_processing.centroid(t.sum(patterns, dim=0))
else:
center = None
# 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,
oversampling=oversampling)
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.])
# 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(surface_normal)
# 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)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# 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, propagation_distance=propagation_distance, oversampling=oversampling)
else:
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
# Now we initialize all the subdominant probe modes
probe_max = t.max(t.abs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
# For a Fourier space probe
if fourier_probe:
probe = tools.propagators.far_field(probe)
probe = t.stack([probe, ] + probe_stack)
obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
if dm_rank is not None and dm_rank != 0:
if dm_rank > n_modes:
raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.')
elif dm_rank == -1:
# dm_rank == -1 is defined to mean full-rank
dm_rank = n_modes
Ws = t.zeros(len(dataset), dm_rank, n_modes, dtype=t.complex64)
# Start with as close to the identity matrix as possible,
# cutting of when we hit the specified maximum rank
for i in range(0, dm_rank):
Ws[:, i, i] = 1
else:
# dm_rank == None or dm_rank = 0 triggers a special case where
# a standard incoherent multi-mode model is used. This is the
# default, because it is so common.
# In this case, we define a set of weights which only has one index
Ws = t.ones(len(dataset))
if hasattr(dataset, 'intensities') and dataset.intensities is not None:
Ws *= (dataset.intensities.to(dtype=Ws.dtype)[:,...]
/ t.mean(dataset.intensities))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
else:
mask = None
if probe_support_radius is not None:
probe_support = t.zeros(probe[0].shape, dtype=t.bool)
xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
Rs = np.sqrt(xs**2 + ys**2)
probe_support[Rs < probe_support_radius] = 1
probe = probe * probe_support[None, :, :]
else:
probe_support = None
return cls(wavelength, det_geo, probe_basis, probe, obj,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets=translation_offsets,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
fourier_probe=fourier_probe,
oversampling=oversampling,
loss=loss, units=units,
simulate_probe_translation=simulate_probe_translation,
simulate_finite_pixels=simulate_finite_pixels)
def interaction(self, index, translations, *args):
# The *args is included so that this can work even when given, say,
# a polarized ptycho dataset that might spit out more inputs.
# Step 1 is to convert the translations for each position into a
# value in pixels
pix_trans = tools.interactions.translations_to_pixel(
self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
# We then add on any recovered translation offset, if they exist
if self.translation_offsets is not None:
pix_trans += (self.translation_scale *
self.translation_offsets[index])
# This restricts the basis probes to stay within the probe support
basis_prs = self.probe * self.probe_support[..., :, :]
# For a Fourier-space probe, we take an IFT
if self.fourier_probe:
basis_prs = tools.propagators.inverse_far_field(basis_prs)
# Now we construct the probes for each shot from the basis probes
if self.weights is not None:
Ws = self.weights[index]
else:
try:
Ws = t.ones(len(index)) # I'm positive this introduced a bug
except:
Ws = 1
if self.weights is None or len(self.weights[0].shape) == 0:
# If a purely stable coherent illumination is defined
prs = Ws[..., None, None, None] * basis_prs
else:
# If a frame-by-frame weight matrix is defined
# This takes the dot product of all the weight matrices with
# the probes. The output has dimensions of translation, then
# coherent mode index, then x,y, and then complex index
# Maybe this can be done with a matmul now?
prs = t.sum(Ws[..., None, None] * basis_prs, axis=-3)
if self.simulate_probe_translation:
det_pix_trans = tools.interactions.translations_to_pixel(
self.detector_geometry['basis'],
translations,
surface_normal=self.surface_normal)
probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] *
self.I_phase[None,...] +
det_pix_trans[:,1,None,None] *
self.J_phase[None,...]))
prs = prs * probe_masks[...,None,:,:]
# Now we actually do the interaction, using the sinc subpixel
# translation model as per usual
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj, pix_trans,
shift_probe=True, multiple_modes=True)
return 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):
measured = tools.measurements.quadratic_background(
wavefields,
self.background,
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling,
simulate_finite_pixels=self.simulate_finite_pixels)
thresholds = t.repeat_interleave(self.threshold, 480,dim=-1)
return t.clamp(measured - thresholds, min=0.001)
# Note: No "loss" function is defined here, because it is added
# dynamically during object creation in __init__
def to(self, *args, **kwargs):
super(FastCCDPtycho, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args, **kwargs)
# move the detector geometry too
det_geo = self.detector_geometry
if 'distance' in det_geo:
det_geo['distance'] = det_geo['distance'].to(*args, **kwargs)
if 'basis' in det_geo:
det_geo['basis'] = det_geo['basis'].to(*args, **kwargs)
if 'corner' in det_geo and det_geo['corner'] is not None:
det_geo['corner'] = det_geo['corner'].to(*args, **kwargs)
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
if self.simulate_probe_translation:
self.I_phase = self.I_phase.to(*args, **kwargs)
self.J_phase = self.J_phase.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.probe_support = self.probe_support.to(*args, **kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list, calculation_width=None):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'cdtools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
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
wavelength = self.wavelength
indices, translations = args_list
data = []
len(indices)
if calculation_width is None:
calculation_width = len(indices)
index_chunks = [indices[i:i + calculation_width]
for i in range(0, len(indices),
calculation_width)]
translation_chunks = [translations[i:i + calculation_width]
for i in range(0, len(indices),
calculation_width)]
# Then we simulate the results
data = [self.forward(idx, trans).detach()
for idx, trans in zip(index_chunks, translation_chunks)]
data = t.cat(data, dim=0)
# And finally, we make the dataset
return Ptycho2DDataset(
translations, data,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def corrected_translations(self, dataset):
translations = dataset.translations.to(
dtype=t.float32, device=self.probe.device)
if (hasattr(self, 'translation_offsets') and
self.translation_offsets is not None):
t_offset = tools.interactions.pixel_to_translations(
self.probe_basis,
self.translation_offsets * self.translation_scale,
surface_normal=self.surface_normal)
return translations + t_offset
else:
return translations
def get_rhos(self):
# If this is the general unified mode model
if self.weights.dim() >= 2:
Ws = self.weights.detach().cpu().numpy()
rhos_out = np.matmul(np.swapaxes(Ws, 1, 2), Ws.conj())
return rhos_out
# This is the purely incoherent case
else:
return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
dtype=np.complex64)
def tidy_probes(self):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
As a first step, we calculate the state of the light field across the
full experiment, using the weight matrices and basis probes. Then, we
use an SVD to update the basis probes so they form an eigenbasis of
the implied density matrix for the full experiment.
Next, the weight matrices for each shot are recalculated so that the
probes generated by weights * basis_probes for each shot are themselves
an eigenbasis for that individual shot's density matrix.
"""
# First we treat the incoherent but stable case, where the weights are
# just one per-shot overall weight
if self.weights.dim() == 1:
probe = self.probe.detach().cpu().numpy()
ortho_probes = analysis.orthogonalize_probes(self.probe.detach())
self.probe.data = ortho_probes
return
# What follows is for the unified OPRP and incoherent multi-mode model,
# where each shot has it's own matrix of weights such that the probe
# state for each shot is self.weights @ self.probe
# We concatenate all the weight matrices, to come up with a state
# corresponding to the summed light field across all the exposures.
# This state will have a large number of modes, but all built from
# the same small number of basis modes
all_weights = t.cat(t.unbind(self.weights.detach(), dim=0), dim=0)
# We generate the orthogonal probes based on this full-experiment
# representation of the light field.
ortho_probes, reexpressed_weights = \
analysis.orthogonalize_probes(
self.probe.detach(),
weight_matrix=all_weights,
return_reexpressed_weights=True
)
# We just orthogonalized the incoherent sum of all the exposures
# across the full experiment, so the output probes are normalized so
# that their intensity matches the summed intensity across the full
# experiment. We divide their amplitudes by the square root of the
# number of shots so that we now have a set of probes corresponding
# to the mean shot
ortho_probes /= np.sqrt(self.weights.shape[0])
reexpressed_weights *= np.sqrt(self.weights.shape[0])
# We now replace the shot-to-shot weights with the versions that have
# been re-expressed in the new basis.
new_weights = t.stack(t.split(reexpressed_weights,
self.weights.shape[1]), dim=0)
# And we save it back to the model
self.probe.data = ortho_probes.to(
device=self.probe.device, dtype=self.probe.dtype)
self.weights.data = new_weights.to(
device=self.weights.device, dtype=self.weights.dtype)
# NOTE: I used to have this part as an option, with "tidy_each_frame",
# because it took such a long time. Now that I've rewritten it properly,
# it's quite fast and so I removed the kwarg because there's really
# no situation where you woudn't want to do this.
# Now, we seek to edit the shot-to-shot weight matrices such that
# self.weights[i] @ self.probes will be properly orthogonalized for
# all i.
# All we need to know about the probes is that they are orthogonalized
# and the intensity within each probe mode
probe_sqrt_intensities = t.linalg.norm(self.probe.data, dim=(-2,-1))
# This does a super fast batched computation
U, S, Vh = t.linalg.svd(self.weights.data * probe_sqrt_intensities,
full_matrices=False)
# We discard the U matrix and re-multiply S & Vh
self.weights.data = S[:,:,None] * (Vh / probe_sqrt_intensities)
def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs):
def get_probes(idx):
basis_prs = self.probe * self.probe_support[..., :, :]
prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs,
axis=-3)
ortho_probes = analysis.orthogonalize_probes(prs)
if mode.lower() == 'amplitude':
return np.abs(ortho_probes.detach().cpu().numpy())
if mode.lower() == 'root_sum_intensity':
return np.sum(np.abs(ortho_probes.detach().cpu().numpy())**2,
axis=0)
if mode.lower() == 'phase':
return np.angle(ortho_probes.detach().cpu().numpy())
probe_matrix = np.zeros([self.probe.shape[0]]*2,
dtype=np.complex64)
np_probes = self.probe.detach().cpu().numpy()
for i in range(probe_matrix.shape[0]):
for j in range(probe_matrix.shape[0]):
probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj())
weights = self.weights.detach().cpu().numpy()
probe_intensities = np.sum(np.tensordot(weights, probe_matrix, axes=1)
* weights.conj(), axis=2)
# Imaginary part is already essentially zero up to rounding error
probe_intensities = np.real(probe_intensities)
values = np.sum(probe_intensities, axis=1)
if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity':
cmap = 'viridis'
else:
cmap = 'twilight'
p.plot_nanomap_with_images(self.corrected_translations(dataset), get_probes, values=values, fig=fig, units=self.units, basis=self.probe_basis, nanomap_colorbar_title='Total Probe Intensity', cmap=cmap, **kwargs),
plot_list = [
('',
lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', image_colorbar_title='Square Root of Intensity'),
lambda self: len(self.weights.shape) >= 2),
('',
lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', image_colorbar_title='Probe Amplitude'),
lambda self: len(self.weights.shape) >= 2),
('',
lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', image_colorbar_title='Probe Phase'),
lambda self: len(self.weights.shape) >= 2),
('Basis Probe Fourier Space Amplitudes',
lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Basis Probe Fourier Space Phases',
lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Basis Probe Real Space Amplitudes',
lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Basis Probe Real Space Phases',
lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Average Weight Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(
np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
fig=fig),
lambda self: len(self.weights.shape) >= 2),
('% of Power in Top Mode',
lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
100 * t.stack([
analysis.calc_mode_power_fractions(
self.probe.data,
weight_matrix=self.weights.data[i])[0]
for i in range(self.weights.shape[0])
], dim=0),
fig=fig,
units=self.units),
lambda self: len(self.weights.shape) >= 2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)),
('Thresholds',
lambda self, fig: p.plot_real(t.repeat_interleave(self.threshold, 480,dim=-1), fig=fig))
]
# def plot_errors(self, dataset):
def save_results(self, dataset):
thresholds = self.thresholds.detach().cpu().numpy()
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
oversampling = self.oversampling
return {'basis': basis, 'translation': translations,
'probe': probe, 'obj': obj,
'background': background,
'oversampling': oversampling,
'weights': weights, 'thresholds':thresholds}
src/cdtools/models/multimode_rpi.py
-354
View File
@@ -1,354 +0,0 @@
import torch as t
from cdtools.models import CDIModel
from cdtools import tools
from cdtools.tools import plotting as p
from cdtools.tools.interactions import RPI_interaction
from cdtools.tools import initializers
from scipy.ndimage import binary_dilation
import numpy as np
from copy import copy
__all__ = ['MultimodeRPI']
class MultimodeRPI(CDIModel):
@property
def obj(self):
return t.complex(self.obj_real, self.obj_imag)
@property
def weights(self):
ws = t.complex(self.weights_real, self.weights_imag)
return ws / 10# / self.obj_real.size().numel()
def __init__(self, wavelength, detector_geometry, probe_basis,
probe, obj_guess, detector_slice=None,
background=None, mask=None, saturation=None,
obj_support=None, oversampling=1, weight_matrix=False):
super(MultimodeRPI, 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.probe_basis = t.tensor(probe_basis)
scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1],
probe.shape[-2]/obj_guess.shape[-2]])
self.obj_basis = self.probe_basis * scale_factor
self.detector_slice = detector_slice
# Maybe something to include in a bit
# self.surface_normal = t.tensor(surface_normal)
self.saturation = saturation
if mask is None:
self.mask = mask
else:
self.mask = t.tensor(mask, dtype=t.bool)
self.probe = t.tensor(probe, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
self.obj_real = t.nn.Parameter(obj_guess.real)
self.obj_imag = t.nn.Parameter(obj_guess.imag)
self.weights_real = t.nn.Parameter(t.eye(probe.shape[0])* 10)# * self.obj_real.size().numel())
self.weights_imag = t.nn.Parameter(t.zeros(probe.shape[0]))
if not weight_matrix:
self.weights_real.requires_grad=False
self.weights_imag.requires_grad=False
# Wait for LBFGS to be updated for complex-valued parameters
# self.obj = t.nn.Parameter(obj_guess.to(t.float32))
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(
self.probe[0][self.detector_slice].shape,
dtype=t.float32)
else:
background = 1e-6 * t.ones(self.probe[0].shape,
dtype=t.float32)
self.background = t.tensor(background, dtype=t.float32)
if obj_support is not None:
self.obj_support = obj_support
self.obj.data = self.obj * obj_support[None, ...]
else:
self.obj_support = t.ones_like(self.obj[0, ...])
self.oversampling = oversampling
@classmethod
def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', weight_matrix=False, probe_threshold=0):
raise NotImplementedError()
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')
# We only need the patterns here, not the inputs associated with them.
_, patterns = dataset[:]
dataset.get_as(*get_as_args[0],**get_as_args[1])
# Set to none to avoid issues with things outside the detector
if auto_center:
center = tools.image_processing.centroid(t.sum(patterns,dim=0))
else:
center = None
# 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,
oversampling=oversampling)
if not isinstance(probe,t.Tensor):
probe = t.as_tensor(probe)
# Potentially need all of this orientation stuff later
#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.])
# 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(surface_normal)
if background is None and hasattr(dataset, 'background') \
and dataset.background is not None:
background = t.sqrt(dataset.background)
elif background is not None:
background = t.sqrt(t.Tensor(background).to(dtype=t.float32))
det_geo = dataset.detector_geometry
# If no mask is given, but one exists in the dataset, load it.
if mask is None and hasattr(dataset, 'mask') \
and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
# Now we initialize the object
if obj_size is None:
# This is a standard size for a well-matched probe and detector
obj_size = (np.array(probe_shape) // 2).astype(int)
if initialization.lower().strip() == 'random':
# I think something to do with the fact that the object is defined
# on a coarser grid needs to be accounted for here that is not
# accounted for yet
scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2)
obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
elif initialization.lower().strip() == 'spectral':
if background is not None:
obj_guess = initializers.RPI_spectral_init(
patterns[0], probe, obj_size, mask=mask,
background=background**2, n_modes=n_modes)
else:
obj_guess = initializers.RPI_spectral_init(
patterns[0], probe, obj_size, mask=mask,
n_modes=n_modes)
else:
raise KeyError('Initialization "' + str(initialization) + \
'" invalid - use "spectral" or "random"')
probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0))
probe_fft = tools.propagators.far_field(probe_intensity)
pad0l = (probe.shape[-2] - obj_size[-2])//2
pad0r = probe.shape[-2] - obj_size[-2] - pad0l
pad1l = (probe.shape[-1] - obj_size[-1])//2
pad1r = probe.shape[-1] - obj_size[-1] - pad1l
probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r]
probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft))
obj_support = probe_lr > t.max(probe_lr) * probe_threshold
obj_support = t.as_tensor(binary_dilation(obj_support))
return cls(wavelength, det_geo, probe_basis,
probe, obj_guess, detector_slice=det_slice,
background=background, mask=mask, saturation=saturation,
obj_support=obj_support, oversampling=oversampling,
weight_matrix=weight_matrix)
def random_init(self, pattern):
scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2)
self.obj.data = scale * t.exp(
2j * np.pi * t.rand(self.obj.shape)).to(
dtype=self.obj.dtype, device=self.obj.device)
def spectral_init(self, pattern):
if self.background is not None:
self.obj.data = initializers.RPI_spectral_init(
pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask,
background=self.background**2, n_modes=self.obj.shape[0]).to(
dtype=self.obj.dtype, device=self.obj.device)
else:
self.obj.data = initializers.RPI_spectral_init(
pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask,
n_modes=self.obj.shape[0]).to(
dtype=self.obj.dtype, device=self.obj.device)
# Needs work
def interaction(self, index, *args):
# including *args allows this to work with all sorts of datasets
# that might include other information in with the index in their
# "input" parameters (such as translations for a ptychography dataset).
# This makes it seamless to use such a dataset even though those
# extra arguments will not be used.
all_exit_waves = []
# Mix the probes with the weight matrix
prs = t.sum(self.weights[..., None, None] * self.probe, axis=-3)
for i in range(self.probe.shape[0]):
pr = prs[i]
# Here we have a 3D probe (one single mode)
# and a 4D object (multiple modes mixing incoherently)
exit_waves = RPI_interaction(pr,
self.obj_support * self.obj[i])
all_exit_waves.append(exit_waves.unsqueeze(0))
# This creates a bunch of modes generated from all possible combos
# of the probe and object modes all strung out along the first index
output = t.cat(all_exit_waves)
# If we have multiple indexes input, we unsqueeze and repeat the stack
# of wavefields enough times to simulate each requested index. This
# seems silly, but it enables (for example) one to do a reconstruction
# from a set of diffraction patterns that are all known to be from the
# same object.
try:
# will fail if index has no length, for example when index
# is just an int. In this case, we just do nothing instead
output = output.unsqueeze(0).repeat(1,len(index),1,1,1)
except TypeError:
pass
return output
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):
# Here I'm taking advantage of an undocumented feature in the
# incoherent_sum measurement function where it will work with
# a 4D wavefield array as well as a 5D array.
return tools.measurements.quadratic_background(wavefields,
self.background,
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
def regularizer(self, factors):
return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \
+ factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2)
def to(self, *args, **kwargs):
super(MultimodeRPI, 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.probe = self.probe.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.obj_basis = self.obj_basis.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
self.background = self.background.to(*args, **kwargs)
# Maybe include in a bit
#self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
raise NotImplementedError('No sim to dataset yet, sorry!')
plot_list = [
('Root Sum Squared Amplitude of all Probes',
lambda self, fig: p.plot_amplitude(
np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)),
fig=fig, basis=self.probe_basis)),
('Object Amplitudes',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig,
basis=self.obj_basis)),
('Object Phases',
lambda self, fig: p.plot_phase(self.obj, fig=fig,
basis=self.obj_basis))
]
def save_results(self, dataset=None, full_obj=False):
# dataset is set as a kwarg here because it isn't needed, but the
# common pattern is to pass a dataset. This makes it okay if one
# continues to use that standard pattern
probe_basis = self.probe_basis.detach().cpu().numpy()
obj_basis = self.obj_basis.detach().cpu().numpy()
probe = self.probe.detach().cpu().numpy()
# Provide the option to save out the subdominant objects or
# just the dominant one
if full_obj:
obj = self.obj.detach().cpu().numpy()
else:
obj = self.obj[0].detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
'probe': probe,'obj': obj,
'background': background}
src/cdtools/models/multislice_ptycho.py
+24 -22
View File
@@ -44,40 +44,40 @@ class MultislicePtycho(CDIModel):
super(MultislicePtycho, self).__init__()
self.register_buffer('wavelength',
t.tensor(wavelength, dtype=dtype))
t.as_tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
dtype=dtype)
self.register_buffer('min_translation',
t.tensor(min_translation, dtype=dtype))
t.as_tensor(min_translation, dtype=dtype))
self.register_buffer('obj_basis',
t.tensor(obj_basis, dtype=dtype))
t.as_tensor(obj_basis, dtype=dtype))
self.register_buffer('exponentiate_obj',
t.tensor(exponentiate_obj, dtype=bool))
t.as_tensor(exponentiate_obj, dtype=bool))
self.register_buffer('interslice_propagator',
t.tensor(interslice_propagator, dtype=t.complex64))
t.as_tensor(interslice_propagator, dtype=t.complex64))
if probe_basis is None:
self.register_buffer('probe_basis',
t.tensor(obj_basis, dtype=dtype))
t.as_tensor(obj_basis, dtype=dtype))
else:
self.register_buffer('probe_basis',
t.tensor(probe_basis, dtype=dtype))
t.as_tensor(probe_basis, dtype=dtype))
self.register_buffer('surface_normal',
t.tensor(surface_normal, dtype=dtype))
t.as_tensor(surface_normal, dtype=dtype))
if saturation is None:
self.saturation = None
else:
self.register_buffer('saturation',
t.tensor(saturation, dtype=dtype))
t.as_tensor(saturation, dtype=dtype))
self.register_buffer('fourier_probe',
t.tensor(fourier_probe, dtype=bool))
t.as_tensor(fourier_probe, dtype=bool))
# Not sure how to make this a buffer...
self.units = units
@@ -86,10 +86,10 @@ class MultislicePtycho(CDIModel):
self.mask = None
else:
self.register_buffer('mask',
t.tensor(mask, dtype=t.bool))
t.as_tensor(mask, dtype=t.bool))
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
# We rescale the probe here so it learns at the same rate as the
# object
@@ -121,34 +121,36 @@ class MultislicePtycho(CDIModel):
# weights and complex-valued per-mode weight matrices
if len(weights.shape) == 1:
# This is if it's just a list of numbers
self.weights = t.nn.Parameter(t.tensor(weights,
self.weights = t.nn.Parameter(t.as_tensor(weights,
dtype=t.float32))
else:
# Now this is a matrix of weights, so it needs to be complex
self.weights = t.nn.Parameter(t.tensor(weights,
self.weights = t.nn.Parameter(t.as_tensor(weights,
dtype=t.complex64))
if translation_offsets is None:
self.translation_offsets = None
else:
t_o = t.tensor(translation_offsets, dtype=t.float32)
t_o = t.as_tensor(translation_offsets, dtype=t.float32)
t_o = t_o / translation_scale
self.translation_offsets = t.nn.Parameter(t_o)
self.register_buffer('translation_scale',
t.tensor(translation_scale, dtype=dtype))
t.as_tensor(translation_scale, dtype=dtype))
if probe_support is None:
probe_support = t.ones_like(self.probe[0], dtype=t.bool)
self.register_buffer('probe_support',
t.tensor(probe_support, dtype=t.bool))
t.as_tensor(probe_support, dtype=t.bool))
self.probe.data *= self.probe_support
self.register_buffer('oversampling',
t.tensor(oversampling, dtype=int))
t.as_tensor(oversampling, dtype=int))
self.register_buffer('simulate_probe_translation',
t.tensor(simulate_probe_translation, dtype=bool))
self.register_buffer(
'simulate_probe_translation',
t.as_tensor(simulate_probe_translation, dtype=bool)
)
if simulate_probe_translation:
Is = t.arange(self.probe.shape[-2], dtype=dtype)
@@ -161,7 +163,7 @@ class MultislicePtycho(CDIModel):
self.register_buffer('J_phase', J_phase)
self.register_buffer('simulate_finite_pixels',
t.tensor(simulate_finite_pixels, dtype=bool))
t.as_tensor(simulate_finite_pixels, dtype=bool))
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
src/cdtools/models/polarization_swept_ptycho.py
-605
View File
@@ -1,605 +0,0 @@
import torch as t
from cdtools.models import CDIModel
from cdtools.datasets import Ptycho2DDataset
from cdtools import tools
from cdtools.tools import plotting as p
from cdtools.tools import analysis
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from scipy import linalg as sla
from copy import copy
__all__ = ['FancyPtycho']
class PolarizationSweptPtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess,
obj_guess,
polarization_states,
detector_slice=None,
surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
min_translation=t.tensor([0, 0], dtype=t.float32),
background=None,
translation_offsets=None,
mask=None,
weights=None,
translation_scale=1,
saturation=None,
probe_support=None,
oversampling=1,
fourier_probe=False,
loss='amplitude mse',
units='um',
simulate_probe_translation=False,
simulate_finite_pixels=False,
dtype=t.float32,
obj_view_crop=0
):
super(PolarizationSweptPtycho, self).__init__()
self.register_buffer('wavelength',
t.tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
dtype=dtype)
self.register_buffer('min_translation',
t.tensor(min_translation, dtype=dtype))
self.register_buffer('probe_basis',
t.tensor(probe_basis, dtype=dtype))
self.detector_slice = copy(detector_slice)
self.register_buffer('surface_normal',
t.tensor(surface_normal, dtype=dtype))
if saturation is None:
self.saturation = None
else:
self.register_buffer('saturation',
t.tensor(saturation, dtype=dtype))
# Not sure how to make this a buffer...
self.units = units
self.fourier_probe = fourier_probe
if mask is None:
self.mask = None
else:
self.register_buffer('mask',
t.tensor(mask, dtype=t.bool))
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 2:
probe_norm = 1 * t.max(t.abs(probe_guess[0]))
else:
probe_norm = 1 * t.max(t.abs(probe_guess))
self.register_buffer('probe_norm', probe_norm.to(dtype))
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
self.obj_view_slice = np.s_[obj_view_crop:-obj_view_crop,
obj_view_crop:-obj_view_crop]
if background is None:
if detector_slice is not None:
dummy_det = t.empty([s//oversampling
for s in self.probe.shape[-2:]])
shape = dummy_det[self.detector_slice].shape
#shape = self.probe[0][self.detector_slice].shape
else:
shape = [s//oversampling for s in self.probe.shape[-2:]]
background = 1e-6 * t.ones(shape, dtype=t.float32)
self.background = t.nn.Parameter(background)
if weights is None:
self.weights = None
else:
self.weights = t.nn.Parameter(t.tensor(weights,
dtype=t.float32))
if translation_offsets is None:
self.translation_offsets = None
else:
t_o = t.tensor(translation_offsets, dtype=t.float32)
t_o = t_o / translation_scale
self.translation_offsets = t.nn.Parameter(t_o)
self.register_buffer('translation_scale',
t.tensor(translation_scale, dtype=dtype))
if probe_support is None:
probe_support = t.ones(self.probe.shape[-2:], dtype=t.bool)
self.register_buffer('probe_support',
t.tensor(probe_support, dtype=t.bool))
self.oversampling = oversampling
self.simulate_probe_translation = simulate_probe_translation
if simulate_probe_translation:
Is = t.arange(self.probe.shape[-2], dtype=dtype)
Js = t.arange(self.probe.shape[-1], dtype=dtype)
Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
I_phase = 2 * np.pi* Is * self.oversampling
J_phase = 2 * np.pi* Js * self.oversampling
self.register_buffer('I_phase', I_phase)
self.register_buffer('J_phase', J_phase)
self.simulate_finite_pixels = simulate_finite_pixels
self.polarization_states = t.nn.Parameter(
t.tensor(polarization_states, dtype=t.complex64))
# by default, don't optimize this, but I think it might be
# interesting to try it because the polarization states
# are not very pure
self.polarization_states.requires_grad = False
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
or loss.lower().strip() == 'amplitude_mse'):
self.loss = tools.losses.amplitude_mse
elif (loss.lower().strip() == 'poisson nll'
or loss.lower().strip() == 'poisson_nll'):
self.loss = tools.losses.poisson_nll
else:
raise KeyError('Specified loss function not supported')
@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,
probe_fourier_crop=None,
propagation_distance=None,
scattering_mode=None,
oversampling=1,
auto_center=False,
fourier_probe=False,
loss='amplitude mse',
units='um',
simulate_probe_translation=False,
simulate_finite_pixels=False,
obj_view_crop=None,
obj_padding=200
):
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')
# We include the *extras to make this work even with datasets, like
# polarization dependent datasets, that might toss out extra inputs
((indices, translations, *extras), patterns) = dataset[:]
polarization_states = dataset.polarization_states
dataset.get_as(*get_as_args[0], **get_as_args[1])
# Set to none to avoid issues with things outside the detector
if auto_center:
center = tools.image_processing.centroid(t.sum(patterns, dim=0))
else:
center = None
# 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,
oversampling=oversampling)
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.])
# 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(surface_normal)
# 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)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=obj_padding)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# 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, propagation_distance=propagation_distance, oversampling=oversampling)
else:
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
if probe_fourier_crop is not None:
probe = tools.propagators.far_field(probe)
probe = probe[probe_fourier_crop:-probe_fourier_crop,
probe_fourier_crop:-probe_fourier_crop]
probe = tools.propagators.inverse_far_field(probe)
# Now we initialize all the subdominant probe modes
probe_max = t.max(t.abs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
# For a Fourier space probe
if fourier_probe:
probe = tools.propagators.far_field(probe)
probe = t.stack([probe, ] + probe_stack)
# The probe gets one extra dimension, so we can simulate one probe
# per polarization state
n_states = polarization_states.shape[0]
probe = t.stack([probe] * n_states, dim=0)
# Looks like an identity matrix
obj_base = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
obj_top = t.stack([obj_base, obj_base*0], dim=0)
obj_bottom = t.stack([obj_base*0, obj_base], dim=0)
obj = t.stack([obj_top, obj_bottom], dim=0)
if obj_view_crop is None:
obj_view_crop = min(probe.shape[-2], probe.shape[-1]) // 2
if obj_view_crop < 0:
obj_view_crop += min(probe.shape[-2], probe.shape[-1]) // 2
obj_view_crop += obj_padding
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
Ws = t.ones(len(dataset))
if hasattr(dataset, 'intensities') and dataset.intensities is not None:
Ws *= (dataset.intensities.to(dtype=Ws.dtype)[:,...]
/ t.mean(dataset.intensities))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
else:
mask = None
if probe_support_radius is not None:
probe_support = t.zeros(probe.shape[-2:], dtype=t.bool)
xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
Rs = np.sqrt(xs**2 + ys**2)
probe_support[Rs < probe_support_radius] = 1
probe = probe * probe_support[None, :, :]
else:
probe_support = None
return cls(wavelength, det_geo, probe_basis, probe, obj,
polarization_states=polarization_states,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets=translation_offsets,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
fourier_probe=fourier_probe,
oversampling=oversampling,
loss=loss, units=units,
simulate_probe_translation=simulate_probe_translation,
simulate_finite_pixels=simulate_finite_pixels,
obj_view_crop=obj_view_crop)
def interaction(self, index, translations, polarization_indices, *args):
# The *args is included so that this can work even when given, say,
# a polarized ptycho dataset that might spit out more inputs.
# Step 1 is to convert the translations for each position into a
# value in pixels
pix_trans = tools.interactions.translations_to_pixel(
self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
# We then add on any recovered translation offset, if they exist
if self.translation_offsets is not None:
pix_trans += (self.translation_scale *
self.translation_offsets[index])
# This restricts the basis probes to stay within the probe support
basis_prs = self.probe * self.probe_support[..., :, :]
# Now we expand the polarization states explicitly.
# pol_basis_prs has dimensions n_states x 2 x n_modes x m x l
pol_basis_prs = (
self.polarization_states[...,None,None,None]
* basis_prs[...,None,:,:,:])
# For a Fourier-space probe, we take an IFT
if self.fourier_probe:
basis_prs = tools.propagators.inverse_far_field(basis_prs)
# So, first we get the appropriate probe for each shot.
# prs is now n_frames x 2 x n_modes x m x l
prs = pol_basis_prs[polarization_indices]
# Now we construct the probes for each shot from the basis probes
if self.weights is not None:
Ws = self.weights[index]
# And then we multiply by the weights, along the 0th dimension
prs = Ws[..., None, None, None, None] * prs#basis_prs
if self.simulate_probe_translation:
det_pix_trans = tools.interactions.translations_to_pixel(
self.det_basis,
translations,
surface_normal=self.surface_normal)
probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] *
self.I_phase[None,...] +
det_pix_trans[:,1,None,None] *
self.J_phase[None,...]))
prs = prs * probe_masks[...,None,None,:,:]
# We automatically rescale the probe to match the background size,
# which allows us to do stuff like let the object be super-resolution,
# while restricting the probe to the detector resolution but still
# doing an explicit real-space limitation of the probe
padding = [self.oversampling * self.background.shape[-2] - prs.shape[-2],
self.oversampling * self.background.shape[-1] - prs.shape[-1]]
if any([p != 0 for p in padding]): # For probe_fourier_crop != 0.
padding = [padding[-1]//2, padding[-1]-padding[-1]//2,
padding[-2]//2, padding[-2]-padding[-2]//2]
prs = tools.propagators.far_field(prs)
prs = t.nn.functional.pad(prs, padding)
prs = tools.propagators.inverse_far_field(prs)
# Now we actually do the interaction, using the sinc subpixel
# translation model as per usual
# After the object, no point in treating the polarization modes
# any differently from the object/probe modes
# So, this stacks up all the output modes in lexicographic order
exit_waves = t.cat(
[self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,0,:,:,:], self.obj[0,0], pix_trans,
shift_probe=True, multiple_modes=True),
self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,0,:,:,:], self.obj[1,0], pix_trans,
shift_probe=True, multiple_modes=True)],
dim=1)
exit_waves = exit_waves + t.cat(
[self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,1,:,:,:], self.obj[0,1], pix_trans,
shift_probe=True, multiple_modes=True),
self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs[...,1,:,:,:], self.obj[1,1], pix_trans,
shift_probe=True, multiple_modes=True)],
dim=1)
return 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,
saturation=self.saturation,
oversampling=self.oversampling,
simulate_finite_pixels=self.simulate_finite_pixels)
# Note: No "loss" function is defined here, because it is added
# dynamically during object creation in __init__
def sim_to_dataset(self, args_list, calculation_width=None):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'cdtools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
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}
mask = self.mask
wavelength = self.wavelength
indices, translations = args_list
data = []
len(indices)
if calculation_width is None:
calculation_width = len(indices)
index_chunks = [indices[i:i + calculation_width]
for i in range(0, len(indices),
calculation_width)]
translation_chunks = [translations[i:i + calculation_width]
for i in range(0, len(indices),
calculation_width)]
# Then we simulate the results
data = [self.forward(idx, trans).detach()
for idx, trans in zip(index_chunks, translation_chunks)]
data = t.cat(data, dim=0)
# And finally, we make the dataset
return Ptycho2DDataset(
translations, data,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=self.get_detector_geometry(),
mask=mask)
def corrected_translations(self, dataset):
translations = dataset.translations.to(
dtype=t.float32, device=self.probe.device)
if (hasattr(self, 'translation_offsets') and
self.translation_offsets is not None):
t_offset = tools.interactions.pixel_to_translations(
self.probe_basis,
self.translation_offsets * self.translation_scale,
surface_normal=self.surface_normal)
return translations + t_offset
else:
return translations
def get_rhos(self):
# If this is the general unified mode model
if self.weights.dim() >= 2:
Ws = self.weights.detach().cpu().numpy()
rhos_out = np.matmul(np.swapaxes(Ws, 1, 2), Ws.conj())
return rhos_out
# This is the purely incoherent case
else:
return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1, normalize=False):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
# NOTE: untested, edited when ortho_probes was updated
for idx in range(self.probe.shape[0]):
oself.probe.data[idx] = analysis.orthogonalize_probes(
self.probe.data[idx])
plot_list = [
('Basis Probe Fourier Space Amplitudes',
lambda self, fig: p.plot_amplitude(
(self.probe if self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig)),
('Basis Probe Fourier Space Phases',
lambda self, fig: p.plot_phase(
(self.probe if self.fourier_probe
else tools.propagators.inverse_far_field(self.probe))
, fig=fig)),
('Basis Probe Real Space Amplitudes',
lambda self, fig: p.plot_amplitude(
(self.probe if not self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig,
basis=self.probe_basis,
units=self.units)),
('Basis Probe Real Space Phases',
lambda self, fig: p.plot_phase(
(self.probe if not self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig,
basis=self.probe_basis,
units=self.units)),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(
np.nanmean(np.abs(self.get_rhos()), axis=0),
fig=fig),
lambda self: len(self.weights.shape) >= 2),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
analysis.calc_top_mode_fraction(self.get_rhos()),
fig=fig,
units=self.units),
lambda self: len(self.weights.shape) >= 2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(
self.obj[np.s_[:,:] + self.obj_view_slice],
fig=fig,
basis=self.probe_basis,
units=self.units)),
('Object Phase',
lambda self, fig: p.plot_phase(
self.obj[np.s_[:,:] + self.obj_view_slice],
fig=fig,
basis=self.probe_basis,
units=self.units)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
]
# def plot_errors(self, dataset):
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
oversampling = self.oversampling
wavelength = self.wavelength.cpu().numpy()
return {'basis': basis, 'translation': translations,
'probe': probe, 'obj': obj,
'background': background,
'oversampling': oversampling,
'weights': weights, 'wavelength': wavelength}
src/cdtools/models/polarized_fancy_ptycho.py
-556
View File
@@ -1,556 +0,0 @@
import torch as t
from cdtools.models import CDIModel, FancyPtycho
from cdtools.datasets import Ptycho2DDataset
from cdtools import tools
from cdtools.tools import plotting as p
# from cdtools.tools import polarized_plotting as pp
from cdtools.tools import analysis
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from scipy import linalg as sla
from copy import copy
from cdtools.tools import polarization
__all__ = ['PolarizedFancyPtycho']
class PolarizedFancyPtycho(FancyPtycho):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess, polarizer, analyzer,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None,
polarizer_offsets=None, analyzer_offsets=None,
polarizer_scale=1, analyzer_scale=1, mask=None,
weights = None, translation_scale = 1, saturation=None,
probe_support = None, oversampling=1,
loss='amplitude mse',units='um'):
super(PolarizedFancyPtycho, self).__init__(wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
weights = weights, translation_scale = 1, saturation=None,
probe_support = None, oversampling=1,
loss='amplitude mse',units='um')
if polarizer_offsets is None:
self.polarizer_offsets = None
else:
self.polarizer_offsets = t.nn.Parameter(t.tensor(polarizer_offsets).to(dtype=t.float32)) / polarizer_scale
if analyzer_offsets is None:
self.analyzer_offsets = None
else:
self.analyzer_offsets = t.nn.Parameter(t.tensor(analyzer_offsets).to(dtype=t.float32)) / analyzer_scale
self.polarizer = polarizer
self.analyzer = analyzer
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
if probe_guess.dim() > 4:
self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
else:
self.probe_norm = 1 * t.max(t.abs(probe_guess))
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
@classmethod
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um', left_polarized=True):
# When using this method, remember to pass through the inputs
model = FancyPtycho.from_dataset(
dataset,
probe_size=probe_size,
randomize_ang=randomize_ang,
padding=padding,
n_modes=n_modes,
dm_rank=dm_rank,
translation_scale=translation_scale,
saturation=saturation,
probe_support_radius=probe_support_radius,
propagation_distance=propagation_distance,
scattering_mode=scattering_mode,
oversampling=oversampling,
auto_center=auto_center,
loss=loss,
units=units)
# Mutate the class to its subclass
model.__class__ = cls
if left_polarized:
x = 1j
else:
x = -1j
probe = model.probe.detach()
probe = t.cat((probe, probe * x), dim=-3)
probe_max = t.max(t.abs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
probe = t.stack([probe, ] + probe_stack)
#print('probe', type(probe), probe.shape)
model.probe.data = probe
#print(model.probe.shape)
# obj = t.stack((model.obj.data, model.obj.data), dim=-3)
# model.obj.data = t.stack((obj.data, obj.data), dim=-4)
# obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
obj = model.obj.detach()
# Abe - Probably something identity matrix-like would be a better
# initialization (e.g. ((obj,0*obj),(0*obj,obj))
obj = t.stack((obj, obj), dim=-3)
obj = t.stack((obj, obj), dim=-4)
#print('object', type(obj), obj.shape)
model.obj.data = obj
#print('polarized fancy ptycho from datset obj')
a = obj.detach()
#plt.imshow(np.real(a[0, 0, :, :]))
#plt.figure()
#plt.imshow(np.real(a[0, 1, :, :]))
#plt.show()
# tensor vs tensor.data
return model
polarizers = [tools.polarization.generate_linear_polarizer(i * 45) for i in range(3)]
@classmethod
def from_dataset2(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um'):
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')
# We include the *extras to make this work even with datasets, like
# polarization dependent datasets, that might toss out extra inputs
(indices, translations, polarizer, analyzer), patterns = dataset[:]
dataset.get_as(*get_as_args[0], **get_as_args[1])
# Set to none to avoid issues with things outside the detector
if auto_center:
center = tools.image_processing.centroid(t.sum(patterns, dim=0))
else:
center = None
if left_polarized:
x = 1j
else:
x = -1j
# 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,
oversampling=oversampling)
probe_shape = t.stack((2, probe_shape), dim=-3)
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.])
# 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(surface_normal)
# 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)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# 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, propagation_distance=propagation_distance, oversampling=oversampling)
else:
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
# Now we initialize all the subdominant probe modes
probe_max = t.max(t.abs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
probe = t.stack([probe, ] + probe_stack)
# probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
probe_x, probe_y = probe, probe * x
probe = t.stact((probe_x, probe_y), dim=-3)
a = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
b = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
c = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
d = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
ab = t.stack((a, b), dim=-3)
cd = t.stack((c, d), dim=-3)
obj = t.stack((ab, cd), dim=-4)
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
if dm_rank is not None and dm_rank != 0:
if dm_rank > n_modes:
raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.')
elif dm_rank == -1:
# dm_rank == -1 is defined to mean full-rank
dm_rank = n_modes
Ws = t.zeros(len(dataset), dm_rank, n_modes, dtype=t.complex64)
# Start with as close to the identity matrix as possible,
# cutting of when we hit the specified maximum rank
for i in range(0, dm_rank):
Ws[:, i, i] = 1
else:
# dm_rank == None or dm_rank = 0 triggers a special case where
# a standard incoherent multi-mode model is used. This is the
# default, because it is so common.
# In this case, we define a set of weights which only has one index
Ws = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
else:
mask = None
if probe_support_radius is not None:
probe_support = t.zeros(probe[0].shape, dtype=t.bool)
xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
Rs = np.sqrt(xs**2 + ys**2)
probe_support[Rs < probe_support_radius] = 1
probe = probe * probe_support[None, :, :]
else:
probe_support = None
return cls(wavelength, det_geo, probe_basis, probe, obj,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets=translation_offsets,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
oversampling=oversampling,
loss=loss, units=units)
def interaction(self, index, translations, polarizer, analyzer, test=False):
# Step 1 is to convert the translations for each position into a
# value in pixels
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
# We then add on any recovered translation offset, if they exist
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
# This restricts the basis probes to stay within the probe support
basis_prs = self.probe * self.probe_support[...,:,:] # This makes no sense
# self.probe is an Nx2xXxY stach of probes
# Now we construct the probes for each shot from the basis probes
Ws = self.weights[index]
if len(self.weights[0].shape) == 0:
# If a purely stable coherent illumination is defined
# Ws is a tensor of length M, M is the number of frames to be processed
prs = Ws[...,None,None,None,None] * basis_prs
else:
raise NotImplementedError('Unstable Modes not Implemented for polarized light')
pol_probes = polarization.apply_linear_polarizer(prs, polarizer)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
pol_probes, self.obj, pix_trans,
shift_probe=True, multiple_modes=True, polarized=True)
# We're losing some efficiency here, because we only need to keep
# around the scalar wavefield after analyzing the waves.
# But I think it's not a huge issue - Abe
analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer)
return analyzed_exit_waves
def vectorial_wavefields(wavefields, func, *args, **kwargs):
wavefields_x = wavefields[..., 0, :, :, :]
wavefields_y = wavefields[..., 1, :, :, :]
out_x = func(wavefields_x, *args, **kwargs)
out_y = func(wavefields_y, *args, **kwargs)
out = t.stack((out_x, out_y), dim=-4)
return out[..., None, :, :]
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):
wavefields_x = wavefields[..., 0, :, :]
wavefields_y = wavefields[..., 1, :, :]
out_x = tools.measurements.quadratic_background(wavefields_x,
self.background,
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
# now, set bckgr to 0 since t shouldn't be calculated twice
out_y = tools.measurements.quadratic_background(wavefields_y,
0,
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
return out_x + out_y
# Note: No "loss" function is defined here, because it is added
# dynamically during object creation in __init__
def to(self, *args, **kwargs):
super(PolarizedFancyPtycho, self).to(*args, **kwargs)
def sim_to_dataset(self, args_list):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'cdtools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
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
wavelength = self.wavelength
indices, translations = args_list
# Then we simulate the results
data = self.forward(indices, translations)
# And finally, we make the dataset
return Ptycho2DDataset(translations, data,
entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def corrected_translations(self, dataset):
translations = dataset.translations.to(dtype=t.float32,device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
def get_rhos(self):
# If this is the general unified mode model
if self.weights.dim() >= 2:
Ws = self.weights.detach().cpu().numpy()
rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
return rhos_out
# This is the purely incoherent case
else:
return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1, normalize=False):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
NOTE: Not updated with the update to orthogonalize_probes
"""
raise NotImplementedError('Function is known to be incorrect! Contact Abe Levitan at abraham.levitan@psi.ch to ask him to fix it.')
# First we treat the purely incoherent case
# I don't love this pattern of using an if statement with a return
# to catch this case, but because it's so much simpler than the
# unified mode case I think it's appropriate
if self.weights.dim() == 1:
probe = self.probe.detach().cpu().numpy()
ortho_probes = analysis.orthogonalize_probes(probe)
self.probe.data = t.as_tensor(ortho_probes,
device=self.probe.device,dtype=self.probe.dtype)
return
# This is for the unified mode case
# Note to future: We could probably do this more cleanly with an
# SVD directly on the Ws matrix, instead of an eigendecomposition
# of the rho matrix.
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = self.probe.detach().cpu().numpy()
ortho_probes, A = analysis.orthogonalize_probes(probe,
density_matrix=overall_rho,
keep_transform=True,
normalize=normalize)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
dm_rank = self.weights.shape[1]
new_Ws = []
for rho in new_rhos:
# These are returned from smallest to largest - we want to keep
# the largest ones
w,v = sla.eigh(rho)
w = w[::-1][:dm_rank]
v = v[:,::-1][:,:dm_rank]
# For situations where the rank of the density matrix is not
# full in reality, but we keep more modes around than needed,
# some ws can go negative due to numerical error! This is
# extremely rare, but comon enough to cause crashes occasionally
# when there are thousands of individual matrices to transform
# every time this is called.
w = np.maximum(w,0)
new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
new_Ws = np.array(new_Ws)
self.weights.data = t.as_tensor(new_Ws,
dtype=self.weights.dtype,device=self.weights.device)
self.probe.data = t.as_tensor(ortho_probes,
device=self.probe.device,dtype=self.probe.dtype)
def plot_wavefront_variation(self, dataset,fig=None,mode='amplitude',**kwargs):
def get_probes(idx):
basis_prs = self.probe * self.probe_support[...,:,:]
prs = t.sum(self.weights[idx,:,:,None,None] * basis_prs, axis=-4)
ortho_probes = analysis.orthogonalize_probes(prs)
if mode.lower() == 'amplitude':
return np.abs(ortho_probes.detach().cpu().numpy())
if mode.lower() == 'root_sum_intensity':
return np.sum(np.abs(ortho_probes.detach().cpu().numpy())**2,axis=0)
if mode.lower() == 'phase':
return np.angle(ortho_probes.detach().cpu().numpy())
probe_matrix = np.zeros([self.probe.shape[0]]*2,
dtype=np.complex64)
np_probes = self.probe.detach().cpu().numpy()
for i in range(probe_matrix.shape[0]):
for j in range(probe_matrix.shape[0]):
probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj())
weights = self.weights.detach().cpu().numpy()
probe_intensities = np.sum(np.tensordot(weights,probe_matrix,axes=1)*
weights.conj(),axis=2)
# Imaginary part is already essentially zero up to rounding error
probe_intensities = np.real(probe_intensities)
values = np.sum(probe_intensities,axis=1)
if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity':
cmap = 'viridis'
else:
cmap = 'twilight'
p.plot_nanomap_with_images(self.corrected_translations(dataset), get_probes, values=values, fig=fig, units=self.units, basis=self.probe_basis, nanomap_colorbar_title='Total Probe Intensity',cmap=cmap,**kwargs),
plot_list = [
('',
lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', image_colorbar_title='Square Root of Intensity'),
lambda self: len(self.weights.shape) >= 2),
('',
lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', image_colorbar_title='Probe Amplitude'),
lambda self: len(self.weights.shape) >= 2),
('',
lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', image_colorbar_title='Probe Phase'),
lambda self: len(self.weights.shape) >= 2),
('Basis Probe Amplitudes (scroll to view modes)',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
('Basis Probe Phases (scroll to view modes)',
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()), axis=0), fig=fig),
lambda self: len(self.weights.shape) >= 2),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig, units=self.units),
lambda self: len(self.weights.shape) >= 2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
]
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
return {'basis':basis, 'translation':translations,
'probe':probe,'obj':obj,
'background':background,
'weights':weights}
src/cdtools/models/rpi.py
+92 -51
View File
@@ -53,6 +53,8 @@ class RPI(CDIModel):
obj_support=None,
oversampling=1,
weight_matrix=False,
exponentiate_obj=False,
phase_only=False,
propagation_distance=0,
units='um',
dtype=t.float32,
@@ -64,7 +66,7 @@ class RPI(CDIModel):
1j * t.ones([1], dtype=dtype)).dtype
self.register_buffer('wavelength',
t.tensor(wavelength, dtype=dtype))
t.as_tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
dtype=dtype)
@@ -73,9 +75,9 @@ class RPI(CDIModel):
# used a bandlimiting constraint and had a larger basis, the user is
# expected to upsample it explicitly before doing RPI.
self.register_buffer('probe_basis',
t.tensor(probe_basis, dtype=dtype))
t.as_tensor(probe_basis, dtype=dtype))
scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1],
scale_factor = t.as_tensor([probe.shape[-1]/obj_guess.shape[-1],
probe.shape[-2]/obj_guess.shape[-2]])
self.register_buffer('obj_basis',
(self.probe_basis * scale_factor).to(dtype=dtype))
@@ -84,7 +86,7 @@ class RPI(CDIModel):
self.saturation = None
else:
self.register_buffer('saturation',
t.tensor(saturation, dtype=dtype))
t.as_tensor(saturation, dtype=dtype))
# not sure how to make this a buffer, or if I have to
self.units = units
@@ -93,19 +95,26 @@ class RPI(CDIModel):
self.mask = None
else:
self.register_buffer('mask',
t.tensor(mask, dtype=t.bool))
t.as_tensor(mask, dtype=t.bool))
self.register_buffer('probe', t.tensor(probe, dtype=complex_dtype))
self.register_buffer('probe', t.as_tensor(probe, dtype=complex_dtype))
self.register_buffer('exponentiate_obj',
t.as_tensor(exponentiate_obj, dtype=bool))
self.register_buffer('phase_only',
t.as_tensor(phase_only, dtype=bool))
# We always use multi-modes to store the object, so we convert it
# if we just get a single 2D array as an input
if obj_guess.dim() == 2:
obj_guess = obj_guess[None, :, :]
self.obj = t.nn.Parameter(t.tensor(obj_guess, dtype=complex_dtype))
self.obj = t.nn.Parameter(t.as_tensor(obj_guess, dtype=complex_dtype))
self.weights = t.nn.Parameter(
t.eye(probe.shape[0], dtype=complex_dtype)* 10)
t.eye(probe.shape[0], dtype=complex_dtype))
if not weight_matrix:
self.weights.requires_grad=False
@@ -115,40 +124,26 @@ class RPI(CDIModel):
dtype=dtype)
self.register_buffer('background',
t.tensor(background, dtype=t.float32))
t.as_tensor(background, dtype=t.float32))
if obj_support is None:
obj_support = t.ones_like(self.obj[0, ...], dtype=int)
self.register_buffer('obj_support',
t.tensor(obj_support, dtype=int))
t.as_tensor(obj_support, dtype=int))
self.obj.data = self.obj * self.obj_support[None, ...]
self.register_buffer('oversampling',
t.tensor(oversampling, dtype=int))
t.as_tensor(oversampling, dtype=int))
self.register_buffer('propagation_distance',
t.tensor(propagation_distance, dtype=dtype))
t.as_tensor(propagation_distance, dtype=dtype))
# The propagation direction of the probe. For now it's fixed,
# but perhaps it would need to be updated in the future
self.register_buffer('prop_dir',
t.tensor([0, 0, 1], dtype=dtype))
# This propagator should be able to be multiplied by the propagation
# distance each time to get a propagator
#universal_propagator = t.angle(ggasp(
# self.probe.shape[-2:],
# self.probe_basis, self.wavelength,
# t.tensor([0, 0, self.wavelength/(2*np.pi)], dtype=dtype),
# propagation_vector=self.prop_dir,
# dtype=complex_dtype,
# propagate_along_offset=True))
# TODO: probably doesn't support non-float-32 dtypes
#self.register_buffer('universal_propagator',
# universal_propagator)
t.as_tensor([0, 0, 1], dtype=dtype))
@classmethod
@@ -165,6 +160,8 @@ class RPI(CDIModel):
oversampling=1,
initialization='random',
weight_matrix=False,
exponentiate_obj=False,
phase_only=False,
probe_threshold=0,
dtype=t.float32,
):
@@ -218,7 +215,7 @@ class RPI(CDIModel):
and dataset.background is not None:
background = t.sqrt(dataset.background)
elif background is not None:
background = t.sqrt(t.Tensor(background).to(dtype=t.float32))
background = t.sqrt(t.as_tensor(background).to(dtype=t.float32))
det_geo = dataset.detector_geometry
@@ -253,7 +250,10 @@ class RPI(CDIModel):
probe, dummy_init_obj,
background=background, mask=mask,
saturation=saturation,
obj_support=obj_support, oversampling=oversampling,
obj_support=obj_support,
oversampling=oversampling,
exponentiate_obj=exponentiate_obj,
phase_only=phase_only,
weight_matrix=weight_matrix)
# I don't love this pattern, where I do the "real" obj initialization
@@ -263,6 +263,13 @@ class RPI(CDIModel):
# is a better pattern for doing this.
rpi_object.init_obj(initialization,
pattern=dataset.patterns[0])
if exponentiate_obj:
rpi_object.obj.data = -1j * t.log(rpi_object.obj.data)
if phase_only:
rpi_object.obj.data.imag[:] = 0
return rpi_object
@@ -273,6 +280,8 @@ class RPI(CDIModel):
obj_size=None,
n_modes=1,
saturation=None, # TODO can we get this from the calibration?
exponentiate_obj=False,
phase_only=False,
initialization='random',
dtype=t.float32
):
@@ -316,10 +325,18 @@ class RPI(CDIModel):
dummy_init_obj,
background=background,
mask=mask,
exponentiate_obj=exponentiate_obj,
phase_only=phase_only,
)
rpi_object.init_obj(initialization)
if exponentiate_obj:
rpi_object.obj.data = -1j * t.log(rpi_object.obj.data)
if phase_only:
rpi_object.obj.data.imag[:] = 0
return rpi_object
@@ -443,20 +460,28 @@ class RPI(CDIModel):
# Mix the probes with the weight matrix
prs = t.sum(self.weights[..., None, None] * self.probe, axis=-3)
if self.exponentiate_obj:
if self.phase_only:
obj = t.exp(1j*self.obj.real)
else:
obj = t.exp(1j*self.obj)
else:
obj = self.obj
for i in range(self.probe.shape[0]):
pr = prs[i]
# Here we have a 3D probe (one single mode)
# and a 4D object (multiple modes mixing incoherently)
exit_waves = RPI_interaction(pr,
self.obj_support * self.obj)
self.obj_support * obj)
all_exit_waves.append(exit_waves)
# This creates a bunch of modes generated from all possible combos
# of the probe and object modes all strung out along the first index
output = t.cat(all_exit_waves)
# If we have multiple indexes input, we unsqueeze and repeat the stack
# of wavefields enough times to simulate each requested index. This
# seems silly, but it enables (for example) one to do a reconstruction
@@ -484,11 +509,12 @@ class RPI(CDIModel):
# Here I'm taking advantage of an undocumented feature in the
# incoherent_sum measurement function where it will work with
# a 4D wavefield array as well as a 5D array.
m = tools.measurements.quadratic_background(wavefields,
self.background,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
m = tools.measurements.quadratic_background(
wavefields,
self.background,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
return m
def loss(self, sim_data, real_data, mask=None):
@@ -510,20 +536,35 @@ class RPI(CDIModel):
lambda self, fig: p.plot_amplitude(
np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)),
fig=fig, basis=self.probe_basis)),
('Dominant Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj[0], fig=fig,
basis=self.obj_basis)),
('Dominant Object Phase',
lambda self, fig: p.plot_phase(self.obj[0], fig=fig,
basis=self.obj_basis)),
('Subdominant Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj[1], fig=fig,
basis=self.obj_basis),
lambda self: len(self.obj) >=2),
('Subdominant Object Phase',
lambda self, fig: p.plot_phase(self.obj[1], fig=fig,
basis=self.obj_basis),
lambda self: len(self.obj) >=2)
('Object Amplitude',
lambda self, fig: p.plot_amplitude(
self.obj,
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object Phase',
lambda self, fig: p.plot_phase(
self.obj,
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Real Part of T',
lambda self, fig: p.plot_real(
self.obj,
fig=fig,
basis=self.obj_basis,
units=self.units,
cmap='cividis'),
lambda self: self.exponentiate_obj),
('Imaginary Part of T',
lambda self, fig: p.plot_imag(
self.obj,
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: self.exponentiate_obj),
]
src/cdtools/models/s_matrix_ptycho.py
-429
View File
@@ -1,429 +0,0 @@
import torch as t
from cdtools.models import CDIModel
from cdtools.datasets import Ptycho2DDataset
from cdtools import tools
from cdtools.tools import cmath
from cdtools.tools import plotting as p
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from copy import copy
__all__ = ['SMatrixPtycho']
class SMatrixPtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis, probe_guess, probe_fourier_support,
s_matrix_guess,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None,
probe_planes = None, mask=None,
weights = None, translation_scale = 1, saturation=None,
oversampling=1):
super(SMatrixPtycho,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
self.surface_normal = t.Tensor(surface_normal)
self.saturation = saturation
if mask is None:
self.mask = mask
else:
self.mask = t.BoolTensor(mask)
# We rescale the probe here so it learns at the same rate as the
# object
# Remember that for S-matrix we have several probes for different
# planes
if probe_guess.dim() > 4:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0,0].to(t.float32)))
else:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
/ self.probe_norm)
self.s_matrix = t.nn.Parameter(s_matrix_guess.to(t.float32))
if background is None:
ew_shape = [s_matrix_guess.shape[0] - 1 + probe_guess.shape[-3],
s_matrix_guess.shape[1] - 1 + probe_guess.shape[-2]]
if detector_slice is not None:
background = 1e-6 * t.ones(t.ones(ew_shape)[self.detector_slice].shape).to(t.float32)
else:
background = 1e-6 * t.ones(ew_shape).to(t.float32)
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)/ translation_scale)
# This maps indices to probe planes to be used. If none, it defaults
# to always being plane 0
if probe_planes is None:
self.probe_planes = None
else:
self.probe_planes = t.LongTensor(probe_planes)
self.translation_scale = translation_scale
self.probe_fourier_support = t.Tensor(probe_fourier_support).to(t.float32)
self.oversampling = oversampling
@classmethod
def from_dataset(cls, dataset, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1):
datasets = [dataset]
propagation_distances = [propagation_distance]
# We only return the 0th element because in the general case, the
# constructor needs to return a stacked datset in addition to
# a model, but for the case of one dataset we only need to return
# the model.
return cls.from_datasets(datasets, probe_convergence_radius,
locality_radius=locality_radius,
probe_size=probe_size,
randomize_ang=randomize_ang,
padding=padding,
n_modes=n_modes,
translation_scale=translation_scale,
saturation=saturation,
propagation_distances=propagation_distances,
scattering_mode=scattering_mode,
oversampling=oversampling)[0]
# This is for the multi-focal-plane case, where each dataset will correspond
# to a different focal plane. The guess propagation distance for each
# dataset can be set individually but otherwise the probes are
# reconstructed entirely separately. All datasets are assumed to have
# the same basic parameters (wavelength, detector geometry, etc) and share
# the same origin in the x-y plane.
@classmethod
def from_datasets(cls, datasets, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distances=None, scattering_mode=None, oversampling=1):
wavelength = datasets[0].wavelength
det_basis = datasets[0].detector_geometry['basis']
det_shape = datasets[0][0][1].shape
distance = datasets[0].detector_geometry['distance']
# Then, generate the probe geometry from the dataset
ewg = tools.initializers.exit_wave_geometry
probe_basis, ew_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
padding=padding,
oversampling=oversampling)
if propagation_distances is None:
propagation_distances = [None] * len(datasets)
# This shrinks the probe to ensure that the output wavefield
# is the correct shape
probe_shape = t.Size(np.array(ew_shape) - (2*locality_radius))
# always do this on the cpu
probe_planes = []
translations = []
patterns = []
probes = []
for i, dataset in enumerate(datasets):
get_as_args = dataset.get_as_args
dataset.get_as(device='cpu')
(indices, tx), pats = dataset[:]
dataset.get_as(*get_as_args[0],**get_as_args[1])
translations.append(tx)
probe_planes.extend([i]*tx.shape[0])
patterns.append(pats)
# Finally, initialize the probe and object using this information
if locality_radius != 0:
probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distances[i], oversampling=oversampling)[locality_radius:-locality_radius,locality_radius:-locality_radius]
else:
probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distances[i], oversampling=oversampling)
# Now we initialize all the subdominant probe modes
probe_max = t.max(cmath.cabs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
probe = t.stack([tools.propagators.inverse_far_field(probe),] + probe_stack)
probes.append(probe)
translations = t.cat(translations)
patterns = t.cat(patterns)
probes = t.stack(probes)
if hasattr(datasets[0], 'sample_info') and \
datasets[0].sample_info is not None and \
'orientation' in datasets[0].sample_info:
surface_normal = datasets[0].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)
# The locality radius correction is probably not needed because
# it will always be way less than 200, but it ensures that there
# is no wrapping in the s-matrix
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200+2*locality_radius)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(datasets[0].background)
else:
background = None
s_matrix = t.zeros([2*locality_radius+1,2*locality_radius+1,obj_size[0],
obj_size[1],2])
s_matrix[locality_radius,locality_radius,:,:,:] = \
tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((translations.shape[0],2)) - 0.5)
weights = t.ones(translations.shape[0])
if hasattr(datasets[0], 'mask') and datasets[0].mask is not None:
mask = datasets[0].mask.to(t.bool)
else:
mask = None
probe_support = t.zeros_like(probes[0,0].to(dtype=t.float32))
xs, ys = np.mgrid[:probes.shape[-3],:probes.shape[-2]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
Rs = np.sqrt(xs**2 + ys**2)
probe_support[Rs<probe_convergence_radius] = 1
probes = probes * probe_support[None,None,:,:]
model = cls(wavelength, det_geo, probe_basis, probes, probe_support,
s_matrix,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets = translation_offsets,
probe_planes = probe_planes,
weights=weights, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
oversampling=oversampling)
# Now we need to produce a concatenated dataset to be used to
# train the model
dataset = Ptycho2DDataset(translations, patterns)
return model, dataset
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
if self.probe_planes is not None:
probes_set = self.probe[self.probe_planes[index]]
else:
probes_set = self.probe[[0]*translations.shape[0]]
all_exit_waves = []
for i in range(probes_set.shape[1]):
exit_waves = []
for j in range(probes_set.shape[0]):
pr = tools.propagators.inverse_far_field(probes_set[j,i] * self.probe_fourier_support)
exit_wave = self.probe_norm * tools.interactions.ptycho_2D_sinc_s_matrix(pr, self.s_matrix, pix_trans[j], shift_probe=True)
exit_waves.append(exit_wave)
exit_waves = t.stack(exit_waves)
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,
saturation=self.saturation,
oversampling=self.oversampling)
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(SMatrixPtycho, 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)
if self.probe_planes is not None:
self.probe_planes = self.probe_planes.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.probe_fourier_support = self.probe_fourier_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.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'cdtools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
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
wavelength = self.wavelength
indices, translations = args_list
# Then we simulate the results
data = self.forward(indices, translations)
# And finally, we make the dataset
return Ptycho2DDataset(translations, data,
entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('First Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0,0], fig=fig, basis=self.probe_basis)),
('First Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0,0], fig=fig, basis=self.probe_basis)),
('Second Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1,0], fig=fig, basis=self.probe_basis),
lambda self: self.probe.shape[0] >=2),
('Second Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1,0], fig=fig, basis=self.probe_basis),
lambda self: self.probe.shape[0] >=2),
('Exit Wave Amplitude under Uniform Illumination',
lambda self, fig: p.plot_amplitude(t.sum(self.s_matrix.data,dim=(0,1)), fig=fig, basis=self.probe_basis)),
('Exit Wave Phase under Uniform Illumination',
lambda self, fig: p.plot_phase(t.sum(self.s_matrix.data,dim=(0,1)), fig=fig, basis=self.probe_basis)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
]
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = probe * self.probe_norm.detach().cpu().numpy()
s_matrix = cmath.torch_to_complex(self.s_matrix.detach().cpu())
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
wavelength = self.wavelength.cpu().numpy()
return {'basis':basis, 'translation':translations,
'probe':probe,'s_matrix':s_matrix,
'background':background,
'weights':weights, 'wavelength':wavelength}
src/cdtools/models/simple_ptycho.py
+34 -35
View File
@@ -1,26 +1,16 @@
import torch as t
from cdtools.models import CDIModel
from cdtools.datasets import Ptycho2DDataset
from cdtools import tools
from cdtools.tools import plotting as p
from copy import copy
from torch.utils import data as torchdata
from datetime import datetime
import numpy as np
__all__ = ['SimplePtycho']
class SimplePtycho(CDIModel):
"""A simple ptychography model for exploring ideas and extensions
"""A simple ptychography model to demonstrate the structure of a model
"""
def __init__(
self,
wavelength,
detector_geometry,
probe_basis,
probe_guess,
obj_guess,
@@ -28,23 +18,22 @@ class SimplePtycho(CDIModel):
):
# We initialize the superclass
super(SimplePtycho,self).__init__()
super().__init__()
# We register all the constants, like wavelength, as buffers.
# This lets the model hook into some nice pytorch features, like
# using model.to, and in principle should make easier to broadcast
# across multiple GPUs
# We register all the constants, like wavelength, as buffers. This
# lets the model hook into some nice pytorch features, like using
# model.to, and broadcasting the model state across multiple GPUs
self.register_buffer('wavelength', t.as_tensor(wavelength))
self.store_detector_geometry(detector_geometry)
self.register_buffer('min_translation', t.as_tensor(min_translation))
self.register_buffer('probe_basis', t.as_tensor(probe_basis))
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
# We cast the probe and object to 64-bit complex tensors
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
# We rescale the probe here so it learns at the same rate as the
# object
# object when using optimizers, like Adam, which set the stepsize
# to a fixed maximum
self.register_buffer('probe_norm', t.max(t.abs(probe_guess)))
# And we store the probe and object guesses as parameters, so
@@ -62,15 +51,13 @@ class SimplePtycho(CDIModel):
det_shape = dataset[0][1].shape
distance = dataset.detector_geometry['distance']
# We extract all the diffraction patterns and translations
(indices, translations), patterns = dataset[:]
# Then, we generate the probe geometry
ewg = tools.initializers.exit_wave_geometry
probe_basis = ewg(det_basis, det_shape, wavelength, distance)
# Next generate the object geometry from the probe geometry and
# the translations
(indices, translations), patterns = dataset[:]
pix_translations = tools.interactions.translations_to_pixel(
probe_basis,
translations,
@@ -86,7 +73,6 @@ class SimplePtycho(CDIModel):
return cls(
wavelength,
dataset.detector_geometry,
probe_basis,
probe,
obj,
@@ -101,10 +87,9 @@ class SimplePtycho(CDIModel):
self.probe_basis,
translations)
pix_trans -= self.min_translation
# This function gets the proper, pixel-accurate crops of the
# object array, and uses a sinc-interpolated subpixel shift of
# the probe to reach subpixel accuracy
# This function extracts the appropriate window from the object and
# multiplies the object and probe functions
return tools.interactions.ptycho_2D_round(
self.probe_norm * self.probe,
self.obj,
@@ -114,15 +99,9 @@ class SimplePtycho(CDIModel):
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)
def loss(self, real_data, sim_data):
return tools.losses.amplitude_mse(real_data, sim_data)
@@ -139,3 +118,23 @@ class SimplePtycho(CDIModel):
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis))
]
def save_results(self, dataset):
# This will save out everything needed to recreate the object
# in the same state, but it's not the best formatted.
base_results = super().save_results()
# So we also save out the main results in a more useable format
probe_basis = self.probe_basis.detach().cpu().numpy()
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = self.obj.detach().cpu().numpy()
wavelength = self.wavelength.cpu().numpy()
results = {
'probe_basis': probe_basis,
'probe': probe,
'obj': obj,
'wavelength': wavelength,
}
return {**base_results, **results}
src/cdtools/models/time_resolved_ptycho_calibration.py
-503
View File
@@ -1,503 +0,0 @@
import torch as t
from cdtools.models import CDIModel
from cdtools.datasets import Ptycho2DDataset
from cdtools import tools
from cdtools.tools import plotting as p
from cdtools.tools import analysis
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from scipy import linalg as sla
from copy import copy
#
# Basic points:
# Just one probe mode, no need to overcomplicate things.
# Weights is only a list of numbers, no matrices or anything like that
# Mandatory probe support in Fourier space
#
# When loading from a dataset, we need information on the zone plate geometry
#
__all__ = ['TimeResolvedPtychoCalibration']
class TimeResolvedPtychoCalibration(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess,
fourier_times, probe_fourier_support,
times, time_dependence, frame_delays,
detector_slice=None,
surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
min_translation=t.tensor([0, 0], dtype=t.float32),
background=None, translation_offsets=None, mask=None,
weights=None, translation_scale=1, saturation=None,
oversampling=1,
loss='amplitude mse', units='um',
simulate_probe_translation=False):
super(TimeResolvedPtychoCalibration, self).__init__()
self.wavelength = t.tensor(wavelength)
self.detector_geometry = copy(detector_geometry)
det_geo = self.detector_geometry
if 'distance' in det_geo:
det_geo['distance'] = t.tensor(det_geo['distance'], dtype=t.float32)
if 'basis' in det_geo:
det_geo['basis'] = t.tensor(det_geo['basis'], dtype=t.float32)
if 'corner' in det_geo:
det_geo['corner'] = t.tensor(det_geo['corner'], dtype=t.float32)
self.min_translation = t.tensor(min_translation)
self.probe_basis = t.tensor(probe_basis)
self.detector_slice = copy(detector_slice)
self.surface_normal = t.tensor(surface_normal)
self.saturation = saturation
self.units = units
if mask is None:
self.mask = mask
else:
self.mask = t.tensor(mask, dtype=t.bool)
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
self.probe_norm = 1 * t.max(t.abs(probe_guess))
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(
self.probe[self.detector_slice].shape,
dtype=t.float32)
else:
background = 1e-6 * t.ones(self.probe[0].shape,
dtype=t.float32)
self.background = t.nn.Parameter(background)
if weights is None:
self.weights = None
else:
self.weights = t.nn.Parameter(t.tensor(weights,
dtype=t.float32))
if translation_offsets is None:
self.translation_offsets = None
else:
t_o = t.tensor(translation_offsets, dtype=t.float32)
t_o = t_o / translation_scale
self.translation_offsets = t.nn.Parameter(t_o)
self.translation_scale = translation_scale
self.probe_fourier_support = probe_fourier_support
self.fourier_times = fourier_times
self.times = times
self.time_dependence = t.nn.Parameter(time_dependence)
self.frame_delays = frame_delays
self.oversampling = oversampling
self.simulate_probe_translation = simulate_probe_translation
if simulate_probe_translation:
Is = t.arange(self.probe.shape[-2], dtype=t.float32)
Js = t.arange(self.probe.shape[-1], dtype=t.float32)
Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
self.I_phase = 2 * np.pi* Is
self.J_phase = 2 * np.pi* Js
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
or loss.lower().strip() == 'amplitude_mse'):
self.loss = tools.losses.amplitude_mse
elif (loss.lower().strip() == 'poisson nll'
or loss.lower().strip() == 'poisson_nll'):
self.loss = tools.losses.poisson_nll
else:
raise KeyError('Specified loss function not supported')
@classmethod
def from_dataset(cls, dataset, zp_geometry, time_window, n_times, n_frames, randomize_ang=0, padding=0, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um', simulate_probe_translation=False):
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')
# We include the *extras to make this work even with datasets, like
# polarization dependent datasets, that might toss out extra inputs
(indices, translations, *extras), patterns = dataset[:]
dataset.get_as(*get_as_args[0], **get_as_args[1])
# Set to none to avoid issues with things outside the detector
if auto_center:
center = tools.image_processing.centroid(t.sum(patterns, dim=0))
else:
center = None
# 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,
oversampling=oversampling)
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.])
# 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(surface_normal)
# 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)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# Finally, initialize the probe and object using this information
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
probe = tools.propagators.far_field(probe)
obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
# we define a set of weights which only has one index
Ws = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
else:
mask = None
# What do we know about the zp?
# delta_r, N, and beamstop ratio. It's probably best, though,
# to just read diameter, beamstop_diameter and focal_length directly
# because that is the most general even if the optic isn't truly
# a zone plate.
zp_distance = zp_geometry['focal_length']
probe_size = probe_basis * t.as_tensor(probe_shape, dtype=t.float32)
pinv_basis = t.tensor(np.linalg.pinv(probe_size).transpose()).to(t.float32)
zone_plate_basis = pinv_basis * wavelength * zp_distance
zone_plate_steps = t.sum(zone_plate_basis,axis=1)
# This may very well mix up x & y and fail on non-square detectors
probe_fourier_support = t.zeros(probe.shape, dtype=t.bool)
xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
xs = zone_plate_steps[1] * (xs - np.mean(xs))
ys = zone_plate_steps[0] * (ys - np.mean(ys))
Rs = np.sqrt(xs**2 + ys**2)
distances = np.sqrt(zp_distance**2 + xs**2 + ys**2)
times = (distances - t.min(distances)) / 2.99792e8
# This sets the support of the probe and also restricts the
# timing matrix so it only considers times that are actually
# in the window defined by the probe fourier support
probe_fourier_support[Rs < zp_geometry['diameter']/2] = 1
times[Rs > zp_geometry['diameter']/2] = 0
times[Rs > zp_geometry['diameter']/2] = t.max(times)
probe_fourier_support[Rs < zp_geometry['beamstop_diameter']/2] = 0
times[Rs < zp_geometry['beamstop_diameter']/2] = t.max(times)
times[Rs < zp_geometry['beamstop_diameter']/2] = t.min(times)
fourier_times = times - t.min(times)
probe = probe * probe_fourier_support
# This is now the time axis for the probe's envelope
times = t.linspace(0, time_window, n_times+1)
time_dependence = t.ones(n_times, dtype=t.complex64)
frame_delays = t.linspace(0, t.max(fourier_times) + time_window, n_frames+2)
frame_delays -= time_window
frame_delays = frame_delays[1:-1]
return cls(wavelength, det_geo, probe_basis, probe, obj,
fourier_times, probe_fourier_support,
times, time_dependence, frame_delays,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets=translation_offsets,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
oversampling=oversampling,
loss=loss, units=units,
simulate_probe_translation=simulate_probe_translation)
def interaction(self, index, translations, *args):
# The *args is included so that this can work even when given, say,
# a polarized ptycho dataset that might spit out more inputs.
# Step 1 is to convert the translations for each position into a
# value in pixels
pix_trans = tools.interactions.translations_to_pixel(
self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
# We then add on any recovered translation offset, if they exist
if self.translation_offsets is not None:
pix_trans += (self.translation_scale *
self.translation_offsets[index])
probes = self.get_probes(space='real')
Ws = self.weights[index]
# This might not work well
prs = Ws[...,None,None,None] * probes
#prs = t.sum(Ws[..., None, None, None] * probes, axis=-3)
#print(prs.shape)
if self.simulate_probe_translation:
det_pix_trans = tools.interactions.translations_to_pixel(
self.detector_geometry['basis'],
translations,
surface_normal=self.surface_normal)
probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] *
self.I_phase[None,...] +
det_pix_trans[:,1,None,None] *
self.J_phase[None,...]))
prs = prs * probe_masks[...,None,:,:]
# Now we actually do the interaction, using the sinc subpixel
# translation model as per usual
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj, pix_trans,
shift_probe=True, multiple_modes=True)
return 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,
saturation=self.saturation,
oversampling=self.oversampling)
# Note: No "loss" function is defined here, because it is added
# dynamically during object creation in __init__
def to(self, *args, **kwargs):
super(TimeResolvedPtychoCalibration, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args, **kwargs)
# move the detector geometry too
det_geo = self.detector_geometry
if 'distance' in det_geo:
det_geo['distance'] = det_geo['distance'].to(*args, **kwargs)
if 'basis' in det_geo:
det_geo['basis'] = det_geo['basis'].to(*args, **kwargs)
if 'corner' in det_geo:
det_geo['corner'] = det_geo['corner'].to(*args, **kwargs)
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
if self.simulate_probe_translation:
self.I_phase = self.I_phase.to(*args, **kwargs)
self.J_phase = self.J_phase.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.probe_fourier_support = self.probe_fourier_support.to(*args,
**kwargs)
self.fourier_times = self.fourier_times.to(*args, **kwargs)
self.times = self.times.to(*args, **kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list, calculation_width=None):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'cdtools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
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
wavelength = self.wavelength
indices, translations = args_list
data = []
len(indices)
if calculation_width is None:
calculation_width = len(indices)
index_chunks = [indices[i:i + calculation_width]
for i in range(0, len(indices),
calculation_width)]
translation_chunks = [translations[i:i + calculation_width]
for i in range(0, len(indices),
calculation_width)]
# Then we simulate the results
data = [self.forward(idx, trans).detach()
for idx, trans in zip(index_chunks, translation_chunks)]
data = t.cat(data, dim=0)
# And finally, we make the dataset
return Ptycho2DDataset(
translations, data,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def corrected_translations(self, dataset):
translations = dataset.translations.to(
dtype=t.float32, device=self.probe.device)
if (hasattr(self, 'translation_offsets') and
self.translation_offsets is not None):
t_offset = tools.interactions.pixel_to_translations(
self.probe_basis,
self.translation_offsets * self.translation_scale,
surface_normal=self.surface_normal)
return translations + t_offset
else:
return translations
def get_probes(self, space='real'):
# This is the part where I need to create the probe modes using the
# time-dependent stuff
# This restricts the basis probes to stay within the probe support
optic_mask = self.probe * self.probe_fourier_support
probes = t.zeros([len(self.frame_delays)] + list(self.probe.shape),
dtype=self.probe.dtype, device=self.probe.device)
for i, delay in enumerate(self.frame_delays):
indices = t.bucketize(self.fourier_times, self.times + delay)
clamped_indices = t.clamp(indices-1, max=len(self.time_dependence)-1)
illumination = t.take(self.time_dependence, clamped_indices)
illumination[indices==0] = 0
illumination[indices==len(self.time_dependence)+1] = 0
probes[i] = illumination * optic_mask
if space.lower()=='real':
return tools.propagators.inverse_far_field(probes)
elif space.lower()=='fourier' or space.lower()=='reciprocal':
return probes
plot_list = [
('Probe Amplitudes (scroll to view modes)',
lambda self, fig: p.plot_amplitude(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)),
('Probe Phases (scroll to view modes)',
lambda self, fig: p.plot_phase(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)),
('Fourier Probe Amplitudes (scroll to view modes)',
lambda self, fig: p.plot_amplitude(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)),
('Fourier Probe Phases (scroll to view modes)',
lambda self, fig: p.plot_phase(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)),
('Optic Amplitude',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
('Optic Phase',
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()), axis=0), fig=fig),
lambda self: len(self.weights.shape) >= 2),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig, units=self.units),
lambda self: len(self.weights.shape) >= 2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis, units=self.units)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)),
('Time Structure',
lambda self, fig: (plt.figure(fig.number) and (plt.clf() or True) and plt.plot(self.time_dependence.real.detach().cpu().numpy()) and plt.plot(self.time_dependence.imag.detach().cpu().numpy())))
]
# def plot_errors(self, dataset):
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
optic = self.probe.detach().cpu().numpy()
optic = optic * self.probe_norm.detach().cpu().numpy()
time_dependence = self.time_dependence.detach().cpu().numpy()
times = self.times.detach().cpu().numpy()
fourier_times = self.fourier_times.detach().cpu().numpy()
probes = self.get_probes().detach().cpu().numpy()
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
return {'basis': basis, 'translation': translations,
'probes': probes, 'optic': optic,
'times': times, 'time_dependence': time_dependence,
'fourier_times': fourier_times,
'obj': obj,
'background': background,
'weights': weights,
}
src/cdtools/models/time_resolved_rpi.py
-298
View File
@@ -1,298 +0,0 @@
import torch as t
from cdtools.models import CDIModel
from cdtools import tools
from cdtools.tools import plotting as p
from cdtools.tools.interactions import RPI_interaction
from cdtools.tools import initializers
from scipy.ndimage import binary_dilation
import numpy as np
from copy import copy
__all__ = ['TimeResolvedRPI']
class TimeResolvedRPI(CDIModel):
@property
def obj(self):
return t.complex(self.obj_real, self.obj_imag)
def __init__(self, wavelength, detector_geometry, probe_basis,
probe, obj_guess, framerate, detector_slice=None,
background=None, mask=None, saturation=None,
obj_support=None, oversampling=1):
super(TimeResolvedRPI, self).__init__()
self.wavelength = t.tensor(wavelength)
self.framerate = framerate
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.probe_basis = t.tensor(probe_basis)
scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1],
probe.shape[-2]/obj_guess.shape[-2]])
self.obj_basis = self.probe_basis * scale_factor
self.detector_slice = detector_slice
# Maybe something to include in a bit
# self.surface_normal = t.tensor(surface_normal)
self.saturation = saturation
if mask is None:
self.mask = mask
else:
self.mask = t.tensor(mask, dtype=t.bool)
self.probe = t.tensor(probe, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
self.obj_real = t.nn.Parameter(obj_guess.real)
self.obj_imag = t.nn.Parameter(obj_guess.imag)
# Wait for LBFGS to be updated for complex-valued parameters
# self.obj = t.nn.Parameter(obj_guess.to(t.float32))
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(
self.probe[0][self.detector_slice].shape,
dtype=t.float32)
else:
background = 1e-6 * t.ones(self.probe[0].shape,
dtype=t.float32)
self.background = t.tensor(background, dtype=t.float32)
if obj_support is not None:
self.obj_support = obj_support
self.obj.data = self.obj * obj_support[None, ...]
else:
self.obj_support = t.ones_like(self.obj[0, ...])
self.oversampling = oversampling
@classmethod
def from_dataset(cls, dataset, probe, framerate, obj_size=None, background=None, mask=None, padding=0, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', probe_threshold=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')
# We only need the patterns here, not the inputs associated with them.
_, patterns = dataset[:]
dataset.get_as(*get_as_args[0],**get_as_args[1])
# Set to none to avoid issues with things outside the detector
if auto_center:
center = tools.image_processing.centroid(t.sum(patterns,dim=0))
else:
center = None
# 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,
oversampling=oversampling)
if not isinstance(probe,t.Tensor):
probe = t.as_tensor(probe)
if background is None and hasattr(dataset, 'background') \
and dataset.background is not None:
background = t.sqrt(dataset.background)
elif background is not None:
background = t.sqrt(t.Tensor(background).to(dtype=t.float32))
det_geo = dataset.detector_geometry
# If no mask is given, but one exists in the dataset, load it.
if mask is None and hasattr(dataset, 'mask') \
and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
# Now we initialize the object
if obj_size is None:
# This is a standard size for a well-matched probe and detector
obj_size = list((np.array(probe_shape) // 2).astype(int))
# I think something to do with the fact that the object is defined
# on a coarser grid needs to be accounted for here that is not
# accounted for yet
scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2)
n_modes = (probe.shape[0] - 1) // framerate + 1
obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0))
probe_fft = tools.propagators.far_field(probe_intensity)
pad0l = (probe.shape[-2] - obj_size[-2])//2
pad0r = probe.shape[-2] - obj_size[-2] - pad0l
pad1l = (probe.shape[-1] - obj_size[-1])//2
pad1r = probe.shape[-1] - obj_size[-1] - pad1l
probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r]
probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft))
obj_support = probe_lr > t.max(probe_lr) * probe_threshold
obj_support = t.as_tensor(binary_dilation(obj_support))
return cls(wavelength, det_geo, probe_basis,
probe, obj_guess, framerate, detector_slice=det_slice,
background=background, mask=mask, saturation=saturation,
obj_support=obj_support, oversampling=oversampling)
def random_init(self, pattern):
scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2)
self.obj.data = scale * t.exp(
2j * np.pi * t.rand(self.obj.shape)).to(
dtype=self.obj.dtype, device=self.obj.device)
# Needs work
def interaction(self, index, *args):
# including *args allows this to work with all sorts of datasets
# that might include other information in with the index in their
# "input" parameters (such as translations for a ptychography dataset).
# This makes it seamless to use such a dataset even though those
# extra arguments will not be used.
all_exit_waves = []
# Mix the probes with the weight matrix
prs = self.probe
for i in range(self.probe.shape[0]):
obj_frame = i // self.framerate
pr = prs[i]
exit_waves = RPI_interaction(pr,
self.obj_support * self.obj[obj_frame])
all_exit_waves.append(exit_waves.unsqueeze(0))
# This creates a bunch of modes generated from all possible combos
# of the probe and object modes all strung out along the first index
output = t.cat(all_exit_waves)
# If we have multiple indexes input, we unsqueeze and repeat the stack
# of wavefields enough times to simulate each requested index. This
# seems silly, but it enables (for example) one to do a reconstruction
# from a set of diffraction patterns that are all known to be from the
# same object.
try:
# will fail if index has no length, for example when index
# is just an int. In this case, we just do nothing instead
output = output.unsqueeze(0).repeat(1,len(index),1,1,1)
except TypeError:
pass
return output
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):
# Here I'm taking advantage of an undocumented feature in the
# incoherent_sum measurement function where it will work with
# a 4D wavefield array as well as a 5D array.
return tools.measurements.quadratic_background(wavefields,
self.background,
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
def regularizer(self, factors):
return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \
+ factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2)
def to(self, *args, **kwargs):
super(TimeResolvedRPI, 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.probe = self.probe.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.obj_basis = self.obj_basis.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
self.background = self.background.to(*args, **kwargs)
# Maybe include in a bit
#self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
raise NotImplementedError('No sim to dataset yet, sorry!')
plot_list = [
('Probe Amplitudes',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
('Object Amplitudes',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig,
basis=self.obj_basis)),
('Object Phases',
lambda self, fig: p.plot_phase(self.obj, fig=fig,
basis=self.obj_basis))
]
def save_results(self, dataset=None, full_obj=False):
# dataset is set as a kwarg here because it isn't needed, but the
# common pattern is to pass a dataset. This makes it okay if one
# continues to use that standard pattern
probe_basis = self.probe_basis.detach().cpu().numpy()
obj_basis = self.obj_basis.detach().cpu().numpy()
probe = self.probe.detach().cpu().numpy()
# Provide the option to save out the subdominant objects or
# just the dominant one
if full_obj:
obj = self.obj.detach().cpu().numpy()
else:
obj = self.obj[0].detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
'probe': probe,'obj': obj,
'background': background}
src/cdtools/tools/__init__.py
+1 -2
View File
@@ -23,5 +23,4 @@ from cdtools.tools import interactions
from cdtools.tools import propagators
from cdtools.tools import measurements
from cdtools.tools import analysis
from cdtools.tools import atoms
from cdtools.tools import polarization
src/cdtools/tools/analysis/__init__.py
+1
View File
@@ -1 +1,2 @@
from cdtools.tools.analysis.analysis import *
from cdtools.tools.analysis.analysis import __all__, __doc__
src/cdtools/tools/analysis/analysis.py
+184 -15
View File
@@ -28,7 +28,8 @@ __all__ = [
'calc_generalized_rms_error',
'remove_phase_ramp',
'remove_amplitude_exponent',
'standardize_reconstruction_set'
'standardize_reconstruction_set',
'standardize_reconstruction_pair',
]
@@ -569,6 +570,8 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
cor_fft = f1 * t.conj(f2)
#from cdtools.tools import plotting as p
#from matplotlib import pyplot as plt
#p.plot_phase(im1[im_slice], cmap='cividis')
#p.plot_phase(im2[im_slice], cmap='cividis')
#p.plot_amplitude(t.log(t.abs(cor_fft)))
#p.plot_amplitude(t.log(t.abs(f1)))
#p.plot_phase(cor_fft)
@@ -576,7 +579,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
F1 = t.abs(f1)**2
F2 = t.abs(f2)**2
# TODO this is still incorrect if the two bases arent equal
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
@@ -590,6 +593,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
max_i = np.max(i_freqs)
max_j = np.max(j_freqs)
frc_range = [0, max(max_i,max_j)]
elif limit.lower().strip() == 'corner':
frc_range = [0, np.max(Rs)]
else:
@@ -607,7 +611,13 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
#n_pix = n_pix / 4 # This is for an apodized image, apodized with a hann window
frc = numerator / np.sqrt(denominator_F1*denominator_F2)
#plt.figure()
#plt.plot(np.real(frc))
#plt.title('real')
#plt.figure()
#plt.plot(np.imag(frc))
#plt.title('imag')
#plt.show()
# This moves from combined-image SNR to single-image SNR
snr /= 2
@@ -765,6 +775,8 @@ def calc_rms_error(field_1, field_2, align_phases=True, normalize=False,
The RMS error, or tensor of RMS errors, depending on the dim argument
"""
field_1 = t.as_tensor(field_1)
field_2 = t.as_tensor(field_2)
sumdims = tuple(d - dims for d in range(dims))
@@ -910,6 +922,8 @@ def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2):
The generalized RMS error, or tensor of generalized RMS errors, depending on the dim argument
"""
# TODO either make this work correctly or just make everything work
# only for tensors
fields_1 = t.as_tensor(fields_1)
fields_2 = t.as_tensor(fields_2)
@@ -928,7 +942,7 @@ def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2):
return t.sqrt(result)
def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, snr=1.):
def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, snr=1., limit='side'):
"""Calculates a Fourier ring correlation between two images
This function requires an input of a basis to allow for FRC calculations
@@ -965,7 +979,7 @@ def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, s
im_np = False
if isinstance(fields_1, np.ndarray):
fields_1 = t.as_tensor(fields_)
fields_1 = t.as_tensor(fields_1)
im_np = True
if isinstance(fields_2, np.ndarray):
fields_2 = t.as_tensor(fields_2)
@@ -975,17 +989,15 @@ def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, s
basis = t.tensor(basis)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
im_slice = np.s_[...,:,:]
if nbins is None:
nbins = np.max(fields_1[...,im_slice].shape[-2:]) // 4
f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2))
f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2))
f1 = t.fft.fftshift(t.fft.fft2(fields_1[im_slice]),dim=(-1,-2))
f2 = t.fft.fftshift(t.fft.fft2(fields_2[im_slice]),dim=(-1,-2))
cor_fft = f1 * t.conj(f2)
F1 = t.abs(f1)**2
F2 = t.abs(f2)**2
@@ -994,18 +1006,30 @@ def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, s
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[0],d=di))
j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[1],d=dj))
i_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[-2],d=di))
j_freqs = np.fft.fftshift(np.fft.fftfreq(cor_fft.shape[-1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
if limit.lower().strip() == 'side':
max_i = np.max(i_freqs)
max_j = np.max(j_freqs)
frc_range = [0, max(max_i,max_j)]
elif limit.lower().strip() == 'corner':
frc_range = [0, np.max(Rs)]
else:
raise ValueError('Invalid FRC limit: choose "side" or "corner"')
# This line is used to get a set of bins that matches the logic
# used by np.histogram, so that this function will match the choices
# of bin edges that comes from the non-generalized version. This also
# gets us the count on the number of pixels per bin so we can calculate
# the threshold curve
n_pix, bins = np.histogram(Rs,bins=nbins)
n_pix, bins = np.histogram(Rs,bins=nbins, range=frc_range)
bins = t.as_tensor(bins)
Rs = t.as_tensor(Rs)
frc = []
for i in range(len(bins)-1):
@@ -1202,8 +1226,8 @@ def standardize_reconstruction_set(
if correct_phase_offset:
obj_1 = np.exp(-1j* np.angle(np.sum(obj_1[window]))) * obj_1
obj_2 = np.exp(-1j* np.angle(np.sum(obj_2))) * obj_2
obj = np.exp(-1j* np.angle(np.sum(obj))) * obj
obj_2 = np.exp(-1j* np.angle(np.sum(obj_2[window]))) * obj_2
obj = np.exp(-1j* np.angle(np.sum(obj[window]))) * obj
# Todo update the translations to account for the determined shift
@@ -1255,3 +1279,148 @@ def standardize_reconstruction_set(
return results
def standardize_reconstruction_pair(
half_1,
half_2,
correct_phase_offset=True,
correct_phase_ramp=True,
correct_amplitude_exponent=False,
window=np.s_[:,:],
nbins=50,
probe_nbins=50,
frc_limit='side',
):
"""Standardizes and analyses a set of two repeat
It's very common to run two subsequent ptycho reconstructions, so that the
effect of sample damage during the first reconstruction can be used in the
estimate of thefinal quality. The difference between the two datasets
datasets can be used to estimate the quality and resolution of each one.
But to do that analysis, first the reconstructions need to be aligned
with respect to each other and normalized in a few ways.
This function takes the results (as output by model.save_results) of
a pair of reconstructions and:
- Aligns the object reconstructions with one another
- Corrects for the global phase offset (by default)
- Sets a sensible value for the object/probe phase ramp (by default)
- Sets a sensible value for the object/probe exponential decay (off by
default)
- Calculates the FRC and derived SSNR.
Then, these results are packaged into an output dictionary. The output
does not retain all the information from the inputs, so if full traceability
is desired, do not delete the files containing the individual
reconstructions
Parameters
----------
half_1 : dict
The result of the first half dataset, as returned by model.save_results
half_2 : dict
The result of the second half dataset, as returned by model.save_results
Returns
-------
results : dict
A dictionary containing the synthesized results
"""
# We get the two half-data reconstructions
obj_1, probe_1, weights_1 = half_1['obj'],half_1['probe'],half_1['weights']
obj_2, probe_2, weights_2 = half_2['obj'],half_2['probe'],half_2['weights']
if correct_phase_ramp:
obj_1, probe_1 = remove_phase_ramp(
half_1['obj'], window, probe=half_1['probe'])
obj_2, probe_2 = remove_phase_ramp(
half_2['obj'], window, probe=half_2['probe'])
if correct_amplitude_exponent:
obj_1, probe_1, weights_1 = remove_amplitude_exponent(
obj_1, window, probe=probe_1,
weights=half_1['weights'],
basis=half_1['basis'],
translations=half_1['translations'])
obj_2, probe_2, weights_2 = remove_amplitude_exponent(
obj_2, window, probe=probe_2,
weights=half_2['weights'],
basis=half_2['basis'],
translations=half_2['translations'])
if correct_phase_offset:
obj_1 = np.exp(-1j* np.angle(np.sum(obj_1[window]))) * obj_1
obj_2 = np.exp(-1j* np.angle(np.sum(obj_2[window]))) * obj_2
# Todo update the translations to account for the determined shift
shift = ip.find_shift(
t.as_tensor(ip.hann_window(obj_1[window])),
t.as_tensor(ip.hann_window(obj_2[window])))
obj_2 = ip.sinc_subpixel_shift(
t.as_tensor(obj_2), shift).numpy()
probe_shift = ip.find_shift(
t.as_tensor(probe_1[0]),
t.as_tensor(probe_2[0]),
)
for idx in range(probe_2.shape[0]):
probe_2[idx] = ip.sinc_subpixel_shift(
t.as_tensor(probe_2[idx]), probe_shift).numpy()
# TODO I'm not sure if the default threshold treats this case right tbh
freqs, frc, threshold = calc_frc(
ip.hann_window(obj_1[window]),
ip.hann_window(obj_2[window]),
half_1['obj_basis'],
nbins=nbins,
limit=frc_limit,
)
probe_freqs, probe_frc, probe_frc_threshold = calc_generalized_frc(
probe_1,
probe_2,
half_1['probe_basis'],
nbins=probe_nbins,
limit=frc_limit,
)
probe_nrms_error = calc_generalized_rms_error(
probe_1,
probe_2,
normalize=True
)
# The correct formulation when the final output is one of the two
# reconstructions
ssnr = np.abs(frc) / (1 - np.abs(frc))
results = {
'obj_1': obj_1,
'probe_1': probe_1,
'weights_1': weights_1,
'translations_1': half_1['translations'],
'background_1': half_1['background'],
'obj_2': obj_2,
'probe_2': probe_2,
'weights_2': weights_2,
'translations_2': half_2['translations'],
'background_2': half_2['background'],
'wavelength': half_1['wavelength'],
'obj_basis': half_1['obj_basis'],
'probe_basis': half_1['probe_basis'],
'frc_freqs': freqs,
'frc': frc,
'frc_threshold': threshold,
'ssnr': ssnr,
'probe_freqs': probe_freqs,
'probe_frc': probe_frc,
'probe_frc_threshold': probe_frc_threshold,
'probe_nrms_error': probe_nrms_error,
}
return results
src/cdtools/tools/atoms/__init__.py
-1
View File
@@ -1 +0,0 @@
from cdtools.tools.atoms.atoms import *
src/cdtools/tools/atoms/atoms.py
-1291
View File
File diff suppressed because it is too large. Load diff
src/cdtools/tools/data/__init__.py
+1
View File
@@ -1 +1,2 @@
from cdtools.tools.data.data import *
from cdtools.tools.data.data import __all__, __doc__
src/cdtools/tools/image_processing/__init__.py
+1
View File
@@ -1 +1,2 @@
from cdtools.tools.image_processing.image_processing import *
from cdtools.tools.image_processing.image_processing import __all__, __doc__
src/cdtools/tools/initializers/__init__.py
+1
View File
@@ -1 +1,2 @@
from cdtools.tools.initializers.initializers import *
from cdtools.tools.initializers.initializers import __all__, __doc__
src/cdtools/tools/interactions/__init__.py
+1
View File
@@ -1 +1,2 @@
from cdtools.tools.interactions.interactions import *
from cdtools.tools.interactions.interactions import __all__, __doc__
src/cdtools/tools/interactions/interactions.py
+30 -33
View File
@@ -7,7 +7,7 @@ for ptychographic reconstruction.
import torch as t
import numpy as np
from cdtools.tools import propagators, image_processing, polarization
from cdtools.tools import propagators, image_processing
__all__ = ['translations_to_pixel', 'pixel_to_translations',
'project_translations_to_sample',
@@ -389,7 +389,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
return t.stack(exit_waves)
def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, polarized=False, polarizer=None, analyzer=None):
def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, probe_support=None):
"""Returns a stack of exit waves accounting for subpixel shifts
This function returns a collection of exit waves, with the first
@@ -438,15 +438,10 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
integer_translations = t.floor(translations)
subpixel_translations = translations - integer_translations
integer_translations = integer_translations.to(dtype=t.int32)
if not polarized:
selections = t.stack([obj[..., tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
else:
selections = t.stack([obj[:, :, tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
# Nx2x2xMxL tensor
selections = t.stack([obj[..., tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
exit_waves = []
if shift_probe:
@@ -459,35 +454,29 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
J = 2 * np.pi * J / probe.shape[-1]
phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I
-subpixel_translations[:,1,None,None]*J))
if polarized:
phase_masks = phase_masks[..., None, :, :]
# Nx2x1xMxL tensor
# probe is (N)(P)x2xMxL tensor
fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2))
if multiple_modes: # Multi-mode probe
if polarized:
shifted_fft_probe = fft_probe * phase_masks[...,None,:,:,:]
else:
shifted_fft_probe = fft_probe * phase_masks[...,None,:,:]
shifted_fft_probe = fft_probe * phase_masks[...,None,:,:]
else:
shifted_fft_probe = fft_probe * phase_masks
shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
dim=(-1,-2)))
if not polarized:
# TODO This is a kludge, I will fix this. I need to handle
# multiple incoherently mixing polarized objects
if multiple_modes and len(selections.shape) == 3: # Multi-mode probe
output = shifted_probe * selections[...,None,:,:]
else:
# This will only work if the
output = shifted_probe * selections
# selections: Nx2x2xMxL
# probe: Nx(P)x2x1xMxL
else:
output = polarization.apply_jones_matrix(shifted_probe, selections, multiple_modes=multiple_modes)
# Note: resist the temptation to remultiply by the probe support here,
# it will fail if you have a probe which is restricted in Fourier space
# TODO This is a kludge, I will fix this.
if multiple_modes and len(selections.shape) == 3: # Multi-mode probe
output = shifted_probe * selections[...,None,:,:]
else:
# This will only work if the
output = shifted_probe * selections
else:
raise NotImplementedError('Object shift not yet implemented')
if single_translation:
return output[0]
else:
@@ -620,6 +609,8 @@ def RPI_interaction(probe, obj):
# The far-field propagator is just a 2D FFT but with an fftshift
fftobj = propagators.far_field(obj)
fftobj_npix = fftobj.shape[-2] * fftobj.shape[-1]
# We calculate the padding that we need to do the upsampling
# This is carefully set up to keep the zero-frequency pixel in the correct
# location as the overall shape changes. Don't mess with this without
@@ -631,9 +622,15 @@ def RPI_interaction(probe, obj):
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad2l, pad2r))
# This keeps the mean intensity equal, instead of spreading out the
# intensity over the upsampled region
# TODO: Keeping this in probably isn't the most efficient
fftobj_npix_new = fftobj.shape[-2] * fftobj.shape[-1]
scale_factor = np.sqrt(fftobj_npix_new / fftobj_npix)
# Again, just an inverse FFT but with an fftshift
upsampled_obj = propagators.inverse_far_field(fftobj)
upsampled_obj = scale_factor * propagators.inverse_far_field(fftobj)
if obj.dim() >= 3:
return probe[None,...] * upsampled_obj
else:
src/cdtools/tools/losses/__init__.py
+1
View File
@@ -1 +1,2 @@
from cdtools.tools.losses.losses import *
from cdtools.tools.losses.losses import __all__, __doc__
src/cdtools/tools/measurements/__init__.py
+1
View File
@@ -1 +1,2 @@
from cdtools.tools.measurements.measurements import *
from cdtools.tools.measurements.measurements import __all__, __doc__
src/cdtools/tools/measurements/measurements.py
+1 -1
View File
@@ -140,7 +140,7 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
if wavefields.dim() == 3:
output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
else:
output = avg_pool2d(output, oversampling)
output = avg_pool2d(output, int(oversampling))
# Then we grab the detector slice
if detector_slice is not None:
src/cdtools/tools/plotting/__init__.py
+1
View File
@@ -1,2 +1,3 @@
from cdtools.tools.plotting.plotting import *
from cdtools.tools.plotting.plotting import __all__, __doc__
src/cdtools/tools/polarization/__init__.py
-1
View File
@@ -1 +0,0 @@
from cdtools.tools.polarization.polarization import *
src/cdtools/tools/polarization/polarization.py
-217
View File
@@ -1,217 +0,0 @@
import numpy as numpy
import torch as t
# Abe - again, we don't need math here. replace with torch-native functions
# in all the definitions here. We have to use torch here if we want to be
# able to calculate derivatives w.r.t. the polarizer angle, say.
import math
from math import sin
from math import cos
__all__ = ['apply_linear_polarizer',
'apply_phase_retardance',
'apply_half_wave_plate',
'apply_quarter_wave_plate',
'apply_circular_polarizer',
'apply_jones_matrix',
'generate_linear_polarizer',
'generate_birefringent_obj']
def generate_linear_polarizer(pol_angle):
single_angle = False
pol_angle = t.as_tensor(pol_angle).to(dtype=t.float32)
if pol_angle.dim() == 0:
pol_angle = t.unsqueeze(pol_angle,0)
single_angle = True
pol_angle_rad = t.deg2rad(pol_angle)
a = t.cos(pol_angle_rad) ** 2
b = t.sin(pol_angle_rad) * t.cos(pol_angle_rad)
c = b
d = t.sin(pol_angle_rad) ** 2
ab = t.stack((a, b), dim=-1)
cd = t.stack((c, d), dim=-1)
jones_matrices = t.stack((ab, cd), dim=-2)
if single_angle:
return jones_matrices[0].to(dtype=t.cfloat)
else:
return jones_matrices.to(dtype=t.cfloat)
def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True):
"""
Applies a linear polarizer to the probe
Parameters:
----------
probe: t.Tensor
A (N)(P)x2xMxL tensor representing the probe, MxL - the size of the probe
The angle between the fast-axis of the linear polarizer and the horizontal axis
polarizer: t.Tensor
A 1D tensor (N) representing the polarizer angles for each of the patterns (or a single tensor of shape (1))
Returns:
--------
linearly polarized probe: t.Tensor
(N)(P)x2x1xMxL
"""
jones_matrices = generate_linear_polarizer(polarizer)
return apply_jones_matrix(probe, jones_matrices, transpose=transpose, multiple_modes=multiple_modes)
def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True):
# print('probe', probe.shape, 'jones matrix', jones_matrix.shape)
# if jones_matrix.shape == t.Size([5, 2, 2, 2, 2]):
# print('probe', probe.shape, 'jones', jones_matrix.shape)
# print('JONES', jones_matrix)
"""
Applies a given Jones matrix to the probe
Parameters:
----------
probe: t.Tensor
A (N)(P)x2xMxL tensor representing the probe
jones_matrix: t.tensor
(N)x2x2x(M)x(L)
Returns:
--------
a probe with the jones matrix applied: t.Tensor
(N)(P)x2xMxL
"""
if transpose:
if jones_matrix.dim() < 4:
jones_matrix = jones_matrix[..., None, None]
if multiple_modes:
jones_matrix = jones_matrix.unsqueeze(-5)
probe = probe[..., None, :, :]
# if jones matrices do not differ from pattern to pattern
if probe.dim() > jones_matrix.dim():
jones_matrix = jones_matrix.unsqueeze(0)
# vice versa
elif jones_matrix.dim() > probe.dim():
probe = probe.unsqueeze(0)
# print('apply jonesmatrix: probe', probe.shape, 'matrix:', jones_matrix)
jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4)
probe = probe.transpose(-1, -3).transpose(-2, -4)
output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3)
else:
raise NotImplementedError
return output
def apply_phase_retardance(probe, phase_shift, multiple_modes=True):
"""
Shifts the y-component of the field wrt the x-component by a given phase shift
Parameters:
----------
probe: t.Tensor
A (...)x2x1xMxL tensor representing the probe
phase_shift: float
phase shift in degrees
Returns:
--------
probe: t.Tensor
(...)x2x1xMxL
"""
theta = t.as_tensor(phase_shift, dtype=t.float32)
theta = t.deg2rad(theta)
probe = probe.to(dtype=t.cfloat)
jones_matrix = t.tensor([[1, 0], [0, t.exp(phase_shift)]]).to(dtype=t.cfloat)
polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
return polarized
def apply_circular_polarizer(probe, left_polarized=True, multiple_modes=True):
"""
Applies a circular polarizer to the probe
Parameters:
----------
probe: t.Tensor
A (...)x2xMxL tensor representing the probe
left_polarizd: bool
True for the left-polarization, False for the right
Returns:
--------
circularly polarized probe: t.Tensor
(...)x2xMxL
"""
probe = probe.to(dtype=t.cfloat)
if left_polarized:
jones_matrix = (1/2 * t.tensor([[1, -1j], [1j, 1]])).to(dtype=t.cfloat)
else:
jones_matrix = 1/2 * t.tensor([[1, 1j], [-1j, 1]]).to(dtype=t.cfloat)
polarized = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
return polarized
def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True):
"""
Parameters:
----------
probe: t.Tensor
A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe
fast_axis_angle: float
The angle between the fast-axis of the polarizer and the horizontal axis
Returns:
--------
polarized probe: t.Tensor
(...)x2x1xMxL
"""
probe = probe.to(dtype=t.cfloat)
theta = math.radians(fast_axis_angle)
exponent = t.exp(-1j * math.pi / 4 * t.ones(2, 2))
jones_matrix = exponent* t.tensor([[(cos(theta))**2 + 1j * (sin(theta))**2, (1 - 1j) * sin(theta) * cos(theta)], [(1 - 1j) * sin(theta) * cos(theta), (sin(theta))**2 + 1j * (cos(theta))**2]]).to(dtype=t.cfloat)
out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
return out
def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True):
"""
Parameters:
----------
probe: t.Tensor
A (...)x2x1xMxL tensor representing the probe, MxL - the size of the probe
fast_axis_angle: float
The angle between the fast-axis of the polarizer and the horizontal axis
Returns:
--------
polarized probe: t.Tensor
(...)x2x1xMxL
"""
probe = probe.to(dtype=t.cfloat)
theta = math.radians(fast_axis_angle)
exponent = t.exp(-1j * math.pi / 2 * t.ones(2, 2))
jones_matrix = exponent * t.tensor([[(cos(theta))**2 - (sin(theta))**2, 2 * sin(theta) * cos(theta)], [2 * sin(theta) * cos(theta), (sin(theta))**2 - (cos(theta))**2]]).to(dtype=t.cfloat)
out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes)
return out
def generate_birefringent_obj(fast_axis=90, phase_ret=10, atten_fast=1, atten_ret=1, global_phase=0):
def to_rad(angle):
angle = t.as_tensor(angle, dtype=t.float32)
angle = t.deg2rad(angle)
return angle
fast_axis = to_rad(fast_axis)
phase_ret = to_rad(phase_ret)
global_phase = to_rad(global_phase)
def coord_rot(angle):
a = t.stack((t.cos(angle), t.sin(angle)), dim=-1)
b = t.stack((-t.sin(angle), t.cos(angle)), dim=-1)
return t.stack((a, b), dim=-2).to(dtype=t.cfloat)
r1 = coord_rot(-fast_axis)
r2 = coord_rot(fast_axis)
p = t.exp(global_phase * 1j) * t.as_tensor([[atten_fast, 0], [0, atten_ret * t.exp(phase_ret*1j)]], dtype=t.cfloat)
return t.matmul(r1, t.matmul(p, r2))
src/cdtools/tools/propagators/__init__.py
+1
View File
@@ -1 +1,2 @@
from cdtools.tools.propagators.propagators import *
from cdtools.tools.propagators.propagators import __all__, __doc__
tests/conftest.py
+48 -1
View File
@@ -14,8 +14,41 @@ import datetime
def pytest_addoption(parser):
parser.addoption(
"--plot", action="store", default=False, help="plot: True to show test plots"
"--plot",
action="store_true",
default=False,
help="when set, shows the test plots"
)
parser.addoption(
"--reconstruction_device",
action="store",
default="cuda",
help="What device to run reconstructions on, if they are being run"
)
parser.addoption(
"--runslow",
action="store_true",
default=False,
help="run slow tests, primarily full reconstruction tests."
)
def pytest_configure(config):
config.addinivalue_line("markers", "slow: mark test as slow to run")
def pytest_collection_modifyitems(config, items):
if config.getoption("--runslow"):
# --runslow given in cli: do not skip slow tests
return
skip_slow = pytest.mark.skip(reason="need --runslow option to run")
for item in items:
if "slow" in item.keywords:
item.add_marker(skip_slow)
@pytest.fixture
def reconstruction_device(request):
return request.config.getoption("--reconstruction_device")
@pytest.fixture
@@ -319,3 +352,17 @@ def test_ptycho_cxis(ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3):
on the cxi files.
"""
return [ptycho_cxi_1, ptycho_cxi_2, ptycho_cxi_3]
@pytest.fixture(scope='module')
def gold_ball_cxi(pytestconfig):
return str(pytestconfig.rootpath) + \
'/examples/example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
@pytest.fixture(scope='module')
def lab_ptycho_cxi(pytestconfig):
return str(pytestconfig.rootpath) + \
'/examples/example_data/lab_ptycho_data.cxi'
tests/models/test_fancy_ptycho.py
+108
View File
@@ -0,0 +1,108 @@
import pytest
import cdtools
import torch as t
import cdtools
from matplotlib import pyplot as plt
@pytest.mark.slow
def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
print('\nTesting performance on the standard transmission ptycho dataset')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
oversampling=2,
exponentiate_obj=True,
dm_rank=2,
probe_support_radius=120,
propagation_distance=5e-3,
units='mm',
obj_view_crop=-50,
)
print('Running reconstruction on provided reconstruction_device,',
reconstruction_device)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
model.tidy_probes()
if show_plot:
model.inspect(dataset)
model.compare(dataset)
# If this fails, the reconstruction has gotten worse
assert model.loss_history[-1] < 0.001
@pytest.mark.slow
def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
print('\nTesting performance on the standard gold balls dataset')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
pad = 10
dataset.pad(pad)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
probe_support_radius=50,
propagation_distance=2e-6,
units='um',
probe_fourier_crop=pad
)
model.translation_offsets.data += \
0.7 * t.randn_like(model.translation_offsets)
# Not much probe intensity instability in this dataset, no need for this
model.weights.requires_grad = False
print('Running reconstruction on provided --reconstruction_device,',
reconstruction_device)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100,
schedule=True):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
model.tidy_probes()
if show_plot:
model.inspect(dataset)
model.compare(dataset)
# This just comes from running a reconstruction when it was working well
# and choosing a rough value. If it triggers this assertion error,
# something changed to make the final quality worse!
assert model.loss_history[-1] < 0.0001
tests/models/test_simple_ptycho.py
+24
View File
@@ -0,0 +1,24 @@
import pytest
import cdtools
from matplotlib import pyplot as plt
@pytest.mark.slow
def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
model = cdtools.models.SimplePtycho.from_dataset(dataset)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
for loss in model.Adam_optimize(100, dataset, batch_size=10):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
if show_plot:
model.inspect(dataset)
model.compare(dataset)
# If this fails, the reconstruction got worse
assert model.loss_history[-1] < 0.013
tests/tools/test_interactions.py
+7 -1
View File
@@ -253,7 +253,6 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
def test_RPI_interaction(random_probe, random_obj):
random_obj1 = random_obj[:79,:68] * 0 + 1
random_probe1 = random_probe * 0 + 1
t_random_obj1 = t.as_tensor(random_obj1)
@@ -269,6 +268,9 @@ def test_RPI_interaction(random_probe, random_obj):
output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups),
norm='ortho'))
output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1]
/ (random_obj1.shape[-2] * random_obj1.shape[-1]))
assert np.allclose(t_output1, output1)
random_obj2 = np.stack([random_obj[:64,:89]]*3)
@@ -286,5 +288,9 @@ def test_RPI_interaction(random_probe, random_obj):
output2 = random_probe2 * fftshift(fft.ifft2(ifftshift(obj2_ups),
norm='ortho'))
output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1]
/ (random_obj2.shape[-2] * random_obj2.shape[-1]))
assert np.allclose(t_output2, output2)
tests/tools/test_polarization.py
-624
View File
@@ -1,624 +0,0 @@
import numpy as np
import torch as t
from cdtools.tools.polarization import apply_linear_polarizer, generate_linear_polarizer
from cdtools.tools.polarization import apply_jones_matrix as jones
# Abe - I removed all the imports that didn't need to be here.
# Abe - A few issues. First, you could just write "from math import cos, sin"
# Second, we already have numpy imported, so better to use np.cos and np.sin
# from math import cos as cos, sin as sin
# numpy also has np.pi and np.deg2rad
from numpy import cos, sin, deg2rad
# Further comments:
#
# 1) You should also introduce an assert statement for the output shape
# checking, e.g. assert out.shape == t.Size(<shape>). See example in the
# first test.
#
# 2) Instead of using "assert a and b and c and d...", use separate assert
# statements for each condition. This is both more readable, and it allows the
# testing system to pinpoint exactly which of the assert statements failed.
#
# 3) USE SPACES INSTEAD OF TABS!!! 4 spaces per indent is the convention for
# this project.
#
# 4) Docstrings should go after the function definition, rather than before
#
# 5) Still missing cases where Jones matrix has multiple entries (N) but probe
# doesn't. Low priority.
angle = 87
angle_2 = angle - 45
def polarizer(angle):
theta = deg2rad(angle)
polarizer = t.tensor([[(cos(theta)) ** 2, sin(2 * theta) / 2], [sin(2 * theta) / 2, sin(theta) ** 2]]).to(dtype=t.cfloat)
return polarizer
exponent = t.exp(-1j * np.pi / 4 * t.ones(2, 2)).to(dtype=t.cfloat)
theta2 = deg2rad(angle_2)
quarter_plate = t.tensor([[(cos(theta2))**2 + 1j * (sin(theta2))**2, (1 - 1j) * sin(theta2) * cos(theta2)],
[(1 - 1j) * sin(theta2) * cos(theta2), (sin(theta2))**2 + 1j * (cos(theta2))**2]]).to(dtype=t.cfloat)
def build_from_quarters(jones1, jones2, jones3, jones4):
x = t.cat((t.stack((jones1, jones1), dim=-1), t.stack((jones2, jones2), dim=-1)), dim=-1)
y = t.cat((t.stack((jones3, jones3), dim=-1), t.stack((jones4, jones4), dim=-1)), dim=-1)
x = t.stack((x, x), dim=-2)
y = t.stack((y, y), dim=-2)
return t.cat((x, y), dim=-2).to(dtype=t.cfloat)
jones_plate = t.matmul(quarter_plate, polarizer(angle))
jones0 = polarizer(0)
jones90 = polarizer(90)
jones45 = polarizer(45)
transpose = True
def test_apply_jones_matrix_no_modes_no_mult_patterns_one_jones_matr():
'''
after applying the polarizer and the quarter_plate, the probe should get circularly polarized
probe: no multiple modes, 1 diffr pattern
2xMxL
jones_matrix: same jones matrix applied to all the pixels
2x2
'''
probe = t.rand(2, 3, 4, dtype=t.cfloat)
print(polarizer)
out = jones(jones(probe, polarizer(angle), multiple_modes=False, transpose=transpose),
quarter_plate, multiple_modes=False, transpose=transpose)
print('expected shape:(2, 3, 4)')
print('actual:', out.shape)
print('simulated:', out)
assert np.allclose(np.real(out[0]), np.imag(out[1]))
assert out.shape == t.Size((2, 3, 4))
def test_generate_linear_polarizer():
pol_angles = [0, 45, 90]
pol_angle1 = 45
pol_angle2 = t.tensor(90)
pol_angle3 = t.tensor([0])
pols = generate_linear_polarizer(pol_angles)
pol1 = generate_linear_polarizer(pol_angle1)
pol2 = generate_linear_polarizer(pol_angle2)
pol3 = generate_linear_polarizer(pol_angle3)
print('polarizers 0, 45, 90 (1D tensor) shape:', pols.shape)
print('shape of the polarizer generated from int:', pol1.shape)
print('shape of the polarizer generated from 0D tensor:', pol2.shape)
print('shape of the linear polarizer generated from t.Size(0) tensor:', pol3.shape)
print('90', pol2)
print(jones90)
probe = t.ones(4, 4)
jones_m = [jones0, jones45, jones90]
jones_m = t.stack([matr for matr in jones_m])
assert pols.shape == t.Size((3, 2, 2))
assert pol1.shape == t.Size((2, 2))
assert pol2.shape == t.Size((2, 2))
assert pol3.shape == t.Size((1, 2, 2))
assert t.allclose(pol1, jones45)
def test_apply_jones_matrix_no_modes_no_mult_patterns_diff_jones_matr():
'''probe: no multiple modes, 1 diffr pattern
2xMxL = 2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2x2xMxL = 2x2x4x4
4 quarters:
1: [:, :, :-2, :-2] - circular_polarizer,
2: [:, :, :-2, -2:] - 0
3: [:, :, -2:, :-2] - 90
4: [:, :, -2:, -2:] - 45
'''
jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
probe = t.ones(2, 4, 4).to(dtype=t.cfloat)
print('jones:', jones_matr)
# print('jones:', jones_matr)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
# jones -> (2,2,x,y), probe, output probe
# interaction ->
print('expected shape:(2, 4, 4)')
print('simulated:', out.shape)
print('simulated:', out)
# Example of using multiple asserts
assert np.allclose(np.real(out[0, :-2, :-2]), np.imag(out[1, :-2, :-2]))
assert t.allclose(out[0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, -2:, -2:], out[0, -2:, -2:])
assert out.shape == t.Size((2, 4, 4))
def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr():
'''
probe: no multiple modes, multiple diffr patterns
Nx2xMxL = 3x2x4x4
jones_matrix: same jones matrix applied to all the pixels
Nx2x2 = 3x2x2
3 different matrices for each probe:
1: 0
2: 45
3: 90
'''
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual:', out.shape)
print('simulated:', out)
assert t.allclose(out[0, 0, :, :], t.ones(4, 4, dtype=t.cfloat))
assert t.allclose(out[0, 1, :, :], t.zeros(4, 4, dtype=t.cfloat))
assert t.allclose(out[1, 0, :, :], out[1, 1])
assert t.allclose(out[2, 0, :, :], 3* t.zeros(4, 4, dtype=t.cfloat))
assert t.allclose(out[2, 1, :, :], 3 * t.ones(4, 4, dtype=t.cfloat))
assert out.shape == t.Size((3, 2, 4, 4))
def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr():
'''
probe: no multiple modes, multiple diffr pattern
Nx2xMxL = 3x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
assert np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
assert t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
assert t.allclose(out[1, 0, :-2, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:])
assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
assert np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:]))
assert t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2])
assert np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:]))
assert out.shape == t.Size((3, 2, 4, 4))
def test_apply_jones_matrix_mult_modes_1_pattern_one_jones_matr():
'''
probe: multiple modes, 1 diffr pattern
Px2xMxL = 2x2x3x4
jones_matrix: same jones matrix applied to all the pixels
2x2 - quarter plate
'''
probe = t.rand(2, 2, 3, 4, dtype=t.cfloat)
out = jones(jones(probe, polarizer(angle), multiple_modes=True, transpose=transpose), quarter_plate, multiple_modes=True, transpose=transpose)
print('expected shape: (2, 2, 3, 4)')
print('actual:', out.shape)
print('simulated:', out)
assert np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :]))
assert np.allclose(np.real(out[1, 0, :, :]), np.imag(out[1, 1, :, :]))
assert out.shape == t.Size((2, 2, 3, 4))
def test_apply_jones_matrix_mult_modes_1_pattern_diff_jones_matr():
'''
probe: multiple modes, 1 diffr pattern
Px2xMxL = 3x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2xMxL = 2x4x4
4 quarters:
1: [:, :, :-2, :-2] - circular_polarizer,
2: [:, :, :-2, -2:] - 0
3: [:, :, -2:, :-2] - 90
4: [:, :, -2:, -2:] - 45
'''
jones_matr = build_from_quarters(jones_plate, jones0, jones90, jones45)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
assert np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
assert t.allclose(out[0, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
assert np.allclose(np.real(out[1, 0, :-2, :-2]), np.imag(out[1, 1, :-2, :-2]))
assert t.allclose(out[1, 0, :-2, -2:], 2 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, -2:, :-2], 2 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 0, -2:, -2:], out[1, 1, -2:, -2:])
assert np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
assert t.allclose(out[2, 0, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:])
assert out.shape == t.Size((3, 2, 4, 4))
def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr():
'''
probe: multiple modes, multiple diffr patterns
NxPx2xMxL = 3x7x2x3x4
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
Nx2x2 = 3x2x2
3 different matrices for each probe in one mode:
1: 0
2: 45
3: 90
'''
probe = t.ones(2, 3, 4, dtype=t.cfloat)
# 1st mode
probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
# 2nd mode
probe_mode2 = 10 * t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 7, 2, 3, 4)')
print('actual:', out.shape)
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
# we'll be checking only one mode
assert t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
assert t.allclose(out[0, 6, 1, :, :], t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :])
assert t.allclose(out[2, 6, 0, :, :], 3 * (10**6) * t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[2, 6, 1, :, :], 3 * (10**6) * t.ones(3, 4, dtype=t.cfloat))
assert out.shape == t.Size((3, 7, 2, 3, 4))
def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr():
'''
probe: multiple modes, multiple diffr pattern
NxPx2xMxL = 3x4x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
(differs across the patterns in each mode)
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
# one mode
probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
print('probe:', probe.shape)
print('jones:', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 4, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated patterns in one mode:', out[:, 3, :, :, :])
o = 10 ** 3
# we'll be checking only one mode (4th)
assert np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2]))
assert t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:])
assert t.allclose(out[1, 3, 0, :-2, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:])
assert t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 1, -2:, :-2], 2 * o * t.ones(2, 2, dtype=t.cfloat))
assert np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:]))
assert t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 1, :-2, :-2], 3 * o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 0, :-2, -2:], 3 * o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2])
assert np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:]))
assert out.shape == t.Size((3, 4, 2, 4, 4))
def test_apply_jones_matrix_no_modes_mult_patterns_one_jones_matr_1():
'''
probe: no multiple modes, multiple diffr patterns
Nx2xMxL = 3x2x4x4
jones_matrix: same jones matrix applied to all the pixels
2x2 = 2x2 - a quarter waveplate
'''
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
jones_matr = jones_plate
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual:', out.shape)
print('simulated:', out)
assert np.allclose(np.real(out[:, 0, :, :]), np.imag(out[:, 1, :, :]))
assert out.shape == t.Size((3, 2, 4, 4))
def test_apply_jones_matrix_no_modes_mult_patterns_diff_jones_matr_1():
'''
probe: no multiple modes, multiple diffr pattern
Nx2xMxL = 3x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2x2xMxL = 2x2x4x4
1: quat plate, 2: 90, 3: 0, 4: 45
'''
jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
# check the 3rd pattern
assert np.allclose(np.real(out[2, 0, :-2, :-2]), np.imag(out[2, 1, :-2, :-2]))
assert t.allclose(out[2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, -2:], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, :-2], 3 * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, -2:], out[2, 1, -2:, -2:])
assert out.shape == t.Size((3, 2, 4, 4))
def test_apply_jones_matrix_mult_modes_mult_pattern_one_jones_matr_1():
'''
probe: multiple modes, multiple diffr patterns
NxPx2xMxL = 3x7x2x3x4
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
2x2 = 2x2
quarter wave plate
'''
probe = t.ones(2, 3, 4, dtype=t.cfloat)
# 1st mode
probe_mode1 = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode1 * 10 ** i for i in range(7)]), dim=1)
jones_matr = jones_plate
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 7, 2, 3, 4)')
print('actual:', out.shape)
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
# we'll be checking only one mode
assert np.allclose(np.real(out[:, :, 0, :, :]), np.imag(out[:, :, 1, :, :]))
assert out.shape == t.Size((3, 7, 2, 3, 4))
def test_apply_jones_matrix_mult_modes_mult_patterns_diff_jones_matr_1():
'''
probe: multiple modes, multiple diffr pattern
NxPx2xMxL = 3x4x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
2x2xMxL = 2x2x4x4
jones matrix:
1: quat plate, 2: 90, 3: 0, 4: 45
'''
jones_matr = build_from_quarters(jones_plate, jones90, jones0, jones45)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
# one mode
probe_mode = t.stack(([probe * (i + 1) for i in range(3)]), dim=0)
probe = t.stack(([probe_mode * (10 ** i) for i in range(4)]), dim=1)
print('probe:', probe.shape)
print('jones:', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 4, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated patterns in one mode:', out[:, 3, :, :, :])
o = 10 ** 2
# we'll be checking only one mode (3th)
assert np.allclose(np.real(out[0, 2, 0, :-2, :-2]), np.imag(out[0, 2, 1, :-2, :-2]))
assert t.allclose(out[0, 2, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 2, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 2, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 2, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 2, 0, -2:, -2:], out[0, 2, 1, -2:, -2:])
assert out.shape == t.Size((3, 4, 2, 4, 4))
return None
def test_apply_jones_matrix_no_mult_modes_one_pattern_probe_mult_patterns_jones_1():
'''
probe:
2xMxL = 2x3x4
jones_matrix:
Nx2x2 = 3x2x2
jones matrices:
1: quat plate, 2: 90, 3: 0
'''
jones_matr = t.stack((jones_plate, jones90, jones0))
probe = t.ones(2, 3, 4, dtype=t.cfloat)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 3, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
# we'll be checking only one mode (3th)
assert np.allclose(np.real(out[0, 0, :, :]), np.imag(out[0, 1, :, :]))
assert t.allclose(out[1, 0, :, :], t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[1, 1, :, :], t.ones(3, 4, dtype=t.cfloat))
assert t.allclose(out[2, 0, :, :], t.ones(3, 4, dtype=t.cfloat))
assert t.allclose(out[2, 1, :, :], t.zeros(3, 4, dtype=t.cfloat))
assert out.shape == t.Size((3, 2, 3, 4))
return None
def test_apply_jones_matrix_no_mult_modes_one_pattern_probe_mult_patterns_jones_2():
'''
probe:
2xMxL = 2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
out = jones(probe, jones_matr, multiple_modes=False, transpose=transpose)
print('expected shape: (3, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated:', out)
assert np.allclose(np.real(out[0, 0, :-2, :-2]), np.imag(out[0, 1, :-2, :-2]))
assert t.allclose(out[0, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 0, -2:, -2:], out[0, 1, -2:, -2:])
assert t.allclose(out[1, 0, :-2, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 0, :-2, -2:], out[1, 1, :-2, -2:])
assert t.allclose(out[1, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 1, -2:, :-2], t.ones(2, 2, dtype=t.cfloat))
assert np.allclose(np.real(out[1, 0, -2:, -2:]), np.real(out[1, 0, -2:, -2:]))
assert t.allclose(out[2, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, :-2], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, :-2, -2:], t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 0, -2:, :-2], out[2, 1, -2:, :-2])
assert np.allclose(np.real(out[2, 0, -2:, -2:]), np.real(out[2, 0, -2:, -2:]))
assert out.shape == t.Size((3, 2, 4, 4))
def test_apply_jones_matrix_mult_modes_one_pattern_probe_mult_patterns_jones_1():
'''
probe: multiple modes, multiple diffr patterns
Px2xMxL = 7x2x3x4
jones_matrix: same jones matrix applied to all the pixels (although differs from pattern to pattern)
Nx2x2 = 3x2x2
3 different matrices for each pattern:
1: 0
2: 45
3: 90
'''
probe = t.ones(2, 3, 4, dtype=t.cfloat)
probe = t.stack(([probe * 10 ** i for i in range(7)]), dim=0)
jones_matr = t.stack(([polarizer(angle) for angle in [0, 45, 90]]), dim=0)
print('probe shape:', probe.shape, 'jones shape', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 7, 2, 3, 4)')
print('actual:', out.shape)
print('simulated (patterns in one mode):', out[:, 6, :, :, :])
# we'll be checking only one mode
assert t.allclose(out[0, 6, 0, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
assert t.allclose(out[0, 6, 1, :, :], t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[1, 6, 0, :, :], out[1, 6, 1, :, :])
assert t.allclose(out[2, 6, 0, :, :], (10**6) * t.zeros(3, 4, dtype=t.cfloat))
assert t.allclose(out[2, 6, 1, :, :], (10**6) * t.ones(3, 4, dtype=t.cfloat))
assert out.shape == t.Size((3, 7, 2, 3, 4))
def test_apply_jones_matrix_mult_modes_one_pattern_probe_mult_patterns_jones_2():
'''
probe: multiple modes, multiple diffr pattern
Px2xMxL = 4x2x4x4
jones_matrix: jones matrices differ from pixel to pixel
Nx2x2xMxL = 3x2x2x4x4
(differs across the patterns in each mode)
jones matrix for the 1st pattern:
1: quat plate, 2: 90, 3: 0, 4: 45
jones matrix for the 2nd pattern:
1: 0, 2: 45, 3: 90, 4: quat plate
jones matrix for the 3rd pattern:
1: 90, 2: 0, 3: 45, 4: quat plate
'''
jones_m = [build_from_quarters(jones_plate, jones90, jones0, jones45),
build_from_quarters(jones0, jones45, jones90, jones_plate),
build_from_quarters(jones90, jones0, jones45, jones_plate)]
jones_matr = t.stack(([i for i in jones_m]), dim=0)
probe = t.ones(2, 4, 4, dtype=t.cfloat)
probe = t.stack(([probe * (10 ** i) for i in range(4)]), dim=0)
print('probe:', probe.shape)
print('jones:', jones_matr.shape)
out = jones(probe, jones_matr, multiple_modes=True, transpose=transpose)
print('expected shape: (3, 4, 2, 4, 4)')
print('actual shape:', out.shape)
print('simulated patterns in one mode:', out[:, 3, :, :, :])
o = 10 ** 3
# we'll be checking only one mode (4th)
assert np.allclose(np.real(out[0, 3, 0, :-2, :-2]), np.imag(out[0, 3, 1, :-2, :-2]))
assert t.allclose(out[0, 3, 0, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 1, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 0, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 1, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[0, 3, 0, -2:, -2:], out[0, 3, 1, -2:, -2:])
assert t.allclose(out[1, 3, 0, :-2, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 1, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 0, :-2, -2:], out[1, 3, 1, :-2, -2:])
assert t.allclose(out[1, 3, 0, -2:, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[1, 3, 1, -2:, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert np.allclose(np.real(out[1, 3, 0, -2:, -2:]), np.real(out[1, 3, 0, -2:, -2:]))
assert t.allclose(out[2, 3, 0, :-2, :-2], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 1, :-2, :-2], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 0, :-2, -2:], o * t.ones(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 1, :-2, -2:], t.zeros(2, 2, dtype=t.cfloat))
assert t.allclose(out[2, 3, 0, -2:, :-2], out[2, 3, 1, -2:, :-2])
assert np.allclose(np.real(out[2, 3, 0, -2:, -2:]), np.real(out[2, 3, 0, -2:, -2:]))
assert out.shape == t.Size((3, 4, 2, 4, 4))