Geometry.plot() and Structure().plot() accept the "transpose" argument which swaps the horizontal and vertical axes"

Author: jewettaijfc

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- `Geometry.plot()` and `Structure().plot()` now accept the `transpose=True` argument which swaps the horizontal and vertical axes of the plot.
+
 ### Changed
 - By default, batch downloads will skip files that already exist locally. To force re-downloading and replace existing files, pass the `replace_existing=True` argument to `Batch.load()`, `Batch.download()`, or `BatchData.load()`.
 - The `BatchData.load_sim_data()` function now overwrites any previously downloaded simulation files (instead of skipping them).

tests/test_components/test_geometry.py

@@ -87,14 +87,15 @@
 _, AX = plt.subplots()
 
 
-@pytest.mark.parametrize("component", GEO_TYPES)
-def test_plot(component):
-    _ = component.plot(z=0, ax=AX)
+@pytest.mark.parametrize("component, transpose", zip(GEO_TYPES, [True, False]))
+def test_plot(component, transpose):
+    _ = component.plot(z=0, ax=AX, transpose=transpose)
     plt.close()
 
 
-def test_plot_with_units():
-    _ = BOX.plot(z=0, ax=AX, plot_length_units="nm")
+@pytest.mark.parametrize("transpose", [True, False])
+def test_plot_with_units(transpose):
+    _ = BOX.plot(z=0, ax=AX, plot_length_units="nm", transpose=transpose)
     plt.close()
 
 
@@ -768,13 +769,12 @@ def test_geometry_touching_intersections_plane(x0):
 
 
 def test_pop_axis():
-    b = td.Box(size=(1, 1, 1))
     for axis in range(3):
         coords = (1, 2, 3)
-        Lz, (Lx, Ly) = b.pop_axis(coords, axis=axis)
-        _coords = b.unpop_axis(Lz, (Lx, Ly), axis=axis)
+        Lz, (Lx, Ly) = td.Box.pop_axis(coords, axis=axis)
+        _coords = td.Box.unpop_axis(Lz, (Lx, Ly), axis=axis)
         assert all(c == _c for (c, _c) in zip(coords, _coords))
-        _Lz, (_Lx, _Ly) = b.pop_axis(_coords, axis=axis)
+        _Lz, (_Lx, _Ly) = td.Box.pop_axis(_coords, axis=axis)
         assert Lz == _Lz
         assert Lx == _Lx
         assert Ly == _Ly
@@ -939,7 +939,8 @@ def test_to_gds(geometry, tmp_path):
     assert len(cell.polygons) == 0
 
 
-def test_custom_surface_geometry(tmp_path):
+@pytest.mark.parametrize("transpose", [True, False])
+def test_custom_surface_geometry(transpose, tmp_path):
     # create tetrahedron STL
     vertices = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]])
     faces = np.array([[1, 2, 3], [0, 3, 2], [0, 1, 3], [0, 2, 1]])
@@ -972,9 +973,13 @@ def test_custom_surface_geometry(tmp_path):
     assert np.isclose(geom.volume(), 1 / 6)
 
     # test intersections
-    assert shapely.equals(geom.intersections_plane(x=0), shapely.Polygon([[0, 0], [0, 1], [1, 0]]))
     assert shapely.equals(
-        geom.intersections_plane(z=0.5), shapely.Polygon([[0, 0], [0, 0.5], [0.5, 0]])
+        geom.intersections_plane(x=0),
+        shapely.Polygon([[0, 0], [0, 1], [1, 0]]),
+    )
+    assert shapely.equals(
+        geom.intersections_plane(z=0.5),
+        shapely.Polygon([[0, 0], [0, 0.5], [0.5, 0]]),
     )
 
     # test inside
@@ -983,7 +988,7 @@ def test_custom_surface_geometry(tmp_path):
 
     # test plot
     _, ax = plt.subplots()
-    _ = geom.plot(z=0.1, ax=ax)
+    _ = geom.plot(z=0.1, ax=ax, transpose=transpose)
     plt.close()
 
     # test inconsistent winding
@@ -1033,7 +1038,7 @@ def test_custom_surface_geometry(tmp_path):
         boundary_spec=td.BoundarySpec.all_sides(td.PML()),
     )
     _, ax = plt.subplots()
-    _ = sim.plot(y=0, ax=ax)
+    _ = sim.plot(y=0, ax=ax, transpose=transpose)
     plt.close()
 
     # allow small triangles

tests/test_components/test_utils.py

@@ -0,0 +1,29 @@
+"""Tests objects shared by multiple components."""
+
+from __future__ import annotations
+
+import random
+
+import pytest
+from shapely.geometry import Point
+
+from tidy3d.components.utils import pop_axis_and_swap, shape_swap_xy, unpop_axis_and_swap
+
+
+@pytest.mark.parametrize("transpose", [True, False])
+def test_pop_axis_and_swap(transpose):
+    for axis in range(3):
+        coords = (1, 2, 3)
+        Lz, (Lx, Ly) = pop_axis_and_swap(coords, axis=axis, transpose=transpose)
+        _coords = unpop_axis_and_swap(Lz, (Lx, Ly), axis=axis, transpose=transpose)
+        assert all(c == _c for (c, _c) in zip(coords, _coords))
+        _Lz, (_Lx, _Ly) = pop_axis_and_swap(_coords, axis=axis, transpose=transpose)
+        assert Lz == _Lz
+        assert Lx == _Lx
+        assert Ly == _Ly
+
+
+def test_shape_swap_xy():
+    p_orig = Point(random.random(), random.random())
+    p_new = shape_swap_xy(p_orig)
+    assert (p_new.coords[0][0], p_new.coords[0][1]) == (p_orig.coords[0][1], p_orig.coords[0][0])

tidy3d/components/data/unstructured/triangular.py

@@ -23,6 +23,7 @@
     SpatialDataArray,
 )
 from tidy3d.components.types import ArrayLike, Ax, Axis, Bound
+from tidy3d.components.utils import pop_axis_and_swap
 from tidy3d.components.viz import add_ax_if_none, equal_aspect, plot_params_grid
 from tidy3d.constants import inf
 from tidy3d.exceptions import DataError
@@ -570,9 +571,10 @@ def does_cover(self, bounds: Bound) -> bool:
 
     """ Plotting """
 
-    @property
-    def _triangulation_obj(self) -> Triangulation:
+    def _triangulation_obj(self, transpose: bool = False) -> Triangulation:
         """Matplotlib triangular representation of the grid to use in plotting."""
+        if transpose:
+            return Triangulation(self.points[:, 1], self.points[:, 0], self.cells)
         return Triangulation(self.points[:, 0], self.points[:, 1], self.cells)
 
     @equal_aspect
@@ -589,6 +591,7 @@ def plot(
         shading: Literal["gourand", "flat"] = "gouraud",
         cbar_kwargs: Optional[dict] = None,
         pcolor_kwargs: Optional[dict] = None,
+        transpose: bool = False,
     ) -> Ax:
         """Plot the data field and/or the unstructured grid.
 
@@ -616,6 +619,8 @@ def plot(
             Additional parameters passed to colorbar object.
         pcolor_kwargs: Dict = {}
             Additional parameters passed to ax.tripcolor()
+        transpose : bool = False
+            Swap horizontal and vertical axes. (This overrides the default ascending axis order)
 
         Returns
         -------
@@ -639,7 +644,7 @@ def plot(
                     f"{self._values_coords_dict} before plotting."
                 )
             plot_obj = ax.tripcolor(
-                self._triangulation_obj,
+                self._triangulation_obj(transpose=transpose),
                 self.values.data.ravel(),
                 shading=shading,
                 cmap=cmap,
@@ -657,14 +662,15 @@ def plot(
         # plot grid if requested
         if grid:
             ax.triplot(
-                self._triangulation_obj,
+                self._triangulation_obj(transpose=transpose),
                 color=plot_params_grid.edgecolor,
                 linewidth=plot_params_grid.linewidth,
             )
 
         # set labels and titles
-        ax_labels = ["x", "y", "z"]
-        normal_axis_name = ax_labels.pop(self.normal_axis)
+        normal_axis_name, ax_labels = pop_axis_and_swap(
+            "xyz", self.normal_axis, transpose=transpose
+        )
         ax.set_xlabel(ax_labels[0])
         ax.set_ylabel(ax_labels[1])
         ax.set_title(f"{normal_axis_name} = {self.normal_pos}")