Fix an issue with the test for the poisson NLL function

CDTools/datasets/__init__.py

@@ -27,6 +27,8 @@ dataset before attempting to do so
 
 from __future__ import division, print_function, absolute_import
 
+# I don't believe that __all__ really needed, but it's nice to define it
+# to be explicit that import * is safe
 __all__ = ['CDataset','Ptycho2DDataset']
 
 from CDTools.datasets.base import CDataset

CDTools/datasets/base.py

@@ -11,11 +11,6 @@ of the following functions:
 * to_cxi
 * inspect
 
-Example implementations of all these functions can be found in the code
-for the Ptycho2DDataset class.
-
-In addition, it is recommended to read through the tutorial section on
-defining a new ptychography dataset before attempting to do so
 """
 
 from __future__ import division, print_function, absolute_import

CDTools/tools/losses/losses.py

@@ -100,7 +100,7 @@ def intensity_mse(intensities, sim_intensities, mask=None):
 
 
     
-def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-4):
+def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-6):
     """ Returns the Poisson negative log likelihood for a simulated dataset's intensities
 
     Calculates the overall Poisson maximum likelihood metric using
@@ -117,6 +117,9 @@ def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-4):
     as long as their shapes match, and the provided mask array can be
     broadcast correctly along them.
 
+    The default value of eps is 1e-6 - a nonzero value here helps avoid
+    divergence of the log function near zero.
+
     Parameters
     ----------
     intensities : torch.Tensor
@@ -125,7 +128,9 @@ def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-4):
         A tensor of simulated detector intensities
     mask : torch.Tensor
         A mask with ones for pixels to include and zeros for pixels to exclude
-
+    eps : float
+        Optional, a small number to add to the simulated intensities
+    
     Returns
     -------
     loss : torch.Tensor

tests/tools/test_losses.py

@@ -55,7 +55,7 @@ def test_intensity_mse():
     assert np.isclose(np_result, np.take(torch_result.numpy(),0))
     
 
-def test_poisson_ml():
+def test_poisson_nll():
     # Make some fake data
     data = np.random.rand(10,100,100)
     # And add some noise to it
@@ -67,14 +67,14 @@ def test_poisson_ml():
     # 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))
+    torch_result = losses.poisson_nll(t.from_numpy(data),t.from_numpy(sim), eps=0)
     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))
+                                      mask = t.from_numpy(mask), eps=0)
     assert np.isclose(np_result, np.take(torch_result.numpy(),0))