Finish the tests for the losses

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]

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))
+
+
+