Coverage for yaptide/converter/converter/shieldhit/beam.py: 91%
110 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-07-01 12:55 +0000
1import math as m
2from dataclasses import dataclass
3from enum import IntEnum, unique
4from typing import Optional, TypeVar, Type
7T = TypeVar("T", bound="LabelledEnum")
10class LabelledEnum(IntEnum):
11 """Base class for enums with a label attribute"""
13 label: str
15 def __new__(cls, value, label):
17 if not isinstance(value, int):
18 raise TypeError("Value must be an integer")
19 obj = int.__new__(cls, value)
20 obj._value_ = value
21 obj.label = label
22 return obj
24 @classmethod
25 def from_str(cls: Type[T], value: str) -> T:
26 """Converts a string to a LabelledEnum"""
27 for method in cls:
28 if method.label == value:
29 return method
31 raise ValueError(f"{cls.__name__} has no value matching {value}")
34@unique
35class BeamSourceType(LabelledEnum):
36 """Beam source type"""
38 SIMPLE = (0, "simple")
39 FILE = (1, "file")
42@unique
43class ModulatorSimulationMethod(LabelledEnum):
44 """Modulator simulation method for beam.dat file"""
46 MODULUS = (0, "modulus")
47 SAMPLING = (1, "sampling")
50@unique
51class ModulatorInterpretationMode(LabelledEnum):
52 """Modulator interpretation mode of data in the input files loaded with the USEBMOD card"""
54 MATERIAL = (0, "material")
55 VACUMM = (1, "vacumm")
58@unique
59class StragglingModel(LabelledEnum):
60 """Straggle model"""
62 GAUSSIAN = (1, "Gaussian")
63 VAVILOV = (2, "Vavilov")
64 NO_STRAGGLING = (0, "no straggling")
67@unique
68class MultipleScatteringMode(LabelledEnum):
69 """Multiple scattering mode"""
71 GAUSSIAN = (1, "Gaussian")
72 MOLIERE = (2, "Moliere")
73 NO_SCATTERING = (0, "no scattering")
76@dataclass(frozen=True)
77class BeamModulator():
78 """Beam modulator card dataclass used in BeamConfig."""
80 filename: str
81 file_content: str
82 zone_id: int
83 simulation: ModulatorSimulationMethod = ModulatorSimulationMethod.MODULUS
84 mode: ModulatorInterpretationMode = ModulatorInterpretationMode.MATERIAL
86 def __str__(self) -> str:
87 """Returns the string representation of the beam modulator card"""
88 modulator_template = """USEBMOD {zone} {filename} ! Zone# and file name for beam modulator
89BMODMC {simulation} ! Simulation method for beam modulator (0-Modulus, 1-Monte Carlo sampling)
90BMODTRANS {mode} ! Interpretation of thicknesses data in the config file (0-Material, 1-Vacuum)"""
91 return modulator_template.format(
92 zone=self.zone_id,
93 filename=self.filename,
94 simulation=self.simulation,
95 mode=self.mode)
98@dataclass
99class BeamConfig:
100 """Class mapping of the beam.dat config file."""
102 particle: int = 2
103 particle_name: Optional[str] = None
104 heavy_ion_a: int = 1
105 heavy_ion_z: int = 1
106 energy: float = 150. # [MeV]
107 energy_spread: float = 1.5 # [MeV]
108 energy_low_cutoff: Optional[float] = None # [MeV]
109 energy_high_cutoff: Optional[float] = None # [MeV]
110 beam_ext_x: float = -0.1 # [cm]
111 beam_ext_y: float = 0.1 # [cm]
112 sad_x: Optional[float] = None # [cm]
113 sad_y: Optional[float] = None # [cm]
114 n_stat: int = 10000
115 beam_pos: tuple[float, float, float] = (0, 0, 0) # [cm]
116 beam_dir: tuple[float, float, float] = (0, 0, 1) # [cm]
117 delta_e: float = 0.03 # [a.u.]
118 nuclear_reactions: bool = True
120 modulator: Optional[BeamModulator] = None
122 straggling: StragglingModel = StragglingModel.VAVILOV
123 multiple_scattering: MultipleScatteringMode = MultipleScatteringMode.MOLIERE
125 heavy_ion_template = "HIPROJ {a} {z} ! A and Z of the heavy ion"
126 energy_cutoff_template = "TCUT0 {energy_low_cutoff} {energy_high_cutoff} ! energy cutoffs [MeV]"
127 sad_template = "BEAMSAD {sad_x} {sad_y} ! BEAMSAD value [cm]"
128 beam_source_type: BeamSourceType = BeamSourceType.SIMPLE
129 beam_source_filename: Optional[str] = None
130 beam_source_file_content: Optional[str] = None
132 beam_dat_template: str = """
133RNDSEED 89736501 ! Random seed
134JPART0 {particle} ! Incident particle type{particle_optional_comment}
135{optional_heavy_ion_line}
136TMAX0 {energy} {energy_spread} ! Incident energy and energy spread; both in (MeV/nucl)
137{optional_energy_cut_off_line}
138NSTAT {n_stat:d} 0 ! NSTAT, Step of saving
139STRAGG {straggling} ! Straggling: 0-Off 1-Gauss, 2-Vavilov
140MSCAT {multiple_scattering} ! Mult. scatt 0-Off 1-Gauss, 2-Moliere
141NUCRE {nuclear_reactions} ! Nucl.Reac. switcher: 1-ON, 0-OFF
142{optional_beam_modulator_lines}
143BEAMPOS {pos_x} {pos_y} {pos_z} ! Position of the beam
144BEAMDIR {theta} {phi} ! Direction of the beam
145BEAMSIGMA {beam_ext_x} {beam_ext_y} ! Beam extension
146{optional_sad_parameter_line}
147DELTAE {delta_e} ! relative mean energy loss per transportation step
148"""
150 @staticmethod
151 def cartesian2spherical(vector: tuple[float, float, float]) -> tuple[float, float, float]:
152 """
153 Transform cartesian coordinates to spherical coordinates.
155 :param vector: cartesian coordinates
156 :return: spherical coordinates
157 """
158 x, y, z = vector
159 r = m.sqrt(x**2 + y**2 + z**2)
160 # acos returns the angle in radians between 0 and pi
161 theta = m.degrees(m.acos(z / r))
162 # atan2 returns the angle in radians between -pi and pi
163 phi = m.degrees(m.atan2(y, x))
164 # lets ensure the angle in degrees is always between 0 and 360, as SHIELD-HIT12A requires
165 if phi < 0.:
166 phi += 360.
167 return theta, phi, r
169 def __str__(self) -> str:
170 """Return the beam.dat config file as a string."""
171 theta, phi, _ = BeamConfig.cartesian2spherical(self.beam_dir)
173 # if particle name is defined, add the comment to the template
174 particle_optional_comment = ""
175 if self.particle_name:
176 particle_optional_comment = f" ({self.particle_name})"
178 # if particle is heavy ion, add the heavy ion line to the template
179 heavy_ion_line = "! no heavy ion"
180 if self.particle == 25:
181 heavy_ion_line = BeamConfig.heavy_ion_template.format(a=self.heavy_ion_a, z=self.heavy_ion_z)
183 # if energy cutoffs are defined, add them to the template
184 cutoff_line = "! no energy cutoffs"
185 if self.energy_low_cutoff is not None and self.energy_high_cutoff is not None:
186 cutoff_line = BeamConfig.energy_cutoff_template.format(
187 energy_low_cutoff=self.energy_low_cutoff,
188 energy_high_cutoff=self.energy_high_cutoff
189 )
191 # if sad was defined, add it to the template
192 sad_line = "! no BEAMSAD value"
193 if self.sad_x is not None or self.sad_y is not None:
194 sad_y_value = self.sad_y if self.sad_y is not None else ""
195 sad_line = BeamConfig.sad_template.format(
196 sad_x=self.sad_x,
197 sad_y=sad_y_value)
199 # if beam modulator was defined, add it to the template
200 mod_lines = str(self.modulator) if self.modulator is not None else '! no beam modulator'
202 # prepare main template
203 result = self.beam_dat_template.format(
204 particle=self.particle,
205 particle_optional_comment=particle_optional_comment,
206 optional_heavy_ion_line=heavy_ion_line,
207 energy=float(self.energy),
208 energy_spread=float(self.energy_spread),
209 optional_energy_cut_off_line=cutoff_line,
210 optional_sad_parameter_line=sad_line,
211 optional_beam_modulator_lines=mod_lines,
212 n_stat=self.n_stat,
213 pos_x=self.beam_pos[0],
214 pos_y=self.beam_pos[1],
215 pos_z=self.beam_pos[2],
216 beam_ext_x=self.beam_ext_x,
217 beam_ext_y=self.beam_ext_y,
218 theta=theta,
219 phi=phi,
220 delta_e=self.delta_e,
221 nuclear_reactions=1 if self.nuclear_reactions else 0,
222 straggling=self.straggling.value,
223 multiple_scattering=self.multiple_scattering.value
224 )
226 # if beam source type is file, add the file name to the template
227 if self.beam_source_type == BeamSourceType.FILE:
228 result += "USECBEAM sobp.dat ! Use custom beam source file"
230 return result