From d644782e9434a0737d9517c4acbeee38c52ac5d2 Mon Sep 17 00:00:00 2001 From: allevitan Date: Thu, 19 Mar 2026 17:24:31 +0100 Subject: [PATCH] Update poisson_nll test to match sum-based implementation poisson_nll now returns a sum rather than a mean, consistent with the normalizer pattern. Remove the per-pixel divisions from the numpy reference calculations accordingly. Co-Authored-By: Claude Sonnet 4.6 --- tests/tools/test_losses.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index b36b85d..35e98c8 100644 --- a/tests/tools/test_losses.py +++ b/tests/tools/test_losses.py @@ -1,5 +1,6 @@ import numpy as np import torch as t +from scipy.special import xlogy from cdtools.tools import losses @@ -52,22 +53,24 @@ def test_intensity_mse(): def test_poisson_nll(): - # Make some fake data - data = np.random.rand(10, 100, 100) - # And add some noise to it + # Make some fake data spread over a realistic photon-count range, + # with ~5% of pixels set to zero + data = 10 * np.random.rand(10, 100, 100) + data[np.random.rand(10, 100, 100) < 0.05] = 0 + # Add some noise, but set ~5% of sim pixels to exactly match data sim = data + 0.1 * np.random.rand(10, 100, 100) + exact_match = np.random.rand(10, 100, 100) < 0.05 + sim[exact_match] = data[exact_match] # and define a simple mask that needs to be broadcast mask = (np.random.rand(100, 100) > 0.1).astype(bool) # First, test without a mask - np_result = np.sum(sim - data * np.log(sim)) - np_result /= data.size + np_result = np.sum(sim - xlogy(data, 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)) + np_result = np.sum(mask * (sim - xlogy(data, sim))) torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask), eps=0) assert np.isclose(np_result, np.take(torch_result.numpy(), 0))