from typing import Dict, List, Optional, Set, Tuple, Union
import numpy as np
import pandas as pd
__all__ = [
"to_cartesian",
"to_pt_eta_phi",
"delta_phi",
"twist",
"add_abs_mom",
"add_mass",
"add_energy",
"add_mt",
"get_vecs",
"fix_event_phi",
"fix_event_z",
"fix_event_y",
"event_to_cartesian",
"proc_event",
"calc_pair_mass",
"boost",
"boost2cm",
"get_momentum",
"cos_delta",
"delta_r",
"delta_r_boosted",
]
"""
Todo:
- Add non inplace versions/options
"""
[docs]def to_cartesian(df: pd.DataFrame, vec: str, drop: bool = False) -> None:
r"""
Vectoriesed conversion of 3-momenta to Cartesian coordinates inplace, optionally dropping old pT,eta,phi features
Arguments:
df: DataFrame to alter
vec: column prefix of vector components to alter, e.g. 'muon' for columns ['muon_pt', 'muon_phi', 'muon_eta']
drop: Whether to remove original columns and just keep the new ones
"""
z = f"{vec}_eta" in df.columns
try:
pt = df[f"{vec}_pT"]
pt_name = f"{vec}_pT"
except KeyError:
pt = df[f"{vec}_pt"]
pt_name = f"{vec}_pt"
if z:
eta = df[f"{vec}_eta"]
phi = df[f"{vec}_phi"]
df[f"{vec}_px"] = pt * np.cos(phi)
df[f"{vec}_py"] = pt * np.sin(phi)
if z:
df[f"{vec}_pz"] = pt * np.sinh(eta)
if drop:
df.drop(columns=[pt_name, f"{vec}_phi"], inplace=True)
if z:
df.drop(columns=[f"{vec}_eta"], inplace=True)
[docs]def to_pt_eta_phi(df: pd.DataFrame, vec: str, drop: bool = False) -> None:
r"""
Vectorised conversion of 3-momenta to pT,eta,phi coordinates inplace, optionally dropping old px,py,pz features
Arguments:
df: DataFrame to alter
vec: column prefix of vector components to alter, e.g. 'muon' for columns ['muon_px', 'muon_py', 'muon_pz']
drop: Whether to remove original columns and just keep the new ones
"""
eta = f"{vec}_pz" in df.columns
px = df[f"{vec}_px"]
py = df[f"{vec}_py"]
if eta:
pz = df[f"{vec}_pz"]
df[f"{vec}_pT"] = np.sqrt(np.square(px) + np.square(py))
if eta:
df[f"{vec}_eta"] = np.arcsinh(pz / df[f"{vec}_pT"])
df[f"{vec}_phi"] = np.arcsin(py / df[f"{vec}_pT"])
df.loc[(df[f"{vec}_px"] < 0) & (df[f"{vec}_py"] > 0), f"{vec}_phi"] = (
np.pi - df.loc[(df[f"{vec}_px"] < 0) & (df[f"{vec}_py"] > 0), f"{vec}_phi"]
)
df.loc[(df[f"{vec}_px"] < 0) & (df[f"{vec}_py"] < 0), f"{vec}_phi"] = -(
np.pi + df.loc[(df[f"{vec}_px"] < 0) & (df[f"{vec}_py"] < 0), f"{vec}_phi"]
)
df.loc[(df[f"{vec}_px"] < 0) & (df[f"{vec}_py"] == 0), f"{vec}_phi"] = np.random.choice(
(-np.pi, np.pi), df[(df[f"{vec}_px"] < 0) & (df[f"{vec}_py"] == 0)].shape[0]
)
if drop:
df.drop(columns=[f"{vec}_px", f"{vec}_py"], inplace=True)
if eta:
df.drop(columns=[f"{vec}_pz"], inplace=True)
[docs]def delta_phi(arr_a: Union[float, np.ndarray], arr_b: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
r"""
Vectorised computation of modulo 2pi angular seperation of array of angles b from array of angles a, in range [-pi,pi]
Arguments:
arr_a: reference angles
arr_b: final angles
Returns:
angular separation as float or np.ndarray
"""
df = pd.DataFrame() # Better way to do this without df?
df["dphi"] = arr_b - arr_a
while len(df[df.dphi > np.pi]) > 0:
df.loc[df.dphi > np.pi, "dphi"] -= 2 * np.pi
while len(df[df.dphi < -np.pi]) > 0:
df.loc[df.dphi < -np.pi, "dphi"] += 2 * np.pi
return df.dphi.values
[docs]def delta_r(dphi: Union[float, np.ndarray], deta: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
r"""
Vectorised computation of delta R separation for arrays of delta phi and delta eta (rapidity or pseudorapidity)
Arguments:
dphi: delta phi separations
deta: delta eta separations
Returns:
delta R separation as float or np.ndarray
"""
return np.sqrt(np.square(dphi) + np.square(deta))
[docs]def twist(dphi: Union[float, np.ndarray], deta: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
r"""
Vectorised computation of twist between vectors (https://arxiv.org/abs/1010.3698)
Arguments:
dphi: delta phi separations
deta: delta eta separations
Returns:
angular separation as float or np.ndarray
"""
return np.arctan(np.abs(dphi / deta))
[docs]def add_abs_mom(df: pd.DataFrame, vec: str, z: bool = True) -> None:
r"""
Vectorised computation 3-momenta magnitude, adding new column in place. Currently only works for Cartesian vectors
Arguments:
df: DataFrame to alter
vec: column prefix of vector components, e.g. 'muon' for columns ['muon_px', 'muon_py', 'muon_pz']
z: whether to consider the z-component of the momenta
"""
# TODO extend to work on pT, eta, phi vectors
if z and f"{vec}_pz" in df.columns:
df[f"{vec}_absp"] = np.sqrt(
np.square(df[f"{vec}_px"]) + np.square(df[f"{vec}_py"]) + np.square(df[f"{vec}_pz"])
)
else:
df[f"{vec}_absp"] = np.sqrt(np.square(df[f"{vec}_px"]) + np.square(df[f"{vec}_py"]))
[docs]def add_mass(df: pd.DataFrame, vec: str) -> None:
r"""
Vectorised computation of mass of 4-vector, adding new column in place.
Arguments:
df: DataFrame to alter
vec: column prefix of vector components, e.g. 'muon' for columns ['muon_px', 'muon_py', 'muon_pz']
"""
if f"{vec}_absp" not in df.columns:
add_abs_mom(df, vec)
df[f"{vec}_mass"] = np.sqrt(np.square(df[f"{vec}_E"]) - np.square(df[f"{vec}_absp"]))
[docs]def add_energy(df: pd.DataFrame, vec: str) -> None:
r"""
Vectorised computation of energy of 4-vector, adding new column in place.
Arguments:
df: DataFrame to alter
vec: column prefix of vector components, e.g. 'muon' for columns ['muon_px', 'muon_py', 'muon_pz']
"""
if f"{vec}_absp" not in df.columns:
add_abs_mom(df, vec)
df[f"{vec}_E"] = np.sqrt(
np.square(df[f"{vec}_mass"] if f"{vec}_mass" in df.columns else 0) + np.square(df[f"{vec}_absp"])
)
[docs]def add_mt(df: pd.DataFrame, vec: str, mpt_name: str = "mpt"):
r"""
Vectorised computation of transverse mass of 4-vector with respect to missing transverse momenta, adding new column in place.
Currently only works for pT, eta, phi vectors
Arguments:
df: DataFrame to alter
vec: column prefix of vector components, e.g. 'muon' for columns ['muon_px', 'muon_py', 'muon_pz']
mpt_name: column prefix of vector of missing transverse momenta components, e.g. 'mpt' for columns ['mpt_pT', 'mpt_phi']
"""
# TODO: extend to work on Cartesian coordinates
try:
df[f"{vec}_mT"] = np.sqrt(
2
* df[f"{vec}_pT"]
* df[f"{mpt_name}_pT"]
* (1 - np.cos(delta_phi(df[f"{vec}_phi"], df[f"{mpt_name}_phi"])))
)
except KeyError:
df[f"{vec}_mt"] = np.sqrt(
2
* df[f"{vec}_pt"]
* df[f"{mpt_name}_pt"]
* (1 - np.cos(delta_phi(df[f"{vec}_phi"], df[f"{mpt_name}_phi"])))
)
[docs]def get_vecs(feats: List[str], strict: bool = True) -> Set[str]:
r"""
Filter list of features to get list of 3-momenta defined in the list. Works for both pT, eta, phi and Cartesian coordinates.
If strict, return only vectors with all coordinates present in feature list.
Arguments:
feats: list of features to filter
strict: whether to require all 3-momenta components to be present in the list
Returns:
set of unique 3-momneta prefixes
"""
low = [f.lower() for f in feats]
all_vecs = [
f
for f in feats
if (f.lower().endswith("_pt") or f.lower().endswith("_phi") or f.lower().endswith("_eta"))
or (f.lower().endswith("_px") or f.lower().endswith("_py") or f.lower().endswith("_pz"))
]
if not strict:
return set([v[: v.rfind("_")] for v in all_vecs])
vecs = [
v[: v.rfind("_")]
for v in all_vecs
if (f'{v[:v.rfind("_")]}_pt'.lower() in low and f'{v[:v.rfind("_")]}_phi'.lower() in low)
or (f'{v[:v.rfind("_")]}_px'.lower() in low and f'{v[:v.rfind("_")]}_py'.lower() in low)
]
return set(vecs)
[docs]def fix_event_phi(df: pd.DataFrame, ref_vec: str) -> None:
r"""
Rotate event in phi such that ref_vec is at phi == 0. Performed inplace. Currently only works on vectors defined in pT, eta, phi
Arguments:
df: DataFrame to alter
ref_vec: column prefix of vector components to use as reference, e.g. 'muon' for columns ['muon_pT', 'muon_eta', 'muon_phi']
"""
# TODO: extend to work on Cartesian coordinates
for v in get_vecs(df.columns):
if v != ref_vec:
df[f"{v}_phi"] = delta_phi(df[f"{ref_vec}_phi"], df[f"{v}_phi"])
df[f"{ref_vec}_phi"] = 0
[docs]def fix_event_z(df: pd.DataFrame, ref_vec: str) -> None:
r"""
Flip event in z-axis such that ref_vec is in positive z-direction. Performed inplace. Works for both pT, eta, phi and Cartesian coordinates.
Arguments:
df: DataFrame to alter
ref_vec: column prefix of vector components to use as reference, e.g. 'muon' for columns ['muon_pT', 'muon_eta', 'muon_phi']
"""
if f"{ref_vec}_eta" in df.columns:
cut = df[f"{ref_vec}_eta"] < 0
for v in get_vecs(df.columns):
try:
df.loc[cut, f"{v}_eta"] = -df.loc[cut, f"{v}_eta"]
except KeyError:
print(f"eta component of {v} not found")
else:
cut = cut = df[f"{ref_vec}_pz"] < 0
for v in get_vecs(df.columns):
try:
df.loc[cut, f"{v}_pz"] = -df.loc[cut, f"{v}_pz"]
except KeyError:
print(f"pz component of {v} not found")
[docs]def fix_event_y(df: pd.DataFrame, ref_vec_0: str, ref_vec_1: str) -> None:
r"""
Flip event in y-axis such that ref_vec_1 has a higher py than ref_vec_0. Performed in place. Works for both pT, eta, phi and Cartesian coordinates.
Arguments:
df: DataFrame to alter
ref_vec_0: column prefix of vector components to use as reference 0, e.g. 'muon' for columns ['muon_pT', 'muon_eta', 'muon_phi']
ref_vec_1: column prefix of vector components to use as reference 1, e.g. 'muon' for columns ['muon_pT', 'muon_eta', 'muon_phi']
"""
if f"{ref_vec_1}_phi" in df.columns:
cut = df[f"{ref_vec_1}_phi"] < 0
for v in get_vecs(df.columns):
if v != ref_vec_0:
df.loc[cut, f"{v}_phi"] = -df.loc[cut, f"{v}_phi"]
else:
cut = df[f"{ref_vec_1}_py"] < 0
for v in get_vecs(df.columns):
if v != ref_vec_0:
df.loc[cut, f"{v}_py"] = -df.loc[cut, f"{v}_py"]
[docs]def event_to_cartesian(df: pd.DataFrame, drop: bool = False, ignore: Optional[List[str]] = None) -> None:
r"""
Convert entire event to Cartesian coordinates, except vectors listed in ignore. Optionally, drop old pT,eta,phi features. Perfomed inplace.
Arguments:
df: DataFrame to alter
drop: whether to drop old coordinates
ignore: vectors to ignore when converting
"""
for v in get_vecs(df.columns):
if ignore is None or v not in ignore:
to_cartesian(df, v, drop=drop)
[docs]def proc_event(
df: pd.DataFrame,
fix_phi: bool = False,
fix_y=False,
fix_z=False,
use_cartesian=False,
ref_vec_0: str = None,
ref_vec_1: str = None,
keep_feats: Optional[List[str]] = None,
default_vals: Optional[List[str]] = None,
) -> None:
r"""
Process event: Pass data through inplace various conversions and drop uneeded columns. Data expected to consist of vectors defined in pT, eta, phi.
Arguments:
df: DataFrame to alter
fix_phi: whether to rotate events using :meth:`~lumin.data_prcoessing.hep_proc.fix_event_phi`
fix_y: whether to flip events using :meth:`~lumin.data_prcoessing.hep_proc.fix_event_y`
fix_z: whether to flip events using :meth:`~lumin.data_prcoessing.hep_proc.fix_event_z`
use_cartesian: wether to convert vectors to Cartesian coordinates
ref_vec_0: column prefix of vector components to use as reference (0) for :meth:~lumin.data_prcoessing.hep_proc.fix_event_phi`,
:meth:`~lumin.data_prcoessing.hep_proc.fix_event_y`, and :meth:`~lumin.data_prcoessing.hep_proc.fix_event_z`
e.g. 'muon' for columns ['muon_pT', 'muon_eta', 'muon_phi']
ref_vec_1: column prefix of vector components to use as reference (1) for :meth:`~lumin.data_prcoessing.hep_proc.fix_event_y`,
e.g. 'muon' for columns ['muon_pT', 'muon_eta', 'muon_phi']
keep_feats: columns to keep which would otherwise be dropped
default_vals: list of default values which might be used to represent missing vector components. These will be replaced with np.nan.
"""
df.replace(
[np.inf, -np.inf] + default_vals if default_vals is not None else [np.inf, -np.inf], np.nan, inplace=True
)
if keep_feats is not None:
for f in keep_feats:
df[f"{f}keep"] = df[f"{f}"]
if fix_phi:
print(f"Setting {ref_vec_0} to phi = 0")
fix_event_phi(df, ref_vec_0)
if fix_y:
print(f"Setting {ref_vec_1} to positve phi")
fix_event_y(df, ref_vec_0, ref_vec_1)
if fix_z:
print(f"Setting {ref_vec_0} to positive eta")
fix_event_z(df, ref_vec_0)
if use_cartesian:
print("Converting to use Cartesian coordinates")
event_to_cartesian(df, drop=True)
if fix_phi and not use_cartesian:
df.drop(columns=[f"{ref_vec_0}_phi"], inplace=True)
elif fix_phi and use_cartesian:
df.drop(columns=[f"{ref_vec_0}_py"], inplace=True)
if keep_feats is not None:
for f in keep_feats:
df[f"{f}"] = df[f"{f}keep"]
df.drop(columns=[f"{f}keep"], inplace=True)
[docs]def calc_pair_mass(
df: pd.DataFrame, masses: Union[Tuple[float, float], Tuple[np.ndarray, np.ndarray]], feat_map: Dict[str, str]
) -> np.ndarray:
r"""
Vectorised computation of invarient mass o
f pair of particles with given masses, using 3-momenta. Only works for vectors defined in Cartesian coordinates.
Arguments:
df: DataFrame vector components
masses: tuple of masses of particles (either constant or different pair of masses per pair of particles)
feat_map: dictionary mapping of requested momentum components to the features in df
Returns:
np.ndarray of invarient masses
"""
# TODO: rewrite to not use a DataFrame for holding parent vector
# TODO: add inplace option
# TODO: extend to work on pT, eta, phi coordinates
tmp = pd.DataFrame()
tmp["0_E"] = np.sqrt(
(masses[0] ** 2)
+ np.square(df.loc[:, feat_map["0_px"]])
+ np.square(df.loc[:, feat_map["0_py"]])
+ np.square(df.loc[:, feat_map["0_pz"]])
)
tmp["1_E"] = np.sqrt(
(masses[1] ** 2)
+ np.square(df.loc[:, feat_map["1_px"]])
+ np.square(df.loc[:, feat_map["1_py"]])
+ np.square(df.loc[:, feat_map["1_pz"]])
)
tmp["p_px"] = df.loc[:, feat_map["0_px"]] + df.loc[:, feat_map["1_px"]]
tmp["p_py"] = df.loc[:, feat_map["0_py"]] + df.loc[:, feat_map["1_py"]]
tmp["p_pz"] = df.loc[:, feat_map["0_pz"]] + df.loc[:, feat_map["1_pz"]]
tmp["p_E"] = tmp.loc[:, "0_E"] + tmp.loc[:, "1_E"]
tmp["p_p2"] = np.square(tmp.loc[:, "p_px"]) + np.square(tmp.loc[:, "p_py"]) + np.square(tmp.loc[:, "p_pz"])
tmp["p_mass"] = np.sqrt(np.square(tmp.loc[:, "p_E"]) - tmp.loc[:, "p_p2"])
return tmp.p_mass.values
[docs]def boost(
ref_vec: Union[np.ndarray, str],
boost_vec: Union[np.ndarray, str],
df: Optional[pd.DataFrame] = None,
rescale_boost: bool = False,
) -> np.ndarray:
r"""
Vectorised boosting of reference vectors along boosting vectors.
N.B. Implementation adapted from ROOT (https://root.cern/)
Arguments:
vec_0: either (N,4) array of 4-momenta coordinates for starting vector,
or prefix name for starting vector, i.e. columns should have names of the form [vec_0]_px, etc.
vec_1: either (N,4) array of 4-momenta coordinates for boosting vector,
or prefix name for boosting vector, i.e. columns should have names of the form [vec_1]_px, etc.
df: DataFrame with data
rescale_boost: whether to divide the boost vector by its energy
Returns:
(N,4) array of boosted vector in Cartesian coordinates
"""
v = get_momentum(df, ref_vec, include_E=True, as_cart=True) if isinstance(ref_vec, str) else ref_vec
b = get_momentum(df, boost_vec, include_E=rescale_boost, as_cart=True) if isinstance(boost_vec, str) else boost_vec
if rescale_boost:
b = (b / b[:, 3:4])[:, :3]
b2 = np.square(np.linalg.norm(b, axis=1))
if b2.max() > 1:
raise ValueError("Boosting vector implies speed greater than c")
g = 1.0 / np.sqrt(1 - b2)
bp = np.sum(v[:, :3] * b, axis=1)
g2 = (g - 1) / b2
g2[b2 <= 0] = 0
bv = v.copy()
bv[:, :3] += (g2[:, None] * bp[:, None] * b) + (g[:, None] * b * v[:, 3][:, None])
bv[:, 3] += bp
bv[:, 3] *= g
return bv
[docs]def boost2cm(vec: Union[np.ndarray, str], df: Optional[pd.DataFrame] = None) -> np.ndarray:
r"""
Vectorised computation of boosting vector required to boost a vector to its centre-of-mass frame
Arguments:
vec: either (N,4) array of 4-momenta coordinates for starting vector,
or prefix name for starting vector, i.e. columns should have names of the form [vec]_px, etc.
df: DataFrame with data is supplying a string `vec`
Returns:
(N,3) array of boosting vector in Cartesian coordinates
"""
v = get_momentum(df, vec, include_E=True, as_cart=True) if isinstance(vec, str) else vec
return -(v / v[:, 3:4])[:, :3]
[docs]def get_momentum(df: pd.DataFrame, vec: str, include_E: bool = False, as_cart: bool = False) -> np.ndarray:
r"""
Extracts array of 3- or 4-momenta coordinates from DataFrame columns
Arguments:
df: DataFrame with data
vec: prefix name for vector, i.e. columns should have names of the form [vec]_px, etc.
as_cart: if True will return momenta in Cartesian coordinates
Returns:
(N, 3|4) array with columns: (px, py, pz, (E)) or (pT, phi, eta, (E))
"""
if f"{vec}_px" in df.columns and f"{vec}_py" in df.columns:
v = df[[f"{vec}_px", f"{vec}_py"]].values
v = (
np.hstack((v, df[f"{vec}_pz"].values[:, None]))
if f"{vec}_pz" in df.columns
else np.hstack((v, np.zeros_like(df.index.values[:, None])))
)
else:
pt = "pT" if f"{vec}_pT" in df.columns else "pt"
v = df[[f"{vec}_{pt}", f"{vec}_phi"]].values
v = (
np.hstack((v[:, 0], df[f"{vec}_eta"].values[:, None]), v[:, 1])
if f"{vec}_eta" in df.columns
else np.hstack((v[:, 0], np.zeros_like(df.index.values[:, None]), v[:, 1]))
)
if as_cart:
v = to_cartesian(pd.DataFrame(v, columns=["vec_pT", "vec_phi", "vec_eta"]), vec="vec", drop=True).values
if include_E:
if f"{vec}_E" not in df.columns:
add_energy(df, vec)
v = np.hstack((v, df[f"{vec}_E"].values[:, None]))
return v
[docs]def cos_delta(
vec_0: Union[np.ndarray, str],
vec_1: Union[np.ndarray, str],
df: Optional[pd.DataFrame] = None,
name: Optional[str] = None,
inplace: bool = False,
) -> Union[None, np.ndarray]:
r"""
Vectorised compututation of the cosine of the angular seperation of `vec_1` from `vec_0`
If `vec_*` are strings, then columns are extracted from DataFrame `df`.
If inplace is True Cosine angle is added a new column to the DataFrame with name `cosdelta_[vec_0]_[vec_1]` or `cosdelta`, unless `name` is set
Arguments:
vec_0: either (N,3) array of 3-momenta coordinates for vector 0,
or prefix name for vector zero, i.e. columns should have names of the form [vec_0]_px, etc.
vec_1: either (N,3) array of 3-momenta coordinates for vector 1,
or prefix name for vector one, i.e. columns should have names of the form [vec_1]_px, etc.
df: DataFrame with data
name: if set, will create a new column in df for cosdelta with given name, otherwise will generate a name
inplace: if True will add new column to df, otherwise will return array of cos_deltas
Returns:
array of cos deltas in not inplace
"""
v0 = get_momentum(df, vec_0) if isinstance(vec_0, str) else vec_0
v1 = get_momentum(df, vec_1) if isinstance(vec_1, str) else vec_1
if name is None:
name = f"cosdelta_{vec_0}_{vec_1}" if isinstance(vec_0, str) and isinstance(vec_1, str) else "cosdelta"
d = np.sum(v0 * v1, axis=1)
mag = np.linalg.norm(v0, axis=1) * np.linalg.norm(v1, axis=1)
if inplace:
df[name] = d / mag
else:
return d / mag
[docs]def delta_r_boosted(
vec_0: Union[np.ndarray, str],
vec_1: Union[np.ndarray, str],
ref_vec: Union[np.ndarray, str],
df: Optional[pd.DataFrame] = None,
name: Optional[str] = None,
inplace: bool = False,
) -> Union[None, np.ndarray]:
r"""
Vectorised compututation of the deltaR seperation of `vec_1` from `vec_0` in the rest-frame of another vector
If `vec_*` are strings, then columns are extracted from DataFrame `df`.
If inplace is True deltaR is added a new column to the DataFrame with name `dR_[vec_0]_[vec_1]_boosted_[ref_vec]` or `dR_boosted`, unless `name` is set
Arguments:
vec_0: either (N,4) array of 4-momenta coordinates for vector 0, in Cartesian coordinates
or prefix name for vector zero, i.e. columns should have names of the form [vec_0]_px, etc.
vec_1: either (N,4) array of 4-momenta coordinates for vector 1, in Cartesian coordinates
or prefix name for vector one, i.e. columns should have names of the form [vec_1]_px, etc.
ref_vec: either (N,4) array of 4-momenta coordinates for the vector in whos rest-frame deltaR should be computed, in Cartesian coordinates
or prefix name for reference vector, i.e. columns should have names of the form [ref_vec]_px, etc.
df: DataFrame with data
name: if set, will create a new column in df for cosdelta with given name, otherwise will generate a name
inplace: if True will add new column to df, otherwise will return array of cos_deltas
Returns:
array of boosted deltaR in not inplace
"""
br = boost2cm(ref_vec, df)[:, :3]
b0 = boost(vec_0, br, df)[:, :3]
b1 = boost(vec_1, br, df)[:, :3]
if name is None:
name = (
f"dR_{vec_0}_{vec_1}_boosted_{ref_vec}"
if isinstance(vec_0, str) and isinstance(vec_1, str) and isinstance(ref_vec, str)
else "dR_boosted"
)
tmp_df = pd.DataFrame(np.hstack((b0, b1)), columns=["0_px", "0_py", "0_pz", "1_px", "1_py", "1_pz"])
for v in range(2):
to_pt_eta_phi(tmp_df, str(v), drop=True)
dphi = delta_phi(tmp_df["0_phi"], tmp_df["1_phi"])
deta = tmp_df["0_eta"] - tmp_df["1_eta"]
dr = delta_r(dphi, deta)
if inplace:
df[name] = dr
else:
return dr