replaced (un)pop_axis() with (un)pop_axis_and_swap() everwhere where "swap_axis" is used

Author: jewettaijfc

tidy3d/components/eme/simulation.py

@@ -316,15 +316,15 @@ def plot_eme_ports(
         rmax = self.geometry.bounds[1][self.axis]
         ports = np.array([rmin + self.port_offsets[0], rmax - self.port_offsets[1]])
         axis, _ = self.parse_xyz_kwargs(x=x, y=y, z=z)
-        _, (axis_x, axis_y) = self.pop_axis([0, 1, 2], axis=axis, swap_axes=swap_axes)
+        _, (axis_x, axis_y) = self.pop_axis_and_swap([0, 1, 2], axis=axis, swap_axes=swap_axes)
         boundaries_x = []
         boundaries_y = []
         if axis_x == self.axis:
             boundaries_x = ports
         if axis_y == self.axis:
             boundaries_y = ports
-        _, (xmin, ymin) = self.pop_axis(self.simulation_bounds[0], axis=axis, swap_axes=swap_axes)
-        _, (xmax, ymax) = self.pop_axis(self.simulation_bounds[1], axis=axis, swap_axes=swap_axes)
+        _, (xmin, ymin) = self.pop_axis_and_swap(self.simulation_bounds[0], axis=axis, swap_axes=swap_axes)
+        _, (xmax, ymax) = self.pop_axis_and_swap(self.simulation_bounds[1], axis=axis, swap_axes=swap_axes)
         segs_x = [((bound, ymin), (bound, ymax)) for bound in boundaries_x]
         line_segments_x = mpl.collections.LineCollection(segs_x, **kwargs)
         segs_y = [((xmin, bound), (xmax, bound)) for bound in boundaries_y]
@@ -365,15 +365,15 @@ def plot_eme_subgrid_boundaries(
         subgrid_boundaries = np.array(eme_grid_spec.subgrid_boundaries)
         subgrids = eme_grid_spec.subgrids
         axis, _ = self.parse_xyz_kwargs(x=x, y=y, z=z)
-        _, (axis_x, axis_y) = self.pop_axis([0, 1, 2], axis=axis, swap_axes=swap_axes)
+        _, (axis_x, axis_y) = self.pop_axis_and_swap([0, 1, 2], axis=axis, swap_axes=swap_axes)
         boundaries_x = []
         boundaries_y = []
         if axis_x == self.axis:
             boundaries_x = subgrid_boundaries
         if axis_y == self.axis:
             boundaries_y = subgrid_boundaries
-        _, (xmin, ymin) = self.pop_axis(self.simulation_bounds[0], axis=axis, swap_axes=swap_axes)
-        _, (xmax, ymax) = self.pop_axis(self.simulation_bounds[1], axis=axis, swap_axes=swap_axes)
+        _, (xmin, ymin) = self.pop_axis_and_swap(self.simulation_bounds[0], axis=axis, swap_axes=swap_axes)
+        _, (xmax, ymax) = self.pop_axis_and_swap(self.simulation_bounds[1], axis=axis, swap_axes=swap_axes)
         segs_x = [((bound, ymin), (bound, ymax)) for bound in boundaries_x]
         line_segments_x = mpl.collections.LineCollection(segs_x, **kwargs)
         segs_y = [((xmin, bound), (xmax, bound)) for bound in boundaries_y]
@@ -412,15 +412,15 @@ def plot_eme_grid(
         kwargs.setdefault("colors", "black")
         cell_boundaries = self.eme_grid.boundaries
         axis, _ = self.parse_xyz_kwargs(x=x, y=y, z=z)
-        _, (axis_x, axis_y) = self.pop_axis([0, 1, 2], axis=axis, swap_axes=swap_axes)
+        _, (axis_x, axis_y) = self.pop_axis_and_swap([0, 1, 2], axis=axis, swap_axes=swap_axes)
         boundaries_x = []
         boundaries_y = []
         if axis_x == self.axis:
             boundaries_x = cell_boundaries
         if axis_y == self.axis:
             boundaries_y = cell_boundaries
-        _, (xmin, ymin) = self.pop_axis(self.simulation_bounds[0], axis=axis, swap_axes=swap_axes)
-        _, (xmax, ymax) = self.pop_axis(self.simulation_bounds[1], axis=axis, swap_axes=swap_axes)
+        _, (xmin, ymin) = self.pop_axis_and_swap(self.simulation_bounds[0], axis=axis, swap_axes=swap_axes)
+        _, (xmax, ymax) = self.pop_axis_and_swap(self.simulation_bounds[1], axis=axis, swap_axes=swap_axes)
         segs_x = [((bound, ymin), (bound, ymax)) for bound in boundaries_x]
         line_segments_x = mpl.collections.LineCollection(segs_x, **kwargs)
         segs_y = [((xmin, bound), (xmax, bound)) for bound in boundaries_y]

tidy3d/components/geometry/base.py

@@ -462,8 +462,8 @@ def _pop_bounds(
             Packed as ``(zmin, zmax), ((xmin, ymin), (xmax, ymax))``.
         """
         b_min, b_max = self.bounds
-        zmin, (xmin, ymin) = self.pop_axis(b_min, axis=axis, swap_axes=swap_axes)
-        zmax, (xmax, ymax) = self.pop_axis(b_max, axis=axis, swap_axes=swap_axes)
+        zmin, (xmin, ymin) = self.pop_axis_and_swap(b_min, axis=axis, swap_axes=swap_axes)
+        zmax, (xmax, ymax) = self.pop_axis_and_swap(b_max, axis=axis, swap_axes=swap_axes)
         return (zmin, zmax), ((xmin, ymin), (xmax, ymax))
 
     @staticmethod
@@ -611,7 +611,7 @@ def _get_plot_labels(axis: Axis, swap_axes: bool = False) -> tuple[str, str]:
         str, str
             Labels of plot, packaged as ``(xlabel, ylabel)``.
         """
-        _, (xlabel, ylabel) = Geometry.pop_axis("xyz", axis=axis, swap_axes=swap_axes)
+        _, (xlabel, ylabel) = Geometry.pop_axis_and_swap("xyz", axis=axis, swap_axes=swap_axes)
         return xlabel, ylabel
 
     def _get_plot_limits(
@@ -734,7 +734,6 @@ def evaluate_inf_shape(shape: Shapely) -> Shapely:
     def pop_axis(
         coord: tuple[Any, Any, Any],
         axis: int,
-        swap_axes: bool = False,
     ) -> tuple[Any, tuple[Any, Any]]:
         """Separates coordinate at ``axis`` index from coordinates on the plane tangent to ``axis``.
 
@@ -744,8 +743,7 @@ def pop_axis(
             Tuple of three values in original coordinate system.
         axis : int
             Integer index into 'xyz' (0,1,2).
-        swap_axes: bool = False
-            Optional: Swap the order of the remaining two axes?
+
 
         Returns
         -------
@@ -756,16 +754,47 @@ def pop_axis(
         """
         plane_vals = list(coord)
         axis_val = plane_vals.pop(axis)
-        if swap_axes:
-            plane_vals = [plane_vals[1], plane_vals[0]]
         return axis_val, tuple(plane_vals)
 
+    @staticmethod
+    def pop_axis_and_swap(
+        coord: tuple[Any, Any, Any],
+        axis: int,
+        swap_axes: bool = False,
+    ) -> tuple[Any, tuple[Any, Any]]:
+        """
+        pop_axis_and_swap() is identical to pop_axis(), except that accepts an
+        additional "swap_axes" argument which reverses the output orer.  Examples:
+
+        pop_axis_and_swap(("x", "y", "z"), 1, swap_axes=False)  ->  "y", ("x", "z")
+        pop_axis_and_swap(("x", "y", "z"), 1, swap_axes=True)   ->  "y", ("z", "x")
+
+        Parameters
+        ----------
+        coord : Tuple[Any, Any, Any]
+            Tuple of three values in original coordinate system.
+        axis : int
+            Integer index into 'xyz' (0,1,2).
+        swap_axes: bool = False
+            Optional: Swap the order of the data from the two remaining axes in the output tuple?
+
+        Returns
+        -------
+        Any, Tuple[Any, Any]
+            The input coordinates are separated into the one along the axis provided
+            and the two on the planar coordinates,
+            like ``axis_coord, (planar_coord1, planar_coord2)``.
+        """
+        axis_val, plane_vals = Geometry.pop_axis(coord, axis)
+        if swap_axes:
+            return axis_val, (plane_vals[1], plane_vals[0])
+        return axis_val, plane_vals
+
     @staticmethod
     def unpop_axis(
         ax_coord: Any,
         plane_coords: tuple[Any, Any],
         axis: int,
-        swap_axes: bool = False,
     ) -> tuple[Any, Any, Any]:
         """Combine coordinate along axis with coordinates on the plane tangent to the axis.
 
@@ -777,23 +806,55 @@ def unpop_axis(
             Values along ordered planar directions.
         axis : int
             Integer index into 'xyz' (0,1,2).
-        swap_axes: bool = False
-            Optional: Swap the order of the entries in plane_coords[]?
-            NOTE: If `axis, plane_coords = pop_axis(coords, axis, transpose)`,
-            then, if you want to get back the original `coords`, you should use:
-            `unpop_axis(ax_coord, plane_coords, axis, transpose)`
 
         Returns
         -------
         Tuple[Any, Any, Any]
             The three values in the xyz coordinate system.
         """
         coords = list(plane_coords)
-        if swap_axes:
-            coords = [coords[1], coords[0]]
         coords.insert(axis, ax_coord)
         return tuple(coords)
 
+    @staticmethod
+    def unpop_axis_and_swap(
+        ax_coord: Any,
+        plane_coords: tuple[Any, Any],
+        axis: int,
+        swap_axes: bool = False,
+    ) -> tuple[Any, Any, Any]:
+        """
+        unpop_axis_and_swap() is identical to unpop_axis(), except that accepts
+        an additional "swap_axes" argument which reverses the order of 
+        plane_coords before sending them to unpop_axis().  For example:
+
+        unpop_axis_and_swap("y", ("x", "z"), 1, swap_axes=False)  -->  ("x", "y", "z")
+        unpop_axis_and_swap("y", ("x", "z"), 1, swap_axes=True)   -->  ("z", "y", "x")
+
+        This function is the inverse of pop_axis_and_swap().  For example:
+        unpop_axis_and_swap("y", ("z", "x"), 1, swap_axes=True)   -->  ("x", "y", "z")
+
+        Parameters
+        ----------
+        ax_coord : Any
+            Value along axis direction.
+        plane_coords : Tuple[Any, Any]
+            Values along ordered planar directions.
+        axis : int
+            Integer index into 'xyz' (0,1,2).
+        swap_axes: bool = False
+            Optional: Swap the order of the entries in plane_coords[]?
+
+        Returns
+        -------
+        Tuple[Any, Any, Any]
+            The three values in the xyz coordinate system.
+        """
+        coords = plane_coords
+        if swap_axes:
+            coords = (coords[1], coords[0])
+        return Geometry.unpop_axis(ax_coord, coords, axis)
+
     @staticmethod
     def parse_xyz_kwargs(**xyz) -> tuple[Axis, float]:
         """Turns x,y,z kwargs into index of the normal axis and position along that axis.
@@ -1813,7 +1874,7 @@ def _order_by_axis(
         """
         vals = 3 * [plane_val]
         vals[self.axis] = axis_val
-        _, (val_x, val_y) = self.pop_axis(vals, axis=axis, swap_axes=swap_axes)
+        _, (val_x, val_y) = self.pop_axis_and_swap(vals, axis=axis, swap_axes=swap_axes)
         return val_x, val_y
 
     @cached_property
@@ -2107,8 +2168,8 @@ def intersections_plane(
         axis, position = self.parse_xyz_kwargs(x=x, y=y, z=z)
         if not self.intersects_axis_position(axis, position):
             return []
-        z0, (x0, y0) = self.pop_axis(self.center, axis=axis, swap_axes=swap_axes)
-        Lz, (Lx, Ly) = self.pop_axis(self.size, axis=axis, swap_axes=swap_axes)
+        z0, (x0, y0) = self.pop_axis_and_swap(self.center, axis=axis, swap_axes=swap_axes)
+        Lz, (Lx, Ly) = self.pop_axis_and_swap(self.size, axis=axis, swap_axes=swap_axes)
         dz = np.abs(z0 - position)
         if dz > Lz / 2 + fp_eps:
             return []
@@ -2183,7 +2244,7 @@ def intersections_with(self, other, swap_axes: bool = False):
         shapes_plane = other.intersections_plane(**xyz_kwargs, swap_axes=swap_axes)
 
         # intersect all shapes with the input self
-        bs_min, bs_max = (self.pop_axis(bounds, axis=normal_ind, swap_axes=swap_axes)[1] for bounds in self.bounds)
+        bs_min, bs_max = (self.pop_axis_and_swap(bounds, axis=normal_ind, swap_axes=swap_axes)[1] for bounds in self.bounds)
 
         shapely_box = self.make_shapely_box(bs_min[0], bs_min[1], bs_max[0], bs_max[1])
         shapely_box = Geometry.evaluate_inf_shape(shapely_box)
@@ -2285,7 +2346,7 @@ def _plot_arrow(
         """
 
         plot_axis, _ = self.parse_xyz_kwargs(x=x, y=y, z=z)
-        _, (dx, dy) = self.pop_axis(direction, axis=plot_axis, swap_axes=swap_axes)
+        _, (dx, dy) = self.pop_axis_and_swap(direction, axis=plot_axis, swap_axes=swap_axes)
 
         # conditions to check to determine whether to plot arrow, taking into account the
         # possibility of a custom arrow base
@@ -2297,12 +2358,12 @@ def _plot_arrow(
             )
             center = arrow_base
 
-        _, (dx, dy) = self.pop_axis(direction, axis=plot_axis, swap_axes=swap_axes)
+        _, (dx, dy) = self.pop_axis_and_swap(direction, axis=plot_axis, swap_axes=swap_axes)
         components_in_plane = any(not np.isclose(component, 0) for component in (dx, dy))
 
         # plot if arrow in plotting plane and some non-zero component can be displayed.
         if arrow_intersecting_plane and components_in_plane:
-            _, (x0, y0) = self.pop_axis(center, axis=plot_axis, swap_axes=swap_axes)
+            _, (x0, y0) = self.pop_axis_and_swap(center, axis=plot_axis, swap_axes=swap_axes)
 
             # Reasonable value for temporary arrow size.  The correct size and direction
             # have to be calculated after all transforms have been set.  That is why we

tidy3d/components/geometry/mesh.py

@@ -604,7 +604,7 @@ def intersections_plane(
             # permute so normal is aligned with z axis
             # and (y, z), (x, z), resp. (x, y) are aligned with (x, y)
             identity = np.eye(3)
-            permutation = self.unpop_axis(identity[2], identity[0:2], axis=axis, swap_axes=swap_axes)
+            permutation = self.unpop_axis_and_swap(identity[2], identity[0:2], axis=axis, swap_axes=swap_axes)
             mapping[:3, :3] = np.array(permutation).T
 
             section2d, _ = section.to_planar(to_2D=mapping)

tidy3d/components/geometry/polyslab.py

@@ -1035,8 +1035,8 @@ def bounds(self, swap_axes: bool = False) -> Bound:
         zmin, zmax = self.slab_bounds
 
         # rearrange axes
-        coords_min = self.unpop_axis(zmin, (xmin, ymin), axis=self.axis, swap_axes=swap_axes)
-        coords_max = self.unpop_axis(zmax, (xmax, ymax), axis=self.axis, swap_axes=swap_axes)
+        coords_min = self.unpop_axis_and_swap(zmin, (xmin, ymin), axis=self.axis, swap_axes=swap_axes)
+        coords_max = self.unpop_axis_and_swap(zmax, (xmax, ymax), axis=self.axis, swap_axes=swap_axes)
         return (tuple(coords_min), tuple(coords_max))
 
     def _extrusion_length_to_offset_distance(self, extrusion: float) -> float:

tidy3d/components/geometry/primitives.py

@@ -141,7 +141,7 @@ def intersections_plane(
         axis, position = self.parse_xyz_kwargs(x=x, y=y, z=z)
         if not self.intersects_axis_position(axis, position):
             return []
-        z0, (x0, y0) = self.pop_axis(self.center, axis=axis, swap_axes=swap_axes)
+        z0, (x0, y0) = self.pop_axis_and_swap(self.center, axis=axis, swap_axes=swap_axes)
         intersect_dist = self._intersect_dist(position, z0)
         if not intersect_dist:
             return []
@@ -403,7 +403,7 @@ def _do_intersections_tilted_plane(
         """
         import trimesh
 
-        z0, (x0, y0) = self.pop_axis(self.center, self.axis, swap_axes=swap_axes)
+        z0, (x0, y0) = self.pop_axis_and_swap(self.center, self.axis, swap_axes=swap_axes)
         half_length = self.finite_length_axis / 2
 
         z_top = z0 + half_length
@@ -445,7 +445,7 @@ def _do_intersections_tilted_plane(
         x = np.hstack((x_bot, x_top))
         y = np.hstack((y_bot, y_top))
         z = np.hstack((np.full_like(x_bot, z_bot), np.full_like(x_top, z_top)))
-        vertices = np.vstack(self.unpop_axis(z, (x, y), self.axis, swap_axes=swap_axes)).T
+        vertices = np.vstack(self.unpop_axis_and_swap(z, (x, y), self.axis, swap_axes=swap_axes)).T
 
         if x_bot.shape[0] == 1:
             m = 1
@@ -501,8 +501,7 @@ def _intersections_normal(self, z: float, swap_axes: bool = False):
 
         if radius_offset <= 0:
             return []
-
-        _, (x0, y0) = self.pop_axis(static_self.center, axis=self.axis, swap_axes=swap_axes)
+        _, (x0, y0) = self.pop_axis_and_swap(static_self.center, axis=self.axis, swap_axes=swap_axes)
         return [shapely.Point(x0, y0).buffer(radius_offset, quad_segs=_N_SHAPELY_QUAD_SEGS)]
 
     def _intersections_side(self, position, axis, swap_axes: bool = False):
@@ -786,5 +785,5 @@ def _local_to_global_side_cross_section(
             axis=axis,
             swap_axes=swap_axes,
         )
-        _, (x_center, y_center) = self.pop_axis(self.center, axis=axis, swap_axes=swap_axes)
+        _, (x_center, y_center) = self.pop_axis_and_swap(self.center, axis=axis, swap_axes=swap_axes)
         return [x_center + lx_offset, y_center + ly_offset]

tidy3d/components/scene.py

@@ -372,8 +372,8 @@ def _get_plot_lims(
     ) -> tuple[tuple[float, float], tuple[float, float]]:
         # if no hlim and/or vlim given, the bounds will then be the usual pml bounds
         axis, _ = Box.parse_xyz_kwargs(x=x, y=y, z=z)
-        _, (hmin, vmin) = Box.pop_axis(bounds[0], axis=axis, swap_axes=swap_axes)
-        _, (hmax, vmax) = Box.pop_axis(bounds[1], axis=axis, swap_axes=swap_axes)
+        _, (hmin, vmin) = Box.pop_axis_and_swap(bounds[0], axis=axis, swap_axes=swap_axes)
+        _, (hmax, vmax) = Box.pop_axis_and_swap(bounds[1], axis=axis, swap_axes=swap_axes)
 
         # account for unordered limits
         if hlim is None:
@@ -647,8 +647,8 @@ def _get_structures_2dbox(
         """
         # if no hlim and/or vlim given, the bounds will then be the usual pml bounds
         axis, _ = Box.parse_xyz_kwargs(x=x, y=y, z=z)
-        _, (hmin, vmin) = Box.pop_axis(self.bounds[0], axis=axis, swap_axes=swap_axes)
-        _, (hmax, vmax) = Box.pop_axis(self.bounds[1], axis=axis, swap_axes=swap_axes)
+        _, (hmin, vmin) = Box.pop_axis_and_swap(self.bounds[0], axis=axis, swap_axes=swap_axes)
+        _, (hmax, vmax) = Box.pop_axis_and_swap(self.bounds[1], axis=axis, swap_axes=swap_axes)
 
         if hlim is not None:
             (hmin, hmax) = hlim
@@ -662,8 +662,8 @@ def _get_structures_2dbox(
         v_size = (vmax - vmin) or inf
 
         axis, center_normal = Box.parse_xyz_kwargs(x=x, y=y, z=z)
-        center = Box.unpop_axis(center_normal, (h_center, v_center), axis=axis, swap_axes=swap_axes)
-        size = Box.unpop_axis(0.0, (h_size, v_size), axis=axis, swap_axes=swap_axes)
+        center = Box.unpop_axis_and_swap(center_normal, (h_center, v_center), axis=axis, swap_axes=swap_axes)
+        size = Box.unpop_axis_and_swap(0.0, (h_size, v_size), axis=axis, swap_axes=swap_axes)
         plane = Box(center=center, size=size)
 
         medium_shapes = []