Figure now accepts complex mosaic layouts

corny/figures.py

@@ -90,6 +90,7 @@ def label_abc(axes, text=[], fontsize=10):
 class Figure:
     """
     Usage:
+        # Map
         with Figure(size=(10,8), projection='flat',
                 tiles='satellite', zoom=16) as ax:
 
@@ -99,6 +100,23 @@ class Figure:
         add_colorbar(ax=ax, points=c, label='Soil moisture (g/g)',
                         bar_kw=dict(shrink=0.5, pad=0.02, aspect=20),
                         ticks=[0.07,0.09,0.12], ticklabels=["1","2","3"])
+        
+        # Layout regular grid
+        with Figure(layout=(2,2)) as axes:
+            axes[0][0].plot(x,y)
+            axes[0][1].plot(x,y)
+            axes[1][0].plot(x,y)
+            axes[1][1].plot(x,y)
+
+        # Layout with complex mosaic
+        with Figure(title="My grid",
+            layout = [[0,0,1],[2,'.',1]],
+            gridspec_kw=dict(width_ratios=[1.4, 1, 1], height_ratios=[1, 2]),
+            ) as axes:
+            axes[0].plot(x,y)
+            axes[1].plot(x,y)
+            axes[2].plot(x,y)
+                        
     """
 
     buff = None
@@ -147,13 +165,41 @@ class Figure:
         
     # Entering `with` statement
     def __enter__(self):
-        self.fig, self.axes = plt.subplots(self.layout[0], self.layout[1],
-                                           figsize=self.size, gridspec_kw=self.gridspec_kw,
-                                           subplot_kw=dict(projection=self.projection))
-        if self.layout[0]==1 and self.layout[1]==1:
-            self.axesflat = [self.axes]
-        else:
-            self.axesflat = self.axes.flatten()
+        
+        return_later = None
+
+        if isinstance(self.layout, tuple):
+            """
+            Regular grids, like (2,4)
+            """
+            self.fig, self.axes = plt.subplots(
+                self.layout[0], self.layout[1],
+                figsize = self.size,
+                gridspec_kw = self.gridspec_kw,
+                subplot_kw = dict(projection=self.projection)
+                )
+            if self.layout[0]==1 and self.layout[1]==1:
+                self.axesflat = [self.axes]
+                return_later = self.axes
+            else:
+                self.axesflat = self.axes.flatten()
+                return_later = self.axesflat
+        
+        elif isinstance(self.layout, list):
+            """
+            Complex mosaic, like [[0,0,1],[2,'.',1]]
+            """
+            self.fig, self.axes = plt.subplot_mosaic(
+                self.layout,
+                layout="constrained",
+                gridspec_kw = self.gridspec_kw,
+                figsize = self.size,
+                subplot_kw = dict(projection=self.projection)
+                )
+            # Convert labeled dict to list
+            self.axesflat = [v for k, v in sorted(self.axes.items(), key=lambda pair: pair[0])]
+            return_later = self.axesflat
+                            
         for ax in self.axesflat:
             if self.extent:
                 ax.set_xlim(self.extent[0], self.extent[1])
@@ -165,7 +211,7 @@ class Figure:
                 print("Figure instance returned. Usage: `with Figure(...) as F:`, `ax = F.ax`, `buff=F.buff`")
             return(self)
         else:
-            return(self.axes) # makes possibe: with Fig() as ax: ax.change
+            return(return_later) # makes possibe: with Fig() as ax: ax.change
 
     # Exiting `with` statement
     def __exit__(self, type, value, traceback):