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

Merge pull request #72 from cdtools-developers/feature/betterplots

Browse Source 
A bunch of long overdue updates to the plotting system
This commit is contained in:
Abe Levitan
2026-04-03 18:23:41 +02:00
committed by GitHub
parent 06105d38c7 a87e3415d2
commit 3f028e2fbf
30 changed files with 2002 additions and 701 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.flake8
+1 -1
View File
@@ -1,2 +1,2 @@
[flake8]
ignore = E501, W503
ignore = E501, W503, E731
.gitignore
+2 -1
View File
@@ -10,4 +10,5 @@ build/*
dist
*/example_data/*
*.h5
.DS_Store
.DS_Store
.ipynb_checkpoints
docs/source/examples.rst
+24 -6
View File
@@ -31,9 +31,9 @@ When reading this script, note the basic workflow. After the data is loaded, a m
Next, the model is moved to the GPU using the :code:`model.to` function. Any device understood by :code:`torch.Tensor.to` can be specified here. The next line is a bit more subtle - the dataset is told to move patterns to the GPU before passing them to the model using the :code:`dataset.get_as` function. This function does not move the stored patterns to the GPU. If there is sufficient GPU memory, the patterns can also be pre-moved to the GPU using :code:`dataset.to`, but the speedup is empirically quite small.
Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at the end of every epoch, to allow some monitoring code to be run.
Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at the end of every epoch, to allow some monitoring code to be run. Inside the loop, :code:`model.inspect(dataset)` is called every epoch to live-update a set of plots showing the current state of the model parameters.
Finally, the results can be studied using :code:`model.inspect(dataset)`, which creates or updates a set of plots showing the current state of the model parameters. :code:`model.compare(dataset)` is also called, which shows how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset.
Finally, :code:`model.compare(dataset)` is called to show how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset.
Fancy Ptycho
@@ -63,11 +63,13 @@ By default, FancyPtycho will also optimize over the following model parameters,
These corrections can be turned off (on) by calling :code:`model.<parameter>.requires_grad = False #(True)`.
Note as well two other changes that are made in this script, when compared to `simple_ptycho.py`. First, a `Reconstructor` object is explicitly created, in this case an `AdamReconstructor`. This object stores a model, dataset, and pytorch optimizer. It is then used to orchestrate the later reconstruction using a call to `Reconstructor.optimize()`.
Note as well two other changes that are made in this script, when compared to :code:`simple_ptycho.py`. First, a :code:`Reconstructor` object is explicitly created, in this case an :code:`AdamReconstructor`. This object stores a model, dataset, and pytorch optimizer. It is then used to orchestrate the later reconstruction using a call to :code:`Reconstructor.optimize()`.
We use this pattern, instead of the simpler call to `model.Adam_optimize()`, because having the reconstructor store the optimizer as well as the model and dataset allows the moment estimates to persist between multiple rounds of optimization. This leads to the second change: In this script, we run two optimization loops. The first loop aggressively refines the probe, with a low minibatch size and a high learning rate. The second loop has a smaller learning rate and a larger batch size, which allow for a more precise final estimation of the object.
We use this pattern, instead of the simpler call to :code:`model.Adam_optimize()`, because having the reconstructor store the optimizer as well as the model and dataset allows the moment estimates to persist between multiple rounds of optimization. This leads to the second change: In this script, we run two optimization loops. The first loop aggressively refines the probe, with a low minibatch size and a high learning rate. The second loop has a smaller learning rate and a larger batch size, which allow for a more precise final estimation of the object.
In this case, we used one reconstructor, but it is possible to create additional reconstructors to zero out all the persistant information in the optimizer, if desired, or even to instantiate multiple reconstructors on the same model with different optimization algorithms (e.g. `model.LBFGS_optimize()`).
In this case, we used one reconstructor, but it is possible to create additional reconstructors to zero out all the persistant information in the optimizer, if desired, or even to instantiate multiple reconstructors on the same model with different optimization algorithms (e.g. :code:`model.LBFGS_optimize()`).
Note also the use of :code:`min_interval=10` in the calls to :code:`model.inspect(dataset)`. Because generating plots can be expensive, passing a minimum interval (in seconds) prevents excessive replots. Finally, the call to :code:`model.inspect(dataset, replot_all=True)` at the end of the script reopens any plot windows that the user may have closed during the reconstruction, so that all results are visible at the end.
Gold Ball Ptycho
@@ -77,11 +79,27 @@ This script shows how the FancyPtycho model might be used in a realistic situati
.. literalinclude:: ../../examples/gold_ball_ptycho.py
Note, in particular, the use of :code:`model.save_on_exception` and :code:`model.save_to_h5` to save the results of the reconstruction. If a different file format is required, :code:`model.save_results` will save to a pure-python dictionary.
Note first the explicit addition of the :code:`plot_level=2` argument in the call to :code:`FancyPtycho.from_dataset`. This value controls which plots are generated. With :code:`plot_level=1`, only the main results are shown - :code:`plot_level=2` shows some more advanced monitoring of the error correction terms (background, position error, etc.), and :code:`plot_level=3` shows all registered plots.
Note also the use of :code:`model.save_on_exception` and :code:`model.save_to_h5` to save the results of the reconstruction. If a different file format is required, :code:`model.save_results` will save to a pure-python dictionary which can be processed further.
Finally, note that there are several small adjustments made to the script to counteract particular sources of error that are present in this dataset, for example the raster grid pathology caused by the scan pattern used. Also note that not every mixin is needed every time - in this case, we turn off optimization of the :code:`weights` parameter.
Near-Field Ptycho
-----------------
This script shows how the FancyPtycho model can be used on a typical near-field ptychography (also known as Fresnel ptychography) dataset.
.. literalinclude:: ../../examples/near_field_ptycho.py
The major change here is the setting of the :code:`near_field=True` argument to :code:`FancyPtycho.from_dataset`. This changes the propagator to a near-field propagator. As noted in the comments, if :code:`propagation_distance` is not set, the model will assume a standard near-field geomtry with flat illumination.
If :code:`propagation_distance` is set, it will assume a Fresnel scaling theorem-type geometry, with :code:`propagation_distance` as the focus-to-sample distance, and the distance set in the dataset object as the sample-to-detector distance.
Finally, note the addition of the :code:`panel_plot_mode=True` argument. This is the default mode, and returns the plots in a panel format, good for easily monitoring the progress of a reconstruction. If individual plots are needed for use in presentations, papers, or otherwise, setting :code:`panel_plot_mode=False` will plot each output in it's own window.
Gold Ball Split
---------------
docs/source/tutorial.rst
+27 -13
View File
@@ -367,28 +367,42 @@ The forward propagator maps the exit wave to the wave at the surface of the dete
Plotting
++++++++
The base CDIModel class has a function, :code:`model.inspect()`, which looks for a class variable called :code:`plot_list` and plots everything contained within. The plot list should be formatted as a list of tuples, with each tuple containing:
The base CDIModel class has a function, :code:`model.inspect()`, which looks for a class variable called :code:`plot_list` and plots everything contained within. The plot list should be formatted as a list of dictionaries, with each dictionary containing:
* :code:`'title'`: the title of the plot
* :code:`'plot_func'`: a function that takes in the model (and optionally a figure) and generates the relevant plot
* :code:`'condition'` (optional): a function that takes in the model and returns whether or not the plot should be generated
* The title of the plot
* A function that takes in the model and generates the relevant plot
* Optional, a function that takes in the model and returns whether or not the plot should be generated
.. code-block:: python
# 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))
{
'title': 'Probe Amplitude',
'plot_func': lambda self, fig:
p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis),
},
{
'title': 'Probe Phase',
'plot_func': lambda self, fig:
p.plot_phase(self.probe, fig=fig, basis=self.probe_basis),
},
{
'title': 'Object Amplitude',
'plot_func': lambda self, fig:
p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis),
},
{
'title': 'Object Phase',
'plot_func': 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`.
More advanced models like :code:`FancyPtycho` also define a :code:`plot_panel_list`, which groups related plots together into multi-subplot figures. The :code:`panel_plot_mode` argument (passed at construction time) controls whether these panels are rendered as combined multi-subplot figures or as individual windows. For a simple model like :code:`SimplePtycho`, :code:`plot_list` is sufficient.
Saving
++++++
examples/fancy_ptycho.py
+10 -9
View File
@@ -1,4 +1,5 @@
import cdtools
import torch as t
from matplotlib import pyplot as plt
filename = 'example_data/lab_ptycho_data.cxi'
@@ -15,9 +16,9 @@ model = cdtools.models.FancyPtycho.from_dataset(
obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix
)
device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)
if t.cuda.is_available():
model.to(device='cuda')
dataset.get_as(device='cuda')
# For this script, we use a slightly different pattern where we explicitly
# create a `Reconstructor` class to orchestrate the reconstruction. The
@@ -31,9 +32,9 @@ recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
# The batch size sets the minibatch size
for loss in recon.optimize(50, lr=0.02, batch_size=10):
print(model.report())
# Plotting is expensive, so we only do it every tenth epoch
if model.epoch % 10 == 0:
model.inspect(dataset)
# Because plotting can be expensive, setting a minimum plotting interval
# (in seconds) can avoid excessive replots.
model.inspect(dataset, min_interval=10)
# It's common to chain several different reconstruction loops. Here, we
# started with an aggressive refinement to find the probe in the previous
@@ -41,12 +42,12 @@ for loss in recon.optimize(50, lr=0.02, batch_size=10):
# and larger minibatch
for loss in recon.optimize(50, lr=0.005, batch_size=50):
print(model.report())
if model.epoch % 10 == 0:
model.inspect(dataset)
model.inspect(dataset, min_interval=10)
# This orthogonalizes the recovered probe modes
model.tidy_probes()
model.inspect(dataset)
# Setting replot_all will reopen any windows which were closed earlier
model.inspect(dataset, replot_all=True)
model.compare(dataset)
plt.show()
examples/fancy_ptycho_inline.ipynb
+151
View File
@@ -0,0 +1,151 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "286054ce",
"metadata": {},
"outputs": [],
"source": [
"%matplotlib inline\n",
"import cdtools\n",
"import torch as t\n",
"from matplotlib import pyplot as plt"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "955bb242-e2ed-47c3-919c-1ea690681445",
"metadata": {},
"outputs": [],
"source": [
"# Load and inspect a dataset\n",
"\n",
"filename = 'example_data/lab_ptycho_data.cxi'\n",
"dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)\n",
"\n",
"dataset.inspect();"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "82f71fb4-f013-46bb-817b-1973ba23336a",
"metadata": {},
"outputs": [],
"source": [
"# Initialize a model from the dataset and move it to the GPU.\n",
"\n",
"model = cdtools.models.FancyPtycho.from_dataset(\n",
" dataset,\n",
" n_modes=3, # Use 3 incoherently mixing probe modes\n",
" oversampling=2, # Simulate the probe on a 2xlarger real-space array\n",
" probe_support_radius=120, # Force the probe to 0 outside a radius of 120 pix\n",
" propagation_distance=5e-3, # Propagate the initial probe guess by 5 mm\n",
" units='mm', # Set the units for the live plots\n",
" obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix,\n",
")\n",
"\n",
"if t.cuda.is_available():\n",
" model.to(device='cuda')\n",
" dataset.get_as(device='cuda')\n",
"\n",
"# Then, create a reconstructor object and view the initialized model\n",
"\n",
"recon = cdtools.reconstructors.AdamReconstructor(model, dataset)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0e22a65e-280b-428d-a6f5-f3206d22110b",
"metadata": {},
"outputs": [],
"source": [
"# Workaround reconstruction pattern for interactive plotting in jupyter:\n",
"# First, a standalone cell to plot the current model state\n",
"\n",
"model.inspect(dataset, replot_all=True);"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "81a768b5",
"metadata": {},
"outputs": [],
"source": [
"# Second, a cell for running the reconstruction. With this pattern, it is safe\n",
"# to interrupt the kernel. Then, the cell above can be re-run to refresh the plots.\n",
"while model.epoch < 50:\n",
" for loss in recon.optimize(1, lr=0.02, batch_size=10):\n",
" print(model.report())\n",
"\n",
"while model.epoch < 100:\n",
" for loss in recon.optimize(1, lr=0.005, batch_size=10):\n",
" print(model.report())"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "237e9286-b6cf-41dc-aeb0-89abfe59b37c",
"metadata": {},
"outputs": [],
"source": [
"# Save out the results\n",
"\n",
"model.save_to_h5('lab_ptycho_reconstruction.h5', dataset)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0f1565b6",
"metadata": {},
"outputs": [],
"source": [
"# Finalize the plotting and create the comparison plot\n",
"\n",
"# This orthogonalizes the recovered probe modes. It is best to do so\n",
"# after saving the results, if you intend to initialize any further\n",
"# reconstructions with the probe.\n",
"model.tidy_probes()\n",
"\n",
"# Final plotting\n",
"model.inspect(dataset)\n",
"model.compare(dataset);"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "63ec3aa4-0d1c-4775-9fb2-3002d404faa4",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
examples/fancy_ptycho_interactive.ipynb
+144
View File
@@ -0,0 +1,144 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "286054ce",
"metadata": {},
"outputs": [],
"source": [
"%matplotlib widget\n",
"import cdtools\n",
"import torch as t\n",
"from matplotlib import pyplot as plt"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "955bb242-e2ed-47c3-919c-1ea690681445",
"metadata": {},
"outputs": [],
"source": [
"# Load and inspect a dataset\n",
"\n",
"filename = 'example_data/lab_ptycho_data.cxi'\n",
"dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)\n",
"\n",
"dataset.inspect();"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "82f71fb4-f013-46bb-817b-1973ba23336a",
"metadata": {},
"outputs": [],
"source": [
"# Initialize a model from the dataset and move it to the GPU.\n",
"\n",
"model = cdtools.models.FancyPtycho.from_dataset(\n",
" dataset,\n",
" n_modes=3, # Use 3 incoherently mixing probe modes\n",
" oversampling=2, # Simulate the probe on a 2xlarger real-space array\n",
" probe_support_radius=120, # Force the probe to 0 outside a radius of 120 pix\n",
" propagation_distance=5e-3, # Propagate the initial probe guess by 5 mm\n",
" units='mm', # Set the units for the live plots\n",
" obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix,\n",
")\n",
"\n",
"if t.cuda.is_available():\n",
" model.to(device='cuda')\n",
" dataset.get_as(device='cuda')\n",
"\n",
"# Then, create a reconstructor object and view the initialized model\n",
"\n",
"recon = cdtools.reconstructors.AdamReconstructor(model, dataset)\n",
"model.inspect(dataset);"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "81a768b5",
"metadata": {},
"outputs": [],
"source": [
"# Workaround reconstruction pattern for interactive plotting in jupyter:\n",
"\n",
"# With this pattern, it is safe to interrupt the kernel. Doing so will\n",
"# trigger an update of the plots, at which point the current state can\n",
"# be viewed. Then this cell can be re-run to continue the reconstruction\n",
"while model.epoch < 50:\n",
" for loss in recon.optimize(1, lr=0.02, batch_size=10):\n",
" print(model.report())\n",
" model.inspect(dataset, min_interval=10)\n",
"\n",
"while model.epoch < 100:\n",
" for loss in recon.optimize(1, lr=0.005, batch_size=10):\n",
" print(model.report())\n",
" model.inspect(dataset, min_interval=10)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "237e9286-b6cf-41dc-aeb0-89abfe59b37c",
"metadata": {},
"outputs": [],
"source": [
"# Save out the results\n",
"\n",
"model.save_to_h5('lab_ptycho_reconstruction.h5', dataset)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0f1565b6",
"metadata": {},
"outputs": [],
"source": [
"# Finalize the plotting and create the comparison plot\n",
"\n",
"# This orthogonalizes the recovered probe modes. It is best to do so\n",
"# after saving the results, if you intend to initialize any further\n",
"# reconstructions with the probe.\n",
"model.tidy_probes()\n",
"\n",
"# Final plotting\n",
"model.inspect(dataset)\n",
"model.compare(dataset);"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "63ec3aa4-0d1c-4775-9fb2-3002d404faa4",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
examples/gold_ball_ptycho.py
+10 -12
View File
@@ -1,6 +1,6 @@
import cdtools
from matplotlib import pyplot as plt
import torch as t
from matplotlib import pyplot as plt
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)
@@ -26,7 +26,8 @@ model = cdtools.models.FancyPtycho.from_dataset(
probe_support_radius=50,
propagation_distance=2e-6,
units='um',
probe_fourier_crop=pad
probe_fourier_crop=pad,
plot_level=2,
)
@@ -39,9 +40,9 @@ 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
device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)
if t.cuda.is_available():
model.to(device='cuda')
dataset.get_as(device='cuda')
# Create the reconstructor
recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
@@ -53,13 +54,11 @@ with model.save_on_exception(
for loss in recon.optimize(20, lr=0.005, batch_size=50):
print(model.report())
if model.epoch % 10 == 0:
model.inspect(dataset)
model.inspect(dataset, min_interval=5)
for loss in recon.optimize(50, lr=0.002, batch_size=100):
print(model.report())
if model.epoch % 10 == 0:
model.inspect(dataset)
model.inspect(dataset, min_interval=5)
# We can often reset our guess of the probe positions once we have a
# good guess of probe and object, but in this case it causes the
@@ -70,8 +69,7 @@ with model.save_on_exception(
# the loss fails to improve after 10 epochs
for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True):
print(model.report())
if model.epoch % 10 == 0:
model.inspect(dataset)
model.inspect(dataset, min_interval=5)
model.tidy_probes()
@@ -79,6 +77,6 @@ model.tidy_probes()
# This saves the final result
model.save_to_h5('example_reconstructions/gold_balls.h5', dataset)
model.inspect(dataset)
model.inspect(dataset, replot_all=True)
model.compare(dataset)
plt.show()
examples/gold_ball_split.py
+3 -3
View File
@@ -32,9 +32,9 @@ for label, dataset in zip(labels, datasets):
model.weights.requires_grad = False
device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)
if t.cuda.is_available():
model.to(device='cuda')
dataset.get_as(device='cuda')
# Create the reconstructor
recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
examples/gold_ball_synthesize.py
+3 -3
View File
@@ -5,11 +5,11 @@ import numpy as np
# We load all three reconstructions
half_1 = cdtools.tools.data.h5_to_nested_dict(
f'example_reconstructions/gold_balls_half_1.h5')
'example_reconstructions/gold_balls_half_1.h5')
half_2 = cdtools.tools.data.h5_to_nested_dict(
f'example_reconstructions/gold_balls_half_2.h5')
'example_reconstructions/gold_balls_half_2.h5')
full = cdtools.tools.data.h5_to_nested_dict(
f'example_reconstructions/gold_balls_full.h5')
'example_reconstructions/gold_balls_full.h5')
# This defines the region of recovered object to use for the analysis.
pad = 260
examples/near_field_ptycho.py
+8 -10
View File
@@ -1,11 +1,11 @@
import cdtools
import torch as t
from matplotlib import pyplot as plt
filename = 'example_data/PETRAIII_P25_Near_Field_Ptycho.cxi'
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)
dataset.inspect()
plt.show()
# Setting near_field equal to True uses an angular spectrum propagator in
# lieu of the default Fourier-transform propagator for far-field ptychography.
@@ -27,11 +27,12 @@ model = cdtools.models.FancyPtycho.from_dataset(
propagation_distance=3.65e-3, # 3.65 downstream from focus
units='um', # Set the units for the live plots
obj_view_crop=-35,
panel_plot_mode=True, # Set to False to get individual figures
)
device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)
if t.cuda.is_available():
model.to(device='cuda')
dataset.get_as(device='cuda')
model.inspect(dataset)
@@ -39,18 +40,15 @@ recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
for loss in recon.optimize(100, lr=0.04, batch_size=10):
print(model.report())
# Plotting is expensive, so we only do it every tenth epoch
if model.epoch % 10 == 0:
model.inspect(dataset)
model.inspect(dataset, min_interval=5)
for loss in recon.optimize(50, lr=0.005, batch_size=50):
print(model.report())
if model.epoch % 10 == 0:
model.inspect(dataset)
model.inspect(dataset, min_interval=5)
# This orthogonalizes the recovered probe modes
model.tidy_probes()
model.inspect(dataset)
model.inspect(dataset, replot_all=True)
model.compare(dataset)
plt.show()
examples/simple_ptycho.py
+8 -6
View File
@@ -8,6 +8,7 @@ more powerful FancyPtycho model and include more information on how to
correct for common sources of error.
"""
import cdtools
import torch as t
from matplotlib import pyplot as plt
# We load an example dataset from a .cxi file
@@ -17,10 +18,12 @@ dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)
# We create a ptychography model from the dataset
model = cdtools.models.SimplePtycho.from_dataset(dataset)
# We move the model to the GPU
device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)
# We move the model to the GPU, if possible
if t.cuda.is_available():
model.to(device='cuda')
dataset.get_as(device='cuda')
model.inspect(dataset)
# We run the reconstruction
for loss in model.Adam_optimize(100, dataset, batch_size=10):
@@ -29,7 +32,6 @@ for loss in model.Adam_optimize(100, dataset, batch_size=10):
# And liveplot the updates to the model as they happen
model.inspect(dataset)
# We study the results
model.inspect(dataset)
# We open a comparison of the simulated and measured data
model.compare(dataset)
plt.show()
examples/transmission_RPI.py
+9 -7
View File
@@ -1,5 +1,6 @@
import cdtools
import pickle
import torch as t
from matplotlib import pyplot as plt
# First, we load an example dataset from a .cxi file
@@ -22,26 +23,27 @@ model = cdtools.models.RPI.from_dataset(dataset, probe, [500,500],
# Let's do this reconstruction on the GPU, shall we?
model.to(device='cuda')
dataset.get_as(device='cuda')
if t.cuda.is_available():
model.to(device='cuda')
dataset.get_as(device='cuda')
# Note that the inspect step takes the vast majority of the time
# The regularization is an L2 regularizer that empirically helps accelerate
# convergence
for loss in model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]):
model.inspect(dataset)
model.inspect(dataset, min_interval=5)
print(model.report())
# Now we use the regularizer to damp all but the top modes
for loss in model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1]):
#model.inspect(dataset)
model.inspect(dataset, min_interval=5)
print(model.report())
# Save results to a python dictionary
results = model.save_results()
# Save results to an h5 file
model.save_to_h5('example_reconstructions/transmission_RPI.h5', dataset)
# Finally, we plot the results
model.inspect(dataset)
model.inspect(dataset, replot_all=True)
model.compare(dataset)
plt.show()
examples/tutorial_finale.py
+4 -2
View File
@@ -1,5 +1,6 @@
from tutorial_basic_ptycho_dataset import BasicPtychoDataset
from tutorial_simple_ptycho import SimplePtycho
import torch as t
from h5py import File
from matplotlib import pyplot as plt
@@ -11,8 +12,9 @@ dataset.inspect()
model = SimplePtycho.from_dataset(dataset)
model.to(device='cuda')
dataset.get_as(device='cuda')
if t.cuda.is_available():
model.to(device='cuda')
dataset.get_as(device='cuda')
for loss in model.Adam_optimize(10, dataset):
model.inspect(dataset)
examples/tutorial_simple_ptycho.py
+20 -8
View File
@@ -108,14 +108,26 @@ class SimplePtycho(CDIModel):
# 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))
{
'title': 'Probe Amplitude',
'plot_func': lambda self, fig:
p.plot_amplitude(self.probe, fig, basis=self.probe_basis),
},
{
'title': 'Probe Phase',
'plot_func': lambda self, fig:
p.plot_phase(self.probe, fig, basis=self.probe_basis)
},
{
'title': 'Object Amplitude',
'plot_func': lambda self, fig:
p.plot_amplitude(self.obj, fig, basis=self.probe_basis)
},
{
'title': 'Object Phase',
'plot_func': lambda self, fig:
p.plot_phase(self.obj, fig, basis=self.probe_basis)
},
]
def save_results(self, dataset):
pyproject.toml
+2 -1
View File
@@ -50,4 +50,5 @@ docs = [
]
[tool.ruff]
line-length = 79
line-length = 79
ignore = ["E501", "E731"]
src/cdtools/datasets/ptycho_2d_dataset.py
-2
View File
@@ -245,8 +245,6 @@ class Ptycho2DDataset(CDataset):
return np.log10((meas_data * mask) + log_offset)
else:
return meas_data * mask
translations = self.translations.detach().cpu().numpy()
# This takes about twice as long as it would to just do it all at
# once, but it avoids creating another self.patterns-sized array
src/cdtools/models/base.py
+342 -73
View File
@@ -29,21 +29,17 @@ loss
"""
import torch as t
from torch.utils import data as torchdata
import matplotlib
from matplotlib import pyplot as plt
from matplotlib.widgets import Slider
from matplotlib import ticker
import numpy as np
import threading
import queue
import time
from scipy import io
from contextlib import contextmanager
from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch
from cdtools.tools.data import nested_dict_to_h5, nested_dict_to_numpy, nested_dict_to_torch
from cdtools.reconstructors import AdamReconstructor, LBFGSReconstructor, SGDReconstructor
from cdtools.datasets import CDataset
from typing import List, Union, Tuple
import os
__all__ = ['CDIModel']
@@ -58,12 +54,25 @@ class CDIModel(t.nn.Module):
functions.
"""
def __init__(self):
def __init__(self, panel_plot_mode=False, plot_level=np.inf):
"""Initializes the CDIModel base class.
Parameters
----------
panel_plot_mode : bool, default: False
If True, plot_panel_list entries are rendered as multi-subplot
figures. If False, each subplot is rendered as its own figure.
plot_level : float, default: np.inf
Only plots whose plot_level <= this value are shown.
"""
super(CDIModel, self).__init__()
self.loss_history = []
self.training_history = ''
self.epoch = 0
self.panel_plot_mode = panel_plot_mode
self.plot_level = plot_level
self.last_inspected_time = None
def from_dataset(self, dataset):
raise NotImplementedError()
@@ -547,98 +556,357 @@ class CDIModel(t.nn.Module):
return msg
# By default, the plot_list is empty
# By default, the plot lists are empty
plot_panel_list = []
plot_list = []
def inspect(self, dataset=None, update=True):
"""Plots all the plots defined in the model's plot_list attribute
def inspect(self, dataset=None, replot_all=False, min_interval=None):
"""Plots all the plots defined in the model's plot_panel_list and plot_list attributes
If update is set to True, it will update any previously plotted set
of plots, if one exists, and then redraw them. Otherwise, it will
plot a new set, and any subsequent updates will update the new set
Updates any previously plotted figures that are still open. Figures
that have been closed are left closed unless replot_all=True.
Optionally, a dataset can be passed, which will allow plotting of any
registered plots which need to incorporate some information from
the dataset (such as geometry or a comparison with measured data).
Plots can be registered in any subclass by defining the plot_list
attribute. This should be a list of tuples in the following format:
( 'Plot Title', function_to_generate_plot(self),
function_to_determine_whether_to_plot(self))
Plots can be registered in any subclass by defining plot_panel_list
and/or plot_list class attributes. See the CDIModel documentation for
the expected dict-based format of each.
Where the third element in the tuple (a function that returns
True if the plot is relevant) is not required.
When panel_plot_mode=True (set in __init__), plot_panel_list entries
are rendered as multi-subplot figures. When False (the default),
each subplot in plot_panel_list is rendered as its own figure,
prepended to any standalone plot_list entries.
The plot_level attribute (set in __init__, default np.inf) controls
which plots are shown: a panel or standalone plot is only shown when
its plot_level <= self.plot_level.
Parameters
----------
dataset : CDataset
Optional, a dataset matched to the model type
update : bool, default: True
Whether to update existing plots or plot new ones
replot_all : bool, default: False
If True, recreate figures that were previously closed by the user.
min_interval : float, optional
If set, skip updating plots if fewer than this many seconds have
elapsed since the last call to inspect().
"""
# We find or create all the figures
first_update = False
if update and hasattr(self, 'figs') and self.figs:
figs = self.figs
elif update:
figs = None
self.figs = []
first_update = True
if (min_interval is not None
and self.last_inspected_time is not None
and time.time() - self.last_inspected_time < min_interval):
return
plot_panel_list = getattr(self, 'plot_panel_list', None) or []
plot_list = getattr(self, 'plot_list', None) or []
if self.panel_plot_mode:
# First we plot all the panels
panel_figs = self._inspect_panel(
plot_panel_list, dataset=dataset, replot_all=replot_all)
# And then we plot all the individual figures
individual_figs = self._inspect_individual_figures(
plot_list, dataset=dataset, replot_all=replot_all
)
self.figs = panel_figs + individual_figs
else:
figs = None
self.figs = []
# If not in panel plot mode, we first flatten the figures
# from the panels
flat = []
for panel in plot_panel_list:
panel_level = panel.get('plot_level', 1)
for plot in panel['plots']:
# We add the plot level from the larger panel
flat.append({**plot, 'plot_level': panel_level})
all_plots = flat + list(plot_list)
idx = 0
for plots in self.plot_list:
# If a conditional is included in the plot, we check whether
# it is True
try:
if len(plots) >=3 and not plots[2](self):
continue
except TypeError as e:
if len(plots) >= 3 and not plots[2](self, dataset):
continue
# We make sure to keep a reference to the open figs around
self.figs = self._inspect_individual_figures(
all_plots, dataset=dataset, replot_all=replot_all)
name = plots[0]
plotter = plots[1]
if self.last_inspected_time is None or replot_all:
# Somehow, this is needed for new figures to appear
if self._is_backend_interactive():
plt.pause(0.05 * len(self.figs))
for fig in self.figs:
fig.canvas.flush_events()
if figs is None:
fig = plt.figure()
self.figs.append(fig)
self.last_inspected_time = time.time()
def _is_backend_interactive(
self
):
"""Returns True if the current matplotlib backend is interactive."""
backend = matplotlib.get_backend().lower()
try:
# matplotlib >= 3.9
interactive_bk = matplotlib.backends.backend_registry.list_builtin(
matplotlib.backends.BackendFilter.INTERACTIVE
)
except AttributeError:
# older matplotlib
interactive_bk = matplotlib.rcsetup.interactive_bk
return backend in [b.lower() for b in interactive_bk]
def _inspect_individual_figures(
self,
plot_list,
dataset=None,
replot_all=False
):
"""Core one-figure-per-plot rendering logic.
This is the function which is called internally by model.inspect()
to plot all the figures registered on plot_list, or to plot all
figures if panel_plot_mode is set to False.
Parameters
----------
plot_list : dict
The list of registerd plots to show.
dataset : CDataset
A CDataset object which this model is reconstructing, which may
store some needed information such as original translations.
replot_all : bool
if set to True, will reopen closed figures. Otherwise, will skip.
Returns
-------
rendered : list
the list of figures object that were rendered this call.
"""
rendered = []
for plot in plot_list:
# Level filter
if plot.get('plot_level', 1) > self.plot_level:
continue
# Condition check
condition = plot.get('condition', None)
if condition is not None:
try:
if not condition(self):
continue
except TypeError:
if not condition(self, dataset):
continue
figsize = plot.get('figure_size', None)
if self.last_inspected_time is not None and \
not replot_all and \
not plt.fignum_exists(plot['title']):
continue
if self.last_inspected_time is None:
fig = plt.figure(plot['title'],
figsize=figsize)
else:
fig = figs[idx]
with plt.rc_context({'figure.raise_window': False}):
fig = plt.figure(plot['title'],
figsize = figsize)
try: # We try just plotting using the simplest allowed signature
plotter(self,fig)
plt.title(name)
except TypeError as e:
# TypeError implies it wanted another argument, i.e. a dataset
try:
plot['plot_func'](self, fig)
if plt.gca().get_title().strip() == '':
plt.title(plot['title'])
except TypeError:
if dataset is not None:
try:
plotter(self, fig, dataset)
plt.title(name)
except Exception as e: # Don't raise errors: it's just plots
plot['plot_func'](self, fig, dataset)
if plt.gca().get_title().strip() == '':
plt.title(plot['title'])
except KeyboardInterrupt:
raise
except Exception:
pass
except Exception as e: # Don't raise errors, it's just a plot
except KeyboardInterrupt:
raise
except Exception:
pass
idx += 1
if update:
# This seems to update the figure without blocking.
rendered.append(fig)
if self._is_backend_interactive():
plt.draw()
fig.canvas.start_event_loop(0.001)
if first_update:
# But this is needed the first time the figures update, or
# they won't get drawn at all
plt.pause(0.05 * len(self.figs))
return rendered
def _inspect_panel(
self,
plot_panel_list,
dataset=None,
replot_all=False,
):
"""Multi-subplot panel rendering.
This is the function which is called internally by model.inspect()
to plot all the figures registered on plot_panel_list, and is only
used when panel_plot_mode is set to True.
Parameters
----------
plot_panel_list : dict
The list of registerd plot panels to show.
dataset : CDataset
A CDataset object which this model is reconstructing, which may
store some needed information such as original translations.
replot_all : bool
if set to True, will reopen closed figures. Otherwise, will skip.
Returns
-------
rendered : list
the list of figures object that were rendered this call.
"""
rendered = []
for panel_def in plot_panel_list[::-1]: # Flip so first ones show on top
panel_level = panel_def.get('plot_level', 1)
if panel_level > self.plot_level:
continue # skip entire panel
panel_condition = panel_def.get('condition', None)
if panel_condition is not None:
try:
if not panel_condition(self):
continue
except TypeError:
if not panel_condition(self, dataset):
continue
nrows, ncols = panel_def['grid']
figsize = panel_def.get('figure_size', None)
if self.last_inspected_time is not None and \
not replot_all and \
not plt.fignum_exists(panel_def['title']):
continue
if self.last_inspected_time is None:
fig = plt.figure(panel_def['title'], figsize=figsize)
else:
with plt.rc_context({'figure.raise_window': False}):
fig = plt.figure(panel_def['title'], figsize=figsize)
for subfig in fig.subfigs:
if hasattr(subfig, '_sliders'):
for slider in subfig._sliders:
slider.disconnect_events()
fig.clear()
gs = fig.add_gridspec(
nrows, ncols,
width_ratios=[1]*ncols,
height_ratios=[1]*nrows,
)
for plot in panel_def['plots']:
condition = plot.get('condition', None)
if condition is not None:
try:
if not condition(self):
continue
except TypeError:
if not condition(self, dataset):
continue
subfig = fig.add_subfigure(gs[plot['subplot'][0],
plot['subplot'][1]])
try:
plot['plot_func'](self, subfig)
if plt.gca().get_title().strip() == '':
plt.title(plot['title'])
except TypeError:
if dataset is not None:
try:
plot['plot_func'](self, subfig, dataset)
if plt.gca().get_title().strip() == '':
plt.title(plot['title'])
except KeyboardInterrupt:
raise
except Exception:
pass
except KeyboardInterrupt:
raise
except Exception:
raise
rendered.append(fig)
if self._is_backend_interactive():
plt.draw()
return rendered
def plot_loss_history(self, fig=None, clear_fig=True):
"""Plots the loss history on a semilogy axis
Parameters
----------
fig : matplotlib.figure.Figure
Default is a new figure, a matplotlib figure to use to plot
clear_fig : bool
Default is True. Whether to clear the figure before plotting.
Returns
-------
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
if clear_fig:
fig.clear()
if len(fig.axes) >= 1:
ax = fig.axes[0]
else:
try:
total_width, total_height = fig.get_size_inches()
except AttributeError:
# Only support one layer of nested subfigures
main_fig = fig.figure # get enclosing figure
fig_w, fig_h = main_fig.get_size_inches()
total_width = fig.bbox.width * fig_w / main_fig.bbox.width
total_height = fig.bbox.height * fig_h / main_fig.bbox.height
except AttributeError:
# Fall back to default figsize
total_width, total_height = (6.4, 4.8)
pad_left = 0.6 / total_height
# De-adjusts for an ad-hoc offset introduced by matplotlib
pad_right = 0.6 / total_width - 0.05
pad_bottom = 0.5 / total_height
pad_top = 0.4 / total_height
im_ax_bottom = pad_bottom
im_ax_height = 1 - pad_top - im_ax_bottom
ax = fig.add_axes(
[pad_left, im_ax_bottom, 1-pad_left-pad_right, im_ax_height]
)
ax.semilogy(self.loss_history)
plt.title('Loss History')
ax.set_xlabel('Epoch')
ax.set_ylabel('Loss Metric')
return fig
def save_figures(self, prefix='', extension='.pdf'):
"""Saves all currently open inspection figures.
@@ -661,14 +929,15 @@ class CDIModel(t.nn.Module):
Default is .eps, the file extension to save with.
"""
if hasattr(self, 'figs') and self.figs:
figs = self.figs
else:
return # No figures to save
if not (hasattr(self, 'figs') and self.figs):
return # No figures to save
for fig in self.figs:
fig.savefig(prefix + fig.axes[0].get_title() + extension,
bbox_inches = 'tight')
if hasattr(fig, '_panel_label') and fig._panel_label:
label = fig._panel_label
else:
label = fig.axes[0].get_title() if fig.axes else 'figure'
fig.savefig(prefix + label + extension, bbox_inches='tight')
def compare(self, dataset, logarithmic=False):
src/cdtools/models/bragg_2d_ptycho.py
+49 -44
View File
@@ -7,7 +7,6 @@ from cdtools.tools.propagators import generate_generalized_angular_spectrum_prop
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from copy import copy
__all__ = ['Bragg2DPtycho']
@@ -80,6 +79,8 @@ class Bragg2DPtycho(CDIModel):
units='um',
dtype=t.float32,
obj_view_crop=0,
panel_plot_mode=False,
plot_level=1,
):
# We need the detector geometry
@@ -91,7 +92,8 @@ class Bragg2DPtycho(CDIModel):
# translation_offsets can stay 2D for now
# propagate_probe and correct_tilt are important!
super(Bragg2DPtycho, self).__init__()
super(Bragg2DPtycho, self).__init__(panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
self.register_buffer('wavelength',
t.as_tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
@@ -258,7 +260,9 @@ class Bragg2DPtycho(CDIModel):
obj_padding=200,
obj_view_crop=None,
units='um',
surface_normal=None
surface_normal=None,
panel_plot_mode=False,
plot_level=1,
):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -436,8 +440,8 @@ class Bragg2DPtycho(CDIModel):
return cls(wavelength, det_geo, obj_basis, probe, obj,
min_translation=min_translation,
probe_basis=probe_basis,
median_propagation =median_propagation,
translation_offsets = translation_offsets,
median_propagation=median_propagation,
translation_offsets=translation_offsets,
weights=weights, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
@@ -448,6 +452,8 @@ class Bragg2DPtycho(CDIModel):
lens=lens,
obj_view_crop=obj_view_crop,
units=units,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
)
@@ -553,7 +559,6 @@ class Bragg2DPtycho(CDIModel):
'orientation': orientation}
detector_geometry = self.detector_geometry
mask = self.mask
wavelength = self.wavelength
indices, translations = args_list
@@ -578,90 +583,90 @@ class Bragg2DPtycho(CDIModel):
plot_list = [
('Basis Probe Fourier Space Amplitudes',
lambda self, fig: p.plot_amplitude(tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Basis Probe Fourier Space Phases',
lambda self, fig: p.plot_phase(tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Basis Probe Real Space Amplitudes, Surface Normal View',
lambda self, fig: p.plot_amplitude(
{'title': 'Basis Probe Fourier Space Amplitudes',
'plot_func': lambda self, fig: p.plot_amplitude(tools.propagators.inverse_far_field(self.probe), fig=fig)},
{'title': 'Basis Probe Fourier Space Phases',
'plot_func': lambda self, fig: p.plot_phase(tools.propagators.inverse_far_field(self.probe), fig=fig)},
{'title': 'Basis Probe Real Space Amplitudes, Surface Normal View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.probe,
fig=fig,
basis=self.probe_basis,
units=self.units,
)),
('Basis Probe Real Space Phases, Surface Normal View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Basis Probe Real Space Phases, Surface Normal View',
'plot_func': lambda self, fig: p.plot_phase(
self.probe,
fig=fig,
basis=self.probe_basis,
units=self.units,
)),
('Basis Probe Real Space Amplitudes, Beam View',
lambda self, fig: p.plot_amplitude(
)},
{'title': 'Basis Probe Real Space Amplitudes, Beam View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.probe,
fig=fig,
basis=self.probe_basis,
view_basis=beam_basis,
units=self.units,
)),
('Basis Probe Real Space Phases, Beam View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Basis Probe Real Space Phases, Beam View',
'plot_func': lambda self, fig: p.plot_phase(
self.probe,
fig=fig,
basis=self.probe_basis,
view_basis=beam_basis,
units=self.units,
)),
('Object Amplitude, Surface Normal View',
lambda self, fig: p.plot_amplitude(
)},
{'title': 'Object Amplitude, Surface Normal View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units,
)),
('Object Phase, Surface Normal View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Object Phase, Surface Normal View',
'plot_func': lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units,
)),
('Object Amplitude, Beam View',
lambda self, fig: p.plot_amplitude(
)},
{'title': 'Object Amplitude, Beam View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
view_basis=beam_basis,
units=self.units,
)),
('Object Phase, Beam View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Object Phase, Beam View',
'plot_func': lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
view_basis=beam_basis,
units=self.units,
)),
('Object Amplitude, Detector View',
lambda self, fig: p.plot_amplitude(
)},
{'title': 'Object Amplitude, Detector View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
view_basis=self.det_basis,
units=self.units,
)),
('Object Phase, Detector View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Object Phase, Detector View',
'plot_func': lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
view_basis=self.det_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))
)},
{'title': 'Corrected Translations',
'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)},
{'title': 'Background',
'plot_func': lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)},
]
src/cdtools/models/fancy_ptycho.py
+262 -110
View File
@@ -7,8 +7,6 @@ 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']
@@ -45,9 +43,12 @@ class FancyPtycho(CDIModel):
near_field=False,
angular_spectrum_propagator=None,
inv_angular_spectrum_propagator=None,
panel_plot_mode=True,
plot_level=2,
):
super(FancyPtycho, self).__init__()
super(FancyPtycho, self).__init__(panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
self.register_buffer('wavelength',
t.as_tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
@@ -252,6 +253,8 @@ class FancyPtycho(CDIModel):
obj_view_crop=None,
obj_padding=200,
near_field=False,
panel_plot_mode=True,
plot_level=2,
):
wavelength = dataset.wavelength
@@ -518,6 +521,8 @@ class FancyPtycho(CDIModel):
near_field=near_field,
angular_spectrum_propagator=angular_spectrum_propagator,
inv_angular_spectrum_propagator=inv_angular_spectrum_propagator,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
)
@@ -551,7 +556,7 @@ class FancyPtycho(CDIModel):
else:
try:
Ws = t.ones(len(index)) # I'm positive this introduced a bug
except:
except TypeError:
Ws = 1
if self.weights is None or len(self.weights[0].shape) == 0:
@@ -776,7 +781,6 @@ class FancyPtycho(CDIModel):
# 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
@@ -840,6 +844,53 @@ class FancyPtycho(CDIModel):
# We discard the U matrix and re-multiply S & Vh
self.weights.data = S[:,:,None] * (Vh / probe_sqrt_intensities)
def get_probe_intensities(self):
"""Returns the effective probe intensity at each scan position.
Handles both the simple (1D weights) and OPRP (2D weights) cases.
Returns
-------
probe_intensities : np.ndarray
Array of probe intensities, one per scan position.
"""
if not hasattr(self, 'weights'):
raise NotImplementedError(
"I don't know how to handle having no weights")
elif self.weights.ndim == 1:
probe_intensities = self.weights.detach().cpu().numpy()**2
else:
# The big case, with OPRP
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()
# The outer one is a sum, because the tensordot is what broadcasts
# the probe matrix along the shot dimension - the second one
# doesn't have to.
weighted_probe_matrices = np.sum(
np.tensordot(weights, probe_matrix, axes=1)[...,None]
* weights.conj().transpose((0,2,1))[...,None,:,:],
axis=-2
)
basis_probe_intensities = np.trace(
probe_matrix, axis1=-2, axis2=-1)
probe_intensities = np.trace(
weighted_probe_matrices, axis1=-2, axis2=-1)
# Imaginary part is already essentially zero up to rounding error
probe_intensities = np.real(
probe_intensities / basis_probe_intensities)
return probe_intensities
def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs):
def get_probes(idx):
@@ -856,22 +907,8 @@ class FancyPtycho(CDIModel):
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)
values = self.get_probe_intensities()
if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity':
cmap = 'viridis'
else:
@@ -888,7 +925,22 @@ class FancyPtycho(CDIModel):
cmap=cmap,
**kwargs),
def plot_illumination_intensity(self, fig, dataset):
"""Plots the probe intensity nanomap. Only used to make a plot for the plot list."""
p.plot_nanomap(
self.corrected_translations(dataset),
self.get_probe_intensities(),
fig=fig,
cmap='viridis',
cmap_label='Intensity (a.u.)',
units=self.units,
convention='probe',
invert_xaxis=True
)
plt.gca().set_aspect('equal')
def plot_translations_and_originals(self, fig, dataset):
"""Only used to make a plot for the plot list."""
p.plot_translations(
@@ -908,112 +960,212 @@ class FancyPtycho(CDIModel):
color='k',
marker='.'
)
plt.legend()
plt.gca().set_aspect('equal')
plt.legend(loc='upper right')
plot_panel_list = [
{
'title': 'Main Results',
'plot_level': 1,
'grid': (2,2),
'figure_size': (8.4,6.8),
'plots': [
{
'title': 'Object Phase',
'subplot': (0,0),
'plot_func': lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
additional_axis_labels=['Mode #',],
units=self.units),
'condition': lambda self: not self.exponentiate_obj,
},
{
'title': 'Object Amplitude',
'subplot': (1,0),
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
additional_axis_labels=['Mode #',],
units=self.units),
'condition': lambda self: not self.exponentiate_obj,
},
{
'title': 'Real Part of T',
'subplot': (0,0),
'plot_func': lambda self, fig: p.plot_real(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
additional_axis_labels=['Mode #',],
units=self.units,
cmap='cividis',
),
'condition': lambda self: self.exponentiate_obj,
},
{
'title': 'Imaginary Part of T',
'subplot': (1,0),
'plot_func': lambda self, fig: p.plot_imag(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
additional_axis_labels=['Mode #',],
units=self.units,
cmap='viridis_r',
),
'condition': lambda self: self.exponentiate_obj,
},
{
'title': 'Probe Modes, Colorized',
'subplot': (0,1),
'plot_func': lambda self, fig: p.plot_colorized(
(self.probe if not self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig,
title='Probe Modes, Real Space',
basis=self.probe_basis,
additional_axis_labels=['Mode #',],
amplitude_scaling=np.sqrt,
units=self.units),
},
{
'title': 'Probe Modes, Amplitude',
'subplot': (1,1),
'plot_func': lambda self, fig: p.plot_amplitude(
(self.probe if not self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig,
title='Probe Modes, Real Space',
basis=self.probe_basis,
additional_axis_labels=['Mode #',],
units=self.units),
},
],
},
{
'title': 'Advanced Monitoring',
'plot_level': 2,
'figure_size': (12.6,6.8),
'grid': (2,3),
'plots': [
{
'title': 'Probe Modes, Fourier Colorized',
'subplot': (0,0),
'plot_func': lambda self, fig: p.plot_colorized(
(self.probe if self.fourier_probe
else tools.propagators.far_field(self.probe)),
fig=fig,
title='Probe Modes, Fourier Space',
additional_axis_labels=['Mode #',],
amplitude_scaling = np.sqrt,
),
},
{
'title': 'Probe Modes, Fourier Amplitude',
'subplot': (1,0),
'plot_func': lambda self, fig: p.plot_amplitude(
(self.probe if self.fourier_probe
else tools.propagators.far_field(self.probe)),
fig=fig,
title='Probe Modes, Fourier Space',
additional_axis_labels=['Mode #',],
),
},
{
'title': 'Illumination Intensity',
'subplot': (0,1),
'plot_func': lambda self, fig, dataset: self.plot_illumination_intensity(fig, dataset),
},
{
'title': 'Detector Background',
'subplot': (1,1),
'plot_func': lambda self, fig: p.plot_amplitude(self.background**2, fig=fig, cmap='viridis', cmap_label='Intensity (detector units)'),
},
{
'title': 'Corrected Translations',
'subplot': (0,2),
'plot_func': lambda self, fig, dataset: self.plot_translations_and_originals(fig, dataset),
},
{
'title': 'Loss History',
'subplot': (1,2),
'plot_func': lambda self, fig: self.plot_loss_history(fig),
},
],
},
{
'title': 'Unstable Probe Refinement Details',
'plot_level': 2,
'figure_size': (8.4,3.4),
'grid': (1,2),
'condition': lambda self: len(self.weights.shape) >= 2,
'plots': [
{
'title': '% of Power in Top Mode',
'subplot': (0,0),
'plot_func': 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),
'condition': lambda self: len(self.weights.shape) >= 2
},
{
'title': 'Mean Weight Matrix Amplitudes',
'subplot': (0,1),
'plot_func': lambda self, fig: p.plot_amplitude(
np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
fig=fig),
'condition': lambda self: len(self.weights.shape) >= 2
},
]
}
]
plot_list = [
('',
lambda self, fig, dataset: self.plot_wavefront_variation(
{'title': 'Quantum Efficiency Mask',
'plot_level': 2,
'plot_func': lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig),
'condition': lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None)},
{'title': 'Per-Exposure Probe Intensity',
'plot_level': 3,
'figure_size': (8,5.3),
'plot_func': 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(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Per-Exposure Probe Amplitudes',
'plot_level': 3,
'figure_size': (8,5.3),
'plot_func': 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(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Per-Exposure Probe Phases',
'plot_level': 3,
'figure_size': (8,5.3),
'plot_func': 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.far_field(self.probe)),
fig=fig)),
('Basis Probe Fourier Space Colorized',
lambda self, fig: p.plot_colorized(
(self.probe if self.fourier_probe
else tools.propagators.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 Colorized',
lambda self, fig: p.plot_colorized(
(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[self.obj_view_slice],
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[self.obj_view_slice],
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[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: self.plot_translations_and_originals(fig, dataset)),
('Background',
lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)),
('Quantum Efficiency Mask',
lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig),
lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None))
'condition': lambda self: len(self.weights.shape) >= 2},
]
src/cdtools/models/multislice_2d_ptycho.py
+52 -47
View File
@@ -47,9 +47,12 @@ class Multislice2DPtycho(CDIModel):
prevent_aliasing=True,
phase_only=False,
units='um',
panel_plot_mode=False,
plot_level=1,
):
super(Multislice2DPtycho, self).__init__()
super(Multislice2DPtycho, self).__init__(panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
self.wavelength = t.tensor(wavelength)
self.detector_geometry = copy(detector_geometry)
self.dz = dz
@@ -154,7 +157,7 @@ class Multislice2DPtycho(CDIModel):
@classmethod
def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None):
def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None, panel_plot_mode=False, plot_level=1):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -286,7 +289,7 @@ class Multislice2DPtycho(CDIModel):
surface_normal=surface_normal,
probe_support=probe_support,
min_translation=min_translation,
translation_offsets = translation_offsets,
translation_offsets=translation_offsets,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
@@ -296,7 +299,9 @@ class Multislice2DPtycho(CDIModel):
exponentiate_obj=exponentiate_obj,
units=units, fourier_probe=fourier_probe,
phase_only=phase_only,
prevent_aliasing=prevent_aliasing)
prevent_aliasing=prevent_aliasing,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
def interaction(self, index, translations):
@@ -581,22 +586,22 @@ class Multislice2DPtycho(CDIModel):
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Probe Fourier Space Amplitude',
lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Probe Fourier Space Phase',
lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Probe Real Space Amplitude',
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)),
('Probe Real Space Phase',
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(
plot_list = [
{'title': 'Probe Fourier Space Amplitude',
'plot_func': lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)},
{'title': 'Probe Fourier Space Phase',
'plot_func': lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)},
{'title': 'Probe Real Space Amplitude',
'plot_func': 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)},
{'title': 'Probe Real Space Phase',
'plot_func': 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)},
{'title': 'Average Weight Matrix Amplitudes',
'plot_func': 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(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': '% of Power in Top Mode',
'plot_func': lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
100 * t.stack([
analysis.calc_mode_power_fractions(
@@ -606,35 +611,35 @@ class Multislice2DPtycho(CDIModel):
], dim=0),
fig=fig,
units=self.units),
lambda self: len(self.weights.shape) >= 2),
('Slice by Slice Real Part of T',
lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
lambda self: self.exponentiate_obj),
('Slice by Slice Imaginary Part of T',
lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: self.exponentiate_obj),
('Integrated Real Part of T',
lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3),
('Integrated Imaginary Part of T',
lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3),
('Slice by Slice Amplitude of Object Function',
lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: not self.exponentiate_obj),
('Slice by Slice Phase of Object Function',
lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units,cmap='cividis'),
lambda self: not self.exponentiate_obj),
('Amplitude of Stacked Object Function',
lambda self, fig: p.plot_amplitude(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: (not self.exponentiate_obj) and self.obj.dim() >=3),
('Phase of Stacked Object Function',
lambda self, fig: p.plot_phase(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3),
('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))
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Slice by Slice Real Part of T',
'plot_func': lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
'condition': lambda self: self.exponentiate_obj},
{'title': 'Slice by Slice Imaginary Part of T',
'plot_func': lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
'condition': lambda self: self.exponentiate_obj},
{'title': 'Integrated Real Part of T',
'plot_func': lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
'condition': lambda self: self.exponentiate_obj and self.obj.dim() >= 3},
{'title': 'Integrated Imaginary Part of T',
'plot_func': lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units),
'condition': lambda self: self.exponentiate_obj and self.obj.dim() >= 3},
{'title': 'Slice by Slice Amplitude of Object Function',
'plot_func': lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Slice by Slice Phase of Object Function',
'plot_func': lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Amplitude of Stacked Object Function',
'plot_func': lambda self, fig: p.plot_amplitude(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
'condition': lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3},
{'title': 'Phase of Stacked Object Function',
'plot_func': lambda self, fig: p.plot_phase(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
'condition': lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3},
{'title': 'Corrected Translations',
'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)},
{'title': 'Background',
'plot_func': lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)},
]
src/cdtools/models/multislice_ptycho.py
+66 -63
View File
@@ -4,11 +4,8 @@ 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__ = ['MultislicePtycho']
@@ -39,10 +36,13 @@ class MultislicePtycho(CDIModel):
simulate_finite_pixels=False,
dtype=t.float32,
exponentiate_obj=False,
obj_view_crop=0
obj_view_crop=0,
panel_plot_mode=False,
plot_level=1,
):
super(MultislicePtycho, self).__init__()
super(MultislicePtycho, self).__init__(panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
self.register_buffer('wavelength',
t.as_tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
@@ -202,6 +202,8 @@ class MultislicePtycho(CDIModel):
obj_view_crop=None,
obj_padding=200,
exponentiate_obj=False,
panel_plot_mode=False,
plot_level=1,
):
wavelength = dataset.wavelength
@@ -409,6 +411,8 @@ class MultislicePtycho(CDIModel):
simulate_finite_pixels=simulate_finite_pixels,
exponentiate_obj=exponentiate_obj,
obj_view_crop=obj_view_crop,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
)
@@ -442,7 +446,7 @@ class MultislicePtycho(CDIModel):
else:
try:
Ws = t.ones(len(index)) # I'm positive this introduced a bug
except:
except TypeError:
Ws = 1
if self.weights is None or len(self.weights[0].shape) == 0:
@@ -636,7 +640,6 @@ class MultislicePtycho(CDIModel):
# 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
@@ -750,61 +753,61 @@ class MultislicePtycho(CDIModel):
plot_list = [
('',
lambda self, fig, dataset: self.plot_wavefront_variation(
{'title': '',
'plot_func': 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(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': '',
'plot_func': 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(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': '',
'plot_func': 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(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Basis Probe Fourier Space Amplitudes',
'plot_func': 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(
fig=fig)},
{'title': 'Basis Probe Fourier Space Phases',
'plot_func': 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(
else tools.propagators.inverse_far_field(self.probe)),
fig=fig)},
{'title': 'Basis Probe Real Space Amplitudes',
'plot_func': 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(
units=self.units)},
{'title': 'Basis Probe Real Space Phases',
'plot_func': 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(
units=self.units)},
{'title': 'Average Weight Matrix Amplitudes',
'plot_func': 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(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': '% of Power in Top Mode',
'plot_func': lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
100 * t.stack([
analysis.calc_mode_power_fractions(
@@ -814,69 +817,69 @@ class MultislicePtycho(CDIModel):
], dim=0),
fig=fig,
units=self.units),
lambda self: len(self.weights.shape) >= 2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Object Amplitude',
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj[(np.s_[:],) + self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object (T) Imaginary Part',
lambda self, fig: p.plot_imag(
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Object (T) Imaginary Part',
'plot_func': lambda self, fig: p.plot_imag(
self.obj[(np.s_[:],) + self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: self.exponentiate_obj),
('Object Phase',
lambda self, fig: p.plot_phase(
'condition': lambda self: self.exponentiate_obj},
{'title': 'Object Phase',
'plot_func': lambda self, fig: p.plot_phase(
self.obj[(np.s_[:],) + self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object (T) Real Part',
lambda self, fig: p.plot_real(
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Object (T) Real Part',
'plot_func': lambda self, fig: p.plot_real(
self.obj[(np.s_[:],) + self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units,
cmap='cividis'),
lambda self: self.exponentiate_obj),
('Object Product Amplitude',
lambda self, fig: p.plot_amplitude(
'condition': lambda self: self.exponentiate_obj},
{'title': 'Object Product Amplitude',
'plot_func': lambda self, fig: p.plot_amplitude(
t.prod(self.obj, dim=0)[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object (T) Sum Imaginary Part',
lambda self, fig: p.plot_imag(
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Object (T) Sum Imaginary Part',
'plot_func': lambda self, fig: p.plot_imag(
t.sum(self.obj, dim=0)[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: self.exponentiate_obj),
('Object Product Phase',
lambda self, fig: p.plot_phase(
'condition': lambda self: self.exponentiate_obj},
{'title': 'Object Product Phase',
'plot_func': lambda self, fig: p.plot_phase(
t.prod(self.obj, dim=0)[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object (T) Sum Real Part',
lambda self, fig: p.plot_real(
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Object (T) Sum Real Part',
'plot_func': lambda self, fig: p.plot_real(
t.sum(self.obj, dim=0)[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units,
cmap='cividis'),
lambda self: self.exponentiate_obj),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
lambda self, fig: p.plot_amplitude(self.background**2, fig=fig))
'condition': lambda self: self.exponentiate_obj},
{'title': 'Corrected Translations',
'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)},
{'title': 'Background',
'plot_func': lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)},
]
src/cdtools/models/rpi.py
+87 -46
View File
@@ -6,8 +6,6 @@ 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
import time
__all__ = ['RPI']
@@ -58,9 +56,12 @@ class RPI(CDIModel):
propagation_distance=0,
units='um',
dtype=t.float32,
panel_plot_mode=True,
plot_level=1,
):
super(RPI, self).__init__()
super(RPI, self).__init__(panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
complex_dtype = (t.ones([1], dtype=dtype) +
1j * t.ones([1], dtype=dtype)).dtype
@@ -164,6 +165,8 @@ class RPI(CDIModel):
phase_only=False,
probe_threshold=0,
dtype=t.float32,
panel_plot_mode=True,
plot_level=1,
):
complex_dtype = (t.ones([1], dtype=dtype) +
1j * t.ones([1], dtype=dtype)).dtype
@@ -227,7 +230,7 @@ class RPI(CDIModel):
# This will be superceded later by a call to init_obj, but it sets
# the shape
if obj_size is None:
obj_size = (np.array(self.probe.shape[-2:]) // 2).astype(int)
obj_size = (np.array(probe.shape[-2:]) // 2).astype(int)
dummy_init_obj = t.ones([n_modes, obj_size[0], obj_size[1]],
dtype=complex_dtype)
@@ -247,14 +250,16 @@ class RPI(CDIModel):
obj_support = t.as_tensor(binary_dilation(obj_support))
rpi_object = cls(wavelength, det_geo, ew_basis,
probe, dummy_init_obj,
probe, dummy_init_obj,
background=background, mask=mask,
saturation=saturation,
obj_support=obj_support,
oversampling=oversampling,
exponentiate_obj=exponentiate_obj,
phase_only=phase_only,
weight_matrix=weight_matrix)
weight_matrix=weight_matrix,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
# I don't love this pattern, where I do the "real" obj initialization
# after creating the rpi object. But, I chose this so that I could
@@ -283,7 +288,9 @@ class RPI(CDIModel):
exponentiate_obj=False,
phase_only=False,
initialization='random',
dtype=t.float32
dtype=t.float32,
panel_plot_mode=True,
plot_level=1,
):
complex_dtype = (t.ones([1], dtype=dtype) +
@@ -308,7 +315,7 @@ class RPI(CDIModel):
# This will be superceded later by a call to init_obj, but it sets
# the shape
if obj_size is None:
obj_size = (np.array(self.probe.shape[-2:]) // 2).astype(int)
obj_size = (np.array(probe.shape[-2:]) // 2).astype(int)
dummy_init_obj = t.ones([n_modes, obj_size[0], obj_size[1]],
dtype=complex_dtype)
@@ -327,6 +334,8 @@ class RPI(CDIModel):
mask=mask,
exponentiate_obj=exponentiate_obj,
phase_only=phase_only,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
)
rpi_object.init_obj(initialization)
@@ -365,21 +374,21 @@ class RPI(CDIModel):
obj_shape=obj_shape,
n_modes=n_modes)
else:
raise KeyError('Initialization "' + str(initialization) + \
raise KeyError('Initialization "' + str(initialization_type) + \
'" invalid - use "spectral", "uniform", or "random"')
def get_obj_shape_and_n_modes(self, obj_shape=None, n_modes=None):
"""Sets defaults for obj shape and n modes"""
if obj_shape == None:
if obj_shape is None:
if hasattr(self, 'obj'):
obj_shape = self.obj.shape[-2:]
else:
obj_size = (np.array(self.probe.shape[-2:]) // 2).astype(int)
obj_shape = [obj_size, obj_size]
if n_modes == None:
if n_modes is None:
if hasattr(self, 'obj'):
n_modes = self.obj.shape[0]
else:
@@ -531,40 +540,72 @@ class RPI(CDIModel):
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 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),
plot_panel_list = [
{
'title': 'RPI Results',
'plot_level': 1,
'figure_size': (12, 3.5),
'grid': (1, 3),
'plots': [
{
'title': 'Object Phase',
'subplot': (0, 0),
'plot_func': lambda self, fig: p.plot_phase(
self.obj,
fig=fig,
title='Object Phase',
basis=self.obj_basis,
units=self.units),
'condition': lambda self: not self.exponentiate_obj,
},
{
'title': 'Real Part of T',
'subplot': (0, 0),
'plot_func': lambda self, fig: p.plot_real(
self.obj,
fig=fig,
title='Real Part of T',
basis=self.obj_basis,
units=self.units,
cmap='cividis'),
'condition': lambda self: self.exponentiate_obj,
},
{
'title': 'Object Amplitude',
'subplot': (0, 1),
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj,
fig=fig,
title='Object Amplitude',
basis=self.obj_basis,
units=self.units),
'condition': lambda self: not self.exponentiate_obj,
},
{
'title': 'Imaginary Part of T',
'subplot': (0, 1),
'plot_func': lambda self, fig: p.plot_imag(
self.obj,
fig=fig,
title='Imaginary Part of T',
basis=self.obj_basis,
units=self.units),
'condition': lambda self: self.exponentiate_obj,
},
{
'title': 'Root Sum Squared Amplitude of all Probes',
'subplot': (0, 2),
'plot_func': 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,
units=self.units),
},
],
},
]
src/cdtools/models/simple_ptycho.py
+21 -9
View File
@@ -76,7 +76,7 @@ class SimplePtycho(CDIModel):
probe_basis,
probe,
obj,
min_translation=min_translation
min_translation=min_translation,
)
@@ -108,14 +108,26 @@ class SimplePtycho(CDIModel):
# 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))
{
'title': 'Probe Amplitude',
'plot_func': lambda self, fig:
p.plot_amplitude(self.probe, fig, basis=self.probe_basis),
},
{
'title': 'Probe Phase',
'plot_func': lambda self, fig:
p.plot_phase(self.probe, fig, basis=self.probe_basis)
},
{
'title': 'Object Amplitude',
'plot_func': lambda self, fig:
p.plot_amplitude(self.obj, fig, basis=self.probe_basis)
},
{
'title': 'Object Phase',
'plot_func': lambda self, fig:
p.plot_phase(self.obj, fig, basis=self.probe_basis)
},
]
def save_results(self, dataset):
src/cdtools/reconstructors/base.py
+8 -5
View File
@@ -16,6 +16,7 @@ from torch.utils import data as td
import threading
import queue
import time
from matplotlib import pyplot as plt
from typing import List, Union
if TYPE_CHECKING:
@@ -232,7 +233,7 @@ class Reconstructor:
def optimize(self,
iterations: int,
batch_size: int = 1,
custom_data_loader: torch.utils.data.DataLoader = None,
custom_data_loader: t.utils.data.DataLoader = None,
regularization_factor: Union[float, List[float]] = None,
thread: bool = True,
calculation_width: int = 10,
@@ -357,10 +358,12 @@ class Reconstructor:
try:
calc.start()
while calc.is_alive():
if hasattr(self.model, 'figs'):
self.model.figs[0].canvas.start_event_loop(0.01)
else:
calc.join()
open_figs = plt.get_fignums()
with plt.rc_context({'figure.raise_window': False}):
for fignum in open_figs:
plt.figure(fignum).canvas.flush_events()
time.sleep(0.01)
except KeyboardInterrupt as e:
stop_event.set()
src/cdtools/tools/plotting/plotting.py
+557 -182
View File
File diff suppressed because it is too large Load Diff
tests/models/test_fancy_ptycho.py
+19 -10
View File
@@ -1,5 +1,7 @@
import pytest
import time
import torch as t
from matplotlib import pyplot as plt
import cdtools
@@ -89,6 +91,8 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
units='mm',
obj_view_crop=-50,
use_qe_mask=True, # test this in the case where no qe mask is defined
panel_plot_mode=True, # test with panel plot mode,
plot_level=4, # test with all plots
)
print('Running reconstruction on provided reconstruction_device,',
@@ -98,24 +102,26 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
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)
if show_plot:
model.inspect(dataset, min_interval=10)
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)
if show_plot:
model.inspect(dataset, min_interval=10)
for loss in model.Adam_optimize(25, dataset, lr=0.001, batch_size=50):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
if show_plot:
model.inspect(dataset, min_interval=10)
model.tidy_probes()
if show_plot:
model.inspect(dataset)
model.compare(dataset)
time.sleep(3)
plt.close('all')
# If this fails, the reconstruction has gotten worse
assert model.loss_history[-1] < 0.0013
@@ -132,6 +138,7 @@ def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_pl
n_modes=1,
near_field=True,
propagation_distance=3.65e-3, # 3.65 downstream from focus
panel_plot_mode=False, # test without panel plot mode
)
print('Running reconstruction on provided reconstruction_device,',
@@ -141,19 +148,21 @@ def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_pl
for loss in model.Adam_optimize(100, dataset, lr=0.04, batch_size=10):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
if show_plot:
model.inspect(dataset, min_interval=10)
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)
if show_plot:
model.inspect(dataset, min_interval=10)
model.tidy_probes()
if show_plot:
model.inspect(dataset)
model.compare(dataset)
time.sleep(3)
plt.close('all')
# If this fails, the reconstruction has gotten worse
assert model.loss_history[-1] < 0.005
tests/models/test_simple_ptycho.py
+6 -2
View File
@@ -1,5 +1,7 @@
import pytest
import time
import torch as t
from matplotlib import pyplot as plt
import cdtools
@@ -18,12 +20,14 @@ def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
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, min_interval=10)
if show_plot:
model.inspect(dataset)
model.compare(dataset)
time.sleep(3)
plt.close('all')
# If this fails, the reconstruction got worse
assert model.loss_history[-1] < 0.013
tests/test_reconstructors.py
+27 -13
View File
@@ -1,4 +1,5 @@
import pytest
import time
import cdtools
import torch as t
import numpy as np
@@ -36,7 +37,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
probe_support_radius=50,
propagation_distance=2e-6,
units='um',
probe_fourier_crop=pad
probe_fourier_crop=pad,
panel_plot_mode=False, # At least one check without panel plot mode
)
model.translation_offsets.data += 0.7 * \
@@ -67,8 +69,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
lr=lr_tup[i],
batch_size=batch_size_tup[i]):
print(model_recon.report())
if show_plot and model_recon.epoch % 10 == 0:
model_recon.inspect(dataset)
if show_plot:
model_recon.inspect(dataset, min_interval=10)
# Check hyperparameter update
assert recon.optimizer.param_groups[0]['lr'] == lr_tup[i]
@@ -86,6 +88,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
if show_plot:
model_recon.inspect(dataset)
model_recon.compare(dataset)
time.sleep(3)
plt.close('all')
# ******* Reconstructions with CDIModel.Adam_optimize *******
print('Running reconstruction using CDIModel.Adam_optimize on provided' +
@@ -99,14 +103,16 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
lr=lr_tup[i],
batch_size=batch_size_tup[i]):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
if show_plot:
model.inspect(dataset, min_interval=10)
model.tidy_probes()
if show_plot:
model.inspect(dataset)
model.compare(dataset)
time.sleep(3)
plt.close('all')
# Ensure equivalency between the model reconstructions during the first
# pass, where they should be identical
@@ -170,8 +176,8 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
for loss in recon.optimize(iterations,
lr=0.4,
regularization_factor=reg_factor_tup[i]):
if show_plot and i == 0:
model_recon.inspect(dataset)
if show_plot:
model_recon.inspect(dataset, min_interval=10)
print(model_recon.report())
# Check hyperparameter update (or lack thereof)
@@ -180,6 +186,8 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
if show_plot:
model_recon.inspect(dataset)
model_recon.compare(dataset)
time.sleep(3)
plt.close('all')
# Check model pointing
assert id(model_recon) == id(recon.model)
@@ -193,13 +201,15 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
dataset,
lr=0.4,
regularization_factor=reg_factor_tup[i]): # noqa
if show_plot and i == 0:
model.inspect(dataset)
if show_plot:
model.inspect(dataset, min_interval=10)
print(model.report())
if show_plot:
model.inspect(dataset)
model.compare(dataset)
time.sleep(3)
plt.close('all')
# Check loss equivalency between the two reconstructions
assert np.allclose(model.loss_history[:epoch_tup[0]], model_recon.loss_history[:epoch_tup[0]])
@@ -271,8 +281,8 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
lr=lr,
batch_size=batch_size):
print(model_recon.report())
if show_plot and model_recon.epoch % 10 == 0:
model_recon.inspect(dataset)
if show_plot:
model_recon.inspect(dataset, min_interval=10)
# Check hyperparameter update
assert recon.optimizer.param_groups[0]['lr'] == lr
@@ -290,6 +300,8 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
if show_plot:
model_recon.inspect(dataset)
model_recon.compare(dataset)
time.sleep(3)
plt.close('all')
# ******* Reconstructions with cdtools.CDIModel.SGD_optimize *******
print('Running reconstruction using CDIModel.SGD_optimize on provided' +
@@ -301,14 +313,16 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
lr=lr,
batch_size=batch_size):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
if show_plot:
model.inspect(dataset, min_interval=10)
model.tidy_probes()
if show_plot:
model.inspect(dataset)
model.compare(dataset)
time.sleep(3)
plt.close('all')
# Ensure equivalency between the model reconstructions
assert np.allclose(model_recon.loss_history[-1], model.loss_history[-1])
tests/tools/test_plotting.py
+80 -13
View File
@@ -11,28 +11,32 @@ def test_plot_amplitude(show_plot):
# Test with tensor
im = t.as_tensor(scipy.datasets.ascent(), dtype=t.complex128)
plotting.plot_amplitude(im, basis=np.array([[0, -1], [-1, 0], [0, 0]]), title='Test Amplitude')
if show_plot:
plt.show()
# Test with numpy array
im = scipy.datasets.ascent().astype(np.complex128)
# Test with numpy array and an extra dimension
im = np.stack([scipy.datasets.ascent().astype(np.complex128)]*3, axis=0)
plotting.plot_amplitude(im, title='Test Amplitude')
# Test with pytorch tensor and two extra dimensions
im = t.as_tensor(np.stack([im]*5, axis=0))
plotting.plot_amplitude(im, title='Test Amplitude',
additional_axis_labels=['Hi','There'])
if show_plot:
plt.show()
plt.close('all')
def test_plot_phase(show_plot):
# Test with tensor
im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
plotting.plot_phase(im, title='Test Phase')
if show_plot:
plt.show()
# Test with numpy array
im = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1]).numpy()
plotting.plot_phase(im, title='Test Phase', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
if show_plot:
plt.show()
plt.close('all')
def test_plot_colorized(show_plot):
@@ -40,11 +44,74 @@ def test_plot_colorized(show_plot):
gaussian = initializers.gaussian([512, 512], [200, 200], amplitude=100, curvature=[.1, .1])
im = gaussian * t.as_tensor(scipy.datasets.ascent(), dtype=t.complex64)
plotting.plot_colorized(im, title='Test Colorize', basis=np.array([[0, -1], [-1, 0], [0, 0]]))
if show_plot:
plt.show()
# Test with numpy array
# Test with numpy array and hsv
im = im.numpy()
plotting.plot_colorized(im, title='Test Colorize')
plotting.plot_colorized(im, title='Test Colorize', use_cmocean=False)
if show_plot:
plt.show()
plt.close('all')
def test_plot_translations(show_plot):
rng = np.random.default_rng(0)
trans_np = rng.uniform(-5e-6, 5e-6, (20, 2))
trans_t = t.as_tensor(trans_np)
# numpy, defaults
plotting.plot_translations(trans_np)
# torch tensor and reuse figure
fig = plotting.plot_translations(trans_t)
plotting.plot_translations(trans_np, lines=False, color='red', label='scan', fig=fig, clear_fig=False)
if show_plot:
plt.show()
plt.close('all')
def test_plot_nanomap(show_plot):
rng = np.random.default_rng(0)
trans_np = rng.uniform(-5e-6, 5e-6, (20, 2))
values_np = np.random.default_rng(1).uniform(0, 1, 20)
trans_t = t.as_tensor(trans_np)
values_t = t.as_tensor(values_np)
# numpy, defaults
plotting.plot_nanomap(trans_np, values_np)
# torch tensors
plotting.plot_nanomap(trans_t, values_t, units='nm', cmap_label='Intensity', convention='sample')
if show_plot:
plt.show()
plt.close('all')
def test_plot_nanomap_with_images(show_plot):
rng = np.random.default_rng(0)
trans_np = rng.uniform(-5e-6, 5e-6, (20, 2))
values_np = np.random.default_rng(1).uniform(0, 1, 20)
# plot_nanomap_with_images requires tensor translations
trans_t = t.as_tensor(trans_np)
values_t = t.as_tensor(values_np)
def get_image_2d(i):
return np.random.default_rng(i).uniform(0, 1, (32, 32))
def get_image_3d(i):
return np.random.default_rng(i).uniform(0, 1, (4, 32, 32))
# basic call, no values
plotting.plot_nanomap_with_images(trans_np, get_image_2d)
# with explicit values
plotting.plot_nanomap_with_images(trans_t, get_image_2d, values=values_np)
# 3D image stack
fig = plt.figure(figsize=(11,7))
plotting.plot_nanomap_with_images(trans_np, get_image_3d, values=values_t, fig=fig)
if show_plot:
plt.show()
plt.close('all')