diff --git a/mbtrack2/utilities/optics.py b/mbtrack2/utilities/optics.py
index fdbd0091ab9166430c2cf329ed90fbe384931055..4c15073a95e9fef858b595d97717b84044edab58 100644
--- a/mbtrack2/utilities/optics.py
+++ b/mbtrack2/utilities/optics.py
@@ -523,8 +523,9 @@ class PhysicalModel:
return (1/sigma_star, a1_L, a2_L, a3_H, a4_H, a3_V, a4_V)
- def change_values(self, start_position, end_position, x_right, y_top,
- shape, rho, x_left=None, y_bottom=None):
+ def change_values(self, start_position, end_position, x_right=None,
+ y_top=None, shape=None, rho=None, x_left=None,
+ y_bottom=None, sym=True):
"""
Change the physical parameters between start_position and end_position.
@@ -532,41 +533,46 @@ class PhysicalModel:
----------
start_position : float
end_position : float
- x_right : float
+ x_right : float, optional
Right horizontal aperture in [m].
- y_top : float
+ y_top : float, optional
Top vertical aperture in [m].
- shape : str
+ shape : str, optional
Shape of the chamber cross section (elli/circ/rect).
- rho : float
+ rho : float, optional
Resistivity of the chamber material [ohm.m].
x_left : float, optional
Left horizontal aperture in [m].
y_bottom : float, optional
Bottom vertical aperture in [m].
+ sym : bool, optional
+ If True, right/left and top/bottum symmetry is applied.
"""
ind = (self.position > start_position) & (self.position < end_position)
- self.x_right[ind] = x_right
- self.y_top[ind] = y_top
-
- if x_left is None:
- self.x_left[ind] = -1*x_right
- else:
+ if x_right is not None:
+ self.x_right[ind] = x_right
+ if y_top is not None:
+ self.y_top[ind] = y_top
+ if x_left is not None:
self.x_left[ind] = x_left
-
- if y_bottom is None:
- self.y_bottom[ind] = -1*y_top
- else:
+ elif (x_right is not None) and (sym is True):
+ self.x_left[ind] = -1*x_right
+ if y_bottom is not None:
self.y_bottom[ind] = y_bottom
+ elif (y_top is not None) and (sym is True):
+ self.y_bottom[ind] = -1*y_top
ind2 = ((self.position[:-1] > start_position) &
(self.position[1:] < end_position))
- self.rho[ind2] = rho
- self.shape[ind2] = shape
-
- def taper(self, start_position, end_position, x_right_start, x_right_end,
- y_top_start, y_top_end, shape, rho, x_left_start=None,
- x_left_end=None, y_bottom_start=None, y_bottom_end=None):
+ if rho is not None:
+ self.rho[ind2] = rho
+ if shape is not None:
+ self.shape[ind2] = shape
+
+ def taper(self, start_position, end_position, x_right_start=None,
+ x_right_end=None, y_top_start=None, y_top_end=None, shape=None,
+ rho=None, x_left_start=None, x_left_end=None,
+ y_bottom_start=None, y_bottom_end=None, sym=True):
"""
Change the physical parameters to have a tapered transition between
start_position and end_position.
@@ -595,28 +601,31 @@ class PhysicalModel:
Bottom vertical aperture at the start of the taper in [m].
y_bottom_end : float, optional
Bottom vertical aperture at the end of the taper in [m].
+ sym : bool, optional
+ If True, right/left and top/bottum symmetry is applied.
"""
ind = (self.position > start_position) & (self.position < end_position)
- self.x_right[ind] = np.linspace(x_right_start, x_right_end, sum(ind))
- self.y_top[ind] = np.linspace(y_top_start, y_top_end, sum(ind))
-
+ if (x_right_start is not None) and (x_right_end is not None):
+ self.x_right[ind] = np.linspace(x_right_start, x_right_end, sum(ind))
+ if sym is True:
+ self.x_left[ind] = -1*np.linspace(x_right_start, x_right_end, sum(ind))
+
+ if (y_top_start is not None) and (y_top_end is not None):
+ self.y_top[ind] = np.linspace(y_top_start, y_top_end, sum(ind))
+ if sym is True:
+ self.y_bottom[ind] = -1*np.linspace(y_top_start, y_top_end, sum(ind))
+
if (x_left_start is not None) and (x_left_end is not None):
self.x_left[ind] = np.linspace(x_left_start, x_left_end, sum(ind))
- else:
- self.x_left[ind] = -1*np.linspace(x_right_start, x_right_end,
- sum(ind))
-
if (y_bottom_start is not None) and (y_bottom_end is not None):
- self.y_bottom[ind] = np.linspace(y_bottom_start, y_bottom_end,
- sum(ind))
- else:
- self.y_bottom[ind] = -1*np.linspace(y_top_start, y_top_end,
- sum(ind))
+ self.y_bottom[ind] = np.linspace(y_bottom_start, y_bottom_end, sum(ind))
- ind2 = ((self.position[:-1] > start_position)
- & (self.position[1:] < end_position))
- self.rho[ind2] = rho
- self.shape[ind2] = shape
+ ind2 = ((self.position[:-1] > start_position) &
+ (self.position[1:] < end_position))
+ if rho is not None:
+ self.rho[ind2] = rho
+ if shape is not None:
+ self.shape[ind2] = shape
def plot_aperture(self):
"""Plot horizontal and vertical apertures."""
@@ -633,3 +642,15 @@ class PhysicalModel:
ylabel="Vertical aperture [mm]")
axs[1].legend(["Top","Bottom"])
+ return (fig, axs)
+
+ def get_aperture(self, s):
+ self.xp = interp1d(self.position, self.x_right, kind='linear')
+ self.xm = interp1d(self.position, self.x_left, kind='linear')
+ self.yp = interp1d(self.position, self.y_top, kind='linear')
+ self.ym = interp1d(self.position, self.y_bottom, kind='linear')
+ aperture = np.array([self.xp(s),
+ self.xm(s),
+ self.yp(s),
+ self.ym(s)])
+ return aperture
\ No newline at end of file