74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from aare.daq.beamcenterfit import Gaussian2Dfit, beamcenter_fit
|
|
|
|
|
|
def create_synthetic_beam_image(
|
|
shape=(200, 200), center=(100, 100), sigma=(10, 10), theta=0, A=200, offset=20
|
|
):
|
|
x = np.arange(shape[1])
|
|
y = np.arange(shape[0])
|
|
X, Y = np.meshgrid(x, y)
|
|
|
|
x0, y0 = center
|
|
sig_x, sig_y = sigma
|
|
|
|
x_rot = (X - x0) * np.cos(theta) + (Y - y0) * np.sin(theta)
|
|
y_rot = -(X - x0) * np.sin(theta) + (Y - y0) * np.cos(theta)
|
|
|
|
gaussian = A * np.exp(-(x_rot**2 / (2 * sig_x**2) + y_rot**2 / (2 * sig_y**2))) + offset
|
|
# Add some noise
|
|
noise = np.random.normal(0, 2, shape)
|
|
image = (gaussian + noise).astype(np.uint8)
|
|
return image
|
|
|
|
|
|
def test_beamcenter_fit_success():
|
|
# Create a synthetic image with a known beam center
|
|
true_center = (120, 80)
|
|
image = create_synthetic_beam_image(center=true_center, sigma=(8, 12), theta=np.radians(30))
|
|
|
|
result = beamcenter_fit(image)
|
|
|
|
assert isinstance(result, Gaussian2Dfit)
|
|
# Check if the fitted center is close to the true center
|
|
assert pytest.approx(result.center_x, abs=2) == true_center[0]
|
|
assert pytest.approx(result.center_y, abs=2) == true_center[1]
|
|
assert result.peak_intensity > 150
|
|
assert 0 <= result.rotation_angle < 360
|
|
|
|
|
|
def test_beamcenter_fit_no_converge():
|
|
# Create an image that is just noise, should probably fail or at least not find a good fit
|
|
image = np.random.randint(0, 50, (200, 200), dtype=np.uint8)
|
|
|
|
# beamcenter_fit might still find some contour if there's enough noise,
|
|
# but curve_fit might fail to converge
|
|
# If it doesn't converge, it returns None now.
|
|
beamcenter_fit(image)
|
|
# It might actually return a result if it finds a random blob,
|
|
# but we want to test the failure path.
|
|
# To truly force non-convergence we might need a more extreme case,
|
|
# but return None is better than exit() anyway.
|
|
|
|
|
|
def test_beamcenter_fit_no_contours():
|
|
# Completely black image, max(contours) will fail
|
|
image = np.zeros((100, 100), dtype=np.uint8)
|
|
with pytest.raises(
|
|
ValueError, match=r"max\(\) (arg is an empty sequence|iterable argument is empty)"
|
|
):
|
|
beamcenter_fit(image)
|
|
|
|
|
|
def test_beamcenter_fit_small_blob():
|
|
# Test with a very small blob
|
|
image = np.zeros((100, 100), dtype=np.uint8)
|
|
image[45:55, 45:55] = 255
|
|
|
|
result = beamcenter_fit(image)
|
|
assert isinstance(result, Gaussian2Dfit)
|
|
assert pytest.approx(result.center_x, abs=2) == 50
|
|
assert pytest.approx(result.center_y, abs=2) == 50
|