Update README.rst

.gitlab-ci.yml

@@ -1,3 +1,6 @@
+# requires up-to-date conda environment to run the tests as 
+# the gitlab runner does not have access to /sf...
+
 build-job:
   stage: build
   script:
@@ -11,5 +14,7 @@ run_tests:
         - echo `pwd`
         - pip install -e .
     script:
-        - echo "Hello, $GITLAB_USER_LOGIN!"
-        - pytest -v
+        - echo "Starting pytest run"
+        - coverage run -m pytest
+        - coverage report
+    coverage: '/TOTAL.*\s+(\d+\%)/'

README.rst

@@ -27,21 +27,18 @@
     :alt: Project generated with PyScaffold
     :target: https://pyscaffold.org/
 
-|
+.. image:: https://gitlab.psi.ch/sf_cristallina/cristallina/badges/master/coverage.svg
+    :alt: Cristallina coverage report
+    :target: https://gitlab.psi.ch/sf_cristallina/cristallina/-/commits/master
+
 
 ===========
-cristallina
+Cristallina data analysis modules and plotting utilities.
 ===========
 
+This is a python package for data analysis, plotting and utility functions for the Cristallina endstation @ SwissFEL.
 
-    Cristallina data analysis modules and plotting utilities.
-
-
-Here we collect modules for data analysis, plotting and utility functions for the Cristallina endstation. 
-
-The data analysis is based on the common SwissFEL data architecture, with convienent access provided by `sf_datafiles <https://github.com/paulscherrerinstitute/sf_datafiles>`_. 
-
-The test suite (based on pytest) requires access to some data only available on either the cristallina consoles or the RA cluster. 
+The data analysis is based on the common SwissFEL data architecture, with convienent access provided by `sf_datafiles <https://github.com/paulscherrerinstitute/sf_datafiles>`_.  Some parts require access to SwissFEL-specific data storage but the aim is to use this package indepedently for data analysis both online and offline operation.
 
 
 

src/cristallina/analysis.py

@@ -3,6 +3,8 @@ from collections import defaultdict
 from typing import Optional
 
 import numpy as np
+from matplotlib import pyplot as plt
+
 import lmfit
 
 from sfdata import SFDataFiles, sfdatafile, SFScanInfo
@@ -97,7 +99,6 @@ def perform_image_calculations(
         res["roi"] = repr(roi)
 
         for image_slice in Images.in_batches(batch_size):
-
             index_slice, im = image_slice
 
             if roi is None:
@@ -178,7 +179,7 @@ def get_contrast_images(
 
 def fit_2d_gaussian(image, roi: Optional[ROI] = None, plot=False):
     """
-    2D Gaussian fit using LMFit for a given image and an optional region of interest. 
+    2D Gaussian fit using LMFit for a given image and an optional region of interest.
     plot=True as optional argument plots the fit results.
 
     Returns the x, y coordinates of the center and the results object which contains
@@ -209,6 +210,10 @@ def fit_2d_gaussian(image, roi: Optional[ROI] = None, plot=False):
         x=x,
         y=y,
         params=params,
+        method="leastsq",
+        verbose=False,
+        nan_policy=None,
+        max_nfev=None,
     )
 
     if roi is not None:
@@ -219,47 +224,184 @@ def fit_2d_gaussian(image, roi: Optional[ROI] = None, plot=False):
     else:
         center_x = result.params["centerx"].value
         center_y = result.params["centery"].value
-    
+
     if plot == True:
-        from matplotlib import pyplot as plt
-        from scipy.interpolate import griddata
-        X, Y = np.meshgrid(np.arange(len_y), np.arange(len_x))
-        Z = griddata((x, y), z, (X, Y), method='linear', fill_value=0)
+        _plot_2d_gaussian_fit(im, z, model, result)
+
+    return center_x, center_y, result
+
+
+def _plot_2d_gaussian_fit(im, z, model, result):
+    """Plot helper function to use the current image data, model and fit result and
+    plots them together.
+    """
+    from matplotlib import pyplot as plt
+    from scipy.interpolate import griddata
+
+    len_y, len_x = im.shape
+
+    X, Y = np.meshgrid(np.arange(len_x), np.arange(len_y))
+    Z = griddata((X.ravel(), Y.ravel()), z, (X, Y), method="linear", fill_value=np.nan)
+
+    fig, axs = plt.subplots(2, 2, figsize=(10, 10), layout="constrained")
+    for ax in axs.ravel():
+        ax.axis("equal")
+        ax.set_xlabel("x")
+        ax.set_ylabel("y")
+
+    vmax = np.max(Z)
+
+    ax = axs[0, 0]
+    art = ax.pcolormesh(X, Y, Z, vmin=0, vmax=vmax, shading="auto")
+    fig.colorbar(art, ax=ax, label="z", shrink=0.5)
+    ax.set_title("Data")
+
+    ax = axs[0, 1]
+    fit = model.func(X, Y, **result.best_values)
+    art = ax.pcolormesh(X, Y, fit, vmin=0, vmax=vmax, shading="auto")
+    fig.colorbar(art, ax=ax, label="z", shrink=0.5)
+    ax.set_title("Fit")
+
+    ax = axs[1, 0]
+    fit = model.func(X, Y, **result.best_values)
+    art = ax.pcolormesh(
+        X, Y, Z - fit, vmin=-0.05 * vmax, vmax=0.05 * vmax, cmap="gray", shading="auto"
+    )
+    fig.colorbar(art, ax=ax, label="z", shrink=0.5)
+    ax.set_title("Data - Fit")
+
+    ax = axs[1, 1]
+    fit = model.func(X, Y, **result.best_values)
+    art = ax.pcolormesh(X, Y, fit, vmin=0, vmax=vmax, shading="auto")
+    ax.contour(X, Y, fit, 8, colors="r", alpha=0.4)
+    fig.colorbar(art, ax=ax, label="z", shrink=0.5)
+    ax.set_title("Data & Fit")
+
+    fig.suptitle("2D Gaussian fit results")
+
+
+def gaussian2d_rot_model(
+    x,
+    y=0.0,
+    amplitude=1.0,
+    centerx=0.0,
+    centery=0.0,
+    sigmax=1.0,
+    sigmay=1.0,
+    rotation=0,
+    background=0,
+):
+    """Returns a two-dimensional Gaussian model from lmfit with a rotation in radians around the center."""
+    sr = np.sin(rotation)
+    cr = np.cos(rotation)
+
+    center_x_rot = centerx * cr - centery * sr
+    center_y_rot = centerx * sr + centery * cr
+
+    x_rot = x * cr - y * sr
+    y_rot = x * sr + y * cr
+
+    return (
+        lmfit.models.gaussian2d(
+            x_rot,
+            y=y_rot,
+            amplitude=amplitude,
+            centerx=center_x_rot,
+            centery=center_y_rot,
+            sigmax=sigmax,
+            sigmay=sigmay,
+        )
+        + background
+    )
+
+
+def fit_2d_gaussian_rotated(
+    image,
+    roi=None,
+    plot=False,
+    vary_rotation=True,
+    vary_background=False,
+):
+    """
+    2D Gaussian fit with rotation using LMFit for a given image and an optional region of interest.
+
+    As the number of free parameters for this kind of fit is large issues with convergence appear often.
+    Here we first fit without rotation, use the obtained parameters as a starting guess.
+
+    plot = True as optional argument plots the fit results.
+
+    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
+
+    mod = lmfit.Model(gaussian2d_rot_model, independent_vars=["x", "y"])
+
+    # Guess parameters, this is one possible approach
+    mod.set_param_hint("amplitude", value=np.max(z) * 0.75, min=0, vary=True)
+
+    mod.set_param_hint("centerx", value=np.mean(x) / 2, vary=True)
+    mod.set_param_hint("centery", value=np.mean(y) / 2, vary=True)
+
+    mod.set_param_hint("sigmax", value=np.mean(x) / 10, vary=True)
+    mod.set_param_hint("sigmay", value=np.mean(y) / 10, vary=True)
+
+    mod.set_param_hint("rotation", value=0.0, min=-np.pi / 2, max=np.pi / 2, vary=False)
+    mod.set_param_hint("background", value=0.0, vary=vary_background)
+
+    params = mod.make_params(verbose=False)
+    # first fit without rotation
+    result = mod.fit(
+        z,
+        x=x,
+        y=y,
+        params=params,
+        method="leastsq",
+        verbose=False,
+        nan_policy=None,
+        max_nfev=20,
+    )
+
+    # now refining with rotation
+    params = result.params
+    params["rotation"].set(vary=vary_rotation)
+
+    result = mod.fit(
+        z,
+        x=x,
+        y=y,
+        params=params,
+        method="leastsq",
+        verbose=False,
+        nan_policy=None,
+        max_nfev=None,
+    )
+
+    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
+
+    if plot == True:
+        _plot_2d_gaussian_fit(im, z, mod, result)
 
-        fig, axs = plt.subplots(2, 2, figsize=(10, 10))
-
-        # vmax = np.nanpercentile(Z, 99.9)
-        vmax = np.max(Z)
-
-        ax = axs[0, 0]
-        art = ax.pcolor(X, Y, Z, vmin=0, vmax=vmax, shading='auto')
-        plt.colorbar(art, ax=ax, label='z')
-        ax.set_title('Data')
-
-        ax = axs[0, 1]
-        fit = model.func(X, Y, **result.best_values)
-        art = ax.pcolor(X, Y, fit, vmin=0, vmax=vmax, shading='auto')
-        plt.colorbar(art, ax=ax, label='z')
-        ax.set_title('Fit')
-
-        ax = axs[1, 0]
-        fit = model.func(X, Y, **result.best_values)
-        art = ax.pcolor(X, Y, Z-fit, vmin=0, shading='auto')
-        plt.colorbar(art, ax=ax, label='z')
-        ax.set_title('Data - Fit')
-        
-        ax = axs[1, 1]
-        fit = model.func(X, Y, **result.best_values)
-        art = ax.pcolor(X, Y, fit, vmin=0, vmax=vmax, shading='auto')
-        ax.contour(X, Y, fit, 8, colors='r',alpha=0.4)
-        plt.colorbar(art, ax=ax, label='z')
-        ax.set_title('Data & Fit')
-
-        for ax in axs.ravel():
-            ax.set_xlabel('x')
-            ax.set_ylabel('y')
-        plt.suptitle('2D Gaussian fit results')    
-        plt.tight_layout()
-        plt.show()
-    
     return center_x, center_y, result

src/cristallina/utils.py

@@ -22,7 +22,7 @@ def scan_info(run_number, base_path=None, small_data=True):
     try:
         if small_data:
             for i in range(len(scan.readbacks)):
-                sd_path_for_step=heuristic_extract_smalldata_path()+'run'+str(run_number).zfill(4)+'/acq'+str(i+1).zfill(4)+'.smalldata.h5'
+                sd_path_for_step=heuristic_extract_smalldata_path()+'run'+str(run_number).zfill(4)+'/acq'+str(i+1).zfill(4)+'.smalldata*.h5'
                 scan.info['scan_files'][i].append(sd_path_for_step)
     except KeyError:
         logger.warning("Failed to load smalldata.") 

tests/test_analysis.py

@@ -6,7 +6,7 @@ import cristallina.analysis
 __author__ = "Alexander Steppke"
 
 	
-def test_2d_gaussian():
+def test_minimal_2d_gaussian():
 
     image = np.array([[0,0,0,0,0],
             [0,0,0,0,0],
@@ -20,3 +20,47 @@ def test_2d_gaussian():
     assert np.allclose(center_y, 2.0, rtol=1e-04)
 
 
+
+def test_2d_gaussian():
+
+        # define normalized 2D gaussian
+        def gauss2d(x=0, y=0, mx=0, my=0, sx=1, sy=1):
+            return (1 / (2 * np.pi * sx * sy) * np.exp(-((x - mx) ** 2 / (2 * sx**2.0) + (y - my) ** 2 / (2 * sy**2))))
+
+        x = np.arange(0, 150, 1)
+        y = np.arange(0, 100, 1)
+        x, y = np.meshgrid(x, y)  
+
+        z = gauss2d(x, y, mx=40, my=50, sx=20, sy=40)
+        
+        center_x, center_y, result = cristallina.analysis.fit_2d_gaussian(z)
+        assert np.allclose(center_x, 40, rtol=1e-04)
+        assert np.allclose(center_y, 50, rtol=1e-04)
+
+
+def test_2d_gaussian_rotated():
+
+        # define normalized 2D gaussian
+        def gauss2d_rotated(x=0, y=0, center_x=0, center_y=0, sx=1, sy=1, rotation=0):
+
+            sr = np.sin(rotation)
+            cr = np.cos(rotation)
+
+            center_x_rot = center_x * cr - center_y * sr
+            center_y_rot = center_x * sr + center_y * cr
+
+            x_rot = x * cr - y * sr
+            y_rot = x * sr + y * cr
+
+            return (1 / (2 * np.pi * sx * sy) * np.exp(-((x_rot - center_x_rot) ** 2 / (2 * sx**2.0) + (y_rot - center_y_rot) ** 2 / (2 * sy**2))))
+
+        x = np.arange(0, 150, 1)
+        y = np.arange(0, 100, 1)
+        x, y = np.meshgrid(x, y)  
+
+        z = 100*gauss2d_rotated(x, y, center_x=40, center_y=50, sx=10, sy=20, rotation=0.5)
+        
+        center_x, center_y, result = cristallina.analysis.fit_2d_gaussian_rotated(z, vary_rotation=True, plot=False)
+        assert np.allclose(center_x, 40, rtol=1e-04)
+        assert np.allclose(center_y, 50, rtol=1e-04)
+        assert np.allclose(result.params['rotation'].value, 0.5, rtol=1e-02)