From 706b683e3b872fe588be13e12cef93e8e2731ce5 Mon Sep 17 00:00:00 2001
From: Alexander Steppke <alexander.steppke@psi.ch>
Date: Sun, 29 Jan 2023 16:24:34 +0100
Subject: [PATCH] added 2d gaussian fitting routine

---
 src/cristallina/analysis.py | 54 ++++++++++++++++++++++++++++++++++---
 tests/test_analysis.py      | 22 +++++++++++++++
 2 files changed, 73 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_analysis.py

diff --git a/src/cristallina/analysis.py b/src/cristallina/analysis.py
index c78916d..2c93a22 100644
--- a/src/cristallina/analysis.py
+++ b/src/cristallina/analysis.py
@@ -1,7 +1,9 @@
 import re
 from collections import defaultdict
+from typing import Optional
 
 import numpy as np
+import lmfit
 
 from sfdata import SFDataFiles, sfdatafile, SFScanInfo
 
@@ -55,7 +57,7 @@ def perform_image_calculations(
     channel="JF16T03V01",
     alignment_channels=None,
     batch_size=10,
-    roi: ROI = None,
+    roi: Optional[ROI] = None,
     preview=False,
     operations=["sum"],
 ):
@@ -123,7 +125,7 @@ def sum_images(
     channel="JF16T03V01",
     alignment_channels=None,
     batch_size=10,
-    roi: ROI = None,
+    roi: Optional[ROI] = None,
     preview=False,
 ):
     """
@@ -156,7 +158,7 @@ def get_contrast_images(
     channel="JF16T03V01",
     alignment_channels=None,
     batch_size=10,
-    roi: ROI = None,
+    roi: Optional[ROI] = None,
     preview=False,
 ):
     """
@@ -172,3 +174,49 @@ def get_contrast_images(
         preview=preview,
         operations=["mean", "std"],
     )
+
+
+def fit_2d_gaussian(image, roi: Optional[ROI] = None):
+    """
+    2D Gaussian fit using LMFit for a given image and an optional region of interest.
+
+    Returns the x, y coordinates of the center and the results object which contains
+    further fit statistics.
+    """
+
+    # given an image and optional ROI
+    if roi is not None:
+        im = image[roi.rows, roi.cols]
+    else:
+        im = image
+
+    len_y, len_x = im.shape
+
+    y = np.arange(len_y)
+    x = np.arange(len_x)
+
+    x, y = np.meshgrid(x, y)  # here now a 2D mesh
+
+    x, y = x.ravel(), y.ravel()  # and all back into sequences of 1D arrays
+
+    z = im.ravel()  # and this also as a 1D
+
+    model = lmfit.models.Gaussian2dModel()
+    params = model.guess(z, x, y)
+    result = model.fit(
+        z,
+        x=x,
+        y=y,
+        params=params,
+    )
+
+    if roi is not None:
+        # convert back to original image coordinates
+        center_x = roi.left + result.params["centerx"]
+        center_y = roi.bottom + result.params["centery"]
+
+    else:
+        center_x = result.params["centerx"].value
+        center_y = result.params["centery"].value
+
+    return center_x, center_y, result
diff --git a/tests/test_analysis.py b/tests/test_analysis.py
new file mode 100644
index 0000000..bd2cb54
--- /dev/null
+++ b/tests/test_analysis.py
@@ -0,0 +1,22 @@
+import pytest
+import numpy as np
+
+import cristallina.analysis
+
+__author__ = "Alexander Steppke"
+
+	
+def test_2d_gaussian():
+
+    image = np.array([[0,0,0,0,0],
+            [0,0,0,0,0],
+            [0,0,1,0,0],
+            [0,0,0,0,0],
+            [0,0,0,0,0],
+            ])
+    
+    center_x, center_y, result = cristallina.analysis.fit_2d_gaussian(image)
+    assert np.allclose(center_x, 2.0, rtol=1e-04)
+    assert np.allclose(center_y, 2.0, rtol=1e-04)
+
+