pybes3.helix

Helix operations

HelixObject

Source code in src/pybes3/helix.py

class HelixObject:
    def __init__(
        self,
        dr: float,
        phi0: float,
        kappa: float,
        dz: float,
        tanl: float,
        *,
        error: np.ndarray = None,
        pivot: TypeObjPosition = (0, 0, 0),
    ):
        self.dr = float(dr)
        self.phi0 = float(phi0)
        self.kappa = float(kappa)
        self.dz = float(dz)
        self.tanl = float(tanl)
        self.error = error
        self.pivot = _regularize_obj_position(pivot)

    @property
    def radius(self) -> float:
        """
        Circular radius of the helix (in cm).
        """
        return 1000 / 2.99792458 / np.abs(self.kappa)

    @property
    def momentum(self) -> vector.MomentumObject3D:
        """
        Momentum of the helix as a 3D vector.
        Note that the momentum is relative to the pivot, so once the pivot is changed,
        the momentum will also change accordingly.
        """
        pt = 1 / abs(self.kappa)
        pz = pt * self.tanl
        phi = (self.phi0 + np.pi / 2) % (2 * np.pi)
        return vector.obj(pt=pt, phi=phi, pz=pz)

    @property
    def position(self) -> vector.VectorObject3D:
        return vector.VectorObject3D(
            x=self.dr * math.cos(self.phi0),
            y=self.dr * math.sin(self.phi0),
            z=self.dz,
        )

    @property
    def charge(self) -> int:
        """
        Returns the charge of the helix.
        """
        return 1 if self.kappa > 1e-10 else -1 if self.kappa < -1e-10 else 0

    def change_pivot(self, *args):
        """
        Change the pivot point of the helix.
        The transformation refers to `Helix` class in `BOSS`.
        """
        # transform helix parameters
        r = self.radius

        old_dr = self.dr
        old_phi0 = self.phi0
        old_dz = self.dz
        tanl = self.tanl
        kappa = self.kappa

        old_pivot = self.pivot
        new_pivot = _regularize_obj_position(args)

        new_dr, new_phi0, new_dz, new_error = _change_pivot(
            r=r,
            old_dr=old_dr,
            old_phi0=old_phi0,
            old_dz=old_dz,
            kappa=kappa,
            tanl=tanl,
            old_error=self.error,
            old_pivot=old_pivot,
            new_pivot=new_pivot,
        )

        return HelixObject(
            dr=new_dr,
            phi0=new_phi0,
            kappa=kappa,
            dz=new_dz,
            tanl=tanl,
            error=new_error,
            pivot=new_pivot,
        )

    def __repr__(self) -> str:
        return f"Bes3Helix(dr={self.dr:.3f}, phi0={self.phi0:.3f}, kappa={self.kappa:.3f}, tanl={self.tanl:.3f}, dz={self.dz:.3f})"

    def isclose(
        self,
        other: "HelixObject",
        *,
        rtol: float = 1e-5,
        atol: float = 1e-8,
        equal_nan: bool = False,
    ) -> bool:
        if xor(
            (self.error is not None),
            (other.error is not None),
        ):
            warnings.warn(
                "One of the helix records has an error matrix, but the other does not. "
                "Ignoring error matrix for isclose check.",
                UserWarning,
            )

        return _obj_isclose(self, other, rtol=rtol, atol=atol, equal_nan=equal_nan)

radius property

Circular radius of the helix (in cm).

momentum property

Momentum of the helix as a 3D vector. Note that the momentum is relative to the pivot, so once the pivot is changed, the momentum will also change accordingly.

charge property

Returns the charge of the helix.

change_pivot(*args)

Change the pivot point of the helix. The transformation refers to Helix class in BOSS.

Source code in src/pybes3/helix.py

def change_pivot(self, *args):
    """
    Change the pivot point of the helix.
    The transformation refers to `Helix` class in `BOSS`.
    """
    # transform helix parameters
    r = self.radius

    old_dr = self.dr
    old_phi0 = self.phi0
    old_dz = self.dz
    tanl = self.tanl
    kappa = self.kappa

    old_pivot = self.pivot
    new_pivot = _regularize_obj_position(args)

    new_dr, new_phi0, new_dz, new_error = _change_pivot(
        r=r,
        old_dr=old_dr,
        old_phi0=old_phi0,
        old_dz=old_dz,
        kappa=kappa,
        tanl=tanl,
        old_error=self.error,
        old_pivot=old_pivot,
        new_pivot=new_pivot,
    )

    return HelixObject(
        dr=new_dr,
        phi0=new_phi0,
        kappa=kappa,
        dz=new_dz,
        tanl=tanl,
        error=new_error,
        pivot=new_pivot,
    )

---

helix_obj(*args, **kwargs)

helix_obj(
    dr: float,
    phi0: float,
    kappa: float,
    dz: float,
    tanl: float,
    *,
    error: np.ndarray | None = None,
    pivot: TypeObjPosition = (0, 0, 0),
) -> HelixObject

helix_obj(
    *,
    params: tuple[float, float, float, float, float],
    error: np.ndarray | None = None,
    pivot: TypeObjPosition = (0, 0, 0),
) -> HelixObject

helix_obj(
    *,
    momentum: TypeObjMomentum,
    position: TypeObjPosition,
    charge: Literal[-1, 1],
    error: np.ndarray | None = None,
    pivot: TypeObjPosition = (0, 0, 0),
) -> HelixObject

Source code in src/pybes3/helix.py

def helix_obj(*args, **kwargs) -> HelixObject:
    pivot = _regularize_obj_position(kwargs.pop("pivot", (0, 0, 0)))

    error = kwargs.pop("error", None)
    if isinstance(error, ak.Array):
        error = error.to_numpy()

    # given helix parameters as positional arguments
    if len(args) > 0:
        _check_kwargs_used_up(kwargs)
        return HelixObject(*args, error=error, pivot=pivot)

    # given helix parameters as keyword arguments
    if "dr" in kwargs:
        dr = kwargs.pop("dr")
        phi0 = kwargs.pop("phi0")
        kappa = kwargs.pop("kappa")
        dz = kwargs.pop("dz")
        tanl = kwargs.pop("tanl")

        _check_kwargs_used_up(kwargs)
        return HelixObject(
            dr=dr, phi0=phi0, kappa=kappa, dz=dz, tanl=tanl, pivot=pivot, error=error
        )

    # given helix parameters as a tuple
    if "params" in kwargs:
        params = kwargs.pop("params")
        if len(params) != 5:
            raise ValueError("params must be a tuple of 5 elements")

        _check_kwargs_used_up(kwargs)
        return HelixObject(*params, pivot=pivot, error=error)

    # given momentum, position and charge
    charge: Literal[-1, 1] = int(kwargs.pop("charge"))
    assert charge in (-1, 1), "Charge must be either -1 or 1"

    momentum = _regularize_obj_momentum(kwargs.pop("momentum"))
    position = _regularize_obj_position(kwargs.pop("position"))

    # compute helix parameters
    kappa = charge / momentum.pt
    phi0 = (momentum.phi - np.pi / 2) % (2 * np.pi)

    dist = (position - pivot).to_2D()
    dr = dist.rho
    if not np.isclose(dist.phi % (2 * np.pi), phi0):
        dr *= -1

    dz = position.z - pivot.z
    tanl = momentum.pz / momentum.pt

    _check_kwargs_used_up(kwargs)
    return HelixObject(
        dr=dr,
        phi0=phi0,
        kappa=kappa,
        dz=dz,
        tanl=tanl,
        pivot=pivot,
        error=error,
    )

---

helix_awk(*args, **kwargs)

helix_awk(
    helix: ak.Array,
    error: ak.Array | None = None,
    pivot: TypeAwkPosition = (0, 0, 0),
) -> HelixAwkwardArray

helix_awk(
    *,
    dr: ak.Array,
    phi0: ak.Array,
    kappa: ak.Array,
    dz: ak.Array,
    tanl: ak.Array,
    error: ak.Array | None = None,
    pivot: TypeAwkPosition = (0, 0, 0),
) -> HelixAwkwardArray

helix_awk(
    *,
    momentum: ak.Array,
    position: ak.Array,
    charge: Literal[-1, 1] | ak.Array,
    error: ak.Array | None = None,
    pivot: TypeAwkPosition = (0, 0, 0),
) -> HelixAwkwardArray

Source code in src/pybes3/helix.py

def helix_awk(*args, **kwargs) -> HelixAwkwardArray:
    # Use a sentinel to distinguish "not provided" from an explicit None
    error = _SENTINEL
    pivot = _SENTINEL

    if len(args) > 0:
        helix = args[0]

        if len(args) > 1:
            error = args[1]
            if "error" in kwargs:
                raise ValueError("Cannot pass both helix and error as positional arguments.")

        if len(args) > 2:
            pivot = args[2]
            if "pivot" in kwargs:
                raise ValueError("Cannot pass both helix and pivot as positional arguments.")

        dr = helix[..., 0]
        phi0 = helix[..., 1]
        kappa = helix[..., 2]
        dz = helix[..., 3]
        tanl = helix[..., 4]

    elif "helix" in kwargs:
        helix = kwargs.pop("helix")

        dr = helix[..., 0]
        phi0 = helix[..., 1]
        kappa = helix[..., 2]
        dz = helix[..., 3]
        tanl = helix[..., 4]

    elif "dr" in kwargs:
        dr = kwargs.pop("dr")
        phi0 = kwargs.pop("phi0")
        kappa = kwargs.pop("kappa")
        dz = kwargs.pop("dz")
        tanl = kwargs.pop("tanl")

    else:
        momentum = kwargs.pop("momentum")
        position = kwargs.pop("position")
        charge = kwargs.pop("charge")

        pivot = kwargs.pop("pivot", (0, 0, 0))
        if not isinstance(pivot, ak.Array):
            pivot = _regularize_obj_position(pivot)

        # compute helix parameters
        kappa = charge / momentum.pt
        phi0 = (momentum.phi - np.pi / 2) % (2 * np.pi)

        dist = (position - pivot).to_2D()
        dr = _fix_dr_sign(dist.rho, phi0, dist.phi)

        dz = position.z - pivot.z
        tanl = momentum.pz / momentum.pt

    # Only pop from kwargs if not already set via positional args
    if error is _SENTINEL:
        error = kwargs.pop("error", None)
    if pivot is _SENTINEL:
        pivot = kwargs.pop("pivot", (0, 0, 0))

    if not isinstance(pivot, ak.Array):
        pivot = _regularize_obj_position(pivot)

        x0 = ak.ones_like(dr) * pivot.x
        y0 = ak.ones_like(dr) * pivot.y
        z0 = ak.ones_like(dr) * pivot.z
        pivot = ak.zip({"x": x0, "y": y0, "z": z0}, with_name="Vector3D")

    res_dict = {
        "dr": dr,
        "phi0": phi0,
        "kappa": kappa,
        "dz": dz,
        "tanl": tanl,
        "pivot": pivot,
    }

    if error is not None:
        res_dict["error"] = error

    _check_kwargs_used_up(kwargs)

    raw_shape = _extract_index(dr.layout)
    return ak.zip(res_dict, depth_limit=len(raw_shape) + 1, with_name="Bes3Helix")

---

dr_phi0_to_x(dr, phi0)

Convert helix parameters to x location.

Parameters:

Name	Type	Description	Default
dr	float	helix[0] parameter, dr.	required
phi0	float	helix[1] parameter, phi0.	required

Returns:

Type	Description
FloatLike	x location of the helix.

Source code in src/pybes3/helix.py

@nb.vectorize(cache=True)
def dr_phi0_to_x(dr: FloatLike, phi0: FloatLike) -> FloatLike:
    """
    Convert helix parameters to x location.

    Parameters:
        dr (float): helix[0] parameter, dr.
        phi0 (float): helix[1] parameter, phi0.

    Returns:
        x location of the helix.
    """
    return dr * np.cos(phi0)

---

dr_phi0_to_y(dr, phi0)

Convert helix parameters to y location.

Parameters:

Name	Type	Description	Default
dr	FloatLike	helix[0] parameter, dr.	required
phi0	FloatLike	helix[1] parameter, phi0.	required

Returns:

Type	Description
FloatLike	y location of the helix.

Source code in src/pybes3/helix.py

@nb.vectorize(cache=True)
def dr_phi0_to_y(dr: FloatLike, phi0: FloatLike) -> FloatLike:
    """
    Convert helix parameters to y location.

    Parameters:
        dr: helix[0] parameter, dr.
        phi0: helix[1] parameter, phi0.

    Returns:
        y location of the helix.
    """
    return dr * np.sin(phi0)

---

phi0_to_phi(phi0)

Convert helix parameter phi0 to momentum phi.

Parameters:

Name	Type	Description	Default
phi0	FloatLike	helix[1] parameter, phi0.	required

Returns:

Type	Description
FloatLike	phi of the momentum vector.

Source code in src/pybes3/helix.py

@nb.vectorize(cache=True)
def phi0_to_phi(phi0: FloatLike) -> FloatLike:
    """
    Convert helix parameter phi0 to momentum phi.

    Parameters:
        phi0: helix[1] parameter, phi0.

    Returns:
        phi of the momentum vector.
    """
    return (phi0 + np.pi / 2) % (2 * np.pi)

---

kappa_to_pt(kappa)

Convert helix parameter to pt.

Parameters:

Name	Type	Description	Default
kappa	FloatLike	helix[2] parameter, kappa.	required

Returns:

Type	Description
FloatLike	pt of the helix.

Source code in src/pybes3/helix.py

@nb.vectorize(cache=True)
def kappa_to_pt(kappa: FloatLike) -> FloatLike:
    """
    Convert helix parameter to pt.

    Parameters:
        kappa: helix[2] parameter, kappa.

    Returns:
        pt of the helix.
    """
    return 1 / np.abs(kappa)

---

kappa_to_charge(kappa)

Convert helix parameter to charge.

Parameters:

Name	Type	Description	Default
kappa	FloatLike	helix[2] parameter, kappa.	required

Returns:

Type	Description
IntLike	charge of the helix.

Source code in src/pybes3/helix.py

@nb.vectorize(cache=True)
def kappa_to_charge(kappa: FloatLike) -> IntLike:
    """
    Convert helix parameter to charge.

    Parameters:
        kappa: helix[2] parameter, kappa.

    Returns:
        charge of the helix.
    """
    return 1 if kappa > 1e-10 else -1 if kappa < -1e-10 else 0

---

kappa_to_radius(kappa)

Convert helix parameter kappa to circular radius.

Parameters:

Name	Type	Description	Default
kappa	FloatLike	helix[2] parameter, kappa.	required

Returns:

Type	Description
FloatLike	circular radius of the helix in cm.

Source code in src/pybes3/helix.py

@nb.vectorize(cache=True)
def kappa_to_radius(kappa: FloatLike) -> FloatLike:
    """
    Convert helix parameter kappa to circular radius.

    Parameters:
        kappa: helix[2] parameter, kappa.

    Returns:
        circular radius of the helix in cm.
    """
    return 1000 / 2.99792458 / np.abs(kappa)

---

HelixAwkwardRecord

Bases: Record

Source code in src/pybes3/helix.py

class HelixAwkwardRecord(ak.Record):
    @property
    def momentum(self) -> vector.MomentumObject3D:
        """
        Returns the momentum of the helix as a 3D vector.

        Returns:
            vector.MomentumObject3D: The momentum vector of the helix.
        """
        pt, phi, pz = _compute_momentum(self.kappa, self.tanl, self.phi0)
        return ak.zip({"pt": pt, "phi": phi, "pz": pz}, with_name="Momentum3D")

    @property
    def position(self) -> vector.VectorObject3D:
        """
        Returns the position of the helix at a given azimuthal angle.

        Returns:
            vector.VectorObject3D: The position vector of the helix.
        """
        x, y, z = _compute_position(self.dr, self.phi0, self.dz)
        return ak.zip({"x": x, "y": y, "z": z}, with_name="Vector3D")

    @property
    def charge(self) -> int:
        """
        Returns the charge of the helix.

        Returns:
            int: The charge of the helix.
        """
        return kappa_to_charge(self.kappa)

    @property
    def radius(self) -> float:
        """
        Returns the radius of the helix.

        Returns:
            The radius of the helix in mm.
        """
        return kappa_to_radius(self.kappa)

    def change_pivot(self, *args) -> HelixAwkwardRecord:
        multi_trk = isinstance(self.pivot.x, ak.Array)
        res_dict, raw_shape = _awk_change_pivot(self, args, is_multi_trk=multi_trk)
        res = ak.Record(res_dict, with_name="Bes3Helix")
        for count in raw_shape:
            res = ak.unflatten(res, count)
        return res

    def isclose(
        self,
        value: "HelixAwkwardRecord",
        *,
        rtol: float = 1e-5,
        atol: float = 1e-8,
        equal_nan: bool = False,
    ) -> bool:
        """
        Check if two helix records are close to each other.

        Args:
            value (HelixAwkwardRecord): The helix record to compare with.

        Returns:
            bool: True if the records are close, False otherwise.
        """
        _helix_isclose_check_error(self.fields, value.fields)
        multi_trk = isinstance(self.pivot.x, ak.Array)
        if multi_trk:
            return _arr_isclose(self, value, rtol=rtol, atol=atol, equal_nan=equal_nan)
        else:
            return _obj_isclose(self, value, rtol=rtol, atol=atol, equal_nan=equal_nan)

momentum property

Returns the momentum of the helix as a 3D vector.

Returns:

Type	Description
MomentumObject3D	vector.MomentumObject3D: The momentum vector of the helix.

position property

Returns the position of the helix at a given azimuthal angle.

Returns:

Type	Description
VectorObject3D	vector.VectorObject3D: The position vector of the helix.

charge property

Returns the charge of the helix.

Returns:

Name	Type	Description
int	int	The charge of the helix.

radius property

Returns the radius of the helix.

Returns:

Type	Description
float	The radius of the helix in mm.

isclose(value, *, rtol=1e-05, atol=1e-08, equal_nan=False)

Check if two helix records are close to each other.

Parameters:

Name	Type	Description	Default
value	HelixAwkwardRecord	The helix record to compare with.	required

Returns:

Name	Type	Description
bool	bool	True if the records are close, False otherwise.

Source code in src/pybes3/helix.py

def isclose(
    self,
    value: "HelixAwkwardRecord",
    *,
    rtol: float = 1e-5,
    atol: float = 1e-8,
    equal_nan: bool = False,
) -> bool:
    """
    Check if two helix records are close to each other.

    Args:
        value (HelixAwkwardRecord): The helix record to compare with.

    Returns:
        bool: True if the records are close, False otherwise.
    """
    _helix_isclose_check_error(self.fields, value.fields)
    multi_trk = isinstance(self.pivot.x, ak.Array)
    if multi_trk:
        return _arr_isclose(self, value, rtol=rtol, atol=atol, equal_nan=equal_nan)
    else:
        return _obj_isclose(self, value, rtol=rtol, atol=atol, equal_nan=equal_nan)

---

HelixAwkwardArray

Bases: Array

Source code in src/pybes3/helix.py

class HelixAwkwardArray(ak.Array):
    @property
    def momentum(self) -> vec_ak.MomentumAwkward3D:
        """
        Returns the momentum of the helix as an awkward array of 3D vectors.

        Returns:
            vector.MomentumNumpy3D: The momentum vectors of the helix.
        """
        pt, phi, pz = _compute_momentum(self.kappa, self.tanl, self.phi0)
        return ak.zip({"pt": pt, "phi": phi, "pz": pz}, with_name="Momentum3D")

    @property
    def position(self) -> vec_ak.VectorAwkward3D:
        """
        Returns the position of the helix at a given azimuthal angle as an awkward array of 3D vectors.

        Returns:
            vector.VectorNumpy3D: The position vectors of the helix.
        """
        x, y, z = _compute_position(self.dr, self.phi0, self.dz)
        return ak.zip({"x": x, "y": y, "z": z}, with_name="Vector3D")

    @property
    def charge(self) -> ak.Array:
        """
        Returns the charge of the helix as an awkward array.

        Returns:
            ak.Array: The charge of the helix.
        """
        return kappa_to_charge(self.kappa)

    @property
    def radius(self) -> ak.Array:
        """
        Returns the radius of the helix as an awkward array.

        Returns:
            ak.Array: The radius of the helix in mm.
        """
        return kappa_to_radius(self.kappa)

    def change_pivot(self, *args) -> "HelixAwkwardArray":
        """
        Changes the pivot point of the helix.
        """
        res_dict, raw_shape = _awk_change_pivot(self, args, is_multi_trk=True)
        res = ak.Array(res_dict, with_name="Bes3Helix")
        for count in raw_shape:
            res = ak.unflatten(res, count)
        return res

    def isclose(
        self,
        other: "HelixAwkwardRecord",
        *,
        rtol: float = 1e-5,
        atol: float = 1e-8,
        equal_nan: bool = False,
    ) -> ak.Array:
        """
        Check if two helix records are close to each other.

        Args:
            other (HelixAwkwardRecord): The helix record to compare with.

        Returns:
            ak.Array: A boolean array indicating if the records are close.
        """
        _helix_isclose_check_error(self.fields, other.fields)
        return _arr_isclose(self, other, rtol=rtol, atol=atol, equal_nan=equal_nan)

momentum property

Returns the momentum of the helix as an awkward array of 3D vectors.

Returns:

Type	Description
MomentumAwkward3D	vector.MomentumNumpy3D: The momentum vectors of the helix.

position property

Returns the position of the helix at a given azimuthal angle as an awkward array of 3D vectors.

Returns:

Type	Description
VectorAwkward3D	vector.VectorNumpy3D: The position vectors of the helix.

charge property

Returns the charge of the helix as an awkward array.

Returns:

Type	Description
Array	ak.Array: The charge of the helix.

radius property

Returns the radius of the helix as an awkward array.

Returns:

Type	Description
Array	ak.Array: The radius of the helix in mm.

change_pivot(*args)

Changes the pivot point of the helix.

Source code in src/pybes3/helix.py

def change_pivot(self, *args) -> "HelixAwkwardArray":
    """
    Changes the pivot point of the helix.
    """
    res_dict, raw_shape = _awk_change_pivot(self, args, is_multi_trk=True)
    res = ak.Array(res_dict, with_name="Bes3Helix")
    for count in raw_shape:
        res = ak.unflatten(res, count)
    return res

isclose(other, *, rtol=1e-05, atol=1e-08, equal_nan=False)

Check if two helix records are close to each other.

Parameters:

Name	Type	Description	Default
other	HelixAwkwardRecord	The helix record to compare with.	required

Returns:

Type	Description
Array	ak.Array: A boolean array indicating if the records are close.

Source code in src/pybes3/helix.py

def isclose(
    self,
    other: "HelixAwkwardRecord",
    *,
    rtol: float = 1e-5,
    atol: float = 1e-8,
    equal_nan: bool = False,
) -> ak.Array:
    """
    Check if two helix records are close to each other.

    Args:
        other (HelixAwkwardRecord): The helix record to compare with.

    Returns:
        ak.Array: A boolean array indicating if the records are close.
    """
    _helix_isclose_check_error(self.fields, other.fields)
    return _arr_isclose(self, other, rtol=rtol, atol=atol, equal_nan=equal_nan)

---