diff --git a/CDTools/tools/losses.py b/CDTools/tools/losses.py index 43720d4..8255ac8 100644 --- a/CDTools/tools/losses.py +++ b/CDTools/tools/losses.py @@ -9,7 +9,7 @@ from __future__ import division, print_function, absolute_import import torch as t -__all__ = ['amplitude_mse', 'intensity_mse', 'poisson_ml'] +__all__ = ['amplitude_mse', 'intensity_mse', 'poisson_nll'] def amplitude_mse(intensities, sim_intensities, mask=None): @@ -85,8 +85,8 @@ def intensity_mse(intensities, sim_intensities, mask=None): -def poisson_ml(intensities, sim_intensities, mask=None): - """ Returns the Poisson maximum likelihood metric for a simulated dataset's intensities +def poisson_nll(intensities, sim_intensities, mask=None): + """ Returns the Poisson negative log likelihood for a simulated dataset's intensities Calculates the overall Poisson maximum likelihood metric using diffraction intensities - the measured set of detector intensities - @@ -94,6 +94,10 @@ def poisson_ml(intensities, sim_intensities, mask=None): for detectors in a single-photon counting mode, with their output scaled to number of photons + Note that this calculation ignores the log(intensities!) term in the + full expression for Poisson negative log likelihood. This term doesn't + change the calculated gradients so isn't worth taking the time to compute + 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. @@ -108,11 +112,12 @@ def poisson_ml(intensities, sim_intensities, mask=None): """ if mask is None: - t.sum(simulated_intensities - - intensities * t.log(simulated_intensities)) \ - / intensities.view(-1).shape[0] + return t.sum(sim_intensities - + intensities * t.log(sim_intensities)) \ + / intensities.view(-1).shape[0] + else: masked_intensities = intensities.masked_select(mask) - masked_sims = simulated_intensities.masked_select(mask) + masked_sims = sim_intensities.masked_select(mask) return t.sum(masked_sims - masked_intensities * t.log(masked_sims)) / masked_intensities.shape[0] diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index 11c4108..d6f9b71 100644 --- a/tests/tools/test_losses.py +++ b/tests/tools/test_losses.py @@ -6,11 +6,76 @@ import numpy as np import torch as t +# The idea here is to use a simple numpy calculation of the various +# objective functions to check the torch implementations and make sure +# that any optimizations in the future don't change the results + def test_amplitude_mse(): - pass + # Make some fake data + data = np.random.rand(10,100,100) + # And add some noise to it + sim = data + 0.1 * np.random.rand(10,100,100) + # and define a simple mask that needs to be broadcast + mask = (np.random.rand(100,100) > 0.1).astype(np.uint8) + + # First, test without a mask + np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2) + np_result /= data.size + torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim)) + assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + + # Then, test with a mask + np_result = np.sum(mask * (np.sqrt(data) - np.sqrt(sim))**2) + np_result /= np.count_nonzero(mask * np.ones_like(data)) + torch_result = losses.amplitude_mse(t.from_numpy(data),t.from_numpy(sim), + mask = t.from_numpy(mask)) + assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + def test_intensity_mse(): - pass + # Make some fake data + data = np.random.rand(10,100,100) + # And add some noise to it + sim = data + 0.1 * np.random.rand(10,100,100) + # and define a simple mask that needs to be broadcast + mask = (np.random.rand(100,100) > 0.1).astype(np.uint8) + + + # First, test without a mask + np_result = np.sum((data - sim)**2) + np_result /= data.size + torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim)) + assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + + # Then, test with a mask + np_result = np.sum(mask * (data - sim)**2) + np_result /= np.count_nonzero(mask * np.ones_like(data)) + torch_result = losses.intensity_mse(t.from_numpy(data),t.from_numpy(sim), + mask = t.from_numpy(mask)) + assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + def test_poisson_ml(): - pass + # Make some fake data + data = np.random.rand(10,100,100) + # And add some noise to it + sim = data + 0.1 * np.random.rand(10,100,100) + # and define a simple mask that needs to be broadcast + mask = (np.random.rand(100,100) > 0.1).astype(np.uint8) + + + # First, test without a mask + np_result = np.sum(sim - data * np.log(sim)) + np_result /= data.size + torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim)) + assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + + # Then, test with a mask + np_result = np.sum(mask * (sim - data * np.log(sim))) + np_result /= np.count_nonzero(mask * np.ones_like(data)) + torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim), + mask = t.from_numpy(mask)) + assert np.isclose(np_result, np.take(torch_result.numpy(),0)) + + +