def test_mul():
assert v1 * 10 == vector.obj(x=10, y=50)
assert 10 * v1 == vector.obj(x=10, y=50)
assert numpy.allclose(
a1 * 10, vector.array({"x": [10, 20, 30, 40], "y": [50, 60, 70, 80]})
)
assert numpy.allclose(
10 * a1, vector.array({"x": [10, 20, 30, 40], "y": [50, 60, 70, 80]})
)
with pytest.raises(TypeError):
v1 * v2
with pytest.raises(TypeError):
a1 * a2
def test_pass_through():
@numba.njit
def pass_through(obj):
return obj
array = vector.array({"px": [1, 2, 3], "py": [10, 20, 30]})
assert isinstance(array, vector._backends.numpy_.VectorNumpy)
assert isinstance(pass_through(array), vector._backends.numpy_.VectorNumpy)
def test_sub():
assert v1 - v2 == vector.obj(x=-9, y=-15)
assert numpy.allclose(
a1 - a2,
vector.array({"x": [-9, -98, -997, -9996], "y": [-15, -194, -1993, -19992]}),
)
assert numpy.allclose(
v1 - a2,
vector.array({"x": [-9, -99, -999, -9999], "y": [-15, -195, -1995, -19995]}),
)
assert numpy.allclose(
a2 - v1,
vector.array({"x": [9, 99, 999, 9999], "y": [15, 195, 1995, 19995]}),
)
with pytest.raises(TypeError):
v1 - 5
with pytest.raises(TypeError):
5 - v1
def test_add():
assert v1 + v2 == vector.obj(x=11, y=25)
assert numpy.allclose(
a1 + a2,
vector.array({"x": [11, 102, 1003, 10004], "y": [25, 206, 2007, 20008]}),
)
assert numpy.allclose(
v1 + a2,
vector.array({"x": [11, 101, 1001, 10001], "y": [25, 205, 2005, 20005]}),
)
assert numpy.allclose(
a2 + v1,
vector.array({"x": [11, 101, 1001, 10001], "y": [25, 205, 2005, 20005]}),
)
with pytest.raises(TypeError):
v1 + 5
with pytest.raises(TypeError):
5 + v1
def test_truediv():
assert v1 / 10 == vector.obj(x=0.1, y=0.5)
with pytest.raises(TypeError):
10 / v1
assert numpy.allclose(
a1 / 10, vector.array({"x": [0.1, 0.2, 0.3, 0.4], "y": [0.5, 0.6, 0.7, 0.8]})
)
with pytest.raises(TypeError):
10 / a1
with pytest.raises(TypeError):
v1 / v2
with pytest.raises(TypeError):
a1 / a2
def get_good_kaons(self, photons='one'):
"""
photons : None, 'one', 'all' -- как работать с фотонами из калориметра. None -- не добавлять их в данные,
'one' -- только пару с лучшим соответствием pi0, 'all' -- добавить все пары
"""
dat_tracks = self.get_dat_tracks()
dat_kaons = self.get_dat_kaons()
dat_glob = self.get_dat_glob()
dat_goods = dat_tracks.join(dat_kaons, how='inner')
goods = dat_goods.groupby('entry').agg(
num=('tz', 'count')).query('num==2').index
dat_goods = dat_goods.reset_index().set_index(
'entry').loc[goods].reset_index().set_index(['entry', 'subentry'])
dat_goods['tcharge'] = np.where(dat_goods['tcharge'] > 0, 'p', 'n')
dat_goods = pd.pivot_table(dat_goods.reset_index(),
values=[
'trho', 'tz', 'tdedx', 'tptot', 'tth',
'tphi', 'ksminv', 'ksalign', 'ksptot',
'ksdpsi', 'ksz0', 'kslen', 'ksth',
'ksphi'
],
index=['entry'],
columns=['tcharge'])
dat_goods.columns = [
'_'.join(map(lambda x: str(x), col)) for col in dat_goods.columns
]
dat_goods.drop([
'ksalign_n', 'ksminv_n', 'ksptot_n', 'ksdpsi_n', 'ksz0_n',
'kslen_n', 'ksth_n', 'ksphi_n'
],
axis=1,
inplace=True)
#kick badruns
dat_glob = dat_glob.query('badrun==False')
dat_goods = dat_goods.join(dat_glob, how='inner')
#add x1, x2
dat_goods = dat_goods.rename(
{
'ksalign_p': 'ksalign',
'ksminv_p': 'ksminv',
'ksptot_p': 'ksptot',
'ksdpsi_p': 'ksdpsi',
'ksz0_p': 'ksz0',
'kslen_p': 'kslen',
'ksth_p': 'ksth',
'ksphi_p': 'ksphi'
},
axis=1)
dat_goods['x1'], dat_goods['x2'] = get_x(dat_goods)
#calc recoil mass
vec = vector.array({
'pt':
dat_goods['ksptot'] * np.sin(dat_goods['ksth']),
'theta':
dat_goods['ksth'],
'phi':
dat_goods['ksphi'],
'mass':
dat_goods['ksminv'],
})
vec0 = vector.obj(px=0, py=0, pz=0, E=dat_goods['emeas'].mean() * 2)
dat_goods['recoil'] = (vec0 - vec).mass
del vec
#add photons
if photons is not None:
dat_photons = self.get_dat_photons()
dat_goods = pd.merge(dat_goods.reset_index(),
dat_photons.reset_index(),
on='entry',
how='left')
dat_goods['subentry'] = dat_goods['subentry'].fillna(0).astype(int)
dat_goods = dat_goods.set_index(['entry', 'subentry'])
if photons == 'one':
dat_goods = dat_goods.sort_values(
'M', ascending=True,
key=lambda x: np.abs(x - 134.97)).groupby('entry').agg(
'first')
return dat_goods
def test_add():
one = vector.Array([[{
"x": 1,
"y": 1.1
}, {
"x": 2,
"y": 2.2
}], [], [{
"x": 3,
"y": 3.3
}]])
two = vector.Array([{
"x": 10,
"y": 20
}, {
"x": 100,
"y": 200
}, {
"x": 1000,
"y": 2000
}])
assert isinstance(one.add(two), vector._backends.awkward_.VectorArray2D)
assert isinstance(two.add(one), vector._backends.awkward_.VectorArray2D)
assert one.add(two).tolist() == [
[{
"x": 11,
"y": 21.1
}, {
"x": 12,
"y": 22.2
}],
[],
[{
"x": 1003,
"y": 2003.3
}],
]
assert two.add(one).tolist() == [
[{
"x": 11,
"y": 21.1
}, {
"x": 12,
"y": 22.2
}],
[],
[{
"x": 1003,
"y": 2003.3
}],
]
two = vector.array({"x": [10, 100, 1000], "y": [20, 200, 2000]})
assert isinstance(one.add(two), vector._backends.awkward_.VectorArray2D)
assert isinstance(two.add(one), vector._backends.awkward_.VectorArray2D)
assert one.add(two).tolist() == [
[{
"x": 11,
"y": 21.1
}, {
"x": 12,
"y": 22.2
}],
[],
[{
"x": 1003,
"y": 2003.3
}],
]
assert two.add(one).tolist() == [
[{
"x": 11,
"y": 21.1
}, {
"x": 12,
"y": 22.2
}],
[],
[{
"x": 1003,
"y": 2003.3
}],
]
two = vector.obj(x=10, y=20)
assert isinstance(one.add(two), vector._backends.awkward_.VectorArray2D)
assert isinstance(two.add(one), vector._backends.awkward_.VectorArray2D)
assert one.add(two).tolist() == [
[{
"x": 11,
"y": 21.1
}, {
"x": 12,
"y": 22.2
}],
[],
[{
"x": 13,
"y": 23.3
}],
]
assert two.add(one).tolist() == [
[{
"x": 11,
"y": 21.1
}, {
"x": 12,
"y": 22.2
}],
[],
[{
"x": 13,
"y": 23.3
}],
]
def f(t, x, y):
"""vector field from the ODE system"""
h, v = x
dv = -9.81
dh = v
return array([dh, dv])
def RISE(t, x, y):
"""the ball touched the ground and now bounces back"""
h, v = x
v_new = -v * 0.8
return array([h, v_new]), y
def ev_top(t, x, y):
"""event for detecting if the ball reaches the peak"""
h, v = x
return v
graph = { # FSM graph
FALL: {
ev_bottom: RISE
},
RISE: {
ev_top: FALL
},
}
z0 = FALL # Init FSM state
x0 = array([5, 0]) # Init values for h, v
y0 = None # no discrete states
t0 = 0.0
t1 = 8.0
dt = 0.02
eps = 1e-3
if __name__ == '__main__':
# Run the simulation
t, x, y = hyint(f, x0, t0, t1, dt, graph, z0, eps, y0)
# Transpose the results
h, v = zip(*x)
#Plot the results
def test_pos():
assert +v1 == vector.obj(x=1, y=5)
assert numpy.allclose(+a1, vector.array({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]}))
def test_neg():
assert -v1 == vector.obj(x=-1, y=-5)
assert numpy.allclose(
-a1, vector.array({"x": [-1, -2, -3, -4], "y": [-5, -6, -7, -8]})
)
# Copyright (c) 2019-2021, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner.
#
# Distributed under the 3-clause BSD license, see accompanying file LICENSE
# or https://github.com/scikit-hep/vector for details.
import numpy
import pytest
import vector
v1 = vector.obj(x=1, y=5)
a1 = vector.array({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]})
v2 = vector.obj(x=10, y=20)
a2 = vector.array({"x": [10, 100, 1000, 10000], "y": [20, 200, 2000, 20000]})
def test_eq():
assert v1 == v1
assert not v1 == v2
assert (a1 == a1).all()
assert not (a1 == a2).any()
assert (v1 == a1).any()
assert not (v1 == a1).all()
assert (a1 == v1).any()
assert not (a1 == v1).all()
def test_ne():
assert not v1 != v1
assert v1 != v2
