Q3D-Calibration/qdx/utils.py

121 lines
3.4 KiB
Python
Raw Normal View History

import uproot
2022-07-11 16:24:14 +08:00
import numpy as np
import matplotlib.pyplot as plt
from sklearn.mixture import GaussianMixture
def readFileData(file, n=6, m=8, minT=800, maxT=4000):
"""Read whole data from root file
Parameters
----------
file : str
root file path
n : int, optional
number of blocks, default 6
m : int, optional
number of binds, default 8
minT/maxT : int, optional
Filtering data, the sum of the left and right sides needs to be in the interval [minT, maxT]
min / max threshold
"""
data = uproot.open(file)["Tree1"]
ldata, rdata = [], []
for i in range(n):
for j in range(m):
na = i // 2
nc = j + 2 * m * (i % 2)
x = data["adc{:d}ch{:d}".format(na, nc)].array(library="np")
y = data["adc{:d}ch{:d}".format(na, nc + m)].array(library="np")
idx = np.where((x + y >= minT) & (x + y <= maxT))[0]
ldata.append(x[idx])
rdata.append(y[idx])
return ldata, rdata
def readBlockData(file, n, m=8, minT=800, maxT=4000):
"""Read block data from root file
2022-07-11 16:24:14 +08:00
Parameters
----------
file : str
root file path
n : int
No.n block
m : int, optional
number of binds, default 8
minT/maxT : int, optional
Filtering data, the sum of the left and right sides needs to be in the interval [minT, maxT]
min / max threshold
"""
data = uproot.open(file)["Tree1"]
2022-07-11 16:24:14 +08:00
ldata, rdata = [], []
for j in range(m):
na = n // 2
nc = j + 2 * m * (n % 2)
x = data["adc{:d}ch{:d}".format(na, nc)].array(library="np")
y = data["adc{:d}ch{:d}".format(na, nc + m)].array(library="np")
idx = np.where((x + y >= minT) & (x + y <= maxT))[0]
ldata.append(x[idx])
rdata.append(y[idx])
2022-07-11 16:24:14 +08:00
return ldata, rdata
2022-07-11 16:24:14 +08:00
def draw_scatter(data, title, s=0.1):
"""Draw points using scatter
Parameters
----------
s : float, optional
size of scatter point, default 0.1
"""
2022-07-11 16:24:14 +08:00
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(1, 1, 1)
for cluster in data:
ax.scatter(cluster[:, 0], cluster[:, 1], s=s)
fig.savefig(title, facecolor="w", transparent=False)
2022-07-11 16:24:14 +08:00
plt.close()
2022-07-11 16:24:14 +08:00
def get_hist(data, delta=1, maxN=50):
"""Gets the boundary of histogram that the maximum count is bigger than threshold
Parameters
----------
delta : int, optional
Minimum bin width. The bin width is an integer multiple of delta.
maxN : int, optional
Maximum count threshold
"""
2022-07-11 16:24:14 +08:00
step = delta
edge = np.arange(data.min(), data.max() + 1, step)
count, _ = np.histogram(data, bins=edge)
try:
while count.max() <= maxN:
step += delta
edge = np.arange(data.min(), data.max() + 1, step)
count, _ = np.histogram(data, bins=edge)
except:
edge = np.arange(data.min(), data.max() + 1, delta)
2022-07-11 16:24:14 +08:00
count, _ = np.histogram(data, bins=edge)
return count, (edge[1:] + edge[:-1]) / 2
def GMM_slash(data):
"""Using Gaussian Mixture Method (GMM) to decompose the data into noise and slashes"""
fit_data = np.array([])
model = GaussianMixture(n_components=2)
model.fit(data[:, :2])
ny = model.predict(data[:, :2])
for i in np.unique(ny):
idx = np.where(ny == i)[0]
fit_data = idx if len(idx) > len(fit_data) else fit_data
return data[fit_data]