Figure.subplot: Fix strange positioning issues after exiting subplot (#…

…2427)

Co-authored-by: Wei Ji <[email protected]>

examples/gallery/images/grdgradient_shading.py

@@ -57,6 +57,6 @@
                 panel=True,
             )
 
-fig.colorbar(position="x9.0c/-1.1c+w10c/0.25c+h", frame="a2000f500+lElevation (m)")
+fig.colorbar(position="JBC+w10c/0.25c+h", frame="a2000f500+lElevation (m)")
 
 fig.show()

pygmt/src/subplot.py

@@ -154,12 +154,17 @@ def subplot(self, nrows=1, ncols=1, **kwargs):
             "Please provide either one of 'figsize' or 'subsize' only."
         )
 
-    with Session() as lib:
-        try:
+    # Need to use separate sessions for "subplot begin" and "subplot end".
+    # Otherwise, "subplot end" will use the last session, which may cause
+    # strange positioning issues for later plotting calls.
+    # See https://github.com/GenericMappingTools/pygmt/issues/2426.
+    try:
+        with Session() as lib:
             arg_str = " ".join(["begin", f"{nrows}x{ncols}", build_arg_string(kwargs)])
             lib.call_module(module="subplot", args=arg_str)
             yield
-        finally:
+    finally:
+        with Session() as lib:
             v_arg = build_arg_string({"V": kwargs.get("V")})
             lib.call_module(module="subplot", args=f"end {v_arg}")
 

pygmt/tests/baseline/test_subplot_outside_plotting_positioning.png.dvc

@@ -0,0 +1,4 @@
+outs:
+- md5: 976a01a8c9f583918c00b5a81fddaf84
+  size: 12112
+  path: test_subplot_outside_plotting_positioning.png

pygmt/tests/test_subplot.py

@@ -100,3 +100,19 @@ def test_subplot_nrows_ncols_less_than_one_error():
     with pytest.raises(GMTInvalidInput):
         with fig.subplot(nrows=0, ncols=-1, figsize=("2c", "1c")):
             pass
+
+
+@pytest.mark.mpl_image_compare
+def test_subplot_outside_plotting_positioning():
+    """
+    Plotting calls are correctly positioned after exiting subplot.
+
+    This is a regression test for
+    https://github.com/GenericMappingTools/pygmt/issues/2426.
+    """
+    fig = Figure()
+    with fig.subplot(nrows=1, ncols=2, figsize=(10, 5)):
+        fig.basemap(region=[0, 10, 0, 10], projection="X?", frame=True, panel=True)
+        fig.basemap(region=[0, 10, 0, 10], projection="X?", frame=True, panel=True)
+    fig.colorbar(position="JBC+w5c+h", cmap="turbo", frame=True)
+    return fig