142 lines
4.2 KiB
Python
142 lines
4.2 KiB
Python
import uproot
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from sklearn.mixture import GaussianMixture
|
|
|
|
|
|
def readFileData(file, count, n=6, m=8, minT=800, maxT=4000):
|
|
"""Read whole data from root file
|
|
|
|
Parameters
|
|
----------
|
|
file : str
|
|
root file path
|
|
count : int
|
|
count that normalized by counts of Faraday cylinder
|
|
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")[:count]
|
|
y = data["adc{:d}ch{:d}".format(na, nc + m)].array(library="np")[:count]
|
|
idx = np.where((x + y >= minT) & (x + y <= maxT))[0]
|
|
ldata.append(x[idx])
|
|
rdata.append(y[idx])
|
|
|
|
return ldata, rdata
|
|
|
|
|
|
def readBlockData(file, count, n, m=8, minT=800, maxT=4000):
|
|
"""Read block data from root file
|
|
|
|
Parameters
|
|
----------
|
|
file : str
|
|
root file path
|
|
count : int
|
|
count that normalized by counts of Faraday cylinder
|
|
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"]
|
|
|
|
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")[:count]
|
|
y = data["adc{:d}ch{:d}".format(na, nc + m)].array(library="np")[:count]
|
|
idx = np.where((x + y >= minT) & (x + y <= maxT))[0]
|
|
ldata.append(x[idx])
|
|
rdata.append(y[idx])
|
|
|
|
return ldata, rdata
|
|
|
|
|
|
def draw_scatter(data, title, s=0.1):
|
|
"""Draw points using scatter
|
|
|
|
Parameters
|
|
----------
|
|
s : float, optional
|
|
size of scatter point, default 0.1
|
|
"""
|
|
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)
|
|
plt.close()
|
|
|
|
|
|
def get_hist(data, step=1, maxN=50, return_edge=False):
|
|
"""Gets the boundary of histogram that the maximum count is bigger than threshold
|
|
|
|
Parameters
|
|
----------
|
|
step : int, optional
|
|
Minimum bin width. The bin width is an integer multiple of step.
|
|
maxN : int, optional
|
|
Maximum count threshold
|
|
return_edge: bool, optional
|
|
If True, then the bin edges are also returned.
|
|
"""
|
|
delta = step
|
|
edge = np.arange(data.min(), data.max() + 1, delta)
|
|
count, _ = np.histogram(data, bins=edge)
|
|
try:
|
|
while count.max() <= maxN:
|
|
delta += step
|
|
edge = np.arange(data.min(), data.max() + 1, delta)
|
|
count, _ = np.histogram(data, bins=edge)
|
|
except:
|
|
edge = np.arange(data.min(), data.max() + 1, step)
|
|
count, _ = np.histogram(data, bins=edge)
|
|
|
|
if return_edge:
|
|
return count / delta, (edge[1:] + edge[:-1]) / 2, edge
|
|
else:
|
|
return count / delta, (edge[1:] + edge[:-1]) / 2
|
|
|
|
|
|
def GMM_clip(data, return_all=False):
|
|
"""Using Gaussian Mixture Method (GMM) to decompose the data into noise and available data
|
|
|
|
Parameters
|
|
----------
|
|
data : numpy.ndarray
|
|
Data to be clipped
|
|
return_all: bool, optional
|
|
If True, then all data will be returned.
|
|
"""
|
|
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
|
|
|
|
if return_all:
|
|
return ny
|
|
else:
|
|
return data[fit_data]
|