diff --git a/src/cristallina/analysis.py b/src/cristallina/analysis.py
index cbbe4e0..5cad074 100644
--- a/src/cristallina/analysis.py
+++ b/src/cristallina/analysis.py
@@ -200,7 +200,7 @@ def fit_2d_gaussian(image, roi: Optional[ROI] = None, plot=False):
 
     x, y = x.ravel(), y.ravel()  # and all back into sequences of 1D arrays
 
-    z = im.ravel()  # and this also as a 1D
+    z = im.ravel() # and this also as a 1D
 
     model = lmfit.models.Gaussian2dModel()
     params = model.guess(z, x, y)
@@ -221,45 +221,52 @@ def fit_2d_gaussian(image, roi: Optional[ROI] = None, plot=False):
         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)
-
-        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()
+        _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.pcolor(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.pcolor(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.pcolor(X, Y, Z-fit, vmin=0, 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.pcolor(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') 
\ No newline at end of file
diff --git a/tests/test_analysis.py b/tests/test_analysis.py
index bd2cb54..fa3a8c5 100644
--- a/tests/test_analysis.py
+++ b/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,21 @@ 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)
+
+