fixes polarization arrow S-P switching

tidy3d/components/geometry.py

@@ -875,13 +875,13 @@ def _plot_arrow(  # pylint:disable=too-many-arguments, too-many-locals
         """
 
         plot_axis, _ = self.parse_xyz_kwargs(x=x, y=y, z=z)
-        arrow_axis = [component == 0 for component in direction]
+        arrow_axis = [np.isclose(component, 0) for component in direction]
         arrow_length, arrow_width = self._arrow_dims(ax, length_factor, width_factor)
 
         # conditions to check to determine whether to plot arrow
         arrow_intersecting_plane = len(self.intersections(x=x, y=y, z=z)) > 0
-        arrow_perp_to_screen = arrow_axis.index(0.0) != plot_axis
-        arrow_not_cartesian_axis = arrow_axis.count(0.0) > 1
+        arrow_perp_to_screen = arrow_axis.index(False) != plot_axis
+        arrow_not_cartesian_axis = arrow_axis.count(False) > 1
 
         # plot if arrow in plotting plane and some non-zero component can be displayed.
         if arrow_intersecting_plane and (arrow_not_cartesian_axis or arrow_perp_to_screen):

tidy3d/components/source.py

@@ -500,10 +500,12 @@ def _pol_vector(self) -> Tuple[float, float, float]:
         normal_dir = [0.0, 0.0, 0.0]
         normal_dir[self.injection_axis] = 1.0
         propagation_dir = list(self._dir_vector)
-        pol_vector = np.cross(normal_dir, propagation_dir)
-        if np.all(pol_vector == 0.0):
-            pol_vector = np.array((0, 1, 0)) if self.injection_axis == 0 else np.array((1, 0, 0))
-        return self.rotate_points(pol_vector, propagation_dir, angle=self.pol_angle)
+        pol_vector_s = np.cross(normal_dir, propagation_dir)
+        if np.all(pol_vector_s == 0.0):
+            pol_vector_s = np.array((0, 1, 0)) if self.injection_axis == 0 else np.array((1, 0, 0))
+        pol_vector_p = np.cross(propagation_dir, pol_vector_s)
+        pol_vector_p = np.array(pol_vector_p) / np.linalg.norm(pol_vector_p)
+        return self.rotate_points(pol_vector_p, propagation_dir, angle=self.pol_angle)
 
 
 class ModeSource(DirectionalSource, PlanarSource):