diff --git a/.gitignore b/.gitignore index 02307b6f..ccc6ca6c 100644 --- a/.gitignore +++ b/.gitignore @@ -175,3 +175,4 @@ src/_version.py *.stp *.sql *.html +*.png \ No newline at end of file diff --git a/examples/spherical_tokamak_from_plasma_with_color.py b/examples/spherical_tokamak_from_plasma_with_color.py new file mode 100644 index 00000000..bade0de2 --- /dev/null +++ b/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 diff --git a/examples/tokamak_from_plasma_with_color.py b/examples/tokamak_from_plasma_with_color.py new file mode 100644 index 00000000..5d51f09e --- /dev/null +++ b/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 diff --git a/src/paramak/assemblies/spherical_tokamak.py b/src/paramak/assemblies/spherical_tokamak.py index c54b668f..7d8df7fd 100644 --- a/src/paramak/assemblies/spherical_tokamak.py +++ b/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 diff --git a/src/paramak/assemblies/tokamak.py b/src/paramak/assemblies/tokamak.py index c9694843..45c312a9 100644 --- a/src/paramak/assemblies/tokamak.py +++ b/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 diff --git a/tests/test_assemblies/test_spherical_tokamak.py b/tests/test_assemblies/test_spherical_tokamak.py index c02b3215..ac4d0826 100644 --- a/tests/test_assemblies/test_spherical_tokamak.py +++ b/tests/test_assemblies/test_spherical_tokamak.py @@ -107,3 +107,30 @@ def test_transport_without_magnets(): cross_sections_xml="tests/cross_sections.xml", ) assert flux > 0.0 + +def test_colors(): + "passing in the colors dictionary should not raise an error" + 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), + }, + ) \ No newline at end of file diff --git a/tests/test_assemblies/test_tokamak.py b/tests/test_assemblies/test_tokamak.py new file mode 100644 index 00000000..9993a415 --- /dev/null +++ b/tests/test_assemblies/test_tokamak.py @@ -0,0 +1,33 @@ +import paramak + + +def test_colors(): + "passing in the colors dictionary should not raise an error" + + 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), + } + ) \ No newline at end of file