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

Update the docstrings on all the tools files and start laying out the docs for real

Browse Source
This commit is contained in:
Abe Levitan
2019-09-03 18:49:45 -04:00
parent 1e3a98c599
commit de186592cb
42 changed files with 1181 additions and 887 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/datasets.py
+35 -10
View File
@@ -2,6 +2,8 @@ from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from copy import copy
import h5py
import pathlib
from CDTools.tools import data as cdtdata
from CDTools.tools import plotting
@@ -86,6 +88,12 @@ class CDataset(torchdata.Dataset):
@classmethod
def from_cxi(cls, cxi_file):
# If a bare string is passed
if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
with h5py.File(cxi_file,'r') as f:
return cls.from_cxi(f)
entry_info = cdtdata.get_entry_info(cxi_file)
sample_info = cdtdata.get_sample_info(cxi_file)
wavelength = cdtdata.get_wavelength(cxi_file)
@@ -139,6 +147,11 @@ class Ptycho_2D_Dataset(CDataset):
self.axes = copy(axes)
self.translations = t.tensor(translations)
self.patterns = t.tensor(patterns)
if self.mask is None:
self.mask = t.ones(self.patterns.shape[-2:]).to(dtype=t.uint8)
self.mask.masked_fill_(t.isnan(t.sum(self.patterns,dim=(0,))),0)
self.patterns.masked_fill_(t.isnan(self.patterns),0)
def __len__(self):
@@ -158,6 +171,11 @@ class Ptycho_2D_Dataset(CDataset):
# perhaps there is a way but I couldn't figure it out.
@classmethod
def from_cxi(cls, cxi_file):
# If a bare string is passed
if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path):
with h5py.File(cxi_file,'r') as f:
return cls.from_cxi(f)
entry_info = cdtdata.get_entry_info(cxi_file)
sample_info = cdtdata.get_sample_info(cxi_file)
wavelength = cdtdata.get_wavelength(cxi_file)
@@ -166,6 +184,7 @@ class Ptycho_2D_Dataset(CDataset):
'basis' : basis,
'corner' : corner}
mask = cdtdata.get_mask(cxi_file)
dark = cdtdata.get_dark(cxi_file)
patterns, axes = cdtdata.get_data(cxi_file)
@@ -190,7 +209,7 @@ class Ptycho_2D_Dataset(CDataset):
axslider = plt.axes([0.15,0.06,0.75,0.03])
translations = self.translations.detach().cpu().numpy()
nanomap_values = self.patterns.sum(dim=(1,2)).detach().cpu().numpy()
nanomap_values = (self.mask.to(t.float32) * self.patterns).sum(dim=(1,2)).detach().cpu().numpy()
def update_colorbar(im):
# If the update brought the colorbar out of whack
@@ -207,6 +226,9 @@ class Ptycho_2D_Dataset(CDataset):
im.colorbar.set_ticks(ticker.LinearLocator(numticks=5))
im.colorbar.draw_all()
def on_pick(event):
update(event.ind[0])
plt.draw()
def update(idx):
idx = int(idx) % len(self)
@@ -226,14 +248,15 @@ class Ptycho_2D_Dataset(CDataset):
bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted())
s = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s /= 4 # A rough value to make the size work out
s = np.ones(len(nanomap_values)) * s
s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s0 /= 4 # A rough value to make the size work out
s = np.ones(len(nanomap_values)) * s0
s[idx] *= 4
nanomap = axes[0].scatter(1e6 * translations[:,0],1e6 * translations[:,1],s=s,c=nanomap_values)
nanomap = axes[0].scatter(1e6 * translations[:,0],1e6 * translations[:,1],s=s,c=nanomap_values, picker=True)
fig.canvas.mpl_connect('pick_event',on_pick)
axes[0].invert_xaxis()
axes[0].set_facecolor('k')
axes[0].set_xlabel('Translation x (um)')
@@ -251,13 +274,15 @@ class Ptycho_2D_Dataset(CDataset):
axes[0].set_title('Nanomap')
bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted())
s = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s /= 4 # A rough value to make the size work out
s = np.ones(len(nanomap_values)) * s
s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s0 /= 4 # A rough value to make the size work out
s = np.ones(len(nanomap_values)) * s0
s[idx] *= 4
axes[0].clear()
nanomap = axes[0].scatter(1e6 * translations[:,0],1e6 * translations[:,1],s=s,c=nanomap_values)
nanomap = axes[0].scatter(1e6 * translations[:,0],1e6 * translations[:,1],s=s,c=nanomap_values, picker=True)
fig.canvas.mpl_connect('pick_event',on_pick)
axes[0].invert_xaxis()
axes[0].set_facecolor('k')
axes[0].set_xlabel('Translation x (um)')
CDTools/models/__init__.py
+1
View File
@@ -341,3 +341,4 @@ class CDIModel(t.nn.Module):
from CDTools.models.simple_ptycho import SimplePtycho
from CDTools.models.fancy_ptycho import FancyPtycho
from CDTools.models.pinhole_plane_ptycho import PinholePlanePtycho
CDTools/tools/analysis.py
+122 -385
View File
@@ -1,3 +1,10 @@
"""Contains basic functions for analyzing the results of reconstructions
The functions in this module are designed to work either with pytorch tensors
or numpy arrays, so they can be used either directly after reconstructions
on the attributes of the models themselves, or after-the-fact once the
data has been stored in numpy arrays.
"""
from __future__ import division, print_function
import torch as t
@@ -20,11 +27,15 @@ def orthogonalize_probes(probes):
defined by the probes in that basis. After diagonalization, the
eigenvectors can be recast into the original basis and returned
Args:
probes (t.Tensor) : n x (image) size tensor, a stack of probes
Parameters
----------
probes : array
An l x n x m complex array representing a stack of probes
Returns:
(t.Tensor) : n x (image) size tensor, a stack of probes
Returns
-------
ortho_probes: array
An l x n x m complex array representing a stack of probes
"""
try:
@@ -74,12 +85,13 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
these ambiguities for real differences between the reconstructions.
The ambiguities and standardizations are:
* Probe and object can be scaled inversely to one another
* So we set the probe intensity to an average per-pixel value of 1
* The probe and object can aquire equal and opposite phase ramps
* So we set the centroid of the FFT of the probe to zero frequency
* The probe and object can each acquire an arbitrary overall phase
* So we set the phase of the sum of all values of both the probe and object to 0
1) a. Probe and object can be scaled inversely to one another
b. So we set the probe intensity to an average per-pixel value of 1
2) a. The probe and object can aquire equal and opposite phase ramps
b. So we set the centroid of the FFT of the probe to zero frequency
3) a. The probe and object can each acquire an arbitrary overall phase
b. So we set the phase of the sum of all values of both the probe and object to 0
When dealing with the properties of the object, a slice is used by
default as the edges of the object often are dominated by unphysical
@@ -88,15 +100,23 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
dominant probe mode (assumed to be the first in the list) is used, but
all the probes are updated with the same factors.
Args:
probe (t.tensor) : tensor or numpy array storing a retrieved probe or stack of incoherently mixed probes
obj (t.tensor) : tensor or numpy array storing a retrieved probe
obj_slice (slice) : optional, a slice to take from the object for calculating normalizations
correct_ramp (bool) : Default False, whether to correct for the relative phase ramps
Parameters
----------
probe : array
A complex array storing a retrieved probe or stack of incoherently mixed probes
obj : array
A complex array storing a retrieved probe
obj_slice : slice
Optional, a slice to take from the object for calculating normalizations
correct_ramp : bool
Default False, whether to correct for the relative phase ramps
Returns:
(t.tensor) : The standardized probe
(t.tensor) : The standardized object
Returns
-------
standardized_probe : array
The standardized probe
standardized_obj : array
The standardized object
"""
# First, we normalize the probe intensity to a fixed value.
@@ -177,18 +197,27 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
precision and uses a sinc interpolation to shift all the probes and objects
to a common frame. Then the images are summed.
Args:
probes (list) : A list of probes or stacks of probe modes
objects (list) : A list of objects
use_probe (bool) : Default False, whether to use the probe or object for alignment
obj_slice (slice) : Optional, A slice of the object to use for alignment and normalization
correct_ramp (bool) : Default False, whether to correct for a relative phase ramp in the probe and object
Parameters
----------
probes : list(array)
A list of probes or stacks of probe modes
objects : list(array)
A list of objects
use_probe : bool
Default False, whether to use the probe or object for alignment
obj_slice : slice
Optional, A slice of the object to use for alignment and normalization
correct_ramp : bool
Default False, whether to correct for a relative phase ramp in the probe and object
Returns:
(array_like) : The synthesized probe
(array_like) : The synthesized object
(list) : a list of standardized objects, for further processing
Returns
-------
synth_probe : array
The synthesized probe
synth_obj : array
The synthesized object
obj_stack : list(array)
A list of standardized objects, for further processing
"""
probe_np = False
@@ -252,7 +281,7 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
"""Calculates a PRTF between each the individual objects and an averaged one
"""Calculates a PRTF between each the individual objects and a synthesized one
The consistency PRTF at any given spatial frequency is defined as the ratio
between the intensity of any given reconstruction and the intensity
@@ -260,16 +289,25 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
Typically, the PRTF is averaged over spatial frequencies with the same
magnitude.
Args:
synth_obj (t.Tensor) : The synthesized object in the numerator of the PRTF
objects (list): A list of objects or diffraction patterns for the denomenator of the PRTF
basis (array_like) : The basis for the reconstruction array to allow output in physical unit
obj_slice : Optional, a slice of the objects to use for calculating the PRTF
nbinbs (int) : Optional, number of bins to use in the histogram. Defaults to a sensible value
Parameters
----------
synth_obj : array
The synthesized object in the numerator of the PRTF
objects : list(array)
A list of objects or diffraction patterns for the denomenator of the PRTF
basis : array
The basis for the reconstruction array to allow output in physical unit
obj_slice : slice
Optional, a slice of the objects to use for calculating the PRTF
nbins : int
Optional, number of bins to use in the histogram. Defaults to a sensible value
Returns:
(t.Tensor) : The frequencies for the PRTF
(t.Tensor) : The values of the PRTF
Returns
-------
freqs : array
The frequencies for the PRTF
PRTF : array
The values of the PRTF
"""
obj_np = False
@@ -279,7 +317,9 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
if isinstance(synth_obj, np.ndarray):
synth_obj = cmath.complex_to_torch(synth_obj).to(t.float32)
if isinstance(basis, t.Tensor):
basis = basis.detach().cpu().numpy()
if obj_slice is None:
obj_slice = np.s_[(objects[0].shape[0]//8)*3:(objects[0].shape[0]//8)*5,
(objects[0].shape[1]//8)*3:(objects[0].shape[1]//8)*5]
@@ -311,352 +351,40 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
prtfs.append(synth_ints/single_ints)
prtf = np.mean(prtfs,axis=0)
if not obj_np:
bins = t.Tensor(bins)
prtfs = t.Tensor(prtfs)
prtf = t.Tensor(prtf)
return bins[:-1], np.mean(prtfs,axis=0)
def calc_deconvolved_cross_correlation(im1, im2):
"""Calculates a cross-correlation between two images with their autocorrelations deconvolved.
This can also be thought of as the inverse Fourier transform of the
object from which the Fourier Ring Correlation is defined.
Args:
im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
Returns:
(t.Tensor) : The deconvolved cross-correlation, in real space
"""
#
# Here's my approach, perhaps it's a little unconventional. I will first
# calculate the phase correlation function as found in ____ (cite a paper
# defining it). This is strongly peaked, so I can take a small window
# of say, 10x10 pixels, and then do a sinc interpolation of that area
# using an FFT with upsampling by a factor of resolution in reciprocal
# space
#
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = cmath.complex_to_torch(im2)
im_np = True
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2)))
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)
if im_np:
cor = cmath.torch_to_complex(cor)
return cor
def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1):
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = cmath.complex_to_torch(im2)
im_np = True
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
im_slice = np.s_[300:-300,300:-300]
if nbins is None:
nbins = np.max(synth_obj[im_slice].shape) // 4
cor_fft = cmath.cmult(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy()
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[0],d=di))
j_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
synth_ints, bins = np.histogram(Rs,bins=nbins,weights=synth_fft)
prtfs = []
for obj in objects:
obj = obj[obj_slice]
single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy()
single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft)
prtfs.append(synth_ints/single_ints)
if not obj_np:
bins = t.Tensor(bins)
prtfs = t.Tensor(prtfs)
return bins[:-1], np.mean(prtfs,axis=0)
def calc_deconvolved_cross_correlation(im1, im2):
"""Calculates a cross-correlation between two images with their autocorrelations deconvolved.
This can also be thought of as the inverse Fourier transform of the
object from which the Fourier Ring Correlation is defined.
Args:
im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
Returns:
(t.Tensor) : The deconvolved cross-correlation, in real space
"""
#
# Here's my approach, perhaps it's a little unconventional. I will first
# calculate the phase correlation function as found in ____ (cite a paper
# defining it). This is strongly peaked, so I can take a small window
# of say, 10x10 pixels, and then do a sinc interpolation of that area
# using an FFT with upsampling by a factor of resolution in reciprocal
# space
#
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = cmath.complex_to_torch(im2)
im_np = True
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2)))
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)
if im_np:
cor = cmath.torch_to_complex(cor)
return cor
def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1):
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = cmath.complex_to_torch(im2)
im_np = True
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
im_slice = np.s_[300:-300,300:-300]
if nbins is None:
nbins = np.max(synth_obj[im_slice].shape) // 4
cor_fft = cmath.cmult(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy()
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[0],d=di))
j_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
synth_ints, bins = np.histogram(Rs,bins=nbins,weights=synth_fft)
prtfs = []
for obj in objects:
obj = obj[obj_slice]
single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy()
single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft)
prtfs.append(synth_ints/single_ints)
if not obj_np:
bins = t.Tensor(bins)
prtfs = t.Tensor(prtfs)
return bins[:-1], np.mean(prtfs,axis=0)
def calc_deconvolved_cross_correlation(im1, im2):
"""Calculates a cross-correlation between two images with their autocorrelations deconvolved.
This can also be thought of as the inverse Fourier transform of the
object from which the Fourier Ring Correlation is defined.
Args:
im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
Returns:
(t.Tensor) : The deconvolved cross-correlation, in real space
"""
#
# Here's my approach, perhaps it's a little unconventional. I will first
# calculate the phase correlation function as found in ____ (cite a paper
# defining it). This is strongly peaked, so I can take a small window
# of say, 10x10 pixels, and then do a sinc interpolation of that area
# using an FFT with upsampling by a factor of resolution in reciprocal
# space
#
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = cmath.complex_to_torch(im2)
im_np = True
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2)))
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)
if im_np:
cor = cmath.torch_to_complex(cor)
return cor
def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1):
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = cmath.complex_to_torch(im2)
im_np = True
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
im_slice = np.s_[300:-300,300:-300]
if nbins is None:
nbins = np.max(synth_obj[im_slice].shape) // 4
cor_fft = cmath.cmult(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy()
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[0],d=di))
j_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
synth_ints, bins = np.histogram(Rs,bins=nbins,weights=synth_fft)
prtfs = []
for obj in objects:
obj = obj[obj_slice]
single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy()
single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft)
prtfs.append(synth_ints/single_ints)
if not obj_np:
bins = t.Tensor(bins)
prtfs = t.Tensor(prtfs)
return bins[:-1], np.mean(prtfs,axis=0)
return bins[:-1], prtf
def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
"""Calculates a cross-correlation between two images with their autocorrelations deconvolved.
This can also be thought of as the inverse Fourier transform of the
object from which the Fourier Ring Correlation is defined.
This is formally defined as the inverse Fourier transform of the normalized
product of the Fourier transforms of the two images. It results in a
kernel, whose characteristic size is related to the exactness of the
possible alignment between the two images, on top of a random background
Args:
im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
im_slice (slice) : Default is from 3/8 to 5/8 across the image, a slice to use in the processing.
Parameters
----------
im1 : array
The first image, as a complex or real valued array
im2 : array
The first image, as a complex or real valued array
im_slice : slice
Default is from 3/8 to 5/8 across the image, a slice to use in the processing.
Returns:
(t.Tensor) : The deconvolved cross-correlation, in real space
Returns
-------
corr : array
The complex-valued deconvolved cross-correlation, in real space
"""
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
@@ -698,23 +426,32 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
Like other analysis functions, this can take input in numpy or pytorch,
and will return output in the respective format.
Args:
im1 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
im2 (t.Tensor) : The first image, as a complex or real valued pytorch tensor or numpy array
basis (t.Tensor) : The basis for the images, defined as is standard for datasets
im_slice (slice) : Default is from 3/8 to 5/8 across the image, a slice to use in the processing.
nbins (int) : Number of bins to break the FRC up into
snr (float) : The signal to noise ratio (for the combined information in both images) to return a threshold curve for.
Parameters
----------
im1 : array
The first image, a complex or real valued array
im2 : array
The first image, a complex or real valued array
basis : array
The basis for the images, defined as is standard for datasets
im_slice : slice
Default is from 3/8 to 5/8 across the image, a slice to use in the processing.
nbins : int
Number of bins to break the FRC up into
snr : float
The signal to noise ratio (for the combined information in both images) to return a threshold curve for.
Returns:
(t.Tensor) : The frequencies associated with each FRC value
(t.Tensor) : The FRC values
(t.Tensor) : The threshold curve for comparison
Returns
-------
freqs : array
The frequencies associated with each FRC value
FRC : array
The FRC values
threshold : array
The threshold curve for comparison
"""
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
CDTools/tools/cmath.py
+112 -61
View File
@@ -27,11 +27,15 @@ def complex_to_torch(x):
complex numbers. This maps a complex type numpy array to a torch
tensor following this convention
Args:
x (array_like): A numpy array to convert
Parameters
----------
x : np.ndarray
A numpy array to convert
Returns:
torch.Tensor : A torch tensor representation of that array
Returns
-------
torch.Tensor
A torch tensor representation of that array
"""
return t.from_numpy(np.stack((np.real(x),np.imag(x)),axis=-1))
@@ -42,13 +46,19 @@ def torch_to_complex(x):
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This maps a torch tensor following that convention
to the appropriate numpy complex array
to the appropriate numpy complex array. Note that, in order for this
function to work, the tensor must be detached from any parameters and
living on the CPU.
Args:
x (torch.Tensor): A tensor to convert
Parameters
----------
x : torch.Tensor
A tensor to convert
Returns:
np.array : A complex typed numpy array corresponding to the input
Returns
-------
np.array
A complex typed numpy array corresponding to the input
"""
x = np.array(x)
@@ -63,72 +73,88 @@ def torch_to_complex(x):
# and thus doesn't need it's own function
#
def cabssq(a):
def cabssq(x):
"""Returns the square of the absolute value of a complex torch tensor
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise absolute value
squared of any toch tensor following that standard.
Args:
x (torch.Tensor): An input tensor
Parameters
----------
x : torch.Tensor
An input tensor
Returns:
array_like : A tensor storing the elementwise absolute value squared
Returns
-------
torch.Tensor
A tensor storing the elementwise absolute value squared
"""
return a[...,0]**2 + a[...,1]**2
return x[...,0]**2 + x[...,1]**2
def cabs(a):
def cabs(x):
"""Returns the absolute value of a complex torch tensor
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise absolute value
of any torch tensor following that standard.
Args:
x (torch.Tensor): An input tensor
Parameters
----------
x : torch.Tensor
An input tensor
Returns:
array_like : A tensor storing the elementwise absolute value
Returns
-------
torch.Tensor
A tensor storing the elementwise absolute value
"""
return t.sqrt(cabssq(a))
return t.sqrt(cabssq(x))
def cphase(a):
def cphase(x):
"""Returns the phase of a complex torch tensor
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise complex phase
of any torch tensor following that standard.
Args:
x (torch.Tensor): An input tensor
Parameters
----------
x : torch.Tensor
An input tensor
Returns:
array_like : A tensor storing the elementwise phase
Returns
-------
torch.Tensor
A tensor storing the elementwise phase
"""
return t.atan2(a[...,1],a[...,0])
return t.atan2(x[...,1],x[...,0])
def cconj(a):
def cconj(x):
"""Returns the complex conjugate of a complex torch tensor
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise complex conjugate
of any torch tensor following that standard.
Args:
x (torch.Tensor): An input tensor
Parameters
----------
x : torch.Tensor
An input tensor
Returns:
array_like : A tensor storing the elementwise complex conjugate
Returns
-------
torch.Tensor
A tensor storing the elementwise complex conjugate
"""
return t.stack((a[...,0],-a[...,1]),dim=-1)
return t.stack((x[...,0],-x[...,1]),dim=-1)
@@ -139,12 +165,17 @@ def cmult(a,b):
complex numbers. This calculates the elementwise product
of two torch tensors following that standard.
Args:
a (torch.Tensor): An input tensor
b (torch.Tensor): A second input tensor
Parameters
----------
a : torch.Tensor
An input tensor
b : torch.Tensor
A second input tensor
Returns:
torch.Tensor : A tensor storing the elementwise product
Returns
-------
torch.Tensor
A tensor storing the elementwise product
"""
@@ -160,12 +191,17 @@ def cdiv(a,b):
complex numbers. This calculates the elementwise quotient
of two torch tensors following that standard.
Args:
a (torch.Tensor): An input tensor
b (torch.Tensor): A second input tensor
Parameters
----------
a : torch.Tensor
An input tensor
b : torch.Tensor
A second input tensor
Returns:
torch.Tensor : A tensor storing the elementwise complex quotient
Returns
-------
torch.Tensor
A tensor storing the elementwise complex quotient
"""
return cmult(a, cconj(b)) / t.unsqueeze(cabssq(b),-1)
@@ -188,12 +224,17 @@ def fftshift(array,dims=None):
represent the complex number and be of dimension 2), but can shift
any arbitrary set of dimensions.
Args:
array (torch.Tensor) : An array of data to be fftshifted
dims (iterable) : A list of all dimensions to shift
Parameters
----------
array : torch.Tensor
An array of data to be fftshifted
dims : iterable
A list of all dimensions to shift
Returns:
torch.Tensor : fftshifted tensor
Returns
-------
torch.Tensor
The fftshifted tensor
"""
@@ -220,12 +261,17 @@ def ifftshift(array,dims=None):
represent the complex number and be of dimension 2), but can shift
any arbitrary set of dimensions.
Args:
array (torch.Tensor) : An array of data to be ifftshifted
dims (iterable) : A list of all dimensions to shift
Parameters
----------
array : torch.Tensor
An array of data to be ifftshifted
dims : list(int)
A list of all dimensions to shift
Returns:
torch.Tensor : ifftshifted tensor
Returns
-------
torch.Tensor
The ifftshifted tensor
"""
@@ -241,15 +287,20 @@ def ifftshift(array,dims=None):
return array
def expi(array):
"""Returns a complex-format array for exp(i* (real array))
def expi(x):
"""Returns a complex-format tensor for exp(i* (x))
Expects the input to be in the form of a real-valued array
Expects the input to be in the form of a real-valued tensor
Args:
array (torch.Tensor) : An array to be exponentiated
Parameters
----------
x : torch.Tensor
An array to be exponentiated
Returns:
torch.Tensor : A complex-format array
Returns
-------
torch.Tensor
A complex-format tensor
"""
return t.stack((t.cos(array),t.sin(array)),dim=-1)
return t.stack((t.cos(x),t.sin(x)),dim=-1)
CDTools/tools/data.py
+145 -116
View File
@@ -1,3 +1,10 @@
"""Contains the base functions for loading and saving data from/to .cxi files
These functions are used when constructing a new dataset class to pull
specific desired information from a .cxi file. These functions should
handle all the needed conversions between standard formats (for example,
transposes of the basis arrays, shifting from object to probe motion, etc).
"""
from __future__ import division, print_function, absolute_import
import h5py
@@ -26,52 +33,6 @@ __all__ = ['get_entry_info',
'add_data',
'add_ptycho_translations']
#
#
# I will put here some thoughts about how to load data into this program.
#
#
# The reconstructions should have the ability to generate datasets.
# So you could write a reconstruction engine and then it would be
# able to simulate data directly in the engine for you to use as a
# reconstruction
#
# I don't even think there needs to be a loading tool for loading cxi files
# because there isn't really a better method beyond just loading the
# file into an h5py object. This file could host the simple cxi file
# browser, perhaps. But I think the reality is that we need individual
# loaders for each kind of experiment. Perhaps we could put some basic
# reuseable tools for inspecting cxi-type h5 files in this file.
#
#
# Then, there can be some more sophisticated tools that load data for
# specific use cases that are common - loading data for a 2D CDI experiment,
# loading data for a 2D Ptycho experiment, loading data for Bragg Ptycho in
# 3D, loading data for a 3D CDI experiment, etc.
#
#
# Perhaps one good way to package this is for the kind of data associated
# with any particular experiment to have it's own kind of dataset or view.
# So there would be a "2D Ptychography" data viewer, which would contain
# all the measured data that comes from a 2D ptychography experiment.
# The specialized functions would plop out these data viewers, and the
# reconstruction classes could be designed around a particular kind of
# viewer with the most general kind just requiring a generic data viewer.
#
# Data viewers could have simple tools like the ability to send themselves
# to the GPU, CPU, change the datatype, etc. I think the most generic thing
# is as a subclass of the torch Data objects, where they would for each slice
# return the index, a set of defining parameters (translation, angle, energy,
# whatever), and a diffraction pattern. They would also have a "setup"
# attribute, or "metadata", or whatever you'd want to call it, that contain
# the various fixed experimental parameters (energy, distance, etc.)
#
# And I think the cxi visualizer should really go into it's own script,
# because it's not a reuseable component.
#
#
# Functions to inspect the basic attributes of a cxi file represented as an
# h5 file object
@@ -85,11 +46,15 @@ def get_entry_info(cxi_file):
is converted to python datetime objects if the string is properly
formatted.
Args:
cxi_file (h5py.File) : a file object to be read
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns:
dict : A dictionary with basic metadata defined in the cxi file
Returns
-------
entry_info : dict
A dictionary with basic metadata defined in the cxi file
"""
e1 = cxi_file['entry_1']
@@ -113,11 +78,15 @@ def get_entry_info(cxi_file):
def get_sample_info(cxi_file):
"""Returns a dictionary with the basic metadata from the cxi file's entry_1/sample_1 attribute
Args:
cxi_file (h5py.File) : a file object to be read
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns:
dict : A dictionary with basic metadata from the sample defined in the cxi file
Returns
-------
sample_info : dict
A dictionary with basic metadata from the sample defined in the cxi file
"""
if 'entry_1/sample_1' not in cxi_file:
@@ -175,11 +144,15 @@ def get_sample_info(cxi_file):
def get_wavelength(cxi_file):
"""Returns the wavelength of the source defined in the cxi file object, in m
Args:
cxi_file (h5py.File) : a file object to be read
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns:
np.float32 : The wavelength of the source defined in the cxi file
Returns
-------
wavelength: np.float32
The wavelength of the source defined in the cxi file
"""
i1 = cxi_file['entry_1/instrument_1']
if 'source_1/wavelength' in i1:
@@ -208,13 +181,19 @@ def get_detector_geometry(cxi_file):
corner location is not. If the corner location is not reported in
the cxi file, no attempt will be made to calculate it.
Args:
cxi_file (h5py.File) : a file object to be read
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns:
distance (np.float32) : The sample to detector distance, in m
basis_vectors (np.array) : The basis vectors for the detector
corner_location (np.array) : The location of the (0,0) pixel in the detector
Returns
-------
distance : np.float32
The sample to detector distance, in m
basis_vectors : np.array
The basis vectors for the detector
corner_location : np.array
The real-space location of the (0,0) pixel in the detector
"""
i1 = cxi_file['entry_1/instrument_1']
@@ -283,11 +262,15 @@ def get_mask(cxi_file):
which is defined to mean that the pixel has signal above the
background. These pixels are treated as on pixels
Args:
cxi_file (h5py.File) : a file object to be read
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns:
np.array : An array storing the mask from the cxi file
Returns
-------
mask : np.array
An array storing the mask from the cxi file
"""
i1 = cxi_file['entry_1/instrument_1']
@@ -311,12 +294,17 @@ def get_dark(cxi_file):
If the darks do not exist, it will return None
Args:
cxi_file (h5py.File) : a file object to be read
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns:
np.array : An array storing the dark image
Returns
-------
dark : np.array
An array storing the dark image
"""
i1 = cxi_file['entry_1/instrument_1']
if 'detector_1/data_dark' in i1:
darks = np.array(i1['detector_1/data_dark'])
@@ -342,16 +330,20 @@ def get_data(cxi_file, cut_zeroes = True):
It will also read out the axes attribute of the data into a list
of strings
Args:
cxi_file (h5py.File) : a file object to be read
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns:
np.array : An array storing the data defined in the cxi file
list : A list of the axes defined in the axes attribute, if any
Returns
-------
data : np.array
An array storing the data defined in the cxi file
axes : list(str)
A list of the axes defined in the axes attribute, if any
"""
# Possible locations for the data
#
# entry_1/detector_1/data
if 'entry_1/data_1/data' in cxi_file:
pull_from = 'entry_1/data_1/data'
elif 'entry_1/instrument_1/detector_1/data' in cxi_file:
@@ -380,12 +372,18 @@ def get_ptycho_translations(cxi_file):
to specify translations of the samples and the CDTools code specifies
translations of the optics.
Args:
cxi_file (h5py.File) : a file object to be read
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns
-------
translations : np.array
An array storing the translations defined in the cxi file
axes : list(str)
A list of the axes defined in the axes attribute, if any
Returns:
np.array : An array storing the translations defined in the cxi file
list : A list of the axes defined in the axes attribute, if any
"""
if 'entry_1/data_1/translation' in cxi_file:
@@ -410,8 +408,10 @@ def get_ptycho_translations(cxi_file):
def create_cxi(filename):
"""Creates a new cxi file with a single entry group
Args:
filename (str) : The path at which to create the file
Parameters
----------
filename : str
The path at which to create the file
"""
file_obj = h5py.File(filename,'w')
file_obj.create_dataset('cxi_version', data=160)
@@ -423,9 +423,12 @@ def create_cxi(filename):
def add_entry_info(cxi_file, metadata):
"""Adds a dictionary of entry metadata to the entry_1 group of a cxi file object
Args:
cxi_file (h5py.File) : The file to add the info to
metadata (dict) : A dictionary containing all the metadata to be stored
Parameters
----------
cxi_file : h5py.File
The file to add the info to
metadata : dict
A dictionary containing all the metadata to be stored
"""
# Just the string and datetime types should be relevant but all are
# included in case the cxi spec becomes more permissive
@@ -448,9 +451,12 @@ def add_sample_info(cxi_file, metadata):
This function will create the sample_1 attribute if it doesn't already exist
Args:
cxi_file (h5py.File) : The file to add the info to
metadata (dict) : A dictionary containing all the metadata to be stored
Parameters
----------
cxi_file : h5py.File
The file to add the info to
metadata : dict
A dictionary containing all the metadata to be stored
"""
if 'entry_1/sample_1' not in cxi_file:
cxi_file['entry_1'].create_group('sample_1')
@@ -485,9 +491,12 @@ def add_source(cxi_file, wavelength):
It stores the energy and wavelength attributes in the source_1 group,
given a wavelength to define them from.
Args:
cxi_file (h5py.File) : The file to add the source to
wavelength (float) : The wavelength of light
Parameters
----------
cxi_file : h5py.File
The file to add the source to
wavelength : float
The wavelength of light
"""
if 'entry_1/instrument_1' not in cxi_file:
cxi_file['entry_1'].create_group('instrument_1')
@@ -507,11 +516,16 @@ def add_detector(cxi_file, distance, basis, corner=None):
detector basis, and corner position (if relevant) based on the provided
information
Args:
cxi_file (h5py.File) : The file to add the detector to
distance (float) : The sample to detector distance
basis (array_like) : The detector basis
corner (array_like) : Optional, the corner position of the detector
Parameters
----------
cxi_file : h5py.File
The file to add the detector to
distance : float
The sample to detector distance
basis : array
The detector basis
corner : array
Optional, the corner position of the detector
"""
if 'entry_1/instrument_1' not in cxi_file:
@@ -522,11 +536,13 @@ def add_detector(cxi_file, distance, basis, corner=None):
d1 = i1['detector_1']
d1['distance'] = np.float32(distance)
d1['x_pixel_size'] = np.linalg.norm(basis[:,1])
d1['y_pixel_size'] = np.linalg.norm(basis[:,0])
if isinstance(basis, t.Tensor):
basis = basis.detach().cpu().numpy()
d1['x_pixel_size'] = np.linalg.norm(basis[:,1])
d1['y_pixel_size'] = np.linalg.norm(basis[:,0])
d1.create_dataset('basis_vectors', data=basis)
if corner is not None:
if isinstance(corner, t.Tensor):
corner = corner.detach().cpu().numpy()
@@ -543,9 +559,12 @@ def add_mask(cxi_file, mask):
most general mask allowed by the cxi file format but it captures the
distinction between pixels to be used and pixels not to be used.
Args:
cxi_file (h5py.File) : The file to add the mask to
mask (array_like) : The mask to save out to the file
Parameters
----------
cxi_file : h5py.File
The file to add the mask to
mask : array
The mask to save out to the file
"""
if 'entry_1/instrument_1' not in cxi_file:
@@ -568,9 +587,12 @@ def add_dark(cxi_file, dark):
It places the dark image data into the data_dark dataset under
entry_1/instrument_1/detector_1.
Args:
cxi_file (h5py.File) : The file to add the mask to
dark (array_like) : The dark image(s) to save out to the file
Parameters
----------
cxi_file : h5py.File
The file to add the mask to
dark : array
The dark image(s) to save out to the file
"""
if 'entry_1/instrument_1' not in cxi_file:
cxi_file['entry_1'].create_group('instrument_1')
@@ -592,10 +614,14 @@ def add_data(cxi_file, data, axes=None):
1) The entry_1/instrument_1/detector_1/data path
2) A softlink at entry_1/data_1/data
Args:
cxi_file (h5py.File) : The file to add the data to
data (array_like) : The data to be saved
axes (list) : Optional, a list of axis names to be saved in the axes attribute
Parameters
----------
cxi_file : h5py.File
The file to add the data to
data : array
The data to be saved
axes : list(str)
Optional, a list of axis names to be saved in the axes attribute
"""
if 'entry_1/data_1' not in cxi_file:
cxi_file['entry_1'].create_group('data_1')
@@ -629,15 +655,18 @@ def add_ptycho_translations(cxi_file, translations):
the standard in cxi files that the translations refer to the object's
translation.
It will generally store them in 3 places:
It will store them in 3 places:
1) The entry_1/sample_1/geometry_1/translation path
2) A softlink at entry_1/data_1/translation
3) A softlink at entry_1/instrument_1/detector_1/translation
Args:
cxi_file (h5py.File) : The file to add the translations to
translations (array_like) : The translations to be saved
Parameters
----------
cxi_file : h5py.File
The file to add the translations to
translations : array
The translations to be saved
"""
if 'entry_1/sample_1' not in cxi_file:
CDTools/tools/image_processing.py
+98 -39
View File
@@ -1,3 +1,12 @@
"""Contains functions for basic image processing needs
This module contains two kinds of image processing tools. The first type
is specific tools for calculating commonly needed metrics (such as the
centroid of an image), directly on complex-valued torch tensors. The second
kind of tools perform common image manipulations on torch tensors, in such
a way that it is safe to include them in automatic differentiation models.
"""
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
@@ -17,12 +26,17 @@ def centroid(im, dims=2):
Beware that the meaning of the centroid is not well defined if your
image contains values less than 0
Args:
im (t.Tensor) : An image or stack of images to calculate from
dims (int) : Default 2, how many trailing dimensions to calculate for
Returns:
t.Tensor : An (i,j) index or stack of indices
Parameters
----------
im : torch.Tensor
An image or stack of images to calculate from
dims : int
Default 2, how many trailing dimensions to calculate with
Returns
-------
centroid : torch.Tensor
An (i,j) index or stack of indices
"""
# For some reason this needs to be a list
indices = [t.arange(im.shape[-dims+i]).to(t.float32) for i in range(dims)]
@@ -46,12 +60,19 @@ def centroid_sq(im, dims=2, comp=False):
represents the real and imaginary part of a complex number, and the
centroid will be calculated for the magnitude squared of those numbers
Args:
im (t.Tensor) : An image or stack of images to calculate from
dims (int) : Default 2, how many trailing dimensions to calculate for
comp (bool) : Default is False, whether the data represents complex numbers
Returns:
t.Tensor : An (i,j) index or stack of indices
Parameters
----------
im : torch.Tensor
An image or stack of images to calculate from
dims : int
Default 2, how many trailing dimensions to calculate for
comp : bool
Default is False, whether the data represents complex numbers
Returns
-------
centroid: torch.Tensor
An (i,j) index or stack of indices
"""
if comp:
im_sq = cmath.cabssq(im)
@@ -67,12 +88,17 @@ def sinc_subpixel_shift(im, shift):
The subpixel shift is done circularly via a multiplication with a linear
phase mask in Fourier space.
Args:
im (torch.Tensor) : A complex-valued tensor to perform the subpixel shift on
shift (array_like) : A length-2 array_like object describing the shift to perform, in pixels
Parameters
----------
im : torch.Tensor
A complex-valued tensor to perform the subpixel shift on
shift : array
A length-2 array describing the shift to perform, in pixels
Returns:
(torch.Tensor) : The subpixel shifted tensor
Returns
-------
shifted_im : torch.Tensor
The subpixel shifted tensor
"""
i = t.arange(im.shape[0]) - im.shape[0]//2
@@ -97,11 +123,21 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
approach outlined in "Efficient subpixel image registration algorithms",
Optics Express (2008) by Manual Guizar-Sicarios et al.
Args:
im1 (t.Tensor): The first real or complex-valued torch tensor
im2 (t.Tensor): The second real or complex-valued torch tensor
search_around (array_like) : Default (0,0), the shift to search in the vicinity of
resolution (int): Default is 10, the resolution to calculate to in units of 1/n
Parameters
----------
im1 : torch.Tensor
The first real or complex-valued torch tensor
im2 : torch.Tensor
The second real or complex-valued torch tensor
search_around : array
Default (0,0), the shift to search in the vicinity of
resolution : int
Default is 10, the fraction of a pixel to calculate to
Returns
-------
shift : torch.Tensor
The relative shift (i,j) needed to best map im1 onto im2
"""
#
# Here's my approach, perhaps it's a little unconventional. I will first
@@ -111,7 +147,8 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
# using an FFT with upsampling by a factor of resolution in reciprocal
# space
#
# If last dimension is not 2, then convert to a complex tensor now
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
@@ -164,11 +201,17 @@ def find_pixel_shift(im1, im2):
to the amount that im1 would have to be shifted by to line up best with
im2
Args:
im1 (t.Tensor): The first real or complex-valued torch tensor
im2 (t.Tensor): The second real or complex-valued torch tensor
search_around (array_like) : Default (0,0), the shift to search in the vicinity of
resolution (int): Default is 10, the resolution to calculate to in units of 1/n
Parameters
----------
im1 : torch.Tensor
The first real or complex-valued torch tensor
im2 : torch.Tensor
The second real or complex-valued torch tensor
Returns
-------
shift : torch.Tensor
The integer-valued shift (i,j) that best maps im1 onto im2
"""
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
@@ -198,10 +241,19 @@ def find_shift(im1, im2, resolution=10):
pixel resolution, and then searchers the nearby area to calculate a
subpixel shift
Args:
im1 (t.Tensor): The first real or complex-valued torch tensor
im2 (t.Tensor): The second real or complex-valued torch tensor
resolution (int): Default is 10, the resolution to calculate to in units of 1/n
Parameters
----------
im1 : torch.Tensor
The first real or complex-valued torch tensor
im2 : torch.Tensor
The second real or complex-valued torch tensor
resolution : int
Default is 10, the fraction of a pixel to calculate to
Returns
-------
shift : torch.Tensor
The relative shift (i,j) needed to best map im1 onto im2
"""
integer_shift = find_pixel_shift(im1,im2)
subpixel_shift = find_subpixel_shift(im1, im2, search_around=integer_shift,
@@ -221,14 +273,21 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True):
Otherwise, the image is assumed to be real. The image and kernel
must either both be real or both be complex.
Args:
image (torch.Tensor) : The image to convolve
kernel (torch.Tensor) : The 1d kernel to convolve with
dim (int) : Default 0, the dimension to convolve along
fftshift_kernel (bool) : Default True, whether to fftshift the kernel first.
Parameters
----------
image : torch.Tensor
The image to convolve
kernel : torch.Tensor
The 1d kernel to convolve with
dim : int
Default 0, the dimension to convolve along
fftshift_kernel : bool
Default True, whether to fftshift the kernel first.
Returns:
(torch.Tensor) : The convolved image
Returns
-------
convolved_im : torch.Tensor
The convolved image
"""
complex_things = 2
CDTools/tools/initializers.py
+96 -42
View File
@@ -1,3 +1,9 @@
"""Contains functions to sensibly initialize reconstructions
The functions in this module both do the geometric calculations needed to
initialize the reconstrucions, and the heuristic calculations for
geierating sensible initializations for the probe guess.
"""
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
@@ -21,20 +27,33 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
if necessary, define the exit wave basis associated with a far-field
diffraction experiment, and return that basis, shape, and detector slice
Args:
det_basis (torch.Tensor) : The detector basis, as defined elsewhere
det_shape (torch.Size) : The (i,j) shape of the detector
wavelength (float) : The wavelength of light for the experiment, in m
distance (float) : The sample-detector distance, in m
center (torch.Tensor) : If defined, the location of the zero frequency pixel
opt_for_fft (bool) : Default is true, whether to increase detector size to improve fft performance
padding (int) : Default is 0, the size of an extra border of nonphysical pixels around the detector
oversampling (int) : Default is 1, the amount to multiply the exit wave shape by.
Parameters
----------
det_basis : torch.Tensor
The detector basis, as defined elsewhere
det_shape : torch.Size
The (i,j) shape of the detector
wavelength : float)
The wavelength of light for the experiment, in m
distance : float
The sample-detector distance, in m
center : torch.Tensor
If defined, the location of the zero frequency pixel
opt_for_fft : bool
Default is true, whether to increase detector size to improve fft performance
padding : int
Default is 0, the size of an extra border of nonphysical pixels around the detector
oversampling : int
Default is 1, the amount to multiply the exit wave shape by.
Returns:
torch.Tensor : The exit wave basis
torch.Tensor : The exit wave's shape
tuple(slice) : The slice corresponding to the physical detector
Returns
-------
basis : torch.Tensor
The exit wave basis
shape : torch.Tensor
The exit wave's shape
slice : slice
The slice corresponding to the physical detector
"""
det_shape = t.Tensor(tuple(det_shape)).to(t.int32)
@@ -94,13 +113,21 @@ def calc_object_setup(probe_shape, translations, padding=0):
attribute. If this is done, the calculated pixel translation will
correspond to (padding,padding)
Args:
probe_shape (torch.Size) : The size of the probe array
translations (torch.Tensor) : Jx2 stack of pixel-valued (i,j) translations
padding (int) : Optional, the size of an extra border to include
Returns:
torch.Size : required size of object array
torch.Tensor : minimum pixel-valued translation
Parameters
----------
probe_shape : torch.Size
The size of the probe array
translations : torch.Tensor
A Jx2 stack of pixel-valued (i,j) translations
padding : int
Optional, the size of an extra border to include
Returns
-------
obj_shape : torch.Size
The minimum required size of the object array
min_translation : torch.Tensor
The minimum pixel-valued translation
"""
# First we look at the translations to find the minimum translation
@@ -134,15 +161,23 @@ def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]):
from it's focal point. The curvature is implemented by adding a quadratic
phase phi = exp(i*curvature/2 r^2) to the Gaussian
Args:
shape (array_like) : A 1x2 array-like object specifying the dimensions of the output array in the form (i shape, j shape)
sigma (array_like): A 1x2 array-like object specifying the i- and j- standard deviation of the gaussian in the form (i stdev, j stdev)
amplitude (float or int): Default 1, the amplitude the gaussian to simulate
center (array_like) : Optional 1x2 array-like object specifying the location of the center of the gaussian (i center, j center)
curvature (array_like) : Optional complex part to add to the gaussian coefficient
Parameters
----------
shape : array
A 1x2 array specifying the dimensions of the output array in the form (i shape, j shape)
sigma : array
A 1x2 array specifying the i- and j- standard deviation of the gaussian in the form (i stdev, j stdev)
amplitude : float
Default 1, the amplitude the gaussian to simulate
center : array
Optional, a 1x2 array specifying the location of the center of the gaussian (i center, j center)
curvature : array
Optional, a complex part to add to the gaussian coefficient
Returns:
torch.Tensor : The complex-style tensor storing the Gaussian
Returns
-------
torch.Tensor
The complex-style tensor storing the Gaussian
"""
if center is None:
center = ((shape[0]-1)/2, (shape[1]-1)/2)
@@ -175,15 +210,23 @@ def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0):
of sigma in the directions parallel to the i and j basis vectors of
the probe basis
Args:
dataset (Ptycho_2D_Dataset) : The dataset whose intensity we want to match
basis (array_like) : The real space basis for exit waves in our experiment
shape (array_like): The shape of the simulated real space arrays
sigma (array_like): The standard deviation of the probe at it's focus
propagation_distance (float) : Optional, a distance to propagate the gaussian from it's focus
Parameters
----------
dataset : Ptycho_2D_Dataset
The dataset whose intensity we want to match
basis : array
The real space basis for exit waves in our experiment
shape : array
The shape of the simulated real space arrays
sigma : array
The standard deviation of the probe at it's focus
propagation_distance : float
Default 0, a distance to propagate the gaussian from it's focus
Returns:
torch.Tensor : The complex-style tensor storing the Gaussian
Returns
-------
torch.Tensor
The complex-style tensor storing the Gaussian
"""
# First, we want to generate the parameters (sigma and curvature) for the
# propagated gaussian. Ignore the purely z-dependent phases
@@ -230,12 +273,23 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
the probe generated this way, which can often overwhelm the rest of the
probe if there is significant noise on the detector
Args:
dataset (Ptycho_2D_Dataset) : The dataset to work from
shape (torch.Size) : The size of the probe array to simulate
det_slice (slice) : A slice or tuple of slices corresponding to the detector region in Fourier space
propagation_distance (float) : Default is no propagation, an amount to propagate the guessed probe from it's focal point
oversampling (int) : Default 1, the width of the region of pixels in the wavefield to bin into a single detector pixel
Parameters
----------
dataset : Ptycho_2D_Dataset
The dataset to work from
shape : torch.Size
The size of the probe array to simulate
det_slice : slice
A slice or tuple of slices corresponding to the detector region in Fourier space
propagation_distance : float
Default is no propagation, an amount to propagate the guessed probe from it's focal point
oversampling : int
Default 1, the width of the region of pixels in the wavefield to bin into a single detector pixel
Returns
-------
torch.Tensor
The complex-style tensor storing the probe guess
"""
CDTools/tools/interactions.py
+77 -34
View File
@@ -1,3 +1,10 @@
""" This module contains various to simulate stages in the probe-sample interaction
All the tools here are designed to work with automatic differentiation. Each
function simulates some aspect of an interaction model that can be used
for ptychographic reconstruction.
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools.cmath import *
@@ -8,11 +15,6 @@ import numpy as np
__all__ = ['translations_to_pixel', 'pixel_to_translations',
'ptycho_2D_round','ptycho_2D_linear','ptycho_2D_sinc']
#
# This file will host tools to turn various kinds of model information
# (probe, 2D object, 3D object, etc) into exit waves leaving the sample
# area.
#
def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])):
@@ -29,10 +31,19 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.
to the +z axis, [0,0,1]. The default sample orientation has a surface
normal parallel to this direction
Args:
basis (torch.Tensor) : The real space basis the wavefields are defined in
translations (torch.Tensor) : A Jx3 stack of real-space translations
surface_normal (torch.Tensor) : Optional, the sample's surface normal
Parameters
----------
basis : torch.Tensor
The real space basis the wavefields are defined in
translations : torch.Tensor
A Jx3 stack of real-space translations, or a single translation
surface_normal : torch.Tensor
Optional, the sample's surface normal
Returns
-------
pixel_translations : torch.Tensor
A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation
"""
projection_1 = t.Tensor([[1,0,0],
@@ -76,10 +87,19 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,
normal parallel to this direction. Because of this, the z direction
translation is always set to zero in the conversion
Args:
basis (torch.Tensor) : The real space basis the wavefields are defined in
translations (torch.Tensor) : A Jx2 stack of pixel-space translations
surface_normal (torch.Tensor) : Optional, the sample's surface normal
Parameters
----------
basis : torch.Tensor
The real space basis the wavefields are defined in
translations : torch.Tensor
A Jx2 stack of pixel-space translations, or a single translation
surface_normal : torch.Tensor
Optional, the sample's surface normal
Returns
-------
real_translations : torch.Tensor
A Jx3 stack of real-space translations, or a single translation
"""
projection_1 = t.Tensor([[1,0,0],
[0,1,0],
@@ -116,13 +136,19 @@ def ptycho_2D_round(probe, obj, translations):
corresponding to the detector. The exit waves are calculated by
shifting the probe by the rounded value of the translation
Args:
probe (torch.Tensor) : An MxL probe function for the exit waves
object (torch.Tensor) : The object function to be probed
translations (torch.Tensor) : The Nx2 array of (i,j) translations to simulate
Parameters
----------
probe : torch.Tensor
An MxL probe function for the exit waves
object : torch.Tensor
The object function to be probed
translations : torch.Tensor
The Nx2 array of (i,j) translations to simulate
Returns:
torch.Tensor : An NxMxL tensor of the calculated exit waves
Returns
-------
exit_waves : torch.Tensor
An NxMxL tensor of the calculated exit waves
"""
single_translation = False
if translations.dim() == 1:
@@ -153,13 +179,21 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
If shift_probe is True, it applies the subpixel shift to the probe,
otherwise the subpixel shift is applied to the object
Args:
probe (torch.Tensor) : An MxL probe function for the exit waves
object (torch.Tensor) : The object function to be probed
translations (torch.Tensor) : The Nx2 array of translations to simulate
shift_probe (bool) : Whether to subpixel shift the probe or object
Returns:
torch.Tensor : An NxMxL tensor of the calculated exit waves
Parameters
----------
probe : torch.Tensor
An MxL probe function for the exit waves
object : torch.Tensor
The object function to be probed
translations : torch.Tensor
The Nx2 array of translations to simulate
shift_probe : bool
Default True, Whether to subpixel shift the probe or object
Returns
-------
exit_waves : torch.Tensor
An NxMxL tensor of the calculated exit waves
"""
single_translation = False
if translations.dim() == 1:
@@ -243,14 +277,23 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10):
If shift_probe is True, it applies the subpixel shift to the probe,
otherwise the subpixel shift is applied to the object
Args:
probe (torch.Tensor) : An MxL probe function for the exit waves
object (torch.Tensor) : The object function to be probed
translations (torch.Tensor) : The Nx2 array of translations to simulate
shift_probe (bool) : Whether to subpixel shift the probe or object
padding (int) : Default 10, if shifting the object, the padding to apply to the object to avoid circular shift effects
Returns:
torch.Tensor : An NxMxL tensor of the calculated exit waves
Parameters
----------
probe : torch.Tensor
An MxL probe function for the exit waves
object : torch.Tensor
The object function to be probed
translations : torch.Tensor
The Nx2 array of translations to simulate
shift_probe : bool
Default True, Whether to subpixel shift the probe or object
padding : int
Default 10, if shifting the object, the padding to apply to the object to avoid circular shift effects
Returns
-------
exit_waves : torch.Tensor
An NxMxL tensor of the calculated exit waves
"""
single_translation = False
if translations.dim() == 1:
CDTools/tools/losses.py
+37 -20
View File
@@ -27,16 +27,21 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
as long as their shapes match, and the provided mask array can be
broadcast correctly along them.
This is empirically the most useful loss function
This is empirically the most useful loss function for most cases
Args:
intensities (torch.Tensor) : A tensor with measured detector values
sim_intensities (torch.Tensor) : A tensor of simulated detector intensities
mask (torch.Tensor) : A mask with ones for pixels to include and zeros for pixels to exclude
Returns:
loss (torch.Tensor) : A single value for the summed mse
Parameters
----------
intensities : torch.Tensor
A tensor with measured detector values
sim_intensities : torch.Tensor
A tensor of simulated detector intensities
mask : torch.Tensor
A mask with ones for pixels to include and zeros for pixels to exclude
Returns
-------
loss : torch.Tensor
A single value for the mean amplitude mse
"""
# I know it would be more efficient if this function took in the
@@ -65,13 +70,19 @@ def intensity_mse(intensities, sim_intensities, mask=None):
as long as their shapes match, and the provided mask array can be
broadcast correctly along them.
Args:
intensities (torch.Tensor) : A tensor with measured detector intensities.
sim_intensities (torch.Tensor) : A tensor of simulated detector intensities
mask (torch.Tensor) : A mask with ones for pixels to include and zeros for pixels to exclude
Parameters
----------
intensities : torch.Tensor
A tensor with measured detector intensities.
sim_intensities : torch.Tensor
A tensor of simulated detector intensities
mask : torch.Tensor
A mask with ones for pixels to include and zeros for pixels to exclude
Returns:
loss (torch.Tensor) : A single value for the summed mse
Returns
-------
loss : torch.Tensor
A single value for the mean intensity mse
"""
if mask is None:
@@ -102,13 +113,19 @@ def poisson_nll(intensities, sim_intensities, mask=None):
as long as their shapes match, and the provided mask array can be
broadcast correctly along them.
Args:
intensities (torch.Tensor) : A tensor with measured detector intensities.
sim_intensities (torch.Tensor) : A tensor of simulated detector intensities
mask (torch.Tensor) : A mask with ones for pixels to include and zeros for pixels to exclude
Parameters
----------
intensities : torch.Tensor
A tensor with measured detector intensities.
sim_intensities : torch.Tensor
A tensor of simulated detector intensities
mask : torch.Tensor
A mask with ones for pixels to include and zeros for pixels to exclude
Returns:
loss (torch.Tensor) : A single value for the poisson ML metric
Returns
-------
loss : torch.Tensor
A single value for the poisson negative log likelihood
"""
if mask is None:
CDTools/tools/measurements.py
+56 -24
View File
@@ -1,3 +1,10 @@
"""This module contains tools to simulate various measurement models
All the measurements here are safe to use in an automatic differentiation
model. There exist tools to simulate detectors with finite saturation
thresholds, backgrounds, and more.
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
@@ -10,7 +17,7 @@ from torch.nn.functional import avg_pool2d
# intensity pattern on a detector
#
__all__ = ['intensity', 'incoherent sum', 'quadratic_background']
__all__ = ['intensity', 'incoherent_sum', 'quadratic_background']
def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
@@ -20,14 +27,21 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
wavefront. If a detector slice is given, the returned array
will only include that slice from the simulated wavefront.
Args:
wavefield (torch.Tensor) : A JxMxNx2 stack of complex wavefields
detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
oversampling (int) : Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
Parameters
----------
wavefield : torch.Tensor
A JxMxNx2 stack of complex wavefields
detector_slice : slice
Optional, a slice or tuple of slices defining a section of the simulation to return
saturation : float
Optional, a maximum saturation value to clamp the resulting intensities to
oversampling : int
Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
Returns:
torch.Tensor : A real MxN array storing the wavefield's intensities
Returns
-------
sim_patterns : torch.Tensor
A real MxN array storing the wavefield's intensities
"""
output = cmath.cabssq(wavefield) + epsilon
@@ -70,13 +84,21 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
The next two indices index the wavefield. The final index is the complex
index.
Args:
wavefields (torch.Tensor) : An LxJxMxNx2 stack of complex wavefields
detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
oversampling (int) : Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
Returns:
torch.Tensor : A real JXMxN array storing the incoherently summed intensities
Parameters
----------
wavefields : torch.Tensor
An LxJxMxNx2 stack of complex wavefields
detector_slice : slice
Optional, a slice or tuple of slices defining a section of the simulation to return
saturation : float
Optional, a maximum saturation value to clamp the resulting intensities to
oversampling : int
Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
Returns
-------
sim_patterns : torch.Tensor
A real JXMxN array storing the incoherently summed intensities
"""
# This syntax just adds an axis to the slice to preserve the J direction
@@ -116,15 +138,25 @@ def quadratic_background(wavefield, background, detector_slice=None, measurement
of background model is commonly used to enforce positivity of the
background model.
Args:
wavefield (torch.Tensor) : A JxMxNx2 stack of complex wavefields
background (torch.Tensor) : An tensor storing the square root of the detector background
detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
measurement (function) : Optional, the measurement function to use. The default is measurements.intensity
saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
oversampling (int) : Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
Returns:
torch.Tensor : A real MxN array storing the wavefield's intensities
Parameters
----------
wavefield : torch.Tensor
A JxMxNx2 stack of complex wavefields
background : torch.Tensor
An tensor storing the square root of the detector background
detector_slice : slice
Optional, a slice or tuple of slices defining a section of the simulation to return
measurement : function
Default is measurements.intensity, the measurement function to use.
saturation : float
Optional, a maximum saturation value to clamp the resulting intensities to
oversampling : int
Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
Returns
-------
sim_patterns : torch.Tensor
A real MxN array storing the wavefield's intensities
"""
if detector_slice is None:
CDTools/tools/plotting.py
+144 -81
View File
@@ -1,3 +1,11 @@
"""This module contains functions for plotting various important metrics
All the plotting functions here can accept torch input or numpy input,
to facilitate their use both for live inspection of running reconstructions
and for after-the-fact analysis. Utilities for plotting complex valued
images exist, as well as plotting scan patterns and nanomaps
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
@@ -7,8 +15,8 @@ import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb
__all__ = ['colorize','plot_1D','plot_amplitude','plot_phase',
'plot_colorized', 'plot_translations','get_units_factor',
__all__ = ['colorize', 'plot_amplitude', 'plot_phase',
'plot_colorized', 'plot_translations', 'get_units_factor',
'plot_nanomap']
@@ -18,11 +26,14 @@ def colorize(z):
in a call to imshow based on an input complex numpy array (not a
torch tensor representing a complex field)
Args:
z (array_like) : A complex-valued array
Returns:
list : A list of arrays for R,G, and B channels of an image.
Parameters
----------
z : array
A complex-valued array
Returns
-------
rgb : list(array)
A list of arrays for the R,G, and B channels of an image
"""
amp = np.abs(z)
@@ -42,11 +53,15 @@ def colorize(z):
def get_units_factor(units):
"""Gets the multiplicative factor associated with a length unit
Args:
units (str) : The abbreviation for the unit type
Parameters
----------
units : str
The abbreviation for the unit type
Returns:
(float) : The factor meters / (unit)
Returns
-------
factor : float
The factor meters / (unit)
"""
u = units.lower()
@@ -66,38 +81,35 @@ def get_units_factor(units):
factor=1e12
return factor
def plot_1D(arr, fig = None, **kwargs):
"""Simple 1D plotter
Args:
im (numpy array) : A 1D array with dimensions (,N)
fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None,
a new figure is created with an Axes subplot at 111.
**kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class
(see https://matplotlib.org/api/axes_api.html#the-axes-class)
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
plt.scatter(np.arange(arr.shape[-1]), arr)
def plot_amplitude(im, fig = None, basis=None, units='um', cmap='viridis', **kwargs):
""" Plots the amplitude of a complex Tensor or numpy array with dimensions NxMx2.
Args:
im (t.Tensor) : An image with dimensions NxMx2.
fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None,
a new figure is created with an Axes subplot at 111.
basis (numpy array) : Optional, the 3x2 probe basis, used to put the axis labels in real space units.
units (str) : The units to convert the basis to
cmap (str) : Default is 'viridis', the colormap to plot with
**kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class
(see https://matplotlib.org/api/axes_api.html#the-axes-class)
"""Plots the amplitude of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
Parameters
----------
im : array
An complex array with dimensions NxM
fig : matplotlib.figure.Figure
Default is a new figure, a matplotlib figure to use to plot
basis : np.array
Optional, the 3x2 probe basis
units : str
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
**kwargs
All other args are passed to fig.add_subplot(111, **kwargs)
Returns
-------
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
@@ -137,15 +149,33 @@ def plot_amplitude(im, fig = None, basis=None, units='um', cmap='viridis', **kwa
def plot_phase(im, fig=None, basis=None, units='um', cmap='auto', **kwargs):
""" Plots the phase of a complex Tensor or numpy array with dimensions NxMx2.
Args:
im (t.Tensor) : An image with dimensions NxMx2.
fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None,
a new figure is created with an Axes subplot at 111.
basis (numpy array) : Optional, the 3x2 probe basis, used to put the axis labels in real space units.
cmap (str) : Default is 'auto', which chooses between twilight and hsv based on availability.
**kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class
(see https://matplotlib.org/api/axes_api.html#the-axes-class)
""" Plots the phase of a complex array with dimensions NxMx2
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
Parameters
----------
im : array
An complex array with dimensions NxM
fig : matplotlib.figure.Figure
Default is a new figure, a matplotlib figure to use to plot
basis : np.array
Optional, the 3x2 probe basis
units : str
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
**kwargs
All other args are passed to fig.add_subplot(111, **kwargs)
Returns
-------
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
@@ -188,23 +218,39 @@ def plot_phase(im, fig=None, basis=None, units='um', cmap='auto', **kwargs):
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
""" Plots the colorized version of a complex Tensor or numpy array with dimensions NxMx2.
The darkness corresponds to the intensity of the image, and the color corresponds
to the phase.
""" Plots the colorized version of a complex array with dimensions NxM
Args:
im (t.Tensor) : An image with dimensions NxMx2.
fig (matplotlib.figure.Figure) : A matplotlib figure to use to plot. If None,
a new figure is created with an Axes subplot at 111.
basis (numpy array) : Optional, the 3x2 probe basis, used to put the axis labels in real space units.
**kwargs: Can be used to set any keyword arguments for the matplotlib.axes.Axes class
(see https://matplotlib.org/api/axes_api.html#the-axes-class)
The darkness corresponds to the intensity of the image, and the color
corresponds to the phase.
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
Parameters
----------
im : array
An complex array with dimensions NxM
fig : matplotlib.figure.Figure
Default is a new figure, a matplotlib figure to use to plot
basis : np.array
Optional, the 3x2 probe basis
units : str
The length units to mark on the plot, default is um
**kwargs
All other args are passed to fig.add_subplot(111, **kwargs)
Returns
-------
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
@@ -240,24 +286,34 @@ def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
def plot_translations(translations, fig=None, units='um', lines=True):
def plot_translations(translations, fig=None, units='um', lines=True, **kwargs):
"""Plots a set of probe translations in a nicely formatted way
Args:
translations: An Nx2 or Nx3 set of translations in real space
fig : Optional, a figure to plot into
units : Default is um, units to report in (assuming input in m)
lines : Whether to plot the lines indicating the path
Parameters
----------
translations : array
An Nx2 or Nx3 set of translations in real space
fig : matplotlib.figure.Figure
Default is a new figure, a matplotlib figure to use to plot
units : str
Default is um, units to report in (assuming input in m)
lines : bool
Whether to plot lines indicating the path taken
**kwargs
All other args are passed to fig.add_subplot(111, **kwargs)
Returns:
None
Returns
-------
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
factor = get_units_factor(units)
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111)
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
@@ -272,21 +328,29 @@ def plot_translations(translations, fig=None, units='um', lines=True):
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
return fig
def plot_nanomap(translations, values, fig=None, units='um', convention='probe'):
"""Plots a set of nanomap data in a flexible way
Args:
translations : An Nx2 or Nx3 set of translations in real space
values : a length-N object of values associated with the translations
fig : Optional, a figure to plot into
units : Default is um, units to report in (assuming input in m)
lines : Whether to plot the lines indicating the path
convention : 'probe' if the translations refer to probe translations, 'obj' if they refer to object translations
Parameters
----------
translations : array
An Nx2 or Nx3 set of translations in real space
values : array
A length-N object of values associated with the translations
fig : matplotlib.figure.Figure
Default is a new figure, a matplotlib figure to use to plot
units : str
Default is um, units to report in (assuming input in m)
convention : str
Default is 'probe', alternative is 'obj'. Whether the translations refer to the probe or object.
Returns:
None
Returns
-------
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
@@ -316,10 +380,9 @@ def plot_nanomap(translations, values, fig=None, units='um', convention='probe')
plt.scatter(factor * trans[:,0],factor * trans[:,1],s=s,c=values)
plt.gca().invert_xaxis()
plt.gca().set_facecolor('k')
plt.xlabel('Translation x (' + units + ')')
plt.ylabel('Translation y (' + units + ')')
plt.colorbar()
return fig
CDTools/tools/projectors.py
+31 -14
View File
@@ -1,3 +1,9 @@
"""This module contains various projection functions
These functions are useful when defining declarative algorithms to run
alongside the automatic differentiation ones, for comparison or in a
situation where they might be needed.
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools.cmath import *
import torch as t
@@ -8,23 +14,29 @@ __all__ = ['modulus', 'support']
def modulus(wavefront, intensities, mask = None):
"""Implements the modulus constraint in torch
This accepts a torch tensor representing the propagated simulated wavefront(s),
This accepts a tensor representing the propagated simulated wavefront(s),
where the last dimension represents the real and imaginary components of
the propagated wavefield(s). It projects the modulus of the diffraction pattern
onto the modulus of the simulated wavefield.
the propagated wavefield(s). It projects the modulus of the diffraction
pattern onto the modulus of the simulated wavefield.
It assumes that the wavefront is stored in an array
[i,j] where i corresponds to the y-axis and j corresponds to the
x-axis, with the origin following the CS standard of being in the
upper right.
Args:
wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
intensities (torch.Tensor): The measured diffraction pattern(s) stored as an JxNxM stack of real tensors
mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex propagated wavefronts
intensities : torch.Tensor
The measured diffraction pattern(s) stored as an JxNxM stack of real tensors
mask : torch.Tensor
A mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
Returns:
torch.Tensor : The JxNxMx2 propagated wavefield with corrected intensities
Returns
-------
projected : torch.Tensor
The JxNxMx2 projected wavefield with corrected intensities
"""
# Calculate amplitudes from intensities
amplitudes = t.sqrt(intensities)
@@ -53,11 +65,16 @@ def support(wavefront, support):
x-axis, with the origin following the CS standard of being in the
upper right.
Args:
wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
support (torch.Tensor) : An NxM support, with 1s within the support and 0s outside
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex propagated wavefronts
support : torch.Tensor
An NxM support, with 1s within the support and 0s outside
Returns:
torch.Tensor : The JxNxMx2 wavefield with the support mask applied
Returns
-------
projected : torch.Tensor
The JxNxMx2 wavefield with the support mask applied
"""
return wavefront * support.to(wavefront.dtype)[...,None]
CDTools/tools/propagators.py
+58 -26
View File
@@ -1,3 +1,8 @@
"""This module contains various propagators for light fields
All the functions here are designed for use in an automatic differentiation
ptychography model. Each function implements a different propagator.
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools.cmath import *
@@ -27,10 +32,15 @@ def far_field(wavefront):
upper right. The zero frequency component of the propagated wavefield is
shifted to the center of the array.
Args:
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
Returns:
torch.Tensor : The JxNxMx2 propagated wavefield
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex wavefronts to be propagated
Returns
-------
propagated : torch.Tensor
The JxNxMx2 propagated wavefield
"""
return fftshift(t.fft(ifftshift(wavefront), 2, normalized=True))
@@ -49,10 +59,15 @@ def inverse_far_field(wavefront):
upper right. The zero frequency component of the propagated wavefield is
assumed to be the center of the array.
Args:
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts propagated to the far-field
Returns:
torch.Tensor : The JxNxMx2 exit wavefield
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex wavefronts propagated to the far-field
Returns
-------
propagated : torch.Tensor
The JxNxMx2 exit wavefield
"""
return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True))
@@ -71,14 +86,21 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, *
transform of the convolution kernel for light propagation in free
space
Args:
shape (iterable) : The shape of the arrays to be propagated
spacing (iterable) : The pixel size in each dimension of the arrays to be propagated
wavelength (float) : The wavelength of light to simulate propagation of
z (float) : The distance to simulate propagation over
Parameters
----------
shape : array
The shape of the arrays to be propagated
spacing : array
The pixel size in each dimension of the arrays to be propagated
wavelength : float
The wavelength of light to simulate propagation of
z : float
The distance to simulate propagation over
Returns:
torch.Tensor : A phase mask which accounts for the phase change that each plane wave will undergo.
Returns
-------
propagator : torch.Tensor
A phase mask which accounts for the phase change that each plane wave will undergo.
"""
ki = 2 * np.pi * fftpack.fftfreq(shape[0],spacing[0])
@@ -109,19 +131,24 @@ def near_field(wavefront, angular_spectrum_propagator):
phase mask.
Args:
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
angular_spectrum_propagator (torch.Tensor) : The NxM phase mask to be applied during propagation
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex wavefronts to be propagated
angular_spectrum_propagator : torch.Tensor
The NxM phase mask to be applied during propagation
Returns:
torch.Tensor : The propagated wavefront
Returns
-------
propagated : torch.Tensor
The propagated wavefront
"""
return t.ifft(cmult(angular_spectrum_propagator,t.fft(wavefront,2)), 2)
def inverse_near_field(wavefront, angular_spectrum_propagator):
""" Inverse ropagates a wavefront via the angular spectrum method
""" Inverse propagates a wavefront via the angular spectrum method
This function accepts an 3D torch tensor, where the last dimension
represents the real and imaginary components of the wavefield, and
@@ -133,12 +160,17 @@ def inverse_near_field(wavefront, angular_spectrum_propagator):
which corresponds to the inverse propagation problem.
Args:
wavefront (torch.Tensor) : The JxNxMx2 stack of complex wavefronts to be propagated
angular_spectrum_propagator (torch.Tensor) : The NxM phase mask to be applied during propagation
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex wavefronts to be propagated
angular_spectrum_propagator : torch.Tensor
The NxM phase mask to be applied during propagation
Returns:
torch.Tensor : The inverse propagated wavefront
Returns
-------
propagated : torch.Tensor
The inverse propagated wavefront
"""
return t.ifft(cmult(t.fft(wavefront,2), cconj(angular_spectrum_propagator)), 2)
docs/source/_templates/autosummary/module.rst
+5
View File
@@ -0,0 +1,5 @@
{{ fullname }}
{{ underline }}
.. automodule:: {{ fullname }}
:members:
docs/source/conf.py
+3 -3
View File
@@ -69,7 +69,7 @@ language = None
# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path .
exclude_patterns = []
exclude_patterns = ['_build','_templates']
# The name of the Pygments (syntax highlighting) style to use.
pygments_style = 'sphinx'
@@ -80,7 +80,7 @@ pygments_style = 'sphinx'
# The theme to use for HTML and HTML Help pages. See the documentation for
# a list of builtin themes.
#
html_theme = 'alabaster'
html_theme = 'sphinx_rtd_theme'
# Theme options are theme-specific and customize the look and feel of a theme
# further. For a list of options available for each theme, see the
@@ -91,7 +91,7 @@ html_theme = 'alabaster'
# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['_static']
# html_static_path = ['_static']
# Custom sidebar templates, must be a dictionary that maps document names
# to template names.
docs/source/datasets.rst
+7
View File
@@ -0,0 +1,7 @@
Datasets
========
.. automodule:: CDTools.datasets
:members:
docs/source/examples.rst
+3
View File
@@ -0,0 +1,3 @@
Examples
========
docs/source/general.rst
+11
View File
@@ -0,0 +1,11 @@
General Reference
=================
A note about the meaning of "array" as a type
A note about SI units in the package
Notes about special things like the probe-convention translations and the transposed basis arrays
docs/source/index.rst
+37 -2
View File
@@ -2,16 +2,51 @@
Introduction to CDTools
=======================
CDTools is a python library for autodifferentiation-based coherent diffractive imaging reconstructions. The core of the library is a set of simple tools built in pytorch for basic operations relevant to coherent diffraction - math operations on complex numbers, subpixel shifts, propagators, and the like. In addition, a set of database types exist to help with easily loading and saving data to/from .cxi files. Finally, a collection of ptychography models are implemented which allow for a variety of different styles of ptychographic reconstructions.
CDTools is a python library for ptychography and CDI reconstructions, using an Automatic Differentiation based approach.
.. code-block:: python
# imports
from matplotlib import pyplot as plt
from CDTools.datasets import Ptycho_2D_Dataset
from CDTools.models import SimplePtycho
# Load the file
dataset = Ptycho_2D_Dataset.from_cxi('ptycho_data.cxi')
# Generate a model from the data
model = SimplePtycho.from_dataset(dataset)
# Run a reconstruction
for i, loss in enumerate(model.Adam_optimize(10, dataset)):
print(i, loss)
# And look at the results!
model.inspect(dataset)
model.compare(dataset)
plt.show()
CDTools makes it simple to load and inspect from data stored in .cxi files using python scripts. Several reconstruction models for common geometries are included "out of the box". For more advanced users, it is straightforward to program new models for AD ptychography, which can then be used right away from the same scripting framework.
The high-level interface to CDTools is built on a lower level "three-legged stool". This consists of tools to access stored data, tools to visualize data and reconstructions, and tools that implement basic operations relevant to coherent diffraction. All of these tools can be used directly alongside the high-level interface, when needed.
Enough blabber. If you're interested, read the docs!
Documentation Overview
======================
.. toctree::
:maxdepth: 2
:maxdepth: 1
installation
examples
tutorial
general
datasets
models
tools/index
docs/source/installation.rst
+3 -3
View File
@@ -1,5 +1,5 @@
How to Get
==========
Installation
============
CDTools can be downloaded from it's `MIT github page`_, and the relevant prerequisites can be downloaded via pip, conda, or most python package managers.
@@ -20,7 +20,7 @@ CDTools has the following prerequisites:
* `python-dateutil <https://github.com/dateutil/dateutil/>`_
* `h5py <https://www.h5py.org/>`_
All of these can be installed via pip or conda. It is required that pytorch is ilt with MKL, as that enables FFTs. Additionally, CUDA support in pytorch is recommended for running any serious reconstructions with the package. The code is written to be python 2.7+ compatible, although it is only tested in python 3.
All of these can be installed via pip or conda. It is required that pytorch is ilt with MKL, as that enables FFTs. Additionally, installing pytorch with CUDA support is recommended for running any serious reconstructions with the package. The code is written to be python 2.7+ compatible, although it is only tested in python 3.
Finally, to run the tests, pytest is required, and to build the docs, sphinx and sphinx-argparse are required.
docs/source/models.rst
+4
View File
@@ -0,0 +1,4 @@
Models
======
docs/source/tools/analysis.rst
+5
View File
@@ -0,0 +1,5 @@
Analysis
========
.. automodule:: CDTools.tools.analysis
:members:
docs/source/tools/cmath.rst
+5
View File
@@ -0,0 +1,5 @@
Cmath
=====
.. automodule:: CDTools.tools.cmath
:members:
docs/source/tools/data.rst
+5
View File
@@ -0,0 +1,5 @@
Data
====
.. automodule:: CDTools.tools.data
:members:
docs/source/tools/image_processing.rst
+5
View File
@@ -0,0 +1,5 @@
Image Processing
================
.. automodule:: CDTools.tools.image_processing
:members:
docs/source/tools/index.rst
+17
View File
@@ -0,0 +1,17 @@
Tools
=====
.. toctree::
:maxdepth: 1
cmath
image_processing
data
initializers
interactions
propagators
measurements
losses
plotting
analysis
projectors
docs/source/tools/initializers.rst
+5
View File
@@ -0,0 +1,5 @@
Initializers
============
.. automodule:: CDTools.tools.initializers
:members:
docs/source/tools/interactions.rst
+5
View File
@@ -0,0 +1,5 @@
Interactions
============
.. automodule:: CDTools.tools.interactions
:members:
docs/source/tools/losses.rst
+5
View File
@@ -0,0 +1,5 @@
Losses
======
.. automodule:: CDTools.tools.losses
:members:
docs/source/tools/measurements.rst
+5
View File
@@ -0,0 +1,5 @@
Measurements
============
.. automodule:: CDTools.tools.measurements
:members:
docs/source/tools/plotting.rst
+5
View File
@@ -0,0 +1,5 @@
Plotting
========
.. automodule:: CDTools.tools.plotting
:members:
docs/source/tools/projectors.rst
+5
View File
@@ -0,0 +1,5 @@
Projectors
==========
.. automodule:: CDTools.tools.projectors
:members:
docs/source/tools/propagators.rst
+5
View File
@@ -0,0 +1,5 @@
Propagators
===========
.. automodule:: CDTools.tools.propagators
:members:
docs/source/tutorial.rst
+2
View File
@@ -0,0 +1,2 @@
Tutorial
========
examples/MIT_BNL_logo.py
+1 -3
View File
@@ -9,7 +9,6 @@ import pickle
from time import time
import datetime
import h5py
import torch as t
import numpy as np
@@ -18,8 +17,7 @@ import numpy as np
# Please contact Abe Levitan (alevitan@mit) if you would like access
filename = '/media/Data Bank/CSX_6_17/Processed_CXIs/79511_p.cxi'
with h5py.File(filename,'r') as f:
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f)
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
# In this dataset, the edges of the patterns are too noisy and are
# masked off anyway. We can easily just remove this data instead of
examples/gold_ball_ensemble_ptycho.py
+1 -4
View File
@@ -1,16 +1,13 @@
from __future__ import division, print_function, absolute_import
import CDTools
import h5py
import numpy as np
import pickle
from matplotlib import pyplot as plt
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
with h5py.File(filename,'r') as f:
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f)
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
results = []
examples/gold_ball_ptycho.py
+2 -3
View File
@@ -1,15 +1,13 @@
from __future__ import division, print_function, absolute_import
import CDTools
import h5py
import numpy as np
import pickle
from matplotlib import pyplot as plt
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
with h5py.File(filename,'r') as f:
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f)
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2)
@@ -30,6 +28,7 @@ for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)):
with open('example_reconstructions/gold_balls.pickle', 'wb') as f:
pickle.dump(model.save_results(dataset),f)
model.inspect(dataset)
dataset.inspect()
examples/simple_ptycho.py
+1 -7
View File
@@ -3,18 +3,12 @@ from __future__ import division, print_function, absolute_import
import CDTools
from CDTools.tools import cmath
from CDTools.tools.plotting import *
import h5py
import torch as t
import numpy as np
from matplotlib import pyplot as plt
import pickle
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
with h5py.File(filename,'r') as f:
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f)
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
model = CDTools.models.SimplePtycho.from_dataset(dataset)
examples/specular_pinhole_ptycho.py
+1 -3
View File
@@ -18,9 +18,7 @@ import numpy as np
# Please contact Abe Levitan (alevitan@mit) if you would like access
filename = '/media/Data Bank/CSX_10_18/Processed_CXIs/110531_p.cxi'
with h5py.File(filename,'r') as f:
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f)
dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(filename)
model = CDTools.models.FancyPtycho.from_dataset(dataset,
randomize_ang = np.pi/4,
tests/test_datasets.py
+3
View File
@@ -197,7 +197,9 @@ def test_Ptycho_2D_Dataset_from_cxi(test_ptycho_cxis):
def test_Ptycho_2D_Dataset_to_cxi(test_ptycho_cxis, tmp_path):
for cxi, expected in test_ptycho_cxis:
print('loading dataset')
dataset = Ptycho_2D_Dataset.from_cxi(cxi)
print('dataset mask is type', dataset.mask.dtype)
with cdtdata.create_cxi(tmp_path / 'test_Ptycho_2D_Dataset_to_cxi.cxi') as f:
dataset.to_cxi(f)
@@ -223,6 +225,7 @@ def test_Ptycho_2D_Dataset_to_cxi(test_ptycho_cxis, tmp_path):
if dataset.detector_geometry['corner'] is not None:
assert 'corner' in read_dataset.detector_geometry
if dataset.mask is not None:
assert t.all(t.eq(dataset.mask,read_dataset.mask))
tests/tools/test_analysis.py
+13
View File
@@ -173,6 +173,19 @@ def test_calc_consistency_prtf():
freqs, prtf = analysis.calc_consistency_prtf(synth_obj, obj_stack, basis, nbins=30)
assert np.allclose(prtf, 0.7)
# Check that it also works with torch input
t_synth_obj = cmath.complex_to_torch(synth_obj)
t_obj_stack = [cmath.complex_to_torch(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, basis, nbins=30)
assert np.allclose(prtf.numpy(), 0.7)
# And also when the basis is in torch
t_synth_obj = cmath.complex_to_torch(synth_obj)
t_obj_stack = [cmath.complex_to_torch(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30)
assert np.allclose(prtf.numpy(), 0.7)
# Check that is uses the right number of bins
assert len(prtf) == 30
assert len(freqs) == 30
tests/tools/test_plotting.py
-7
View File
@@ -9,13 +9,6 @@ import torch as t
import scipy.misc
import matplotlib.pyplot as plt
def test_plot_1D(show_plot):
# Plot simple linear scatter plot
arr = np.arange(10)
plotting.plot_1D(arr, title = 'Linear Plot')
if show_plot:
plt.show()
def test_plot_amplitude(show_plot):
# Test with tensor
im = cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64))