diff --git a/CDTools/datasets.py b/CDTools/datasets.py
index c18d801..5dd48f4 100644
--- a/CDTools/datasets.py
+++ b/CDTools/datasets.py
@@ -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)')
diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py
index 9b5bf61..252d005 100644
--- a/CDTools/models/__init__.py
+++ b/CDTools/models/__init__.py
@@ -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
diff --git a/CDTools/tools/analysis.py b/CDTools/tools/analysis.py
index e7e7dc4..da24b7a 100644
--- a/CDTools/tools/analysis.py
+++ b/CDTools/tools/analysis.py
@@ -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)
diff --git a/CDTools/tools/cmath.py b/CDTools/tools/cmath.py
index 7a89d4e..cf52757 100644
--- a/CDTools/tools/cmath.py
+++ b/CDTools/tools/cmath.py
@@ -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)
diff --git a/CDTools/tools/data.py b/CDTools/tools/data.py
index 5d70b43..8ee8849 100644
--- a/CDTools/tools/data.py
+++ b/CDTools/tools/data.py
@@ -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:
diff --git a/CDTools/tools/image_processing.py b/CDTools/tools/image_processing.py
index 1dc6025..4ea5d3c 100644
--- a/CDTools/tools/image_processing.py
+++ b/CDTools/tools/image_processing.py
@@ -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
diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py
index dfd9e1e..52326e8 100644
--- a/CDTools/tools/initializers.py
+++ b/CDTools/tools/initializers.py
@@ -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
"""
diff --git a/CDTools/tools/interactions.py b/CDTools/tools/interactions.py
index 0cb4b11..e7365ea 100644
--- a/CDTools/tools/interactions.py
+++ b/CDTools/tools/interactions.py
@@ -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:
diff --git a/CDTools/tools/losses.py b/CDTools/tools/losses.py
index 8255ac8..2f92a45 100644
--- a/CDTools/tools/losses.py
+++ b/CDTools/tools/losses.py
@@ -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:
diff --git a/CDTools/tools/measurements.py b/CDTools/tools/measurements.py
index 524d5aa..da502f5 100644
--- a/CDTools/tools/measurements.py
+++ b/CDTools/tools/measurements.py
@@ -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:
diff --git a/CDTools/tools/plotting.py b/CDTools/tools/plotting.py
index 29fd30d..74a96bd 100644
--- a/CDTools/tools/plotting.py
+++ b/CDTools/tools/plotting.py
@@ -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
diff --git a/CDTools/tools/projectors.py b/CDTools/tools/projectors.py
index b4b2a04..eef5461 100644
--- a/CDTools/tools/projectors.py
+++ b/CDTools/tools/projectors.py
@@ -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]
diff --git a/CDTools/tools/propagators.py b/CDTools/tools/propagators.py
index 23c919a..d66ea4d 100644
--- a/CDTools/tools/propagators.py
+++ b/CDTools/tools/propagators.py
@@ -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)
diff --git a/docs/source/_templates/autosummary/module.rst b/docs/source/_templates/autosummary/module.rst
new file mode 100644
index 0000000..6090b5e
--- /dev/null
+++ b/docs/source/_templates/autosummary/module.rst
@@ -0,0 +1,5 @@
+{{ fullname }}
+{{ underline }}
+
+.. automodule:: {{ fullname }}
+ :members:
diff --git a/docs/source/conf.py b/docs/source/conf.py
index d22dfe8..ac2b4af 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -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.
diff --git a/docs/source/datasets.rst b/docs/source/datasets.rst
new file mode 100644
index 0000000..f8cda23
--- /dev/null
+++ b/docs/source/datasets.rst
@@ -0,0 +1,7 @@
+Datasets
+========
+
+.. automodule:: CDTools.datasets
+ :members:
+
+
diff --git a/docs/source/examples.rst b/docs/source/examples.rst
new file mode 100644
index 0000000..ce68152
--- /dev/null
+++ b/docs/source/examples.rst
@@ -0,0 +1,3 @@
+Examples
+========
+
diff --git a/docs/source/general.rst b/docs/source/general.rst
new file mode 100644
index 0000000..3c0b72d
--- /dev/null
+++ b/docs/source/general.rst
@@ -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
+
+
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 60636c3..42c19cc 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -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
diff --git a/docs/source/installation.rst b/docs/source/installation.rst
index 6bd191d..f7be677 100644
--- a/docs/source/installation.rst
+++ b/docs/source/installation.rst
@@ -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 `_
* `h5py `_
-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.
diff --git a/docs/source/models.rst b/docs/source/models.rst
new file mode 100644
index 0000000..3ed0c8a
--- /dev/null
+++ b/docs/source/models.rst
@@ -0,0 +1,4 @@
+Models
+======
+
+
diff --git a/docs/source/tools/analysis.rst b/docs/source/tools/analysis.rst
new file mode 100644
index 0000000..4afd115
--- /dev/null
+++ b/docs/source/tools/analysis.rst
@@ -0,0 +1,5 @@
+Analysis
+========
+
+.. automodule:: CDTools.tools.analysis
+ :members:
diff --git a/docs/source/tools/cmath.rst b/docs/source/tools/cmath.rst
new file mode 100644
index 0000000..763bb0a
--- /dev/null
+++ b/docs/source/tools/cmath.rst
@@ -0,0 +1,5 @@
+Cmath
+=====
+
+.. automodule:: CDTools.tools.cmath
+ :members:
diff --git a/docs/source/tools/data.rst b/docs/source/tools/data.rst
new file mode 100644
index 0000000..91bee3f
--- /dev/null
+++ b/docs/source/tools/data.rst
@@ -0,0 +1,5 @@
+Data
+====
+
+.. automodule:: CDTools.tools.data
+ :members:
diff --git a/docs/source/tools/image_processing.rst b/docs/source/tools/image_processing.rst
new file mode 100644
index 0000000..e4dec5b
--- /dev/null
+++ b/docs/source/tools/image_processing.rst
@@ -0,0 +1,5 @@
+Image Processing
+================
+
+.. automodule:: CDTools.tools.image_processing
+ :members:
diff --git a/docs/source/tools/index.rst b/docs/source/tools/index.rst
new file mode 100644
index 0000000..5d8bea6
--- /dev/null
+++ b/docs/source/tools/index.rst
@@ -0,0 +1,17 @@
+Tools
+=====
+
+.. toctree::
+ :maxdepth: 1
+
+ cmath
+ image_processing
+ data
+ initializers
+ interactions
+ propagators
+ measurements
+ losses
+ plotting
+ analysis
+ projectors
diff --git a/docs/source/tools/initializers.rst b/docs/source/tools/initializers.rst
new file mode 100644
index 0000000..7be9fef
--- /dev/null
+++ b/docs/source/tools/initializers.rst
@@ -0,0 +1,5 @@
+Initializers
+============
+
+.. automodule:: CDTools.tools.initializers
+ :members:
diff --git a/docs/source/tools/interactions.rst b/docs/source/tools/interactions.rst
new file mode 100644
index 0000000..8e0c235
--- /dev/null
+++ b/docs/source/tools/interactions.rst
@@ -0,0 +1,5 @@
+Interactions
+============
+
+.. automodule:: CDTools.tools.interactions
+ :members:
diff --git a/docs/source/tools/losses.rst b/docs/source/tools/losses.rst
new file mode 100644
index 0000000..0dd25ec
--- /dev/null
+++ b/docs/source/tools/losses.rst
@@ -0,0 +1,5 @@
+Losses
+======
+
+.. automodule:: CDTools.tools.losses
+ :members:
diff --git a/docs/source/tools/measurements.rst b/docs/source/tools/measurements.rst
new file mode 100644
index 0000000..301236f
--- /dev/null
+++ b/docs/source/tools/measurements.rst
@@ -0,0 +1,5 @@
+Measurements
+============
+
+.. automodule:: CDTools.tools.measurements
+ :members:
diff --git a/docs/source/tools/plotting.rst b/docs/source/tools/plotting.rst
new file mode 100644
index 0000000..53512bd
--- /dev/null
+++ b/docs/source/tools/plotting.rst
@@ -0,0 +1,5 @@
+Plotting
+========
+
+.. automodule:: CDTools.tools.plotting
+ :members:
diff --git a/docs/source/tools/projectors.rst b/docs/source/tools/projectors.rst
new file mode 100644
index 0000000..60b042f
--- /dev/null
+++ b/docs/source/tools/projectors.rst
@@ -0,0 +1,5 @@
+Projectors
+==========
+
+.. automodule:: CDTools.tools.projectors
+ :members:
diff --git a/docs/source/tools/propagators.rst b/docs/source/tools/propagators.rst
new file mode 100644
index 0000000..12976bf
--- /dev/null
+++ b/docs/source/tools/propagators.rst
@@ -0,0 +1,5 @@
+Propagators
+===========
+
+.. automodule:: CDTools.tools.propagators
+ :members:
diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst
new file mode 100644
index 0000000..b23b9e5
--- /dev/null
+++ b/docs/source/tutorial.rst
@@ -0,0 +1,2 @@
+Tutorial
+========
diff --git a/examples/MIT_BNL_logo.py b/examples/MIT_BNL_logo.py
index 581b904..4a26abc 100644
--- a/examples/MIT_BNL_logo.py
+++ b/examples/MIT_BNL_logo.py
@@ -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
diff --git a/examples/gold_ball_ensemble_ptycho.py b/examples/gold_ball_ensemble_ptycho.py
index 6f0d803..8dd4a65 100644
--- a/examples/gold_ball_ensemble_ptycho.py
+++ b/examples/gold_ball_ensemble_ptycho.py
@@ -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 = []
diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py
index bdf70f3..eb6c551 100644
--- a/examples/gold_ball_ptycho.py
+++ b/examples/gold_ball_ptycho.py
@@ -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()
diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py
index 664b45b..77cf05a 100644
--- a/examples/simple_ptycho.py
+++ b/examples/simple_ptycho.py
@@ -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)
diff --git a/examples/specular_pinhole_ptycho.py b/examples/specular_pinhole_ptycho.py
index 89708ff..96fcccb 100644
--- a/examples/specular_pinhole_ptycho.py
+++ b/examples/specular_pinhole_ptycho.py
@@ -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,
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 91d38c1..5685546 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -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))
diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py
index c0b5ada..c0feade 100644
--- a/tests/tools/test_analysis.py
+++ b/tests/tools/test_analysis.py
@@ -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
diff --git a/tests/tools/test_plotting.py b/tests/tools/test_plotting.py
index 5fc04ef..a6dbe47 100644
--- a/tests/tools/test_plotting.py
+++ b/tests/tools/test_plotting.py
@@ -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))