Coverage for yaptide/converter/converter/fluka/helper_parsers/beam_parser.py: 80%
76 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-07-01 12:55 +0000
1from math import cos, atan, pi
2from dataclasses import dataclass
3from enum import Enum
6class BeamShape(Enum):
7 """Enum representing beam shape"""
9 GAUSSIAN = 1
10 SQUARE = 2
11 CIRCULAR = 3
13 def __str__(self):
14 if self == BeamShape.GAUSSIAN:
15 return 'gaussian'
16 if self == BeamShape.SQUARE:
17 return 'flat square'
18 if self == BeamShape.CIRCULAR:
19 return 'flat circular'
20 return ''
23@dataclass(frozen=False)
24class FlukaBeam:
25 """Class representing beam config in a FLUKA input file."""
27 energy_MeV: float = 150.
28 beam_pos: tuple[float, float, float] = (0, 0, 0) # [cm]
29 beam_dir: tuple[float, float] = (0, 0) # cosines respective to x and y axes
30 z_negative: bool = False
31 shape: BeamShape = BeamShape.GAUSSIAN
32 shape_x: float = 0
33 shape_y: float = 0
34 particle_name: str = 'PROTON'
35 heavy_ion_a: int = 1
36 heavy_ion_z: int = 1
39particle_dict = {
40 1: {
41 'name': 'NEUTRON',
42 'a': 1
43 },
44 2: {
45 'name': 'PROTON',
46 'a': 1
47 },
48 3: {
49 'name': 'PION-',
50 'a': 1
51 },
52 4: {
53 'name': 'PION+',
54 'a': 1
55 },
56 5: {
57 'name': 'PIZERO',
58 'a': 1
59 },
60 6: {
61 'name': 'ANEUTRON',
62 'a': 1
63 },
64 7: {
65 'name': 'APROTON',
66 'a': 1
67 },
68 8: {
69 'name': 'KAON-',
70 'a': 1
71 },
72 9: {
73 'name': 'KAON+',
74 'a': 1
75 },
76 10: {
77 'name': 'KAONZERO',
78 'a': 1
79 },
80 11: {
81 'name': 'KAONLONG',
82 'a': 1
83 },
84 12: {
85 'name': 'PHOTON',
86 'a': 1
87 },
88 15: {
89 'name': 'MUON-',
90 'a': 1
91 },
92 16: {
93 'name': 'MUON+',
94 'a': 1
95 },
96 21: {
97 'name': 'DEUTERON',
98 'a': 2
99 },
100 22: {
101 'name': 'TRITON',
102 'a': 3
103 },
104 23: {
105 'name': '3-HELIUM',
106 'a': 3
107 },
108 24: {
109 'name': '4-HELIUM',
110 'a': 4
111 },
112 25: {
113 'name': 'HEAVYION',
114 'a': 1
115 },
116 26: {
117 'name': 'ELECTRON',
118 'a': 1
119 }
120}
123def convert_energy(beam_json: dict) -> float:
124 """
125 Extract energy from beam JSON and provide it in Fluka convention.
127 For particles other than HEAVYIONS (i.e. protons, alpha particles, neutrons, electrons) MeV unit is used.
128 For HEAVYIONS (i.e. nuclei heavier than helium, defined by A and Z numbers) MeV/u is used.
129 Note than MeV/u (atomic mass unit) is not the same as MeV/nucl (number of nucleons).
130 For more details see:
131 https://flukafiles.web.cern.ch/manual/chapters/description_input/description_options/beam.html#beam.
132 """
133 energy_MeV_nucl = beam_json['energy']
134 particle = particle_dict[beam_json['particle']['id']]
135 energy_Fluka_standard = energy_MeV_nucl * particle['a']
136 if particle['name'] == 'HEAVYION':
137 # Assuming that MeV/nucl is a good approximation of MeV/u
138 energy_Fluka_standard = energy_MeV_nucl
140 return energy_Fluka_standard
143def parse_particle_name(particle_json: dict):
144 """Parse particle ID to FLUKA particle name."""
145 particle_id = particle_json['id']
146 if particle_id in particle_dict:
147 particle = particle_dict[particle_id]
148 return particle['name']
149 raise ValueError('Particle ID not supported by FLUKA')
152def parse_shape_params(shape_params_json: dict) -> tuple[BeamShape, float, float]:
153 """Parse shape params from JSON to FLUKA shape params."""
154 shape = shape_params_json['type']
155 if shape == 'Flat circular':
156 return BeamShape.CIRCULAR, shape_params_json['x'], shape_params_json['y']
157 if shape == 'Flat square':
158 return BeamShape.SQUARE, shape_params_json['x'], shape_params_json['y']
159 if shape == 'Gaussian':
160 return BeamShape.GAUSSIAN, shape_params_json['x'], shape_params_json['y']
161 raise ValueError('Shape type not supported by FLUKA')
164def cartesian_to_spherical(coords: tuple[float, float, float]):
165 """
166 Convert cartesian coordinates to spherical coordinates
167 and return cosines of angles respective to x and y axes
168 """
169 x, y, z = coords
170 theta = pi / 2
171 if x != 0:
172 theta = atan(z / x)
173 phi = pi / 2
174 if y != 0:
175 phi = atan((x**2 + z**2)**0.5 / y)
176 return cos(theta), cos(phi)
179def parse_beam(beam_json: dict) -> FlukaBeam:
180 """Parse beam from JSON to FLUKA beam."""
181 fluka_beam = FlukaBeam()
182 fluka_beam.energy_MeV = convert_energy(beam_json)
183 fluka_beam.particle_name = parse_particle_name(beam_json['particle'])
184 if fluka_beam.particle_name == 'HEAVYION':
185 fluka_beam.heavy_ion_a = beam_json['particle']['a']
186 fluka_beam.heavy_ion_z = beam_json['particle']['z']
187 fluka_beam.beam_pos = tuple(beam_json['position'])
188 shape, shape_x, shape_y = parse_shape_params(beam_json['sigma'])
189 fluka_beam.shape = shape
190 fluka_beam.shape_x = shape_x
191 fluka_beam.shape_y = shape_y
192 theta, phi = cartesian_to_spherical(beam_json['direction'])
193 fluka_beam.beam_dir = (theta, phi)
194 if beam_json['direction'][2] < 0:
195 fluka_beam.z_negative = True
196 return fluka_beam