Source code for zgoubidoo.commands.contrib.iba

"""Zgoubidoo's

More details here.
TODO
"""
from typing import List, Optional, Union, Iterable, Tuple
import numpy as _np
import pandas as _pd
import lmfit
from ..magnetique import Dipole as _Dipole
from ..magnetique import PolarMagnet as _PolarMagnet
from ..magnetique import Multipole as _Multipole
from ..magnetique import Quadrupole as _Quadrupole
from ..magnetique import Drift as _Drift
from ..commands import Collimator as _Collimator
from ..commands import Marker as _Marker
from ..commands import Ymy as _Ymy
from ..actions import Fit as _Fit
from ..actions import FitType as _FitType
from ..commands import Chamber as _Chamber
from ..particules import Proton as _Proton
from ..commands import ChangeRef as _ChangeRef
from ..beam import Beam as _Beam
from ..beam import BeamInputDistribution as _BeamInputDistribution
from ... import ureg as _ureg
from ...input import Input as _Input
from ...mappings import MappedParametersType as _MappedParametersType
from ... import Kinematics as _Kinematics
from ...zgoubi import Zgoubi as _Zgoubi
from ...zgoubi import ZgoubiResults as _ZgoubiResults
from ...polarity import PolarityType as _PolarityType
from ...polarity import Polarity as _Polarity
from ...polarity import HorizontalPolarity as _HorizontalPolarity
from ...polarity import VerticalPolarity as _VerticalPolarity
from ...fieldmaps import EngeModel as _EngeModel
import zgoubidoo


[docs]class DipoleIBA(_Dipole): """IBA generic dipole. The main feature is to provide a method to compute extra drift length taken by the polar extension of the field map. Examples: >>> d = DipoleIBA() >>> d.extra_drift() """ @property def extra_drift(self) -> _ureg.Quantity: """Compute the extra drift length of the field map extension. See Also: `_PolarMagnet.drift_length_from_polar`. Returns: the drift lenght. """ return _PolarMagnet.drift_length_from_polar(radius=self.radius, magnet_angle=self.AT, poles_angle=self.OMEGA_E - self.OMEGA_S)
[docs] def process_fit_field_profile(self, fit: lmfit.model.ModelResult): """ Args: fit: Returns: """ self.LAM_E = fit.best_values['lam_e'] * _ureg.m self.LAM_S = fit.best_values['lam_s'] * _ureg.m self.C0_E = fit.best_values['ce_0'] self.C1_E = fit.best_values['ce_1'] self.C2_E = fit.best_values['ce_2'] self.C3_E = fit.best_values['ce_3'] self.C4_E = fit.best_values['ce_4'] self.C5_E = fit.best_values['ce_5'] self.C0_S = fit.best_values['cs_0'] self.C1_S = fit.best_values['cs_1'] self.C2_S = fit.best_values['cs_2'] self.C3_S = fit.best_values['cs_3'] self.C4_S = fit.best_values['cs_4'] self.C5_S = fit.best_values['cs_5']
#self.XL = (fit.best_values['offset_s'] - fit.best_values['offset_e']) * _ureg.m
[docs]class B1G(DipoleIBA): """Proteus One 40 degree dipole. Examples: >>> B1G() """ PARAMETERS = { 'LABEL1': 'B1G', 'B0': 14 * _ureg.kilogauss, # It is important to keep it expressed in kilogauss here 'AT': 50 * _ureg.degree, 'ACENT': 25 * _ureg.degree, 'RM': 1500 * _ureg.mm, }
[docs] def post_init(self, magnet_opening=50 * _ureg.degree, poles_opening=40 * _ureg.degree, entrance_pole_trim=0 * _ureg.degree, exit_pole_trim=0 * _ureg.degree, entrance_fringe_lambda=9 * _ureg.cm, exit_fringe_lambda=9 * _ureg.cm, **kwargs, ): """ TODO Args: magnet_opening: total angular opening of the magnet (i.e. of the field map) poles_opening: angular opening of the poles entrance_pole_trim: angular shift of the entrance pole exit_pole_trim: angular shift of the exit pole entrance_fringe_lambda: effective length of the entrance fringe field exit_fringe_lambda: effective length of the exit fringe field Example: >>> b1g = B1G() >>> b1g.fit() """ self.AT = magnet_opening self.ACENT = self.AT / 2 self.OMEGA_E = poles_opening / 2 - entrance_pole_trim self.OMEGA_S = -poles_opening / 2 + entrance_pole_trim self.LAM_E = entrance_fringe_lambda self.LAM_S = exit_fringe_lambda self.RE = _PolarMagnet.efb_offset_from_polar(radius=self.RM, magnet_angle=magnet_opening, poles_angle=poles_opening ) self.RS = _PolarMagnet.efb_offset_from_polar(radius=self.RM, magnet_angle=magnet_opening, poles_angle=poles_opening ) self.TE = _PolarMagnet.efb_angle_from_polar(magnet_angle=magnet_opening, poles_angle=poles_opening ) self.TS = -_PolarMagnet.efb_angle_from_polar(magnet_angle=magnet_opening, poles_angle=poles_opening ) self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=134, min=0, max=400) self._field_profile_model.params['offset_s'].set(value=1302, min=1000, max=1400) self._field_profile_model.params['lam_e'].set(value=99) self._field_profile_model.params['lam_s'].set(value=75) self._field_profile_model.params['amplitude'].set(value=-2.0) self._field_profile_model.params['field_offset'].set(vary=True, value=0.0, min=-1e-3, max=1e-3)
[docs]class B2G(DipoleIBA): """Proteus One 70 degree dipole. Examples: >>> B2G() """ PARAMETERS = { 'LABEL1': 'B2G', 'B0': 14 * _ureg.kilogauss, # It is important to keep it expressed in kilogauss here 'RM': 1500 * _ureg.mm, }
[docs] def post_init(self, magnet_opening=80 * _ureg.degree, poles_opening=70 * _ureg.degree, entrance_pole_trim=0 * _ureg.degree, exit_pole_trim=0 * _ureg.degree, entrance_fringe_lambda=9 * _ureg.cm, exit_fringe_lambda=9 * _ureg.cm, **kwargs, ): """ TODO Args: magnet_opening: total angular opening of the magnet (i.e. of the field map) poles_opening: angular opening of the poles entrance_pole_trim: angular shift of the entrance pole exit_pole_trim: angular shift of the exit pole entrance_fringe_lambda: effective length of the entrance fringe field exit_fringe_lambda: effective length of the exit fringe field Example: >>> b2g = B2G() >>> b2g.fit() """ self.AT = magnet_opening self.ACENT = self.AT / 2 self.OMEGA_E = poles_opening / 2 - entrance_pole_trim self.OMEGA_S = -poles_opening / 2 + entrance_pole_trim self.LAM_E = entrance_fringe_lambda self.LAM_S = exit_fringe_lambda self.RE = _PolarMagnet.efb_offset_from_polar(radius=self.RM, magnet_angle=magnet_opening, poles_angle=poles_opening ) self.RS = _PolarMagnet.efb_offset_from_polar(radius=self.RM, magnet_angle=magnet_opening, poles_angle=poles_opening ) self.TE = _PolarMagnet.efb_angle_from_polar(magnet_angle=magnet_opening, poles_angle=poles_opening ) self.TS = -_PolarMagnet.efb_angle_from_polar(magnet_angle=magnet_opening, poles_angle=poles_opening ) self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=400, min=0, max=500) self._field_profile_model.params['offset_s'].set(value=2000, min=1500, max=2300) self._field_profile_model.params['lam_e'].set(value=50) self._field_profile_model.params['lam_s'].set(value=50) self._field_profile_model.params['amplitude'].set(value=-2.0) self._field_profile_model.params['field_offset'].set(vary=True, value=0.0, min=-1e-3, max=1e-3)
[docs]class B3G(DipoleIBA): """Proteus One 60 degree dipole. Examples: >>> B3G() """ PARAMETERS = { 'LABEL1': 'B3G', 'B0': 14 * _ureg.kilogauss, # It is important to keep it expressed in kilogauss here 'RM': 1600 * _ureg.mm, 'LAM_E': 9 * _ureg.cm, 'C0_E': 0.67634054, # Obtained from field map fit 'C1_E': 1.15776841, # Obtained from field map fit 'C2_E': -0.16937986, # Obtained from field map fit 'C3_E': 0.07696388, # Obtained from field map fit 'SHIFT_E': 0 * _ureg.cm, 'LAM_S': 9 * _ureg.cm, 'C0_S': 0.77890820, # Obtained from field map fit 'C1_S': 0.94490545, # Obtained from field map fit 'C2_S': -0.13034787, # Obtained from field map fit 'C3_S': 0.02948957, # Obtained from field map fit 'SHIFT_S': 0 * _ureg.cm, 'XPAS': 1.0 * _ureg.mm, }
[docs] def post_init(self, magnet_opening=70 * _ureg.degree, poles_opening=60 * _ureg.degree, entrance_pole_trim=1.125 * _ureg.degree, exit_pole_trim=1.125 * _ureg.degree, entrance_fringe_lambda=9 * _ureg.cm, exit_fringe_lambda=9 * _ureg.cm, entrance_pole_curvature=1e9 * _ureg.cm, exit_pole_curvature=1e9 * _ureg.cm, **kwargs, ): """ TODO Args: magnet_opening: total angular opening of the magnet (i.e. of the field map) poles_opening: angular opening of the poles entrance_pole_trim: angular shift of the entrance pole exit_pole_trim: angular shift of the exit pole entrance_fringe_lambda: effective length of the entrance fringe field exit_fringe_lambda: effective length of the exit fringe field entrance_pole_curvature: curvature of the pole at the entrance of the magnet exit_pole_curvature: curvature of the pole at the exit of the magnet Example: >>> b3g = B3G() >>> b3g.fit() """ self.AT = magnet_opening self.ACENT = self.AT / 2 self.OMEGA_E = poles_opening / 2 - entrance_pole_trim self.OMEGA_S = -poles_opening / 2 + entrance_pole_trim self.LAM_E = entrance_fringe_lambda self.LAM_S = exit_fringe_lambda self.RE = _PolarMagnet.efb_offset_from_polar(radius=self.RM, magnet_angle=magnet_opening, poles_angle=poles_opening ) self.RS = _PolarMagnet.efb_offset_from_polar(radius=self.RM, magnet_angle=magnet_opening, poles_angle=poles_opening ) self.TE = _PolarMagnet.efb_angle_from_polar(magnet_angle=magnet_opening, poles_angle=poles_opening ) self.TS = -_PolarMagnet.efb_angle_from_polar(magnet_angle=magnet_opening, poles_angle=poles_opening ) self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=400, min=0, max=500) self._field_profile_model.params['offset_s'].set(value=2000, min=1500, max=2300) self._field_profile_model.params['lam_e'].set(value=50) self._field_profile_model.params['lam_s'].set(value=50) self._field_profile_model.params['amplitude'].set(value=-2.0) self._field_profile_model.params['field_offset'].set(vary=True, value=0.0, min=-1e-3, max=1e-3)
[docs]class SMX(_Multipole): """Proteus One inline (horizontal) scanning magnet. """ PARAMETERS = { 'XL': 0.26317552187769017 * _ureg.m, 'LAM_E': 0.03879980891555799 * _ureg.m, 'C0_E': 2.260586829767431, 'C2_E': -0.08454912439481363, 'C3_E': 0.03165152479873875, 'C4_E': -0.016368674101593452, 'C5_E': 0.002271083461208103, 'LAM_S': 0.019573704193528198 * _ureg.m, 'C0_S': 0.007471653696559507, 'C2_S': -0.1596436867616995, 'C3_S': 0.02917407661831325, 'C4_S': -0.0022865133242955886, 'C5_S': 7.362703656240262e-05, }
[docs] def post_init(self, **kwargs): """ Args: **kwargs: Returns: """ self.LABEL1 = self.__class__.__name__ self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=0.2, min=0.1, max=0.3) self._field_profile_model.params['offset_s'].set(value=0.5, min=0.35, max=0.55) self._field_profile_model.params['lam_e'].set(value=0.03) self._field_profile_model.params['lam_s'].set(value=0.03) self._field_profile_model.params['amplitude'].set(value=-0.82, min=-0.9, max=-0.75) self._field_profile_model.params['field_offset'].set(vary=True, value=0.0, min=-1e-3, max=1e-3)
[docs] def process_fit_field_profile(self, fit: lmfit.model.ModelResult): """ Args: fit: Returns: """ self.LAM_E = fit.best_values['lam_e'] * _ureg.m self.LAM_S = fit.best_values['lam_s'] * _ureg.m self.C0_E = fit.best_values['ce_0'] self.C1_E = fit.best_values['ce_1'] self.C2_E = fit.best_values['ce_2'] self.C3_E = fit.best_values['ce_3'] self.C4_E = fit.best_values['ce_4'] self.C5_E = fit.best_values['ce_5'] self.C0_S = fit.best_values['cs_0'] self.C1_S = fit.best_values['cs_1'] self.C2_S = fit.best_values['cs_2'] self.C3_S = fit.best_values['cs_3'] self.C4_S = fit.best_values['cs_4'] self.C5_S = fit.best_values['cs_5'] self.XL = (fit.best_values['offset_s'] - fit.best_values['offset_e']) * _ureg.m
[docs]class SMY(_Multipole): """Proteus One crossline (vertical) scanning magnet. """ PARAMETERS = { 'XL': 0.15221715277508374 * _ureg.m, 'R1': 90 * _ureg.degree, 'LAM_E': 0.037857895089871904 * _ureg.m, 'C0_E': 0.1999859299335233, 'C2_E': 0.08542911466613756, 'C3_E': -0.028865773164774223, 'C4_E': -0.004292970946705814, 'C5_E': 0.00536700990602016, 'LAM_S': 0.03846389166292385 * _ureg.m, 'C0_S': 0.2482542942709081, 'C2_S': 0.02638766212507734, 'C3_S': 0.008102763875877633, 'C4_S': -0.006877068617401515, 'C5_S': 0.0007046649650343981, }
[docs] def post_init(self, **kwargs): """ Args: **kwargs: Returns: """ self.LABEL1 = self.__class__.__name__ self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=0.1, min=0.1, max=0.4) self._field_profile_model.params['offset_s'].set(value=0.7, min=0.4, max=0.8) self._field_profile_model.params['lam_e'].set(value=0.035) self._field_profile_model.params['lam_s'].set(value=0.035) self._field_profile_model.params['amplitude'].set(value=-0.4, min=-0.6, max=-0.3) self._field_profile_model.params['field_offset'].set(value=0.0, min=-1e-3, max=1e-3)
[docs] def process_fit_field_profile(self, fit: lmfit.model.ModelResult): """ Args: fit: Returns: """ self.LAM_E = fit.best_values['lam_e'] * _ureg.m self.LAM_S = fit.best_values['lam_s'] * _ureg.m self.C0_E = fit.best_values['ce_0'] self.C1_E = fit.best_values['ce_1'] self.C2_E = fit.best_values['ce_2'] self.C3_E = fit.best_values['ce_3'] self.C4_E = fit.best_values['ce_4'] self.C5_E = fit.best_values['ce_5'] self.C0_S = fit.best_values['cs_0'] self.C1_S = fit.best_values['cs_1'] self.C2_S = fit.best_values['cs_2'] self.C3_S = fit.best_values['cs_3'] self.C4_S = fit.best_values['cs_4'] self.C5_S = fit.best_values['cs_5'] self.XL = (fit.best_values['offset_s'] - fit.best_values['offset_e']) * _ureg.m
[docs]class T1G(_Multipole): """Proteus One steering magnet. """ PARAMETERS = { 'XL': 209.1 * _ureg.mm, 'B1': 1e-6 * _ureg.gauss, }
[docs] def post_init(self, **kwargs): """ Args: **kwargs: Returns: """ self.LABEL1 = self.__class__.__name__
[docs]class T2G(_Multipole): """Proteus One steering magnet. """ PARAMETERS = { 'XL': 209.1 * _ureg.mm, 'B1': 1e-6 * _ureg.gauss, }
[docs] def post_init(self, **kwargs): """ Args: **kwargs: Returns: """ self.LABEL1 = self.__class__.__name__
[docs]class QuadrupoleIBA(_Quadrupole): """IBA generic quadrupole. Support for gradient/current conversion and for polarity. """
[docs] def post_init(self, polarity: _PolarityType = _HorizontalPolarity, p=None, l_eff: float = None, **kwargs ): """ Args: polarity: p: l_eff: gradient: **kwargs: Returns: """ self._current = 0.0 self._gradient = 0.0 if self.PARAMETERS.get('POLARITY'): self.polarity = self.POLARITY else: self.polarity = polarity self.__P = p self.__L_EFF = l_eff
@property def polarity(self) -> str: """ Returns: """ return str(self._polarity) @polarity.setter def polarity(self, value: _PolarityType): self._polarity = value @property def current(self): """ Returns: """ return self._current @current.setter def current(self, current): if current < 0.0: raise Exception("Trying to set a quadrupole with a negative current. Change the polarity instead.") self._current = current self._gradient = _np.poly1d(self.p)(current) / self.l_eff @property def gradient(self): """ Returns: """ return int(self._polarity) * self._gradient @gradient.setter def gradient(self, gradient): if gradient < 0.0: raise Exception("Trying to set a quadrupole with a negative gradient. Change the polarity instead.") self._gradient = gradient self._current = _np.roots([self.p[0], self.p[1], self.p[2] - _np.abs(gradient * self.l_eff)])[1] super().gradient = gradient
[docs] def set_value(self, value): """Set the current.""" self.current = value
@property def k1(self): """ Returns: """ return self.gradient @k1.setter def k1(self, value): self._gradient = _np.abs(value) if value > 0: self._polarity = _HorizontalPolarity elif value == 0: self._polarity = _Polarity else: self._polarity = _VerticalPolarity self._current = _np.roots([self.p[0], self.p[1], self.p[2] - _np.abs(self._gradient * self.l_eff)])[1] @property def p(self): """ Returns: """ return self.__P @p.setter def p(self, value): self.__P = value if self._current is not None: self.current = self._current @property def l_eff(self): """ Returns: """ return self.__L_EFF @l_eff.setter def l_eff(self, value): self.__L_EFF = value if self._current is not None: self.current = self._current @property def B0(self): return self.gradient * self.R0.to('m').magnitude * _ureg.tesla
[docs]class QLong(QuadrupoleIBA): """Proteus One long quadrupole. """ PARAMETERS = { 'XL': 490 * _ureg.mm, }
[docs] def post_init(self, polarity: _PolarityType = _HorizontalPolarity, **kwargs): """ Args: polarity: **kwargs: Returns: """ super().post_init(polarity=polarity, p=[-2.45367E-05, 8.09526E-02, -6.91149E-03], l_eff=0.490)
[docs]class QShort(QuadrupoleIBA): """Proteus One short quadrupole. """ PARAMETERS = { 'XL': 290 * _ureg.mm, }
[docs] def post_init(self, polarity: _PolarityType = _HorizontalPolarity, **kwargs): """ Args: polarity: **kwargs: Returns: """ super().post_init(polarity=polarity, p=[-2.27972E-05, 4.98563E-02, -1.58432E-02], l_eff=0.290)
[docs]class QWall(QuadrupoleIBA): """Proteus One 'wall' quadrupole. """ PARAMETERS = { 'XL': 297 * _ureg.mm, }
[docs] def post_init(self, polarity: _PolarityType = _HorizontalPolarity, **kwargs): """ Args: polarity: **kwargs: Returns: """ super().post_init(polarity=polarity, p=[-2.23625E-06, 2.46011E-02, 8.21584E-04], l_eff=0.297)
[docs]class QPMQ(QuadrupoleIBA): """Proteus One 'PMQ' quadrupole. """ PARAMETERS = { 'XL': 297 * _ureg.mm, }
[docs] def post_init(self, polarity: _PolarityType = _HorizontalPolarity, **kwargs): """ Args: polarity: **kwargs: Returns: """ super().post_init(polarity=polarity, gradient=17.5)
[docs]class Q1C(QShort): """Proteus One 'Q1C' gantry quadrupole. """ PARAMETERS = { 'LABEL1': 'Q1C', 'POLARITY': _VerticalPolarity, }
[docs]class Q2C(QShort): """Proteus One 'Q2C' gantry quadrupole. """ PARAMETERS = { 'LABEL1': 'Q2C', 'POLARITY': _VerticalPolarity, }
[docs]class Q1G(QWall): """Proteus One 'Q1G' gantry quadrupole. """ PARAMETERS = { 'LABEL1': 'Q1G', 'POLARITY': _VerticalPolarity, }
[docs]class Q2G(QWall): """Proteus One 'Q2G' gantry quadrupole. """ PARAMETERS = { 'LABEL1': 'Q2G', 'POLARITY': _HorizontalPolarity, }
[docs]class Q3G(QShort): """Proteus One 'Q3G' gantry quadrupole. """ PARAMETERS = { 'LABEL1': 'Q3G', 'POLARITY': _HorizontalPolarity, }
[docs] def post_init(self, **kwargs, ): """ TODO Args: **kwargs: Example: >>> q3g = Q3G() """ self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=100, min=0, max=200) self._field_profile_model.params['offset_s'].set(value=400, min=300, max=500) self._field_profile_model.params['lam_e'].set(value=50) self._field_profile_model.params['lam_s'].set(value=50) self._field_profile_model.params['amplitude'].set(value=0.2) self._field_profile_model.params['field_offset'].set(vary=True, value=0.0, min=-1e-3, max=1e-3)
[docs]class Q4G(QShort): """Proteus One 'Q4G' gantry quadrupole. """ PARAMETERS = { 'LABEL1': 'Q4G', 'POLARITY': _VerticalPolarity, }
[docs] def post_init(self, **kwargs, ): """ TODO Args: **kwargs: Example: >>> q4g = Q4G() """ self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=100, min=0, max=200) self._field_profile_model.params['offset_s'].set(value=400, min=300, max=500) self._field_profile_model.params['lam_e'].set(value=50) self._field_profile_model.params['lam_s'].set(value=50) self._field_profile_model.params['amplitude'].set(value=0.2) self._field_profile_model.params['field_offset'].set(vary=True, value=0.0, min=-1e-3, max=1e-3)
[docs]class Q5G(QLong): """Proteus One 'Q5G' gantry quadrupole. """ PARAMETERS = { 'LABEL1': 'Q5G', 'POLARITY': _HorizontalPolarity, }
[docs] def post_init(self, **kwargs, ): """ TODO Args: **kwargs: Example: >>> q5g = Q5G() """ self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=275, min=100, max=350) self._field_profile_model.params['offset_s'].set(value=883, min=700, max=950) self._field_profile_model.params['lam_e'].set(value=100) self._field_profile_model.params['lam_s'].set(value=100) self._field_profile_model.params['amplitude'].set(value=-0.25) self._field_profile_model.params['field_offset'].set(vary=True, value=0.0, min=-1e-3, max=1e-3)
[docs]class Q6G(QShort): """Proteus One 'Q6G' gantry quadrupole. """ PARAMETERS = { 'LABEL1': 'Q6G', 'POLARITY': _VerticalPolarity, }
[docs] def post_init(self, **kwargs, ): """ TODO Args: **kwargs: Example: >>> q6g = Q6G() """ self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=100, min=0, max=200) self._field_profile_model.params['offset_s'].set(value=400, min=300, max=500) self._field_profile_model.params['lam_e'].set(value=50) self._field_profile_model.params['lam_s'].set(value=50) self._field_profile_model.params['amplitude'].set(value=0.2) self._field_profile_model.params['field_offset'].set(vary=True, value=0.0, min=-1e-3, max=1e-3)
[docs]class Q7G(QShort): """Proteus One 'Q7G' gantry quadrupole. """ PARAMETERS = { 'LABEL1': 'Q7G', 'POLARITY': _HorizontalPolarity, }
[docs] def post_init(self, **kwargs, ): """ TODO Args: **kwargs: Example: >>> q7g = Q7G() """ self._field_profile_model = _EngeModel() self._field_profile_model.params['ce_0'].set(vary=True) self._field_profile_model.params['ce_1'].set(vary=False) self._field_profile_model.params['ce_2'].set(vary=True) self._field_profile_model.params['ce_3'].set(vary=True) self._field_profile_model.params['ce_4'].set(vary=True) self._field_profile_model.params['ce_5'].set(vary=True) self._field_profile_model.params['cs_0'].set(vary=True) self._field_profile_model.params['cs_1'].set(vary=False) self._field_profile_model.params['cs_2'].set(vary=True) self._field_profile_model.params['cs_3'].set(vary=True) self._field_profile_model.params['cs_4'].set(vary=True) self._field_profile_model.params['cs_5'].set(vary=True) self._field_profile_model.params['offset_e'].set(value=100, min=0, max=200) self._field_profile_model.params['offset_s'].set(value=400, min=300, max=500) self._field_profile_model.params['lam_e'].set(value=50) self._field_profile_model.params['lam_s'].set(value=50) self._field_profile_model.params['amplitude'].set(value=0.2) self._field_profile_model.params['field_offset'].set(vary=True, value=0.0, min=-1e-3, max=1e-3)
[docs]class HorizontalSlits(_Collimator): """Proteus One horizontal slits. """ PARAMETERS = { 'J': 0, 'IFORM': 1, 'IA': 1, }
[docs] def post_init(self, **kwargs): """ Args: **kwargs: Returns: """ self.LABEL1 = self.__class__.__name__
[docs]class SL1G(HorizontalSlits): """First vertical collimator """ PARAMETERS = { 'LABEL1': 'SL1G', 'C1': 27.5 * _ureg.mm, 'C2': 1.5 * _ureg.cm, }
[docs]class SL3G(HorizontalSlits): """First vertical collimator """ PARAMETERS = { 'LABEL1': 'SL3G', 'C1': 27.5 * _ureg.mm, 'C2': 0.5 * _ureg.cm, }
[docs]class VerticalSlits(_Collimator): """Proteus One vertical slits. """ PARAMETERS = { 'J': 0, 'IFORM': 1, 'IA': 1, }
[docs] def post_init(self, **kwargs): """ Args: **kwargs: Returns: """ self.LABEL1 = self.__class__.__name__
[docs]class SL2G(VerticalSlits): """First horizontal collimator """ PARAMETERS = { 'LABEL1': 'SL2G', 'C1': 0.6 * _ureg.cm, 'C2': 27.5 * _ureg.mm, }
[docs]class ResearchArea: """Proteus One research area input sequence.""" def __init__(self): """ TODO """ pass
[docs]class CGTR: """Proteus One compact gantry (CGTR) input sequence. TODO Examples: >>> cgtr = CGTR() >>> cgtr.shoot() """ def __init__(self, kinematics: _Kinematics = _Kinematics(230 * _ureg.MeV, kinetic=True), b1g: Optional[B1G] = None, b2g: Optional[B2G] = None, b3g: Optional[B3G] = None, t1g: Optional[T1G] = None, t2g: Optional[T2G] = None, q1g: Optional[Q1G] = None, q2g: Optional[Q2G] = None, q3g: Optional[Q3G] = None, q4g: Optional[Q4G] = None, q5g: Optional[Q5G] = None, q6g: Optional[Q6G] = None, q7g: Optional[Q7G] = None, smx: Optional[SMX] = None, smy: Optional[SMY] = None, sl1g: Optional[SL1G] = None, sl2g: Optional[SL2G] = None, sl3g: Optional[SL3G] = None, beam: Optional[_Beam] = None, with_fit: bool = True, ): """ Args: kinematics: b1g: b2g: b3g: t1g: t2g: q1g: q2g: q3g: q4g: q5g: q6g: q7g: smx: smy: beam: with_fit: if True the dipole magnets of the line will be fit. """ self.b1g: B1G = b1g or B1G() self.b2g: B2G = b2g or B2G() self.b3g: B3G = b3g or B3G() self.t1g: T1G = t1g or T1G() self.t2g: T2G = t2g or T2G() self.q1g: Q1G = q1g or Q1G() self.q2g: Q2G = q2g or Q2G() self.q3g: Q3G = q3g or Q3G() self.q4g: Q4G = q4g or Q4G() self.q5g: Q5G = q5g or Q5G() self.q6g: Q6G = q6g or Q6G() self.q7g: Q7G = q7g or Q7G() self.smx: SMX = smx or SMX() self.smy: SMY = smy or SMY() self.sl1g: SL1G = sl1g or SL1G() self.sl2g: SL2G = sl2g or SL2G() self.sl3g: SL3G = sl3g or SL3G() self.beam: _Beam = beam or _BeamInputDistribution('BUNCH', slices=1, kinematics=kinematics.brho) self.start: _Marker = _Marker('START') self.iso: _Marker = _Marker('ISO') if with_fit: self.fit_dipoles(kinematics=kinematics) self.zi: _Input = _Input('CGTR', line=[ self.beam, self.start, _Collimator('C1G', IA=1, IFORM=2, J=0, C1=5 * _ureg.mm, C2=5 * _ureg.mm), _Chamber('Chamber1', IA=1, IFORM=2, J=0, C1=29.75 * _ureg.mm, C2=29.75 * _ureg.mm), _Drift('C1G_T1G', XL=9.2995 * _ureg.cm), _Ymy(), self.t1g, _Drift('T1G_T2G', XL=2.09 * _ureg.cm), self.t2g, _Drift('T2G_Q1G', XL=38.7855 * _ureg.cm), self.q1g, _Drift('Q1G_Q2G', XL=30.3 * _ureg.cm), self.q2g, _Drift('Q2G_SL1G', XL=19.719 * _ureg.cm + 3 * _ureg.cm), self.sl1g, _Drift('SL1G_SL2G', XL=3 * _ureg.cm + 3 * _ureg.cm + 1 * _ureg.cm), self.sl2g, _Drift('SL2G_B1G', XL=39.734 * _ureg.cm + 3 * _ureg.cm - self.b1g.extra_drift), _Chamber(IA=2), _Chamber('Chamber2', IA=1, IFORM=1, J=0, C1=2.9 * _ureg.cm, C2=1.29 * _ureg.cm, C3=self.b1g.RM), self.b1g, _Chamber(IA=2), _Ymy(), _Chamber('Chamber3', IA=1, IFORM=2, J=0, C1=29.75 * _ureg.mm, C2=29.75 * _ureg.mm), _Drift('B1G_Q3G', XL=26.44 * _ureg.cm - self.b1g.extra_drift), self.q3g, _Drift('Q3G_Q4G', XL=32.6 * _ureg.cm), self.q4g, _Drift('Q4G_Q5G', XL=33.4 * _ureg.cm), self.q5g, _Drift('Q5G_Q6G', XL=33.5 * _ureg.cm), self.q6g, _Drift('Q6G_Q7G', XL=36.0 * _ureg.cm), self.q7g, _Drift('Q7G_SL3G', XL=16.5682507 * _ureg.cm + 3 * _ureg.cm), self.sl3g, _Drift('SL3G_B2G', XL=30.9927502 * _ureg.cm + 3 * _ureg.cm - self.b2g.extra_drift), _Chamber(IA=2), _Chamber('Chamber4', IA=1, IFORM=1, J=0, C1=4.9 * _ureg.cm, C2=2.2 * _ureg.cm, C3=self.b2g.RM), self.b2g, _Chamber(IA=2), _Drift('B2G_SMX', XL=31.77 * _ureg.cm - self.b2g.extra_drift - (self.smx.length - 159 * _ureg.mm)/2), self.smx, _Drift('SMX_SMY', XL=13.04 * _ureg.cm - (self.smx.length - 159 * _ureg.mm) / 2 - ( self.smy.length - 109 * _ureg.mm) / 2), self.smy, _Drift('SMY_B3G', XL=20.39 * _ureg.cm - self.b3g.extra_drift), _Collimator('B3G_ENTRY', IA=1, IFORM=1, J=0, C1=11 / 2 * _ureg.cm, C2=9 / 2 * _ureg.cm), self.b3g, _Collimator('B3G_EXIT', IA=1, IFORM=1, J=0, C1=13.65 / 2 * _ureg.cm, C2=9 / 2 * _ureg.cm), _Drift('FINAL', XL=1101.071 * _ureg.mm - self.b3g.extra_drift), self.iso, ], ) self.kinematics = kinematics self.tracks: Optional[_pd.DataFrame] = None self.results: Optional[_ZgoubiResults] = None @property def line(self) -> _Input: """Provides the full CGTR input sequence. Returns: the CGTR input sequence. """ return self.zi @property def gantry_angle(self): """Gantry angle (in IBA reference system).""" pass @gantry_angle.setter def gantry_angle(self, angle): pass
[docs] def fit_dipoles(self, kinematics: _Kinematics, dipoles: Optional[List[DipoleIBA]] = None): """Adjusts the main field of the dipoles according to the beam energy. The adjustment is made so that the reference trajectory exists the magnet on axis. Args: kinematics: the beam kinematics dipoles: a list of dipoles to fit """ z = _Zgoubi() dipoles = dipoles or [self.b1g, self.b2g, self.b3g] for dipole in dipoles: dipole.fit(kinematics=kinematics, zgoubi=z) z.wait()
[docs] def run(self, zgoubi: zgoubidoo.Zgoubi, identifier: Optional[_MappedParametersType] = None, ): """ Args: zgoubi: TODO identifier: TODO Returns: """ zgoubi(code_input=self.zi, identifier=identifier or {})
[docs] def fit(self, zgoubi: zgoubidoo.Zgoubi, identifier: _MappedParametersType, fit: zgoubidoo.commands.Fit, ) -> zgoubidoo.commands.Fit: """ Args: zgoubi: identifier: fit: Returns: """ zi = _Input(name=f"FIT_{fit.LABEL1}", line=self.zi.line + [fit]) zi.IL = 0 zgoubi(code_input=zi, identifier=identifier) return fit
[docs] def shoot(self, x: float, y: float, zgoubi: zgoubidoo.Zgoubi, fit_type: _FitType = _Fit, ) -> zgoubidoo.commands.Fit: """ Args: x: y: zgoubi: TODO fit_type: Returns: The Fit command (instance object) with the results of the fit. """ fit = fit_type( PENALTY=1e-8, PARAMS=[ _Fit.Parameter(line=self.zi, place='SMX', parameter=SMX.B1_), _Fit.Parameter(line=self.zi, place='SMY', parameter=SMY.B1_), ], CONSTRAINTS=[ _Fit.EqualityConstraint(line=self.zi, place='ISO', variable=_Fit.FitCoordinates.Y, value=x), _Fit.EqualityConstraint(line=self.zi, place='ISO', variable=_Fit.FitCoordinates.Z, value=y), ] ) fit = self.fit(zgoubi=zgoubi or _Zgoubi(), identifier={'SPOT_X': x, 'SPOT_Y': y}, fit=fit, ) return fit
[docs] def spots(self, spots: Iterable[Tuple[float, float]], fit_type: _FitType = _Fit, with_tracks: bool = False, debug_fit: bool = False ) -> Union[_pd.DataFrame, List[_Fit]]: """ Args: spots: fit_type: with_tracks: debug_fit: Returns: """ self.zi.cleanup() self.beam.REFERENCE = 1 z = _Zgoubi() fits = [ self.shoot(x=float(spot[0]), y=float(spot[1]), zgoubi=z, fit_type=fit_type) for spot in spots ] z.cleanup() if debug_fit: return fits self.zi.IL = 2 if with_tracks else 0 self.beam.REFERENCE = 0 for f in fits: for p, r in f.results: self.zi.update(r.results) self.run(zgoubi=z, identifier={**p, **{'SMX.B1': r.results.at[1, 'final'], 'SMY.B1': r.results.at[2, 'final']}}, ) self.results = z.collect() tracks = self.results.tracks if len(tracks) == 0: return tracks self.tracks = tracks return tracks
[docs] def plot(self, artist: zgoubidoo.vis.ZgoubidooPlotlyArtist = None, start: Optional[Union[str, zgoubidoo.commands.Command]] = None, stop: Optional[Union[str, zgoubidoo.commands.Command]] = None, crosshair: bool = True ): """Plot the P1 beamline. TODO Args: artist: an artist object for the rendering start: first element of the beamline to be plotted stop: last element of the beamline to be plotted crosshair: draw a crosshair indicating the isocenter """ self.line.survey() if artist is None: artist = zgoubidoo.vis.ZgoubidooPlotlyArtist() artist.plot_beamline(self.line, start=start, stop=stop) return artist