Merge pull request #356 from fusion-energy/adding-color-option

added option for assembly part colors

.gitignore

@@ -175,3 +175,4 @@ src/_version.py
 *.stp
 *.sql
 *.html
+*.png

examples/spherical_tokamak_from_plasma_with_color.py

@@ -0,0 +1,72 @@
+from pathlib import Path
+
+from example_util_functions import transport_particles_on_h5m_geometry
+
+import paramak
+
+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,
+    colors={
+        "layer_1": (0.4, 0.9, 0.4),
+        "layer_2": (0.6, 0.8, 0.6),
+        "plasma": (1., 0.7, 0.8, 0.6),
+        "layer_3": (0.1, 0.1, 0.9),
+        "layer_4": (0.4, 0.4, 0.8),
+        "layer_5": (0.5, 0.5, 0.8),
+    },
+)
+my_reactor.save(f"spherical_tokamak_from_plasma_with_colors.step")
+
+# show colors with inbuild vtk viewer
+# from cadquery.vis import show
+# show(my_reactor)
+
+# cadquery also supports svg export
+# currently needs converting to compound first as svg export not supported by assembly objects
+# lots of options https://cadquery.readthedocs.io/en/latest/importexport.html#exporting-svg
+my_reactor.toCompound().export("spherical_tokamak_from_plasma_with_colors.svg")
+
+# show colors with png file export
+# first install plugin with
+# pip install git+https://github.com/jmwright/cadquery-png-plugin
+import cadquery_png_plugin.plugin
+# lots of options
+# https://github.com/jmwright/cadquery-png-plugin/blob/d2dd6e8a51b7e165ee80240a701c5b434dfe0733/cadquery_png_plugin/plugin.py#L276-L298
+my_reactor.exportPNG(
+    options={
+        "width":1280,
+        "height":1024,
+        "zoom":1.4,
+    },
+    file_path='spherical_tokamak_from_plasma_with_colors.png'
+)
+
+# from cad_to_dagmc import CadToDagmc
+# my_model = CadToDagmc()
+# material_tags = [
+#     "mat1"
+# ] * 6
+# my_model.add_cadquery_object(cadquery_object=my_reactor, material_tags=material_tags)
+# my_model.export_dagmc_h5m_file(min_mesh_size=3.0, max_mesh_size=20.0)
+
+# h5m_filename = "dagmc.h5m"
+# flux = transport_particles_on_h5m_geometry(
+#     h5m_filename=h5m_filename,
+#     material_tags=material_tags,
+#     nuclides=["H1"] * len(material_tags),
+#     cross_sections_xml="tests/cross_sections.xml",
+# )
+# assert flux > 0.0

examples/tokamak_from_plasma_with_color.py

@@ -0,0 +1,73 @@
+from example_util_functions import transport_particles_on_h5m_geometry
+
+import paramak
+
+my_reactor = paramak.tokamak_from_plasma(
+    radial_build=[
+        (paramak.LayerType.GAP, 10),
+        (paramak.LayerType.SOLID, 30),
+        (paramak.LayerType.SOLID, 50),
+        (paramak.LayerType.SOLID, 10),
+        (paramak.LayerType.SOLID, 120),
+        (paramak.LayerType.SOLID, 20),
+        (paramak.LayerType.GAP, 60),
+        (paramak.LayerType.PLASMA, 300),
+        (paramak.LayerType.GAP, 60),
+        (paramak.LayerType.SOLID, 20),
+        (paramak.LayerType.SOLID, 120),
+        (paramak.LayerType.SOLID, 10),
+    ],
+    elongation=2,
+    triangularity=0.55,
+    rotation_angle=180,
+    colors={
+        "layer_1": (0.4, 0.9, 0.4),
+        "layer_2": (0.6, 0.8, 0.6),
+        "plasma": (1., 0.7, 0.8, 0.6),
+        "layer_3": (0.1, 0.1, 0.9),
+        "layer_4": (0.4, 0.4, 0.8),
+        "layer_5": (0.5, 0.5, 0.8),
+    }
+)
+my_reactor.save(f"tokamak_with_colors.step")
+print(f"Saved as tokamak_with_colors.step")
+
+
+# show colors with inbuild vtk viewer
+# from cadquery.vis import show
+# show(my_reactor)
+
+# cadquery also supports svg export
+# currently needs converting to compound first as svg export not supported by assembly objects
+# lots of options https://cadquery.readthedocs.io/en/latest/importexport.html#exporting-svg
+# my_reactor.toCompound().export("tokamak_from_plasma_with_colors.svg")
+
+# show colors with png file export
+# first install plugin with
+# pip install git+https://github.com/jmwright/cadquery-png-plugin
+import cadquery_png_plugin.plugin
+# lots of options
+# https://github.com/jmwright/cadquery-png-plugin/blob/d2dd6e8a51b7e165ee80240a701c5b434dfe0733/cadquery_png_plugin/plugin.py#L276-L298
+my_reactor.exportPNG(
+    options={
+        "width":1280,
+        "height":1024,
+        "zoom":1.4,
+    },
+    file_path='tokamak_from_plasma_with_colors.png'
+)
+
+# from cad_to_dagmc import CadToDagmc
+# my_model = CadToDagmc()
+# material_tags = ["mat1"] * 6  # as inner and outer layers are one solid there are only 6 solids in model
+# my_model.add_cadquery_object(cadquery_object=my_reactor, material_tags=material_tags)
+# my_model.export_dagmc_h5m_file(min_mesh_size=3.0, max_mesh_size=20.0)
+
+# h5m_filename = "dagmc.h5m"
+# flux = transport_particles_on_h5m_geometry(
+#     h5m_filename=h5m_filename,
+#     material_tags=material_tags,
+#     nuclides=["H1"] * len(material_tags),
+#     cross_sections_xml="tests/cross_sections.xml",
+# )
+# assert flux > 0.0

src/paramak/assemblies/spherical_tokamak.py

@@ -109,15 +109,23 @@ def spherical_tokamak_from_plasma(
     rotation_angle: float = 180.0,
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
+    colors: dict = {},
 ):
-    """_summary_
+    """Creates a spherical tokamak fusion reactor from a radial build and plasma parameters.
+
 
     Args:
 
+        radial_build: sequence of tuples containing the radial build of the
+            reactor. Each tuple should contain a LayerType and a float 
         elongation (float, optional): _description_. Defaults to 2.0.
         triangularity (float, optional): _description_. Defaults to 0.55.
         rotation_angle (Optional[str], optional): _description_. Defaults to 180.0.
         extra_cut_shapes (Sequence, optional): _description_. Defaults to [].
+        colors (dict, optional): the colors to assign to the assembly parts. Defaults to {}.
+            Each dictionary entry should be a key that matches the assembly part name
+            (e.g. 'plasma', or 'layer_1') and a tuple of 3 or 4 floats between 0 and 1
+            representing the RGB or RGBA values.
 
     Returns:
         _type_: _description_
@@ -148,26 +156,33 @@ def spherical_tokamak_from_plasma(
         rotation_angle=rotation_angle,
         extra_cut_shapes=extra_cut_shapes,
         extra_intersect_shapes=extra_intersect_shapes,
+        colors=colors,
     )
 
 
 def spherical_tokamak(
-    radial_build: Union[Sequence[Sequence[Tuple[str, float]]], Sequence[Tuple[str, float]]],
+    radial_build: Sequence[Tuple[LayerType, float]],
     vertical_build: Sequence[Tuple[str, float]],
     triangularity: float = 0.55,
     rotation_angle: Optional[str] = 180.0,
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
+    colors: dict = {},
 ):
-    """_summary_
+    """  Creates a spherical tokamak fusion reactor from a radial build and vertical build.
 
     Args:
 
-        radial_build
+        radial_build: sequence of tuples containing the radial build of the
+            reactor. Each tuple should contain a LayerType and a float 
         elongation (float, optional): _description_. Defaults to 2.0.
         triangularity (float, optional): _description_. Defaults to 0.55.
         rotation_angle (Optional[str], optional): _description_. Defaults to 180.0.
         extra_cut_shapes (Sequence, optional): _description_. Defaults to [].
+        colors (dict, optional): the colors to assign to the assembly parts. Defaults to {}.
+            Each dictionary entry should be a key that matches the assembly part name
+            (e.g. 'plasma', or 'layer_1') and a tuple of 3 or 4 floats between 0 and 1
+            representing the RGB or RGBA values.
 
     Returns:
         _type_: _description_
@@ -225,7 +240,8 @@ def spherical_tokamak(
     for i, entry in enumerate(extra_cut_shapes):
 
         if isinstance(entry, cq.Workplane):
-            my_assembly.add(entry, name=f"add_extra_cut_shape_{i+1}")
+            name = f"add_extra_cut_shape_{i+1}"
+            my_assembly.add(entry, name=name, color=cq.Color(*colors.get(name, (0.5,0.5,0.5))))
         else:
             raise ValueError(f"extra_cut_shapes should only contain cadquery Workplanes, not {type(entry)}")
 
@@ -245,14 +261,14 @@ def spherical_tokamak(
         for i, entry in enumerate(extra_intersect_shapes):
             reactor_entry_intersection = entry.intersect(reactor_compound)
             intersect_shapes_to_cut.append(reactor_entry_intersection)
-            my_assembly.add(reactor_entry_intersection, name=f"extra_intersect_shapes_{i+1}")
+            name = f"extra_intersect_shapes_{i+1}"
+            my_assembly.add(reactor_entry_intersection, name=name, color=cq.Color(*colors.get(name, (0.5,0.5,0.5))))
 
     # builds just the core if there are no extra parts
     if len(extra_cut_shapes) == 0 and len(intersect_shapes_to_cut) == 0:
-        for i, entry in enumerate(inner_radial_build):
-            my_assembly.add(entry, name=f"inboard_layer_{i+1})")
-        for i, entry in enumerate(blanket_layers):
-            my_assembly.add(entry, name=f"outboard_layer_{i+1})")
+        for i, entry in enumerate(inner_radial_build+blanket_layers):
+            name = f"layer_{i+1}"
+            my_assembly.add(entry, name=name, color=cq.Color(*colors.get(name, (0.5,0.5,0.5))))
     else:
         shapes_and_components = []
         for i, entry in enumerate(inner_radial_build + blanket_layers):
@@ -264,10 +280,10 @@ def spherical_tokamak(
             shapes_and_components.append(entry)
 
         for i, entry in enumerate(shapes_and_components):
-            my_assembly.add(
-                entry, name=f"layer_{i+1})"
-            )  # TODO track the names of shapes, even when extra shapes are made due to splitting
+            # TODO track the names of shapes, even when extra shapes are made due to splitting
+            name=f"layer_{i+1}"
+            my_assembly.add(entry, name=name, color=cq.Color(*colors.get(name, (0.5,0.5,0.5))))
 
-    my_assembly.add(plasma, name="plasma")
+    my_assembly.add(plasma, name="plasma", color=cq.Color(*colors.get("plasma", (0.5,0.5,0.5))))
 
     return my_assembly

src/paramak/assemblies/tokamak.py

@@ -156,7 +156,27 @@ def tokamak_from_plasma(
     rotation_angle: float = 180.0,
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
+    colors: dict = {}
 ):
+    """
+    Creates a tokamak fusion reactor from a radial build and plasma parameters.
+
+    Args:
+        radial_build: sequence of tuples containing the radial build of the
+            reactor. Each tuple should contain a LayerType and a float
+        elongation: The elongation of the plasma. Defaults to 2.0.
+        triangularity: The triangularity of the plasma. Defaults to 0.55.
+        rotation_angle: The rotation angle of the plasma. Defaults to 180.0.
+        extra_cut_shapes: A list of extra shapes to cut the reactor with. Defaults to [].
+        extra_intersect_shapes: A list of extra shapes to intersect the reactor with. Defaults to [].
+        colors (dict, optional): the colors to assign to the assembly parts. Defaults to {}.
+            Each dictionary entry should be a key that matches the assembly part name
+            (e.g. 'plasma', or 'layer_1') and a tuple of 3 or 4 floats between 0 and 1
+            representing the RGB or RGBA values.
+
+    Returns:
+        CadQuery.Assembly: A CadQuery Assembly object representing the tokamak fusion reactor.
+    """
 
     inner_equatorial_point = sum_up_to_plasma(radial_build)
     plasma_radial_thickness = get_plasma_value(radial_build)
@@ -184,27 +204,35 @@ def tokamak_from_plasma(
         rotation_angle=rotation_angle,
         extra_cut_shapes=extra_cut_shapes,
         extra_intersect_shapes=extra_intersect_shapes,
+        colors=colors
     )
 
 
 def tokamak(
-    radial_build: Union[Sequence[Sequence[Tuple[str, float]]], Sequence[Tuple[str, float]]],
+    radial_build: Sequence[Tuple[str, float]],
     vertical_build: Sequence[Tuple[str, float]],
     triangularity: float = 0.55,
     rotation_angle: float = 180.0,
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
+    colors: dict = {}
 ):
     """
-    Creates a tokamak fusion reactor from a radial build and plasma parameters.
+    Creates a tokamak fusion reactor from a radial and vertical build.
 
     Args:
-        radial_build: A list of tuples containing the radial build of the reactor.
-        elongation: The elongation of the plasma. Defaults to 2.0.
+        radial_build: sequence of tuples containing the radial build of the
+            reactor. Each tuple should contain a LayerType and a float 
+        vertical_build: sequence of tuples containing the vertical build of the
+            reactor. Each tuple should contain a LayerType and a float 
         triangularity: The triangularity of the plasma. Defaults to 0.55.
         rotation_angle: The rotation angle of the plasma. Defaults to 180.0.
         extra_cut_shapes: A list of extra shapes to cut the reactor with. Defaults to [].
         extra_intersect_shapes: A list of extra shapes to intersect the reactor with. Defaults to [].
+        colors (dict, optional): the colors to assign to the assembly parts. Defaults to {}.
+            Each dictionary entry should be a key that matches the assembly part name
+            (e.g. 'plasma', or 'layer_1') and a tuple of 3 or 4 floats between 0 and 1
+            representing the RGB or RGBA values.
 
     Returns:
         CadQuery.Assembly: A CadQuery Assembly object representing the tokamak fusion reactor.
@@ -248,7 +276,8 @@ def tokamak(
 
     for i, entry in enumerate(extra_cut_shapes):
         if isinstance(entry, cq.Workplane):
-            my_assembly.add(entry, name=f"add_extra_cut_shape_{i+1}")
+            name = f"add_extra_cut_shape_{i+1}"
+            my_assembly.add(entry, name=name, color=cq.Color(*colors.get(name, (0.5,0.5,0.5))))
         else:
             raise ValueError(f"extra_cut_shapes should only contain cadquery Workplanes, not {type(entry)}")
 
@@ -268,14 +297,14 @@ def tokamak(
         for i, entry in enumerate(extra_intersect_shapes):
             reactor_entry_intersection = entry.intersect(reactor_compound)
             intersect_shapes_to_cut.append(reactor_entry_intersection)
-            my_assembly.add(reactor_entry_intersection, name=f"extra_intersect_shapes_{i+1}")
+            name=f"extra_intersect_shapes_{i+1}"
+            my_assembly.add(reactor_entry_intersection, name=name, color=cq.Color(*colors.get(name, (0.5,0.5,0.5))))
 
     # builds just the core if there are no extra parts
     if len(extra_cut_shapes) == 0 and len(intersect_shapes_to_cut) == 0:
-        for i, entry in enumerate(inner_radial_build):
-            my_assembly.add(entry, name=f"inboard_layer_{i+1})")
-        for i, entry in enumerate(blanket_layers):
-            my_assembly.add(entry, name=f"outboard_layer_{i+1})")
+        for i, entry in enumerate(inner_radial_build+blanket_layers):
+            name=f"layer_{i+1}"
+            my_assembly.add(entry, name=name, color=cq.Color(*colors.get(name, (0.5,0.5,0.5))))
     else:
         shapes_and_components = []
         for i, entry in enumerate(inner_radial_build + blanket_layers):
@@ -287,10 +316,10 @@ def tokamak(
             shapes_and_components.append(entry)
 
         for i, entry in enumerate(shapes_and_components):
-            my_assembly.add(
-                entry, name=f"layer_{i+1})"
-            )  # TODO track the names of shapes, even when extra shapes are made due to splitting
+            name=f"layer_{i+1}"
+            # TODO track the names of shapes, even when extra shapes are made due to splitting
+            my_assembly.add(entry, name=name, color=cq.Color(*colors.get(name, (0.5,0.5,0.5))))
 
-    my_assembly.add(plasma, name="plasma")
+    my_assembly.add(plasma, name="plasma", color=cq.Color(*colors.get("plasma", (0.5,0.5,0.5))))
 
     return my_assembly