new method: generate_obj (#78)

Author: fandreuz

Diff for: bladex/propeller.py

@@ -7,6 +7,9 @@
 from OCC.Core.BRepAlgoAPI import BRepAlgoAPI_Fuse
 from OCC.Extend.DataExchange import write_stl_file
 from OCC.Display.SimpleGui import init_display
+from smithers.io.obj import ObjHandler, WavefrontOBJ
+from smithers.io.stlhandler import STLHandler
+
 
 class Propeller(object):
     """
@@ -22,26 +25,32 @@ class Propeller(object):
 
     def __init__(self, shaft, blade, n_blades):
         self.shaft_solid = shaft.generate_solid()
+
         blade.apply_transformations(reflect=True)
-        blade_solid = blade.generate_solid(max_deg=2, 
-                                           display=False, 
-                                           errors=None)
+        blade_solid = blade.generate_solid(
+            max_deg=2, display=False, errors=None
+        )
         blades = []
         blades.append(blade_solid)
-        for i in range(n_blades-1):
-            blade.rotate(rad_angle=1.0*2.0*np.pi/float(n_blades))
-            blade_solid = blade.generate_solid(max_deg=2, display=False, errors=None)
+        for i in range(n_blades - 1):
+            blade.rotate(rad_angle=1.0 * 2.0 * np.pi / float(n_blades))
+            blade_solid = blade.generate_solid(
+                max_deg=2, display=False, errors=None
+            )
             blades.append(blade_solid)
         blades_combined = blades[0]
-        for i in range(len(blades)-1):
-            boolean_union = BRepAlgoAPI_Fuse(blades_combined, blades[i+1])
+        for i in range(len(blades) - 1):
+            boolean_union = BRepAlgoAPI_Fuse(blades_combined, blades[i + 1])
             boolean_union.Build()
             if not boolean_union.IsDone():
-                raise RuntimeError('Unsuccessful assembling of blade')
+                raise RuntimeError("Unsuccessful assembling of blade")
             blades_combined = boolean_union.Shape()
+        self.blades_solid = blades_combined
+
         boolean_union = BRepAlgoAPI_Fuse(self.shaft_solid, blades_combined)
         boolean_union.Build()
         result_compound = boolean_union.Shape()
+
         sewer = BRepBuilderAPI_Sewing(1e-2)
         sewer.Add(result_compound)
         sewer.Perform()
@@ -51,9 +60,9 @@ def generate_iges(self, filename):
         """
         Export the .iges CAD for the propeller with shaft.
 
-        :param string filename: path (with the file extension) where to store 
+        :param string filename: path (with the file extension) where to store
             the .iges CAD for the propeller and shaft
-        :raises RuntimeError: if the solid assembling of blades is not 
+        :raises RuntimeError: if the solid assembling of blades is not
             completed successfully
         """
         iges_writer = IGESControl_Writer()
@@ -64,13 +73,111 @@ def generate_stl(self, filename):
         """
         Export the .stl CAD for the propeller with shaft.
 
-        :param string filename: path (with the file extension) where to store 
+        :param string filename: path (with the file extension) where to store
             the .stl CAD for the propeller and shaft
-        :raises RuntimeError: if the solid assembling of blades is not 
+        :raises RuntimeError: if the solid assembling of blades is not
             completed successfully
         """
         write_stl_file(self.sewed_full_body, filename)
 
+    def generate_obj(self, filename, region_selector="by_coords"):
+        """
+        Export the .obj CAD for the propeller with shaft. The resulting
+        file contains two regions: `propellerTip` and `propellerStem`, selected
+        according to the criteria passed in the parameter `region_selector`.
+
+        :param string filename: path (with the file extension).
+        :param string region_selector: Two selectors available:
+
+            * `by_coords`: We compute :math:`x`, the smallest X coordinate of
+                the solid which represents the blades of the propeller. Then all
+                the polygons (belonging to both blades and shaft) composed of
+                points whose X coordinate is greater than :math:`x` are
+                considered to be part of the region `propellerTip`. The rest
+                belongs to `propellerStem`;
+            * `blades_and_stem`: The two regions are simply given by the two
+                solids which are used in :func:`__init__`.
+        :raises RuntimeError: if the solid assembling of blades is not
+            completed successfully
+        """
+
+        # we write the propeller to STL, then re-open it to obtain the points
+        write_stl_file(self.shaft_solid, "/tmp/temp_shaft.stl")
+        shaft = STLHandler.read("/tmp/temp_shaft.stl")
+        write_stl_file(self.blades_solid, "/tmp/temp_blades.stl")
+        blades = STLHandler.read("/tmp/temp_blades.stl")
+
+        obj_instance = WavefrontOBJ()
+
+        # add vertexes. first of all we check for duplicated vertexes
+        all_vertices = np.concatenate(
+            [shaft["points"], blades["points"]], axis=0
+        )
+
+        # unique_mapping maps items in all_vertices to items in unique_vertices
+        unique_vertices, unique_mapping = np.unique(
+            all_vertices, return_inverse=True, axis=0
+        )
+        obj_instance.vertices = unique_vertices
+
+        def cells_to_np(cells):
+            cells = np.asarray(cells)
+            return unique_mapping[cells.flatten()].reshape(-1, cells.shape[1])
+
+        # append a list of cells to obj_instance.polygons, possibly with a
+        # region name
+        def append_cells(cells, region_name=None):
+            if region_name is not None:
+                obj_instance.regions_change_indexes.append(
+                    (
+                        np.asarray(obj_instance.polygons).shape[0],
+                        len(obj_instance.regions),
+                    )
+                )
+                obj_instance.regions.append(region_name)
+
+            if len(obj_instance.polygons) == 0:
+                obj_instance.polygons = np.array(cells_to_np(cells))
+            else:
+                obj_instance.polygons = np.concatenate(
+                    [obj_instance.polygons, cells_to_np(cells)], axis=0
+                )
+
+        shaft_cells = np.asarray(shaft["cells"])
+        # the 0th point in blades if the last+1 point in shaft
+        blades_cells = np.asarray(blades["cells"]) + len(shaft["points"])
+
+        if region_selector == "blades_and_stem":
+            append_cells(blades_cells, region_name="propellerTip")
+            append_cells(shaft_cells, region_name="propellerStem")
+        elif region_selector == "by_coords":
+            minimum_blades_x = np.min(blades["points"][:, 0])
+            tip_boolean_array = shaft["points"][:, 0] >= minimum_blades_x
+            shaft_cells_tip = np.all(
+                tip_boolean_array[shaft_cells.flatten()].reshape(
+                    -1, shaft_cells.shape[1]
+                ),
+                axis=1,
+            )
+
+            append_cells(
+                np.concatenate(
+                    [blades_cells, shaft_cells[shaft_cells_tip]], axis=0
+                ),
+                region_name="propellerTip",
+            )
+            append_cells(
+                shaft_cells[np.logical_not(shaft_cells_tip)],
+                region_name="propellerStem",
+            )
+        else:
+            raise ValueError("This selector is not supported at the moment")
+
+        # this is needed because indexes start at 1
+        obj_instance.polygons += 1
+
+        ObjHandler.write(filename, obj_instance)
+
     def display(self):
         """
         Display the propeller with shaft.

Diff for: tests/test_propeller.py

@@ -4,22 +4,24 @@
 import bladex.profiles as pr
 import bladex.blade as bl
 from bladex import NacaProfile, Shaft, Propeller
+from smithers.io.obj import ObjHandler
+from smithers.io.stlhandler import STLHandler
 
 
 def create_sample_blade_NACApptc():
     sections = np.asarray([NacaProfile('5407') for i in range(13)])
-    radii=np.array([0.034375, 0.0375, 0.04375, 0.05, 0.0625, 0.075, 0.0875, 
+    radii=np.array([0.034375, 0.0375, 0.04375, 0.05, 0.0625, 0.075, 0.0875,
                     0.1, 0.10625, 0.1125, 0.11875, 0.121875, 0.125])
-    chord_lengths = np.array([0.039, 0.045, 0.05625, 0.06542, 0.08125, 
-                              0.09417, 0.10417, 0.10708, 0.10654, 0.10417, 
+    chord_lengths = np.array([0.039, 0.045, 0.05625, 0.06542, 0.08125,
+                              0.09417, 0.10417, 0.10708, 0.10654, 0.10417,
                               0.09417, 0.07867, 0.025])
-    pitch = np.array([0.35, 0.35, 0.36375, 0.37625, 0.3945, 0.405, 0.40875, 
+    pitch = np.array([0.35, 0.35, 0.36375, 0.37625, 0.3945, 0.405, 0.40875,
                       0.4035, 0.3955, 0.38275, 0.3645, 0.35275, 0.33875])
-    rake=np.array([0.0 ,0.0, 0.0005, 0.00125, 0.00335, 0.005875, 0.0075, 
+    rake=np.array([0.0 ,0.0, 0.0005, 0.00125, 0.00335, 0.005875, 0.0075,
                    0.007375, 0.006625, 0.00545, 0.004033, 0.0033, 0.0025])
-    skew_angles=np.array([6.6262795, 3.6262795, -1.188323, -4.4654502, 
-                          -7.440779, -7.3840979, -5.0367916, -1.3257914, 
-                          1.0856404, 4.1448947, 7.697235, 9.5368917, 
+    skew_angles=np.array([6.6262795, 3.6262795, -1.188323, -4.4654502,
+                          -7.440779, -7.3840979, -5.0367916, -1.3257914,
+                          1.0856404, 4.1448947, 7.697235, 9.5368917,
                           11.397609])
     return bl.Blade(
         sections=sections,
@@ -38,54 +40,54 @@ class TestPropeller(TestCase):
     def test_sections_inheritance_NACApptc(self):
         prop= create_sample_blade_NACApptc()
         self.assertIsInstance(prop.sections[0], pr.NacaProfile)
-        
+
     def test_radii_NACApptc(self):
         prop = create_sample_blade_NACApptc()
-        np.testing.assert_equal(prop.radii, np.array([0.034375, 0.0375, 0.04375, 
-                                                      0.05, 0.0625, 0.075, 
-                                                      0.0875, 0.1, 0.10625, 
-                                                      0.1125, 0.11875, 0.121875, 
+        np.testing.assert_equal(prop.radii, np.array([0.034375, 0.0375, 0.04375,
+                                                      0.05, 0.0625, 0.075,
+                                                      0.0875, 0.1, 0.10625,
+                                                      0.1125, 0.11875, 0.121875,
                                                       0.125]))
 
     def test_chord_NACApptc(self):
         prop = create_sample_blade_NACApptc()
-        np.testing.assert_equal(prop.chord_lengths,np.array([0.039, 0.045, 
-                                                             0.05625, 0.06542, 
-                                                             0.08125, 0.09417, 
-                                                             0.10417, 0.10708, 
-                                                             0.10654, 0.10417, 
-                                                             0.09417, 0.07867, 
+        np.testing.assert_equal(prop.chord_lengths,np.array([0.039, 0.045,
+                                                             0.05625, 0.06542,
+                                                             0.08125, 0.09417,
+                                                             0.10417, 0.10708,
+                                                             0.10654, 0.10417,
+                                                             0.09417, 0.07867,
                                                              0.025]))
 
     def test_pitch_NACApptc(self):
         prop = create_sample_blade_NACApptc()
-        np.testing.assert_equal(prop.pitch, np.array([0.35, 0.35, 0.36375, 
-                                                      0.37625, 0.3945, 0.405, 
-                                                      0.40875, 0.4035, 0.3955, 
-                                                      0.38275, 0.3645, 0.35275, 
+        np.testing.assert_equal(prop.pitch, np.array([0.35, 0.35, 0.36375,
+                                                      0.37625, 0.3945, 0.405,
+                                                      0.40875, 0.4035, 0.3955,
+                                                      0.38275, 0.3645, 0.35275,
                                                       0.33875]))
 
     def test_rake_NACApptc(self):
         prop = create_sample_blade_NACApptc()
-        np.testing.assert_equal(prop.rake, np.array([0.0 ,0.0, 0.0005, 0.00125, 
-                                                     0.00335, 0.005875, 0.0075, 
-                                                     0.007375, 0.006625, 0.00545, 
+        np.testing.assert_equal(prop.rake, np.array([0.0 ,0.0, 0.0005, 0.00125,
+                                                     0.00335, 0.005875, 0.0075,
+                                                     0.007375, 0.006625, 0.00545,
                                                      0.004033, 0.0033, 0.0025]))
 
     def test_skew_NACApptc(self):
         prop = create_sample_blade_NACApptc()
-        np.testing.assert_equal(prop.skew_angles, np.array([6.6262795, 
-                                                            3.6262795, 
-                                                            -1.188323, 
-                                                            -4.4654502, 
-                                                            -7.440779, 
-                                                            -7.3840979, 
-                                                            -5.0367916, 
-                                                            -1.3257914, 
-                                                            1.0856404, 
-                                                            4.1448947, 
-                                                            7.697235, 
-                                                            9.5368917, 
+        np.testing.assert_equal(prop.skew_angles, np.array([6.6262795,
+                                                            3.6262795,
+                                                            -1.188323,
+                                                            -4.4654502,
+                                                            -7.440779,
+                                                            -7.3840979,
+                                                            -5.0367916,
+                                                            -1.3257914,
+                                                            1.0856404,
+                                                            4.1448947,
+                                                            7.697235,
+                                                            9.5368917,
                                                             11.397609]))
 
     def test_sections_array_different_length(self):
@@ -156,5 +158,63 @@ def test_generate_stl(self):
         self.assertTrue(os.path.isfile('tests/test_datasets/propeller_and_shaft.stl'))
         self.addCleanup(os.remove, 'tests/test_datasets/propeller_and_shaft.stl')
 
+    def test_generate_obj_by_coords(self):
+        sh = Shaft("tests/test_datasets/shaft.iges")
+        prop = create_sample_blade_NACApptc()
+        prop = Propeller(sh, prop, 4)
+        prop.generate_obj("tests/test_datasets/propeller_and_shaft.obj", region_selector='by_coords')
+
+        data = ObjHandler.read('tests/test_datasets/propeller_and_shaft.obj')
+        assert data.regions == ['propellerTip','propellerStem']
+
+        # we want 0 to be the first index
+        data.polygons = np.asarray(data.polygons) - 1
+
+        tip_poly = data.polygons[:data.regions_change_indexes[1][0]]
+        stem_poly = data.polygons[data.regions_change_indexes[1][0]:]
+
+        blades_stl = STLHandler.read('/tmp/temp_blades.stl')
+        shaft_stl = STLHandler.read('/tmp/temp_shaft.stl')
+
+        # same vertices
+        all_vertices = np.concatenate(
+            [shaft_stl["points"], blades_stl["points"]], axis=0
+        )
+        unique_vertices = np.unique(all_vertices, axis=0)
+        np.testing.assert_almost_equal(data.vertices, unique_vertices, decimal=3)
+
+        blades_min_x = np.min(blades_stl['points'][:,0])
+
+        assert np.all(data.vertices[np.asarray(tip_poly).flatten()][:,0] >= blades_min_x)
+        assert not any(np.all(data.vertices[np.asarray(stem_poly).flatten()][:,0].reshape(-1,data.polygons.shape[1]) >= blades_min_x, axis=1))
+
+    def test_generate_obj_blades_and_stem(self):
+        sh = Shaft("tests/test_datasets/shaft.iges")
+        prop = create_sample_blade_NACApptc()
+        prop = Propeller(sh, prop, 4)
+        prop.generate_obj("tests/test_datasets/propeller_and_shaft.obj", region_selector='blades_and_stem')
+
+        data = ObjHandler.read('tests/test_datasets/propeller_and_shaft.obj')
+        assert data.regions == ['propellerTip','propellerStem']
+
+        tip_polygons = np.asarray(data.polygons[:data.regions_change_indexes[1][0]]) - 1
+        stem_polygons = np.asarray(data.polygons[data.regions_change_indexes[1][0]:]) - 1
+
+        blades_stl = STLHandler.read('/tmp/temp_blades.stl')
+        shaft_stl = STLHandler.read('/tmp/temp_shaft.stl')
+
+        # same vertices
+        all_vertices = np.concatenate(
+            [shaft_stl["points"], blades_stl["points"]], axis=0
+        )
+
+        unique_vertices, indexing = np.unique(
+            all_vertices, return_index=True, axis=0
+        )
+        np.testing.assert_almost_equal(data.vertices, unique_vertices, decimal=3)
+
+        assert np.all(indexing[stem_polygons.flatten()] < shaft_stl['points'].shape[0])
+        assert np.all(indexing[tip_polygons.flatten()] >= shaft_stl['points'].shape[0])
+
     def test_isdisplay(self):
         assert hasattr(Propeller, "display") == True