improve test coverage and simplify flatten_shapely_geometry

Author: dmarek-flex

tests/test_components/test_geometry.py

@@ -12,6 +12,15 @@
 import pytest
 import shapely
 import trimesh
+from shapely.geometry import (
+    GeometryCollection,
+    LineString,
+    MultiLineString,
+    MultiPoint,
+    MultiPolygon,
+    Point,
+    Polygon,
+)
 
 import tidy3d as td
 from tidy3d.compat import _shapely_is_older_than
@@ -22,6 +31,7 @@
     SnapLocation,
     SnappingSpec,
     flatten_groups,
+    flatten_shapely_geometries,
     snap_box_to_grid,
     traverse_geometries,
 )
@@ -1137,7 +1147,14 @@ def test_subdivide():
 @pytest.mark.parametrize("snap_location", [SnapLocation.Boundary, SnapLocation.Center])
 @pytest.mark.parametrize(
     "snap_behavior",
-    [SnapBehavior.Off, SnapBehavior.Closest, SnapBehavior.Expand, SnapBehavior.Contract],
+    [
+        SnapBehavior.Off,
+        SnapBehavior.Closest,
+        SnapBehavior.Expand,
+        SnapBehavior.Contract,
+        SnapBehavior.StrictExpand,
+        SnapBehavior.StrictContract,
+    ],
 )
 def test_snap_box_to_grid(snap_location, snap_behavior):
     """ "Test that all combinations of SnappingSpec correctly modify a test box without error."""
@@ -1158,12 +1175,78 @@ def test_snap_box_to_grid(snap_location, snap_behavior):
     new_box = snap_box_to_grid(grid, box, snap_spec)
 
     if snap_behavior != SnapBehavior.Off and snap_location == SnapLocation.Boundary:
-        # Check that the box boundary slightly off from 0.1 was correctly snapped to 0.1
-        assert math.isclose(new_box.bounds[0][1], xyz[1])
-        # Check that the box boundary slightly off from 0.3 was correctly snapped to 0.3
-        assert math.isclose(new_box.bounds[1][1], xyz[3])
-        # Check that the box boundary outside the grid was snapped to the smallest grid coordinate
-        assert math.isclose(new_box.bounds[0][2], xyz[0])
+        # Strict behaviors have different snapping rules, so skip these specific assertions
+        if snap_behavior not in (SnapBehavior.StrictExpand, SnapBehavior.StrictContract):
+            # Check that the box boundary slightly off from 0.1 was correctly snapped to 0.1
+            assert math.isclose(new_box.bounds[0][1], xyz[1])
+            # Check that the box boundary slightly off from 0.3 was correctly snapped to 0.3
+            assert math.isclose(new_box.bounds[1][1], xyz[3])
+            # Check that the box boundary outside the grid was snapped to the smallest grid coordinate
+            assert math.isclose(new_box.bounds[0][2], xyz[0])
+
+
+def test_snap_box_to_grid_strict_behaviors():
+    """Test StrictExpand and StrictContract behaviors specifically."""
+    xyz = np.linspace(0, 1, 11)  # Grid points at 0.0, 0.1, 0.2, ..., 1.0
+    coords = td.Coords(x=xyz, y=xyz, z=xyz)
+    grid = td.Grid(boundaries=coords)
+
+    # Test StrictExpand: should always move endpoints outwards, even if coincident
+    box_coincident = td.Box(
+        center=(0.1, 0.2, 0.3), size=(0, 0, 0)
+    )  # Centered exactly on grid points
+    snap_spec_strict_expand = SnappingSpec(
+        location=[SnapLocation.Boundary] * 3, behavior=[SnapBehavior.StrictExpand] * 3
+    )
+
+    expanded_box = snap_box_to_grid(grid, box_coincident, snap_spec_strict_expand)
+
+    # StrictExpand should move bounds outwards even when already on grid
+    assert expanded_box.bounds[0][0] < 0.1  # Left bound moved left from 0.1
+    assert expanded_box.bounds[1][0] > 0.1  # Right bound moved right from 0.1
+    assert expanded_box.bounds[0][1] < 0.2  # Bottom bound moved down from 0.2
+    assert expanded_box.bounds[1][1] > 0.2  # Top bound moved up from 0.2
+
+    # Test StrictContract: should always move endpoints inwards, even if coincident
+    box_large = td.Box(center=(0.5, 0.5, 0.5), size=(0.4, 0.4, 0.4))  # Spans multiple grid cells
+    snap_spec_strict_contract = SnappingSpec(
+        location=[SnapLocation.Boundary] * 3, behavior=[SnapBehavior.StrictContract] * 3
+    )
+
+    contracted_box = snap_box_to_grid(grid, box_large, snap_spec_strict_contract)
+
+    # StrictContract should make the box smaller than the original
+    assert contracted_box.size[0] < box_large.size[0]
+    assert contracted_box.size[1] < box_large.size[1]
+    assert contracted_box.size[2] < box_large.size[2]
+
+    # Test edge case: box coincident with grid boundaries
+    box_on_grid = td.Box(
+        center=(0.15, 0.25, 0.35), size=(0.1, 0.1, 0.1)
+    )  # Boundaries at 0.1,0.2 and 0.2,0.3
+
+    # Regular Expand shouldn't change a box already coincident with grid
+    snap_spec_regular_expand = SnappingSpec(
+        location=[SnapLocation.Boundary] * 3, behavior=[SnapBehavior.Expand] * 3
+    )
+    regular_expanded = snap_box_to_grid(grid, box_on_grid, snap_spec_regular_expand)
+    assert np.allclose(regular_expanded.bounds, box_on_grid.bounds)  # Should be unchanged
+
+    # StrictExpand should still expand even when coincident
+    strict_expanded = snap_box_to_grid(grid, box_on_grid, snap_spec_strict_expand)
+    assert not np.allclose(strict_expanded.bounds, box_on_grid.bounds)  # Should be changed
+    assert strict_expanded.size[0] > box_on_grid.size[0]  # Should be larger
+
+    # Test with margin parameter for strict behaviors
+    snap_spec_strict_expand_margin = SnappingSpec(
+        location=[SnapLocation.Boundary] * 3,
+        behavior=[SnapBehavior.StrictExpand] * 3,
+        margin=(1, 1, 1),  # Consider 1 additional grid point when expanding
+    )
+
+    margin_expanded = snap_box_to_grid(grid, box_coincident, snap_spec_strict_expand_margin)
+    # With margin=1, should expand even further than without margin
+    assert margin_expanded.size[0] >= expanded_box.size[0]
 
 
 def test_triangulation_with_collinear_vertices():
@@ -1431,3 +1514,105 @@ def test_trim_dims_and_bounds_edge():
             assert np.all(np.array(expected_trimmed_bounds) == np.array(trimmed_bounds)), (
                 "Unexpected trimmed bounds"
             )
+
+
+def test_flatten_shapely_geometries():
+    """Test the flatten_shapely_geometries utility function comprehensively."""
+    # Test 1: Single polygon (should be wrapped in list and returned)
+    single_polygon = Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])
+    result = flatten_shapely_geometries(single_polygon)
+    assert len(result) == 1
+    assert result[0] == single_polygon
+
+    # Test 2: List of polygons (should return as-is)
+    poly1 = Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])
+    poly2 = Polygon([(2, 0), (3, 0), (3, 1), (2, 1)])
+    polygon_list = [poly1, poly2]
+    result = flatten_shapely_geometries(polygon_list)
+    assert len(result) == 2
+    assert result == polygon_list
+
+    # Test 3: MultiPolygon (should be flattened)
+    multi_polygon = MultiPolygon([poly1, poly2])
+    result = flatten_shapely_geometries(multi_polygon)
+    assert len(result) == 2
+    assert result[0] == poly1
+    assert result[1] == poly2
+
+    # Test 4: Empty geometries (should be filtered out)
+    empty_polygon = Polygon()
+    mixed_list = [poly1, empty_polygon, poly2]
+    result = flatten_shapely_geometries(mixed_list)
+    assert len(result) == 2
+    assert empty_polygon not in result
+
+    # Test 5: GeometryCollection (should be recursively flattened)
+    line = LineString([(0, 0), (1, 1)])
+    point = Point(0, 0)
+    collection = GeometryCollection([poly1, line, point, poly2])
+    result = flatten_shapely_geometries(collection)
+    assert len(result) == 2  # Only polygons kept by default
+    assert poly1 in result
+    assert poly2 in result
+
+    # Test 6: Custom keep_types parameter
+    result_with_lines = flatten_shapely_geometries(collection, keep_types=(Polygon, LineString))
+    assert len(result_with_lines) == 3  # 2 polygons + 1 line
+    assert poly1 in result_with_lines
+    assert poly2 in result_with_lines
+    assert line in result_with_lines
+
+    # Test 7: Nested collections and multi-geometries
+    line1 = LineString([(0, 0), (1, 1)])
+    line2 = LineString([(2, 2), (3, 3)])
+    multi_line = MultiLineString([line1, line2])
+    nested_collection = GeometryCollection(
+        [
+            collection,  # Contains poly1, line, point, poly2
+            multi_line,
+            poly1,
+        ]
+    )
+    result = flatten_shapely_geometries(nested_collection)
+    assert len(result) == 3  # poly1 (from collection), poly2 (from collection), poly1 (direct)
+
+    # Test 8: MultiPoint (should be handled)
+    point1 = Point(0, 0)
+    point2 = Point(1, 1)
+    multi_point = MultiPoint([point1, point2])
+    result = flatten_shapely_geometries(multi_point, keep_types=(Point,))
+    assert len(result) == 2
+    assert point1 in result
+    assert point2 in result
+
+    # Test 9: MultiLineString (should be handled)
+    result = flatten_shapely_geometries(multi_line, keep_types=(LineString,))
+    assert len(result) == 2
+    assert line1 in result
+    assert line2 in result
+
+    # Test 10: Mixed empty and non-empty geometries
+    empty_multi = MultiPolygon([])
+    mixed_with_empty = [poly1, empty_multi, empty_polygon, poly2]
+    result = flatten_shapely_geometries(mixed_with_empty)
+    assert len(result) == 2
+    assert poly1 in result
+    assert poly2 in result
+
+    # Test 11: Deeply nested structure
+    inner_collection = GeometryCollection([poly1, line])
+    outer_multi = MultiPolygon([poly2])
+    deep_collection = GeometryCollection([inner_collection, outer_multi])
+    result = flatten_shapely_geometries(deep_collection)
+    assert len(result) == 2
+    assert poly1 in result
+    assert poly2 in result
+
+    # Test 12: All geometry types filtered out
+    points_and_lines = GeometryCollection([Point(0, 0), LineString([(0, 0), (1, 1)])])
+    result = flatten_shapely_geometries(points_and_lines)  # Default keeps only Polygons
+    assert len(result) == 0
+
+    # Test 13: Edge case - single empty geometry
+    result = flatten_shapely_geometries(empty_polygon)
+    assert len(result) == 0

tests/test_components/test_microwave.py

@@ -326,8 +326,8 @@ def test_antenna_parameters():
         antenna_params.partial_realized_gain("invalid")
 
 
-def test_path_integral_plotting():
-    """Test that all types of path integrals correctly plot themselves."""
+def test_path_spec_plotting():
+    """Test that all types of path specification correctly plot themselves."""
 
     mean_radius = (COAX_R2 + COAX_R1) * 0.5
     size = [COAX_R2 - COAX_R1, 0, 0]
@@ -372,6 +372,28 @@ def test_path_integral_plotting():
     current_integral.plot(y=2, ax=ax)
     plt.close()
 
+    current_integral = td.CompositeCurrentIntegralSpec(
+        path_specs=(
+            td.CurrentIntegralAxisAlignedSpec(center=(-1, -1, 0), size=(1, 1, 0), sign="-"),
+            td.CurrentIntegralAxisAlignedSpec(center=(1, 1, 0), size=(1, 1, 0), sign="-"),
+        ),
+        sum_spec="sum",
+    )
+    ax = current_integral.plot(z=0)
+    plt.close()
+
+
+@pytest.mark.parametrize("clockwise", [False, True])
+def test_custom_current_specification_sign(clockwise):
+    """Make sure the sign is correctly calculated for custom current specs."""
+    current_integral = td.CustomCurrentIntegral2DSpec.from_circular_path(
+        center=(0, 0, 0), radius=0.4, num_points=31, normal_axis=2, clockwise=clockwise
+    )
+    if clockwise:
+        assert current_integral.sign == "-"
+    else:
+        assert current_integral.sign == "+"
+
 
 def test_composite_current_integral_validation():
     """Ensures that the CompositeCurrentIntegralSpec is validated correctly."""
@@ -704,6 +726,6 @@ def test_mode_solver_with_microwave_mode_spec():
     # mms_data.field_components["Ez"].isel(mode_index=0, f=0).real.plot(ax=ax)
     # ax.set_aspect("equal")
     # plt.show()
-    # print(mms_data.to_dataframe())
+    mms_data.to_dataframe()
 
     assert np.all(np.isclose(mms_data.microwave_data.Z0.real, 28.6, 0.2))

tidy3d/components/geometry/utils.py

@@ -12,13 +12,11 @@
 import pydantic.v1 as pydantic
 import shapely
 from shapely.geometry import (
-    GeometryCollection,
-    MultiLineString,
-    MultiPoint,
-    MultiPolygon,
     Polygon,
 )
-from shapely.geometry.base import BaseGeometry
+from shapely.geometry.base import (
+    BaseMultipartGeometry,
+)
 
 from tidy3d.components.autograd.utils import get_static
 from tidy3d.components.base import Tidy3dBaseModel
@@ -85,9 +83,7 @@ def flatten_shapely_geometries(
             continue
         if isinstance(geom, keep_types):
             flat.append(geom)
-        elif isinstance(geom, (MultiPolygon, MultiLineString, MultiPoint, GeometryCollection)):
-            flat.extend(flatten_shapely_geometries(geom.geoms, keep_types))
-        elif isinstance(geom, BaseGeometry) and hasattr(geom, "geoms"):
+        elif isinstance(geom, BaseMultipartGeometry):
             flat.extend(flatten_shapely_geometries(geom.geoms, keep_types))
     return flat
 
@@ -114,7 +110,7 @@ def merging_geometries_on_plane(
 
     Returns
     -------
-    List[Tuple[Any, shapely]]
+    List[Tuple[Any, Shapely]]
         List of shapes and their property value on the plane after merging.
     """
 

tidy3d/components/microwave/path_integrals/base_spec.py

@@ -16,7 +16,7 @@
 from tidy3d.components.types.base import Axis, Direction
 from tidy3d.components.validators import assert_line
 from tidy3d.constants import MICROMETER, fp_eps
-from tidy3d.exceptions import SetupError, Tidy3dError
+from tidy3d.exceptions import SetupError
 from tidy3d.log import log
 
 
@@ -88,7 +88,6 @@ def main_axis(self) -> Axis:
         for index, value in enumerate(self.size):
             if value != 0:
                 return index
-        raise Tidy3dError("Failed to identify axis.")
 
     def _vertices_2D(self, axis: Axis) -> tuple[Coordinate2D, Coordinate2D]:
         """Returns the two vertices of this path in the plane defined by ``axis``."""

tidy3d/components/microwave/path_integrals/current_spec.py

@@ -20,7 +20,7 @@
 from tidy3d.components.validators import assert_plane
 from tidy3d.components.viz import add_ax_if_none
 from tidy3d.constants import fp_eps
-from tidy3d.exceptions import SetupError, Tidy3dError
+from tidy3d.exceptions import SetupError
 
 
 class CurrentIntegralAxisAlignedSpec(AbstractAxesRH, Box):
@@ -52,7 +52,6 @@ def main_axis(self) -> Axis:
         for index, value in enumerate(self.size):
             if value == 0:
                 return index
-        raise Tidy3dError("Failed to identify axis.")
 
     def _to_path_integral_specs(
         self, h_horizontal=None, h_vertical=None

tidy3d/components/mode/mode_solver.py

@@ -1383,8 +1383,6 @@ def _make_path_integrals(
     def _add_microwave_data(self, mode_solver_data: ModeSolverData) -> ModeSolverData:
         """Calculate and add microwave data to ``mode_solver_data`` which uses the path specifications.
         Note: this modifies ``mode_solver_data`` in-place."""
-        if self.mode_spec.microwave_mode_spec is None:
-            return mode_solver_data
         voltage_integrals, current_integrals = self._make_path_integrals()
         # Need to operate on the full symmetry expanded fields
         mode_solver_data_expanded = mode_solver_data.symmetry_expanded_copy