Merge pull request #362 from fusion-energy/adding_utils_to_assembly

added paramak.Assembly class that inherits from cq.Assembly and adds convenient methods

src/paramak/assemblies/assembly.py

@@ -0,0 +1,37 @@
+# Creates an assembly class that inherits from cadquery's assembly class
+# and adds a few conveniences methods remove() and names()
+
+import warnings
+import cadquery as cq
+
+
+class Assembly(cq.Assembly):
+    """Nested assembly of Workplane and Shape objects defining their relative positions."""
+
+    elongation=None
+    triangularity=None
+    major_radius=None
+    minor_radius=None
+
+    def remove(self, name: str):
+        new_assembly = Assembly()
+        part_found = False
+        for part in self:
+            if part[1].endswith(f'/{name}'):
+                part_found = True
+            else:
+                new_assembly.add(part[0], name=part[1], color=part[3], loc=part[2])
+        if not part_found:
+            warnings.warn(f'Part with name {name} not found')
+
+        new_assembly.elongation = self.elongation
+        new_assembly.triangularity = self.triangularity
+        new_assembly.major_radius = self.major_radius
+        new_assembly.minor_radius = self.minor_radius
+        return new_assembly
+
+    def names(self):
+        names = []
+        for part in self:
+            names.append(part[1].split('/')[-1])
+        return names

src/paramak/assemblies/spherical_tokamak.py

@@ -1,6 +1,7 @@
 from typing import Optional, Sequence, Tuple, Union
 
 import cadquery as cq
+from .assembly import Assembly
 
 from ..utils import (
     get_plasma_index,
@@ -110,7 +111,7 @@ def spherical_tokamak_from_plasma(
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
     colors: dict = {},
-):
+) -> Assembly:
     """Creates a spherical tokamak fusion reactor from a radial build and plasma parameters.
 
 
@@ -168,7 +169,7 @@ def spherical_tokamak(
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
     colors: dict = {},
-):
+) -> Assembly:
     """  Creates a spherical tokamak fusion reactor from a radial build and vertical build.
 
     Args:
@@ -235,7 +236,7 @@ def spherical_tokamak(
         center_column=blanket_cutting_cylinder,
     )
 
-    my_assembly = cq.Assembly()
+    my_assembly = Assembly()
 
     for i, entry in enumerate(extra_cut_shapes):
 
@@ -286,4 +287,9 @@ def spherical_tokamak(
 
     my_assembly.add(plasma, name="plasma", color=cq.Color(*colors.get("plasma", (0.5,0.5,0.5))))
 
+    my_assembly.elongation = elongation
+    my_assembly.triangularity = triangularity
+    my_assembly.major_radius = major_radius
+    my_assembly.minor_radius = minor_radius
+
     return my_assembly

src/paramak/assemblies/tokamak.py

@@ -1,6 +1,7 @@
-from typing import Optional, Sequence, Tuple, Union
+from typing import Sequence, Tuple
 
 import cadquery as cq
+from .assembly import Assembly
 
 from ..utils import get_plasma_index, LayerType
 from ..workplanes.blanket_from_plasma import blanket_from_plasma
@@ -33,7 +34,7 @@ def create_center_column_shield_cylinders(radial_build, rotation_angle, center_c
 
     number_of_cylinder_layers = count_cylinder_layers(radial_build)
 
-    for index, item in enumerate(radial_build):
+    for _, item in enumerate(radial_build):
         if item[0] == LayerType.PLASMA:
             break
 
@@ -157,7 +158,7 @@ def tokamak_from_plasma(
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
     colors: dict = {}
-):
+) -> Assembly:
     """
     Creates a tokamak fusion reactor from a radial build and plasma parameters.
 
@@ -216,7 +217,7 @@ def tokamak(
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
     colors: dict = {}
-):
+) -> Assembly:
     """
     Creates a tokamak fusion reactor from a radial and vertical build.
 
@@ -272,7 +273,7 @@ def tokamak(
         center_column=inner_radial_build[0],  # blanket_cutting_cylinder,
     )
 
-    my_assembly = cq.Assembly()
+    my_assembly = Assembly()
 
     for i, entry in enumerate(extra_cut_shapes):
         if isinstance(entry, cq.Workplane):
@@ -322,4 +323,9 @@ def tokamak(
 
     my_assembly.add(plasma, name="plasma", color=cq.Color(*colors.get("plasma", (0.5,0.5,0.5))))
 
+    my_assembly.elongation = elongation
+    my_assembly.triangularity = triangularity
+    my_assembly.major_radius = major_radius
+    my_assembly.minor_radius = minor_radius
+
     return my_assembly

tests/test_assemblies/test_assembly.py

@@ -0,0 +1,19 @@
+import cadquery as cq
+from paramak.assemblies.assembly import Assembly
+
+def test_remove_and_names():
+
+    sphere1 = cq.Workplane().moveTo(2, 2).sphere(1)
+    box1 = cq.Workplane().box(1, 1, 1)
+    assembly = Assembly()
+    assembly.add(box1, name="box1", color=cq.Color(0.5, 0.5, 0.5))
+    assembly.add(sphere1, name="sphere")
+
+    assembly2 = assembly.remove('sphere')
+    assembly3 = assembly.remove('box1')
+    assembly4 = assembly.remove('bosdfsdf')
+
+    assert assembly.names() == ['box1', 'sphere']
+    assert assembly2.names() == ['box1']
+    assert assembly3.names() == ['sphere']
+    assert assembly4.names() == ['box1', 'sphere']

tests/test_assemblies/test_spherical_tokamak.py

@@ -133,4 +133,28 @@ def test_colors():
             "layer_4": (0.4, 0.4, 0.8),
             "layer_5": (0.5, 0.5, 0.8),
         },
-    )
+    )
+
+def test_attributes():
+    "passing in the colors dictionary should not raise an error"
+    my_reactor = paramak.spherical_tokamak_from_plasma(
+        radial_build=[
+            (paramak.LayerType.GAP, 10),
+            (paramak.LayerType.SOLID, 50),
+            (paramak.LayerType.SOLID, 15),
+            (paramak.LayerType.GAP, 50),
+            (paramak.LayerType.PLASMA, 300),
+            (paramak.LayerType.GAP, 60),
+            (paramak.LayerType.SOLID, 15),
+            (paramak.LayerType.SOLID, 60),
+            (paramak.LayerType.SOLID, 10),
+        ],
+        elongation=2,
+        triangularity=0.55,
+        rotation_angle=180,
+    )
+
+    assert my_reactor.elongation == 2
+    assert my_reactor.triangularity == 0.55
+    assert my_reactor.major_radius == 275
+    assert my_reactor.minor_radius == 150