From 7c302ae363d378e12f9fb9273b0a49b24f278650 Mon Sep 17 00:00:00 2001
From: Chris Mackey <chris@ladybug.tools>
Date: Wed, 16 Oct 2024 16:40:09 -0700
Subject: [PATCH] fix(roof): Add an option for selected_indices in
 Roof.snap_to_grid

---
 dragonfly/roof.py   | 28 +++++++++++++++++-----------
 tests/model_test.py |  5 +++++
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/dragonfly/roof.py b/dragonfly/roof.py
index 5c4c2666..fbbcf090 100644
--- a/dragonfly/roof.py
+++ b/dragonfly/roof.py
@@ -483,7 +483,7 @@ def update_geometry_2d(self, new_polygon_2d, polygon_index):
         new_face_3d = Face3D(roof_verts, plane=roof_plane)
         self.update_geometry_3d(new_face_3d, polygon_index)
 
-    def snap_to_grid(self, grid_increment, tolerance=0.01):
+    def snap_to_grid(self, grid_increment, selected_indices=None, tolerance=0.01):
         """Snap naked roof vertices to the nearest grid node defined by an increment.
 
         This is useful for coordinating the Roof specification with the grid snapping
@@ -499,6 +499,9 @@ def snap_to_grid(self, grid_increment, tolerance=0.01):
                 typically should be equal to the tolerance or larger but should
                 not be larger than the smallest detail of the Room2D that you
                 wish to resolve.
+            selected_indices: An optional list of indices for specific roof
+                geometries to be snapped to the grid. If None, all of the roof
+                geometry will be snapped. (Default: None).
             tolerance: The minimum distance between vertices below which they are
                 considered co-located. (Default: 0.01,
                 suitable for objects in meters).
@@ -509,16 +512,19 @@ def snap_to_grid(self, grid_increment, tolerance=0.01):
 
         # loop through the polygons and snap the vertices
         new_polygons = []
-        for poly, poly_info in zip(polygons, poly_ridge_info):
-            new_poly = []
-            for pt, pt_info in zip(poly, poly_info):
-                if len(pt_info) == 0:  # not on a ridge line; move it anywhere
-                    new_x = grid_increment * round(pt.x / grid_increment)
-                    new_y = grid_increment * round(pt.y / grid_increment)
-                    new_poly.append(Point2D(new_x, new_y, pt.z))
-                else:  # on a ridge line; don't move that point!
-                    new_poly.append(pt)
-            new_polygons.append(new_poly)
+        for i, (poly, poly_info) in enumerate(zip(polygons, poly_ridge_info)):
+            if selected_indices is None or i in selected_indices:
+                new_poly = []
+                for pt, pt_info in zip(poly, poly_info):
+                    if len(pt_info) == 0:  # not on a ridge line; move it anywhere
+                        new_x = grid_increment * round(pt.x / grid_increment)
+                        new_y = grid_increment * round(pt.y / grid_increment)
+                        new_poly.append(Point2D(new_x, new_y))
+                    else:  # on a ridge line; don't move that point!
+                        new_poly.append(pt)
+                new_polygons.append(new_poly)
+            else:
+                new_polygons.append(poly)
 
         # project the points back onto the roof
         proj_dir = Vector3D(0, 0, 1)
diff --git a/tests/model_test.py b/tests/model_test.py
index 5d627c7f..a749b3e6 100644
--- a/tests/model_test.py
+++ b/tests/model_test.py
@@ -680,6 +680,11 @@ def test_building_assigned_roof():
     hb_models = model.to_honeybee('District', None, False, tolerance=0.01)
     assert len(hb_models) == 1
 
+    stories[0].roof.snap_to_grid(0.1)
+    stories[1].roof.snap_to_grid(0.1)
+    assert len(stories[0].roof) == 1
+    assert len(stories[1].roof) == 2
+
 
 def test_building_assigned_roofs_2():
     """Test another serialization of a model with roofs assigned to the building."""