mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-11 14:02:38 +02:00
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 <noreply@anthropic.com>
77 lines
3.1 KiB
Python
77 lines
3.1 KiB
Python
import numpy as np
|
|
import torch as t
|
|
from scipy.special import xlogy
|
|
|
|
from cdtools.tools import losses
|
|
|
|
|
|
# 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():
|
|
# 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(bool)
|
|
|
|
# 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():
|
|
# 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(bool)
|
|
|
|
# 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_nll():
|
|
# 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 - 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 - 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))
|