adding returnfig

cophus · Oct 23, 2024 · b291ecb · b291ecb
1 parent c2eeb59
commit b291ecb
Showing 1 changed file with 15 additions and 5 deletions.
diff --git a/py4DSTEM/process/polar/polar_peaks.py b/py4DSTEM/process/polar/polar_peaks.py
@@ -797,6 +797,7 @@ def model_radial_background(
     refine_model=True,
     plot_result=True,
     figsize=(8, 4),
+    returnfig=False,
 ):
     """
     User provided radial background model, of the form:
@@ -900,11 +901,20 @@ def background_model(q, *coefs):
 
     # plotting
     if plot_result:
-        self.plot_radial_background(
-            q_pixel_units=False,
-            plot_background_model=True,
-            figsize=figsize,
-        )
+        if returnfig:
+            fig,ax = self.plot_radial_background(
+                q_pixel_units=False,
+                plot_background_model=True,
+                figsize=figsize,
+                returnfig=returnfig,
+            )
+            return fig, ax
+        else:
+            self.plot_radial_background(
+                q_pixel_units=False,
+                plot_background_model=True,
+                figsize=figsize,
+            )            
 
 
 def refine_peaks(