Merge branch 'polygon-plots' of https://github.com/boutproject/xBOUT …

…into polygon-plots
boutproject · Mar 25, 2024 · ef5740a · ef5740a
2 parents 59f82e4 + d1fdc43
commit ef5740a
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 30 deletions.
diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py
@@ -1071,7 +1071,7 @@ def pcolormesh(self, ax=None, **kwargs):
         Colour-plot a radial-poloidal slice on the R-Z plane
         """
         return plotfuncs.plot2d_wrapper(self.data, xr.plot.pcolormesh, ax=ax, **kwargs)
-    
+
     def polygon(self, ax=None, **kwargs):
         """
         Colour-plot of a radial-poloidal slice on the R-Z plane using polygons

diff --git a/xbout/geometries.py b/xbout/geometries.py
@@ -381,14 +381,14 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
             "total_poloidal_distance",
             "zShift",
             "zShift_ylow",
-            "Rxy_corners",   # Lower left corners
+            "Rxy_corners",  # Lower left corners
             "Rxy_lower_right_corners",
             "Rxy_upper_left_corners",
             "Rxy_upper_right_corners",
-            "Zxy_corners",   # Lower left corners
+            "Zxy_corners",  # Lower left corners
             "Zxy_lower_right_corners",
             "Zxy_upper_left_corners",
-            "Zxy_upper_right_corners"
+            "Zxy_upper_right_corners",
         ],
     )
 
@@ -427,23 +427,24 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
         ds = ds.set_coords(("R", "Z"))
     else:
         ds = ds.set_coords(("Rxy", "Zxy"))
-    
+
     # Add cell corners as coordinates for polygon plotting
     if "Rxy_lower_right_corners" in ds:
         ds = ds.rename(
-            Rxy_corners = "Rxy_lower_left_corners",
-            Zxy_corners = "Zxy_lower_left_corners"
+            Rxy_corners="Rxy_lower_left_corners", Zxy_corners="Zxy_lower_left_corners"
+        )
+        ds = ds.set_coords(
+            (
+                "Rxy_lower_left_corners",
+                "Rxy_lower_right_corners",
+                "Rxy_upper_left_corners",
+                "Rxy_upper_right_corners",
+                "Zxy_lower_left_corners",
+                "Zxy_lower_right_corners",
+                "Zxy_upper_left_corners",
+                "Zxy_upper_right_corners",
             )
-        ds = ds.set_coords((
-            "Rxy_lower_left_corners",
-            "Rxy_lower_right_corners",
-            "Rxy_upper_left_corners",
-            "Rxy_upper_right_corners",
-            "Zxy_lower_left_corners",
-            "Zxy_lower_right_corners",
-            "Zxy_upper_left_corners",
-            "Zxy_upper_right_corners"
-             ))
+        )
 
     # Rename zShift_ylow if it was added from grid file, to be consistent with name if
     # it was added from dump file

diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py
@@ -120,7 +120,9 @@ def plot_separatrices(da, ax, *, x="R", y="Z", **kwargs):
             )
             default_style = {"color": "black", "linestyle": "--"}
             if any(x for x in kwargs if x in ["c", "ls"]):
-                raise ValueError("When passing separatrix plot style kwargs, use 'color' and 'linestyle' instead lf 'c' and 'ls'")
+                raise ValueError(
+                    "When passing separatrix plot style kwargs, use 'color' and 'linestyle' instead lf 'c' and 'ls'"
+                )
             style = {**default_style, **kwargs}
             ax.plot(x_sep, y_sep, **style)
 

diff --git a/xbout/tests/test_against_collect.py b/xbout/tests/test_against_collect.py
@@ -220,5 +220,4 @@ def test_new_collect_indexing_slice(self, tmp_path_factory):
 
 
 @pytest.mark.skip
-class test_speed_against_old_collect:
-    ...
+class test_speed_against_old_collect: ...
diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py
@@ -472,8 +472,7 @@ def test_combine_along_y(self, tmp_path_factory, bout_xyt_example_files):
         xrt.assert_identical(actual, fake)
 
     @pytest.mark.skip
-    def test_combine_along_t(self):
-        ...
+    def test_combine_along_t(self): ...
 
     @pytest.mark.parametrize(
         "bout_v5,metric_3D", [(False, False), (True, False), (True, True)]
@@ -623,8 +622,7 @@ def test_drop_vars(self, tmp_path_factory, bout_xyt_example_files):
         assert "n" in ds.keys()
 
     @pytest.mark.skip
-    def test_combine_along_tx(self):
-        ...
+    def test_combine_along_tx(self): ...
 
     def test_restarts(self):
         datapath = Path(__file__).parent.joinpath(

diff --git a/xbout/utils.py b/xbout/utils.py
@@ -167,12 +167,16 @@ def _1d_coord_from_spacing(spacing, dim, ds=None, *, origin_at=None):
             )
 
         point_to_use = {
-            spacing.metadata["bout_xdim"]: spacing.metadata.get("MXG", 0)
-            if spacing.metadata["keep_xboundaries"]
-            else 0,
-            spacing.metadata["bout_ydim"]: spacing.metadata.get("MYG", 0)
-            if spacing.metadata["keep_yboundaries"]
-            else 0,
+            spacing.metadata["bout_xdim"]: (
+                spacing.metadata.get("MXG", 0)
+                if spacing.metadata["keep_xboundaries"]
+                else 0
+            ),
+            spacing.metadata["bout_ydim"]: (
+                spacing.metadata.get("MYG", 0)
+                if spacing.metadata["keep_yboundaries"]
+                else 0
+            ),
             spacing.metadata["bout_zdim"]: spacing.metadata.get("MZG", 0),
         }