mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
93 lines
3.9 KiB
Python
93 lines
3.9 KiB
Python
import torch as t
|
|
import numpy as np
|
|
|
|
from cdtools.tools import measurements
|
|
|
|
|
|
def test_intensity():
|
|
wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10))
|
|
epsilon = 1e-6
|
|
np_result = np.abs(wavefields.numpy())**2 + epsilon
|
|
assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon),
|
|
t.as_tensor(np_result))
|
|
|
|
# Test single field case
|
|
assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon),
|
|
t.as_tensor(np_result[0]))
|
|
|
|
det_slice = np.s_[3:, 5:8]
|
|
assert t.allclose(measurements.intensity(wavefields, det_slice, epsilon=epsilon),
|
|
t.as_tensor(np_result[(np.s_[:],) + det_slice]))
|
|
|
|
# Test single field case
|
|
assert t.allclose(measurements.intensity(wavefields[0], det_slice, epsilon=epsilon),
|
|
t.as_tensor(np_result[0][det_slice]))
|
|
|
|
# With oversampling on
|
|
np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4
|
|
|
|
# With multiple fields
|
|
assert t.allclose(measurements.intensity(wavefields, epsilon=epsilon, oversampling=2),
|
|
t.as_tensor(np_oversampling_result,))
|
|
|
|
# With a single field
|
|
assert t.allclose(measurements.intensity(wavefields[0], epsilon=epsilon, oversampling=2),
|
|
t.as_tensor(np_oversampling_result[0],))
|
|
|
|
|
|
def test_incoherent_sum():
|
|
|
|
# With no explicit slice given
|
|
|
|
wavefields = t.rand((5, 4, 10, 10)) + 1j * t.rand((5, 4, 10, 10))
|
|
epsilon = 1e-6
|
|
np_result = np.sum(np.abs(wavefields.numpy())**2, axis=-3) + epsilon
|
|
assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon),
|
|
t.as_tensor(np_result))
|
|
# Test single field case
|
|
assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon),
|
|
t.as_tensor(np_result[0]))
|
|
|
|
# With a slice given
|
|
det_slice = np.s_[3:, 5:8]
|
|
assert t.allclose(measurements.incoherent_sum(wavefields, det_slice, epsilon=epsilon),
|
|
t.as_tensor(np_result[(np.s_[:],) + det_slice]))
|
|
# Test single field case
|
|
assert t.allclose(measurements.incoherent_sum(wavefields[0, :], det_slice, epsilon=epsilon),
|
|
t.as_tensor(np_result[0][det_slice]))
|
|
|
|
# With oversampling on
|
|
np_oversampling_result = (np_result[:, ::2, ::2] + np_result[:, 1::2, ::2] + np_result[:, ::2, 1::2] + np_result[:, 1::2, 1::2]) / 4
|
|
|
|
# With multiple fields
|
|
assert t.allclose(measurements.incoherent_sum(wavefields, epsilon=epsilon, oversampling=2),
|
|
t.as_tensor(np_oversampling_result,))
|
|
|
|
# With a single field
|
|
assert t.allclose(measurements.incoherent_sum(wavefields[0, :], epsilon=epsilon, oversampling=2),
|
|
t.as_tensor(np_oversampling_result[0],))
|
|
|
|
|
|
def test_quadratic_background():
|
|
# test with intensity
|
|
wavefields = t.rand((5, 10, 10)) + 1j * t.rand((5, 10, 10))
|
|
epsilon = 1e-6
|
|
background = t.rand((10, 10))
|
|
np_result = np.abs(wavefields.numpy())**2 + background.numpy()**2 + epsilon
|
|
det_slice = np.s_[3:, 5:8]
|
|
|
|
result = measurements.quadratic_background(wavefields, background[det_slice],
|
|
detector_slice=det_slice,
|
|
epsilon=epsilon,
|
|
measurement=measurements.intensity)
|
|
assert t.allclose(result, t.tensor(np_result[(np.s_[:],) + det_slice]))
|
|
|
|
# test with incoherent sum but no slice and no stack
|
|
wavefields = t.rand((4, 10, 10)) + 1j * t.rand((4, 10, 10))
|
|
np_result = np.sum(np.abs(wavefields.numpy())**2, axis=0)
|
|
np_result += background.numpy()**2
|
|
result = measurements.quadratic_background(wavefields, background,
|
|
epsilon=epsilon,
|
|
measurement=measurements.incoherent_sum)
|
|
assert t.allclose(result, t.tensor(np_result))
|