feat: more modular psf fitting and test

.github/workflows/tests.yml

@@ -38,9 +38,7 @@ jobs:
           python -m pip install -U nox
 
       - name: "Tests: run"
-        run: |
-          python -m nox --non-interactive --error-on-missing-interpreter \
-            --session "${{matrix.session}}-${{matrix.python-version}}"
+        run: python -m nox --non-interactive --session "${{matrix.session}}-${{matrix.python-version}}"
 
   docs-build:
     runs-on: ubuntu-latest

.gitignore

@@ -21,4 +21,8 @@ _docs
 jupyter_execute
 __pycache__
 _docs
-*.f*t*
+*.f*t*
+*.npz
+*.array
+*.log
+*.mp4

noxfile.py

@@ -15,6 +15,7 @@ def test(session):
 def docs(session):
     session.install(".[docs]")
     with session.chdir("docs"):
+        session.run("mkdir", "autoapi", external=True)
         session.run(
             "python",
             "-m",

pyproject.toml

@@ -14,7 +14,7 @@ classifiers = [
     "Programming Language :: Python",
     "Programming Language :: Python :: 3",
 ]
-version = "0.0.11b"
+version = "0.0.12b"
 dependencies = [
     "numpy",
     "scikit-image",

src/eloy/psf.py

@@ -6,6 +6,7 @@
 """
 
 from scipy.optimize import minimize
+from functools import partial
 import numpy as np
 
 gaussian_sigma_to_fwhm = 2.0 * np.sqrt(2.0 * np.log(2.0))
@@ -50,6 +51,31 @@ def moments(data):
     }
 
 
+def gaussian(
+    xs=None,
+    ys=None,
+    amplitude=None,
+    x=None,
+    y=None,
+    sigma_x=None,
+    sigma_y=None,
+    theta=None,
+    background=None,
+    **kwargs,
+):
+    dx = xs - x
+    dy = ys - y
+    a = (np.cos(theta) ** 2) / (2 * sigma_x**2) + (np.sin(theta) ** 2) / (
+        2 * sigma_y**2
+    )
+    b = -(np.sin(2 * theta)) / (4 * sigma_x**2) + (np.sin(2 * theta)) / (4 * sigma_y**2)
+    c = (np.sin(theta) ** 2) / (2 * sigma_x**2) + (np.cos(theta) ** 2) / (
+        2 * sigma_y**2
+    )
+    psf = amplitude * np.exp(-(a * dx**2 + 2 * b * dx * dy + c * dy**2))
+    return psf + background
+
+
 def fit_gaussian(data, init=None):
     """
     Fit a 2D Gaussian to the data.
@@ -68,14 +94,7 @@ def fit_gaussian(data, init=None):
     """
     x, y = np.indices(data.shape)
 
-    def model(height, xo, yo, sx, sy, theta, m):
-        dx = x - xo
-        dy = y - yo
-        a = (np.cos(theta) ** 2) / (2 * sx**2) + (np.sin(theta) ** 2) / (2 * sy**2)
-        b = -(np.sin(2 * theta)) / (4 * sx**2) + (np.sin(2 * theta)) / (4 * sy**2)
-        c = (np.sin(theta) ** 2) / (2 * sx**2) + (np.cos(theta) ** 2) / (2 * sy**2)
-        psf = height * np.exp(-(a * dx**2 + 2 * b * dx * dy + c * dy**2))
-        return psf + m
+    model = partial(gaussian, x, y)
 
     def nll(params):
         ll = np.sum(np.power(model(*params) - data, 2))
@@ -98,4 +117,14 @@ def nll(params):
     ]
 
     opt = minimize(nll, p0, bounds=bounds).x
-    return {k: v for k, v in zip(keys, opt)}
+    result = {k: v for k, v in zip(keys, opt)}
+
+    # Ensure sigma_x is the larger value and sigma_y is the smaller
+    if result["sigma_x"] < result["sigma_y"]:
+        result["sigma_x"], result["sigma_y"] = result["sigma_y"], result["sigma_x"]
+        result["theta"] += np.pi / 2
+
+    # Ensure theta is always positive
+    result["theta"] = result["theta"] % (2 * np.pi)
+
+    return result

tests/test_psf.py

@@ -0,0 +1,43 @@
+import numpy as np
+from eloy import psf
+
+
+def test_fit_gaussian_sigma_with_noise():
+    rng = np.random.default_rng(12345)
+    size = 41
+    xs, ys = np.indices((size, size))
+
+    amplitude = 1.0
+    background = 0.1
+    x0 = 20.3
+    y0 = 19.7
+    sigma_x_true = 3.0
+    sigma_y_true = 2.0
+    theta = 0.3
+
+    img = psf.gaussian(
+        xs=xs,
+        ys=ys,
+        amplitude=amplitude,
+        x=x0,
+        y=y0,
+        sigma_x=sigma_x_true,
+        sigma_y=sigma_y_true,
+        theta=theta,
+        background=background,
+    )
+
+    noise = rng.normal(scale=0.02, size=img.shape)
+    data = img + noise
+
+    res = psf.fit_gaussian(data)
+
+    # fit_gaussian enforces sigma_x >= sigma_y; compare sorted values
+    fitted = sorted([res["sigma_x"], res["sigma_y"]], reverse=True)
+    true = sorted([sigma_x_true, sigma_y_true], reverse=True)
+
+    rel_err_x = abs(fitted[0] - true[0]) / true[0]
+    rel_err_y = abs(fitted[1] - true[1]) / true[1]
+
+    # require both sigmas to be recovered within 10% relative error
+    assert rel_err_x < 0.10 and rel_err_y < 0.10