Add the complex math tools and some basic loss functions

CDTools/tools/__init__.py

@@ -1 +1,4 @@
 from __future__ import division, print_function, absolute_import
+
+from CDTools.tools import cmath
+from CDTools.tools import losses

CDTools/tools/cmath.py

@@ -0,0 +1,243 @@
+"""Contains basic functions for dealing with complex numbers in pytorch.
+
+Since pytorch doesn't have built-in support for complex numbers, but the
+fast fourier transforms in pytorch assume a specific format for complex
+arrays, this module uses that format to store complex numbers. It exposes
+functions for converting between complex numpy arrays and torch tensors
+stored in that format, as well as basic complex math operations implemented
+on the torch tensors
+"""
+from __future__ import division, print_function, absolute_import
+import numpy as np
+import torch as t
+
+
+__all__ = ['complex_to_torch','torch_to_complex','cabssq','cabs','cconj',
+           'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift']
+
+
+#
+# These define the conversions to and from this format
+#
+
+def complex_to_torch(x):
+    """Maps a complex numpy array to a torch tensor
+
+    Pytorch uses tensors with a final dimension of 2 to represent
+    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
+
+    Returns:
+        torch.Tensor : A torch tensor representation of that array
+
+    """
+    return t.from_numpy(np.stack((np.real(x),np.imag(x)),axis=-1))
+
+
+def torch_to_complex(x):
+    """Maps a torch tensor to the a complex numpy array
+
+    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
+
+    Args:
+        x (torch.Tensor): A tensor to convert
+
+    Returns:
+        np.array : A complex typed numpy array corresponding to the input
+
+    """
+    x = np.array(x)
+    x = x[...,0] + x[...,1] * 1j
+    return x
+
+
+#
+# And these define the basic operations on these arrays. Note that
+# multiplication between a complex valued and real valued pytorch
+# tensor will proceed as expected because of torch's broadcasting
+# and thus doesn't need it's own function
+#
+
+def cabssq(a):
+    """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
+
+    Returns:
+        array_like : A tensor storing the elementwise absolute value squared
+
+    """
+    return a[...,0]**2 + a[...,1]**2
+
+
+def cabs(a):
+    """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
+
+    Returns:
+        array_like : A tensor storing the elementwise absolute value
+
+    """
+    return t.sqrt(cabssq(a))
+
+
+def cphase(a):
+    """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 phase
+    of any torch tensor following that standard.
+
+    Args:
+        x (torch.Tensor): An input tensor
+
+    Returns:
+        array_like : A tensor storing the elementwise phase
+
+    """
+    return t.atan2(a[...,1],a[...,0])
+    
+    
+def cconj(a):
+    """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
+
+    Returns:
+        array_like : A tensor storing the elementwise complex conjugate
+
+    """
+    return t.stack((a[...,0],-a[...,1]),dim=-1)
+
+
+
+def cmult(a,b):
+    """Returns the complex product of two torch tensors
+
+    Pytorch uses tensors with a final dimension of 2 to represent
+    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
+
+    Returns:
+        torch.Tensor : A tensor storing the elementwise product
+
+    """
+    
+    real = a[...,0] * b[...,0] - a[...,1] * b[...,1]
+    imag = a[...,0] * b[...,1] + a[...,1] * b[...,0]
+    return t.stack((real,imag),dim=-1)
+
+
+def cdiv(a,b):
+    """Returns the complex quotient of two torch tensors
+
+    Pytorch uses tensors with a final dimension of 2 to represent
+    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
+
+    Returns:
+        torch.Tensor : A tensor storing the elementwise complex quotient
+
+    """
+    return cmult(a, cconj(b)) / cabssq(b)
+
+
+    
+#
+# Not entirely sure if these belong here, but heck with it.
+# We just need the ability to do fftshifts
+#
+
+
+def fftshift(array,dims=None):
+    """Drop-in torch replacement for scipy.fftpack.fftshift
+    
+    This maps a tensor, assumed to be the output of a fast Fourier
+    transform, into a tensor whose zero-frequency element is at the
+    center of the tensor instead of the start. It will by default shift
+    every dimension in the tensor but the last (which is assumed to
+    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
+
+    Returns:
+        torch.Tensor : fftshifted tensor
+    
+    """
+    
+    if dims is None:
+        dims=list(range(array.dim()))[:-1]
+    for dim in dims:
+        length = array.size()[dim]
+        cut_to = (length + 1) // 2
+        cut_len = length - cut_to
+        array = t.cat((array.narrow(dim,cut_to,cut_len),
+                       array.narrow(dim,0,cut_to)), dim)
+    return array
+
+
+
+def ifftshift(array,dims=None):
+    """Drop-in torch replacement for scipy.fftpack.iftshift
+    
+    This maps a tensor, assumed to be the shifted output of a fast
+    Fourier transform, into a tensor whose zero-frequency element is
+    back at the start of the tensor instead of the center. It is the
+    inverse of the fftshift operator. It will by default shift
+    every dimension in the tensor but the last (which is assumed to
+    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
+
+    Returns:
+        torch.Tensor : ifftshifted tensor
+    
+    """
+    
+    if dims is None:
+        dims=list(range(array.dim()))[:-1]
+    for dim in dims:
+        length = array.size()[dim]
+        cut_to = length // 2
+        cut_len = length - cut_to
+        
+        array = t.cat((array.narrow(dim,cut_to,cut_len),
+                       array.narrow(dim,0,cut_to)), dim)
+    return array
+
+

CDTools/tools/losses.py

@@ -0,0 +1,112 @@
+"""Contains various loss functions to be used for optimization
+
+It exposes three losses, one returning the mean squared amplitude error, one
+that returns the mean squared intensity error, and one that returns the
+maximum likelihood metric for a system with Poisson statistics.
+
+"""
+from __future__ import division, print_function, absolute_import
+import torch as t
+
+
+__all__ = ['amplitude_mse', 'intensity_mse', 'poisson_ml']
+
+
+def amplitude_mse(intensities, sim_intensities, mask=None):
+    """ Returns the mean squared error of a simulated dataset's amplitudes
+
+    Calculates the summed mean squared error between a given set of 
+    measured diffraction intensities and a simulated set.
+
+    This function calculates the mean squared error between their
+    associated amplitudes. Because this is not well defined for negative
+    numbers, make sure that all the intensities are >0 before using this
+    loss.
+
+    It can accept intensity and simulated intensity tensors of any shape
+    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
+
+    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
+
+    """
+
+    # I know it would be more efficient if this function took in the
+    # amplitudes instead of the intensities, but I want to be consistent
+    # with all the errors working off of the same inputs
+
+    if mask is None:
+        return t.sum((t.sqrt(sim_intensities) -
+                      t.sqrt(intensities))**2)
+    else:
+        return t.sum((t.sqrt(sim_intensities.masked_select(mask)) -
+                      t.sqrt(intensities.masked_select(mask)))**2)
+
+
+    
+def intensity_mse(intensities, sim_intensities, mask=None):
+    """ Returns the mean squared error of a simulated dataset's intensities
+
+    Calculates the summed mean squared error between a given set of 
+    diffraction intensities - the measured set of detector intensities -
+    and a simulated set of diffraction intensities. This function
+    calculates the mean squared error between the intensities.
+
+    It can accept intensity and simulated intensity tensors of any shape
+    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
+
+    Returns:
+        loss (torch.Tensor) : A single value for the summed mse
+
+    """
+    if mask is None:
+        return t.sum((sim_intensities - intensities)**2)
+    else:
+        return t.sum((sim_intensities.masked_select(mask) -
+                      intensities.masked_select(mask))**2)
+
+
+    
+def poisson_ml(intensities, sim_intensities, mask=None):
+    """ Returns the Poisson maximum likelihood metric for a simulated dataset's intensities
+
+    Calculates the overall Poisson maximum likelihood metric using
+    diffraction intensities - the measured set of detector intensities -
+    and a simulated set of intensities. This loss would be appropriate
+    for detectors in a single-photon counting mode, with their output
+    scaled to number of photons
+
+    It can accept intensity and simulated intensity tensors of any shape
+    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
+
+    Returns:
+        loss (torch.Tensor) : A single value for the poisson ML metric
+
+    """
+    if mask is None:
+        t.sum(simulated_intensities -
+              intensities * t.log(simulated_intensities))
+    else:
+        return t.sum(simulated_intensities.masked_select(mask) -
+                     intensities.masked_select(mask) *
+                     t.log(simulated_intensities.masked_select(mask)))