Allow passing of figures for GUI.

pymagglobal/_commands.py

@@ -121,7 +121,7 @@ def lt2yr(times):
     return -times*1000 + 1950
 
 
-def local_curve(args):
+def local_curve(args, fig=None):
     '''Handle the local command and create a local curve at a given
     location.
 
@@ -129,6 +129,8 @@ def local_curve(args):
     ----------
     args : object
         The SimpleNamespace object returned by ArgumentParser.parse_args().
+    fig : matplotlib.figure.Figure, optional
+        A figure to plot into. This is used for the GUI.
 
     Returns
     -------
@@ -142,7 +144,7 @@ def local_curve(args):
     # create a local curve using the core function, check is performed
     # in args2times
     curves = core.local_curve(times, (args.lat, args.lon), args.model,
-                               field_type=args.type, check=False)
+                              field_type=args.type, check=False)
     # output formats for dif and nez components
     fmts = {'dif': ('%.2f', '%2.6f', '%2.6f', '%1.7e'),
             'nez': ('%.2f', '%1.7e', '%1.7e', '%1.7e')}
@@ -169,7 +171,8 @@ def local_curve(args):
 
     # if the --no-show flag is not set, plot the local curve
     if not args.no_show or args.savefig is not None:
-        fig = plt.figure(figsize=(10, 7))
+        if fig is None:
+            fig = plt.figure(figsize=(10, 7))
         fig.suptitle(f'Local curves at ({args.lat}°, {args.lon}°) '
                      f'for {args.model.name}')
         axs = np.empty(3, dtype=object)
@@ -190,7 +193,7 @@ def local_curve(args):
         return fig
 
 
-def dipole_series(args):
+def dipole_series(args, fig=None):
     '''Handle the dipole command and create a dipole-moment time series for
     the given model.
 
@@ -198,6 +201,8 @@ def dipole_series(args):
     ----------
     args : object
         The object returned by ArgumentParser.parse_args().
+    fig : matplotlib.figure.Figure, optional
+        A figure to plot into. This is used for the GUI.
 
     Returns
     -------
@@ -228,7 +233,8 @@ def dipole_series(args):
 
     # if the --no-show flag is not set, plot the time series
     if not args.no_show or args.savefig is not None:
-        fig = plt.figure(figsize=(10, 7))
+        if fig is None:
+            fig = plt.figure(figsize=(10, 7))
         ax = fig.add_subplot(111)
         ax.plot(times, dip_ser)
         ax.set_title(f'Dipole moment series for {args.model.name}')
@@ -462,7 +468,7 @@ def plot_coeffs(gs, ls, ms, args, unit='nT', name=r'Coefficients $g_\ell^m$'):
     return fig
 
 
-def maps(args):
+def maps(args, fig=None):
     '''Handle the map command and create field map of the model for a given
     epoch.
 
@@ -470,6 +476,8 @@ def maps(args):
     ----------
     args : object
         the object returned by ArgumentParser.parse_args().
+    fig : matplotlib.figure.Figure, optional
+        A figure to plot into. This is used for the GUI.
 
     Returns
     -------
@@ -526,13 +534,38 @@ def maps(args):
             field /= 1000
 
         cbar_hght = 0.06
-        proj = ccrs.Mollweide()
+        if fig is None:
+            proj = ccrs.Mollweide()
+            fig, axs = plt.subplots(1, 3, figsize=(13, 3),
+                                    subplot_kw={'projection': proj})
+            colaxs = []
+            for it in range(3):
+                bnds = axs[it].get_position().bounds
+
+                colaxs.append(fig.add_axes([bnds[0],
+                                            bnds[1]-0.1-cbar_hght,
+                                            bnds[2],
+                                            cbar_hght]))
+            cbar_orientation = 'horizontal'
+        else:
+            axs = fig.get_axes()
+            proj = axs[0].projection
+
+            colaxs = []
+            for it in range(3):
+                bnds = axs[it].get_position().bounds
+
+
+                colaxs.append(fig.add_axes([bnds[0]+bnds[2]+0.1+1.5*cbar_hght,
+                                            bnds[1],
+                                            0.8*cbar_hght,
+                                            bnds[3]]))
+            cbar_orientation = 'vertical'
+
         plt_lat, plt_lon, _ = proj.transform_points(ccrs.Geodetic(),
                                                     z_at[1],
                                                     90 - z_at[0]).T
 
-        fig, axs = plt.subplots(1, 3, figsize=(13, 3),
-                                subplot_kw={'projection': proj})
         fig.suptitle(f'Field maps for {args.model.name} at epoch '
                      f'{args.epoch} {args.t_unit}')
 
@@ -543,13 +576,8 @@ def maps(args):
             axs[it].coastlines(alpha=0.8, lw=0.5)
             axs[it].set_global()
             axs[it].set_title(f'{utils._names[args.type][it]} [{units[it]}]')
-            bnds = axs[it].get_position().bounds
-            # colorbar for the mean
-            colax = fig.add_axes([bnds[0],
-                                  bnds[1]-0.1-cbar_hght,
-                                  bnds[2],
-                                  cbar_hght])
+
             fig.colorbar(mappable,
-                         cax=colax,
-                         orientation='horizontal')
+                         cax=colaxs[it],
+                         orientation=cbar_orientation)
         return fig