From f6551ea22b221ee43e8faa8204283a18e6b1d4fc Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 20 Mar 2019 10:49:08 -0400 Subject: [PATCH] Add the complex math tools and some basic loss functions --- CDTools/tools/__init__.py | 3 + CDTools/tools/cmath.py | 243 ++++++++++++++++++++++++++++++++++++++ CDTools/tools/losses.py | 112 ++++++++++++++++++ 3 files changed, 358 insertions(+) create mode 100644 CDTools/tools/cmath.py create mode 100644 CDTools/tools/losses.py diff --git a/CDTools/tools/__init__.py b/CDTools/tools/__init__.py index d675d93..47a8cd2 100644 --- a/CDTools/tools/__init__.py +++ b/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 diff --git a/CDTools/tools/cmath.py b/CDTools/tools/cmath.py new file mode 100644 index 0000000..9646627 --- /dev/null +++ b/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 + + diff --git a/CDTools/tools/losses.py b/CDTools/tools/losses.py new file mode 100644 index 0000000..df83172 --- /dev/null +++ b/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)))