Coverage for yaptide/converter/converter/shieldhit/geo.py: 84%
96 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-07-01 12:55 +0000
1from typing import Optional
2from converter.common import format_float, rotate
3from converter.solid_figures import SolidFigure, BoxFigure, CylinderFigure, SphereFigure
4from dataclasses import dataclass, field
5from enum import IntEnum
8class DefaultMaterial(IntEnum):
9 """
10 List of special material ids that are not referring to the `mat` and have a separate
11 meaning. E.g. Assigning id=1000 to a zone in `geo` file would mean, that the zone is made out
12 of vacuum, not that it refers to the 1000th item in `mat` file.
13 """
15 BLACK_HOLE = 0
16 VACUUM = 1000
18 @staticmethod
19 def is_default_material(material_value: int) -> bool:
20 """Check if the material is one of the predefined default materials."""
21 return material_value in DefaultMaterial._value2member_map_
24def parse_figure(figure: SolidFigure, number: int) -> str:
25 """Parse a SolidFigure into a string representation of SH12A input file."""
26 if type(figure) is BoxFigure:
27 return _parse_box(figure, number)
28 if type(figure) is CylinderFigure:
29 return _parse_cylinder(figure, number)
30 if type(figure) is SphereFigure:
31 return _parse_sphere(figure, number)
33 raise ValueError(f'Unexpected solid figure type: {figure}')
36def _parse_box(box: BoxFigure, number: int) -> str:
37 """Parse a BoxFigure into a str representation of SH12A input file."""
38 x_vec = rotate([box.x_edge_length, 0, 0], box.rotation)
39 y_vec = rotate([0, box.y_edge_length, 0], box.rotation)
40 z_vec = rotate([0, 0, box.z_edge_length], box.rotation)
41 diagonal_vec = x_vec + y_vec + z_vec
42 start_position = (
43 box.position[0] - diagonal_vec[0] / 2,
44 box.position[1] - diagonal_vec[1] / 2,
45 box.position[2] - diagonal_vec[2] / 2,
46 )
47 box_template = """
48 BOX {number:>4}{p1:>10}{p2:>10}{p3:>10}{p4:>10}{p5:>10}{p6:>10}
49 {p7:>10}{p8:>10}{p9:>10}{p10:>10}{p11:>10}{p12:>10}"""
50 return box_template.format(
51 number=number,
52 p1=format_float(start_position[0], 10),
53 p2=format_float(start_position[1], 10),
54 p3=format_float(start_position[2], 10),
55 p4=format_float(x_vec[0], 10),
56 p5=format_float(x_vec[1], 10),
57 p6=format_float(x_vec[2], 10),
58 p7=format_float(y_vec[0], 10),
59 p8=format_float(y_vec[1], 10),
60 p9=format_float(y_vec[2], 10),
61 p10=format_float(z_vec[0], 10),
62 p11=format_float(z_vec[1], 10),
63 p12=format_float(z_vec[2], 10),
64 )
67def _parse_cylinder(cylinder: CylinderFigure, number: int) -> str:
68 """Parse a CylinderFigure into a str representation of SH12A input file."""
69 height_vect = rotate([0, 0, cylinder.height], cylinder.rotation)
70 lower_base_position = (
71 cylinder.position[0] - height_vect[0] / 2,
72 cylinder.position[1] - height_vect[1] / 2,
73 cylinder.position[2] - height_vect[2] / 2,
74 )
75 rcc_template = """
76 RCC {number:>4}{p1:>10}{p2:>10}{p3:>10}{p4:>10}{p5:>10}{p6:>10}
77 {p7:>10}"""
78 return rcc_template.format(
79 number=number,
80 p1=format_float(lower_base_position[0], 10),
81 p2=format_float(lower_base_position[1], 10),
82 p3=format_float(lower_base_position[2], 10),
83 p4=format_float(height_vect[0], 10),
84 p5=format_float(height_vect[1], 10),
85 p6=format_float(height_vect[2], 10),
86 p7=format_float(cylinder.radius_top, 10),
87 )
90def _parse_sphere(sphere: SphereFigure, number: int) -> str:
91 """Parse a SphereFigure into a str representation of SH12A input file."""
92 sphere_entry_template = """
93 SPH {number:>4}{p1:>10}{p2:>10}{p3:>10}{p4:>10}"""
94 return sphere_entry_template.format(
95 number=number,
96 p1=format_float(sphere.position[0], 10),
97 p2=format_float(sphere.position[1], 10),
98 p3=format_float(sphere.position[2], 10),
99 p4=format_float(sphere.radius, 10),
100 )
103@dataclass
104class Material:
105 """Dataclass mapping for SH12A materials."""
107 name: str
108 sanitized_name: str
109 uuid: str
110 icru: int
111 density: Optional[float] = None
112 custom_stopping_power: bool = False
113 idx: int = 0
115 property_template = """{name} {value}\n"""
116 property_template_name = """{name}\n"""
118 def __str__(self) -> str:
119 result = self.property_template.format(name='MEDIUM', value=self.idx)
120 result += self.property_template.format(name='ICRU', value=self.icru)
121 if self.density is not None:
122 result += self.property_template.format(name='RHO', value=format_float(self.density, 10))
124 if self.custom_stopping_power:
125 result += self.property_template_name.format(name='LOADDEDX')
127 result += 'END\n'
128 return result
131@dataclass
132class Zone:
133 """Dataclass mapping for SH12A zones."""
135 uuid: str
137 id: int = 1
138 figures_operators: list[set[int]] = field(default_factory=lambda: [{1}])
139 material: int = 0
140 material_override: dict[str, str] = field(default_factory=dict)
142 zone_template: str = """
143 {id:03d} {operators}"""
145 def __str__(self) -> str:
146 return self.zone_template.format(
147 id=self.id,
148 operators='OR'.join([' '.join([f'{id:+5}' for id in figure_set])
149 for figure_set in self.figures_operators]),
150 )
153@dataclass
154class StoppingPowerFile:
155 """Dataclass mapping for SH12A stopping power files."""
157 icru: int
158 name: str
159 content: str
162@dataclass
163class GeoMatConfig:
164 """Class mapping of the geo.dat config file."""
166 figures: list[SolidFigure] = field(default_factory=lambda: [
167 CylinderFigure(position=(0., 0., 10.), radius_top=10., radius_bottom=10., height=20., rotation=(90., 0., 0.)),
168 CylinderFigure(position=(0., 0., 7.5), radius_top=15., radius_bottom=15., height=25., rotation=(90., 0., 0.)),
169 CylinderFigure(position=(0., 0., 5.), radius_top=20., radius_bottom=20., height=30., rotation=(90., 0., 0.)),
170 ])
171 zones: list[Zone] = field(default_factory=lambda: [
172 Zone(
173 uuid='',
174 id=1,
175 figures_operators=[{
176 1,
177 }],
178 material=1,
179 ),
180 Zone(
181 uuid='',
182 id=2,
183 figures_operators=[{-1, 2}],
184 material=1000,
185 ),
186 Zone(
187 uuid='',
188 id=3,
189 figures_operators=[{-2, 3}],
190 material=0,
191 ),
192 ])
193 materials: list[Material] = field(default_factory=lambda: [Material('', '', '', 276)])
194 jdbg1: int = 0
195 jdbg2: int = 0
196 title: str = 'Unnamed geometry'
197 available_custom_stopping_power_files: dict[int, StoppingPowerFile] = field(default_factory=lambda: {})
199 geo_template: str = """
200{jdbg1:>5}{jdbg1:>5} {title}
201{figures}
202 END
203{zones_geometries}
204 END
205{zones_materials}
206"""
208 @staticmethod
209 def _split_zones_to_rows(zones: list[int], max_size=14) -> list[list[int]]:
210 """Split list of Zones into rows of not more than 14 elements"""
211 return [zones[i:min(i + max_size, len(zones))] for i in range(0, len(zones), max_size)]
213 def _get_zone_material_string(self) -> str:
214 """Generate material_id, zone_id pairs string (for geo.dat)."""
215 # Cut lists into chunks of max size 14
216 zone_ids = [
217 ''.join([f'{id:>5}' for id in row])
218 for row in GeoMatConfig._split_zones_to_rows([zone.id for zone in self.zones])
219 ]
220 material_ids = [
221 ''.join([f'{mat:>5}' for mat in row])
222 for row in GeoMatConfig._split_zones_to_rows([zone.material for zone in self.zones])
223 ]
224 return '\n'.join([*zone_ids, *material_ids])
226 def get_geo_string(self) -> str:
227 """Generate geo.dat config."""
228 return self.geo_template.format(
229 jdbg1=self.jdbg1,
230 jdbg2=self.jdbg2,
231 title=self.title,
232 # we increment idx because shieldhit indexes from 1 while python indexes lists from 0
233 figures=''.join([parse_figure(figure, idx + 1) for idx, figure in enumerate(self.figures)])[1:],
234 zones_geometries=''.join([str(zone) for zone in self.zones])[1:],
235 zones_materials=self._get_zone_material_string(),
236 )
238 def get_mat_string(self) -> str:
239 """Generate mat.dat config."""
240 # we increment idx because shieldhit indexes from 1 while python indexes lists from 0
241 material_strings = []
242 materials_filtered = filter(lambda x: not DefaultMaterial.is_default_material(x.icru), self.materials)
243 for idx, material in enumerate(materials_filtered):
244 material.idx = idx + 1
245 material_strings.append(str(material))
247 return ''.join(material_strings)