Add plot kwargs to AFMulator (#364)

To enable overlaying the atoms to the generated AFM images.

Author: huangchieh

ppafm/ocl/AFMulator.py

@@ -628,14 +628,15 @@ def check_scan_window(self):
                     "the size of the scan window and the force field grid."
                 )
 
-    def plot_images(self, X, outdir="afm_images", prefix="df"):
+    def plot_images(self, X, outdir="afm_images", prefix="df", **plot_kwargs):
         """
         Plot simulated AFM images and save them to a directory.
 
         Arguments:
             X: np.ndarray. AFM images to plot.
             outdir: str. Path to directory where files are saved.
             prefix: str. Prefix string for saved files.
+            plot_kwargs: Additional keyword arguments passed to :func:`.plotImages`.
         """
         if not os.path.exists(outdir):
             os.makedirs(outdir)
@@ -652,14 +653,12 @@ def plot_images(self, X, outdir="afm_images", prefix="df"):
             self.scan_window[0][1],
             self.scan_window[1][1],
         ]
-        plotImages(
-            os.path.join(outdir, prefix),
-            X,
-            slices=list(range(0, len(X))),
-            zs=zTips,
-            extent=extent,
-            cmap=self.colorscale,
-        )
+
+        for key in ["F", "slices", "extent", "zs", "cmap"]:
+            if key in plot_kwargs:
+                warnings.warn(f"\n plot_images(): '{key}' found in plot_kwargs but is  set internally, removing from plot_kwargs.")
+                del plot_kwargs[key]
+        plotImages(os.path.join(outdir, prefix), X, slices=list(range(0, len(X))), zs=zTips, extent=extent, cmap=self.colorscale, **plot_kwargs)
 
 
 def _get_params(file_path):