def reset_cov(self, big: bool = False):
"""
Reset the covariance matrix with initial values
"""
if not big:
self.cov = diag_matrix(
[1 / WX, 1 * VSLP, 1 / WY, 1 * VSLP, 1 / WT, 1 * VSLN])
else:
self.cov = diag_matrix(
[5 / WX, 50 * VSLP, 5 / WY, 50 * VSLP, 5 / WT, 10 * VSLN])
def __init__(self, root_inp: Union[np.array, None] = None,
mdet_inp: Union[np.array, None] = None,
fit_tracks: bool = False):
"""
--- TrackFinding class ---
Finds tracks for hits on any TRASGO-like detector
:param root_inp: ROOT input with values
trbnum, cell, col, row, x, y, z, time, charge
:param mdet_inp: Input with our own format
:param fit_tracks: Check if use or not TrackFitting class by Tim Track method
after use TrackFitting class by Kalman Filter Method
"""
# ------- Kalman Filter variables ------- #
self.r = np.zeros([NPLAN, NPAR]) # Vector (parameters); dimension -> Number of
# Planes x maximum hits x parameters
self.C = np.zeros([NPLAN, NPAR, NPAR]) # Error Matrices
self.rp = np.zeros(self.r.shape) # Projected vector and matrices
self.Cp = np.zeros(self.C.shape)
self.rn = np.zeros(self.r.shape) # UNUSED projected vectors with noises
self.Cn = np.zeros(self.C.shape)
self.C0 = diag_matrix(NPAR, [5 / WX, 50 * VSLP, 5 / WY, 50 * VSLP, 5 / WT, 10 * VSLN]) # Error matrix
self.V = diag_matrix(NDAC, [SIGX ** 2, SIGY ** 2, SIGT ** 2])
reco_dim = 1 + 3 * NPLAN + NPAR + 1 # Dimension (at axis 1) of reconstructed data matrix
self.m_stat = np.zeros([0, reco_dim])
self.mtrec = np.zeros([0, reco_dim])
self.lower_plane_id = NPLAN - 1
# --------------------------------------- #
self.mdet = mdet_inp
self.mdet_xy = None
if fit_tracks:
self.fit_tracks = TrackFitting()
else:
self.fit_tracks = None
def set_transport_func(ks: float, dz: int) -> np.array:
"""
It sets the transport k_mat between both planes separated by dz
:param ks: Slope -> sqrt( 1 + XP**2 + YP**2)
:param dz: Distance between planes
:return: Transport function (k_mat | Numpy array)
"""
F = diag_matrix([1] * NPAR) # Identity 6x6
F[0, 1] = dz
F[2, 3] = dz
F[4, 5] = ks * dz # - ks * dz
return F
def transport(self, dz: float):
"""
Displace the saeta and its covariance matrix.
Move to KFSaeta
"""
self.displace(dz)
transport_matrix = diag_matrix([1] * NPAR) # Identity 6x6
transport_matrix[0, 1] = dz
transport_matrix[2, 3] = dz
transport_matrix[4, 5] = self.ks * dz # - ks * dz
self.cov = transport_matrix @ self.cov @ transport_matrix.T
def cov(self):
"""
Move to KFHit
"""
return diag_matrix([SIGX**2, SIGY**2, SIGT**2])
def tim_track_fit(self, state_vec: Union[np.array, None] = None,
cells_path: Union[np.array, None] = None,
vstat_fcut: Union[np.array, None] = None):
"""
Tim Track Fitting Method
:param state_vec: State vector (SAETA)
:param cells_path: Array with parameters of sorted hits
:param vstat_fcut: Complete array with SAETAs and cell paths.
"""
# Assign saeta and cells safely
if state_vec is None and cells_path is None:
if vstat_fcut is not None:
k_dim = 1 + NDAC * NPLAN
self.saeta = vstat_fcut[k_dim:-1]
self.k_vector = vstat_fcut[:k_dim]
else:
raise Exception("All input values undefined!")
elif state_vec is not None and cells_path is not None:
if vstat_fcut is not None:
print("Variable vstat_cutf defined but unused, used state_vec and cells_path instead")
self.saeta = state_vec
self.k_vector = cells_path
vw = np.array([WX, WY, WT]) # Vector of Weights
mvw = diag_matrix(3, vw) # Wieght Matrix (inverse of the V diagonal k_mat)
# saeta = v_stat[13:-1]
# k_vector = v_stat[:13]
vs = self.saeta # all_reco_saetas[it, 13:-1] # State vector
mK = np.zeros([NPAR, NPAR]) # K k_mat initialization
va = np.zeros(NPAR) # Final state vector initialization
so = 0 # So (Store soi values on loop)
for ip in range(NPLAN): # Loop over hits in each track from TOP
zi = VZ1[ip] # [0, 522, 902, 1739] mm
ii = ip * 3 + 1 # Initial index to search values
dxi = self.k_vector[ii] * WCX
dyi = self.k_vector[ii + 1] * WCY
dti = self.k_vector[ii + 2]
vdat = np.array([dxi, dyi, dti]) # Measured data
mKi = self.set_K(vs, zi, mvw)
mG = self.set_mG(vs, zi)
vg0 = self.set_g0(vs, zi)
vai = self.set_vstat(mG, mvw, vdat, vg0)
mK += mKi
va += vai
so += np.dot((vdat - vg0).T, np.dot(mvw, (vdat - vg0))) # soi values
mK = np.asmatrix(mK) # K k_mat (symmetric)
mErr = mK.I # Error k_mat (symmetric)
va = np.asmatrix(va).T # Vertical Measurement Vector
vsol = np.dot(mErr, va) # SEA equation
sks = float(np.dot(vsol.T, np.dot(mK, vsol))) # s'·K·s
sa = float(np.dot(vsol.T, va)) # s'·a
S = sks - 2 * sa + so # S = s'·K·s - 2s'·a + So
DoF = NPLAN * NDAC - NPAR # Degrees of Freedom
prob = stats.chi2.sf(S, DoF)
vsol = np.asarray(vsol.T)[0] # Make it a normal array again
return vsol, prob