pybes3.tracks

Helix operations

HelixObject

Source code in src/pybes3/tracks/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/tracks/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: Optional[np.ndarray] = None,
    pivot: TypeObjPosition = (0, 0, 0)
) -> HelixObject

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

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

Source code in src/pybes3/tracks/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,
    )

---

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/tracks/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/tracks/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/tracks/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/tracks/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	IntLike	helix[2] parameter, kappa.	required

Returns:

Type	Description
FloatLike	charge of the helix.

Source code in src/pybes3/tracks/helix.py

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

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

    Returns:
        charge of the helix.
    """
    return np.int8(1) if kappa > 1e-10 else np.int8(-1) if kappa < -1e-10 else np.int8(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/tracks/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/tracks/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:
            float: The radius of the helix in mm.
        """
        return kappa_to_radius(self.kappa)

    def change_pivot(self, *args):
        multi_trk = isinstance(self.pivot.x, ak.Array)

        # regularize pivot
        if multi_trk:
            if len(args) == 3:
                x, y, z = args
                new_pivot = None
            elif len(args) == 1:
                new_pivot = args[0]
                if isinstance(new_pivot, (vector.VectorObject3D, ak.Array, ak.Record)):
                    x = new_pivot.x
                    y = new_pivot.y
                    z = new_pivot.z
                    new_pivot = None
                else:
                    x = new_pivot[0]
                    y = new_pivot[1]
                    z = new_pivot[2]
                    new_pivot = None
            else:
                raise ValueError(
                    "change_pivot requires either 3 arguments (x, y, z) or 1 argument (pivot)."
                )

            if new_pivot is None:
                x = (
                    ak.ones_like(self.dr) * x
                    if not isinstance(x, (ak.Array, ak.Record))
                    else x
                )
                y = (
                    ak.ones_like(self.dr) * y
                    if not isinstance(y, (ak.Array, ak.Record))
                    else y
                )
                z = (
                    ak.ones_like(self.dr) * z
                    if not isinstance(z, (ak.Array, ak.Record))
                    else z
                )
                new_pivot = ak.Array({"x": x, "y": y, "z": z}, with_name="Vector3D")

            old_pivot_x = _flat_to_numpy(self.pivot.x)
            old_pivot_y = _flat_to_numpy(self.pivot.y)
            old_pivot_z = _flat_to_numpy(self.pivot.z)
            old_pivot = vector.arr({"x": old_pivot_x, "y": old_pivot_y, "z": old_pivot_z})

            new_pivot_x = _flat_to_numpy(new_pivot.x)
            new_pivot_y = _flat_to_numpy(new_pivot.y)
            new_pivot_z = _flat_to_numpy(new_pivot.z)
            new_pivot = vector.arr({"x": new_pivot_x, "y": new_pivot_y, "z": new_pivot_z})

        else:
            new_pivot = _regularize_obj_position(args)
            old_pivot = vector.obj(
                x=self.pivot.x,
                y=self.pivot.y,
                z=self.pivot.z,
            )

        # do transformation
        r = _flat_to_numpy(self.radius)
        old_dr = _flat_to_numpy(self.dr)
        old_phi0 = _flat_to_numpy(self.phi0)
        old_dz = _flat_to_numpy(self.dz)
        tanl = _flat_to_numpy(self.tanl)
        kappa = _flat_to_numpy(self.kappa)

        if "error" in self.fields:
            old_error = (
                _flat_to_numpy(self.error).reshape(-1, 5, 5)
                if multi_trk
                else _flat_to_numpy(self.error).reshape(5, 5)
            )
        else:
            old_error = None

        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=old_error,
            old_pivot=old_pivot,
            new_pivot=new_pivot,
        )

        res_dict = {
            "dr": new_dr,
            "phi0": new_phi0,
            "kappa": kappa,
            "dz": new_dz,
            "tanl": tanl,
            "pivot": ak.Record(
                {"x": new_pivot.x, "y": new_pivot.y, "z": new_pivot.z},
                with_name="Vector3D",
            ),
        }

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

        res = ak.Record(res_dict, with_name="Bes3Helix")

        if multi_trk:
            raw_shape = _extract_index(self.dr.layout)
            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.
        """
        if xor(
            ("error" in self.fields),
            ("error" in value.fields),
        ):
            warnings.warn(
                "One of the helix records has an error matrix, but the other does not. "
                "Ignoring error matrix for isclose check.",
                UserWarning,
            )

        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:

Name	Type	Description
float	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/tracks/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.
    """
    if xor(
        ("error" in self.fields),
        ("error" in value.fields),
    ):
        warnings.warn(
            "One of the helix records has an error matrix, but the other does not. "
            "Ignoring error matrix for isclose check.",
            UserWarning,
        )

    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/tracks/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.
        """
        if len(args) == 3:
            x, y, z = args
            pivot = None
        elif len(args) == 1:
            pivot = args[0]
            if isinstance(pivot, vector.VectorObject3D):
                x = pivot.x
                y = pivot.y
                z = pivot.z
                pivot = None
            elif not isinstance(pivot, ak.Array):
                x = pivot[0]
                y = pivot[1]
                z = pivot[2]
                pivot = None
        else:
            raise ValueError(
                "change_pivot requires either 3 arguments (x, y, z) or 1 argument (pivot)."
            )

        if pivot is None:
            x = ak.ones_like(self.dr) * x if not isinstance(x, ak.Array) else x
            y = ak.ones_like(self.dr) * y if not isinstance(y, ak.Array) else y
            z = ak.ones_like(self.dr) * z if not isinstance(z, ak.Array) else z
            pivot = ak.Array({"x": x, "y": y, "z": z}, with_name="Vector3D")

        raw_shape = _extract_index(self.dr.layout)

        r = _flat_to_numpy(self.radius)

        old_dr = _flat_to_numpy(self.dr)
        old_phi0 = _flat_to_numpy(self.phi0)
        old_dz = _flat_to_numpy(self.dz)
        tanl = _flat_to_numpy(self.tanl)
        kappa = _flat_to_numpy(self.kappa)

        old_pivot_x = _flat_to_numpy(self.pivot.x)
        old_pivot_y = _flat_to_numpy(self.pivot.y)
        old_pivot_z = _flat_to_numpy(self.pivot.z)
        old_pivot = vector.arr({"x": old_pivot_x, "y": old_pivot_y, "z": old_pivot_z})

        new_pivot_x = _flat_to_numpy(pivot.x)
        new_pivot_y = _flat_to_numpy(pivot.y)
        new_pivot_z = _flat_to_numpy(pivot.z)
        new_pivot = vector.arr({"x": new_pivot_x, "y": new_pivot_y, "z": new_pivot_z})

        old_error = (
            _flat_to_numpy(self.error).reshape(-1, 5, 5) if "error" in self.fields else None
        )

        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=old_error,
            old_pivot=old_pivot,
            new_pivot=new_pivot,
        )

        res_dict = {
            "dr": new_dr,
            "phi0": new_phi0,
            "kappa": kappa,
            "tanl": tanl,
            "dz": new_dz,
            "pivot": ak.zip(
                {"x": new_pivot.x, "y": new_pivot.y, "z": new_pivot.z},
                with_name="Vector3D",
            ),
        }

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

        raw_shape = _extract_index(self.dr.layout)
        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.
        """
        if xor(
            ("error" in self.fields),
            ("error" in other.fields),
        ):
            warnings.warn(
                "One of the helix records has an error matrix, but the other does not. "
                "Ignoring error matrix for isclose check.",
                UserWarning,
            )

        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/tracks/helix.py

def change_pivot(self, *args) -> "HelixAwkwardArray":
    """
    Changes the pivot point of the helix.
    """
    if len(args) == 3:
        x, y, z = args
        pivot = None
    elif len(args) == 1:
        pivot = args[0]
        if isinstance(pivot, vector.VectorObject3D):
            x = pivot.x
            y = pivot.y
            z = pivot.z
            pivot = None
        elif not isinstance(pivot, ak.Array):
            x = pivot[0]
            y = pivot[1]
            z = pivot[2]
            pivot = None
    else:
        raise ValueError(
            "change_pivot requires either 3 arguments (x, y, z) or 1 argument (pivot)."
        )

    if pivot is None:
        x = ak.ones_like(self.dr) * x if not isinstance(x, ak.Array) else x
        y = ak.ones_like(self.dr) * y if not isinstance(y, ak.Array) else y
        z = ak.ones_like(self.dr) * z if not isinstance(z, ak.Array) else z
        pivot = ak.Array({"x": x, "y": y, "z": z}, with_name="Vector3D")

    raw_shape = _extract_index(self.dr.layout)

    r = _flat_to_numpy(self.radius)

    old_dr = _flat_to_numpy(self.dr)
    old_phi0 = _flat_to_numpy(self.phi0)
    old_dz = _flat_to_numpy(self.dz)
    tanl = _flat_to_numpy(self.tanl)
    kappa = _flat_to_numpy(self.kappa)

    old_pivot_x = _flat_to_numpy(self.pivot.x)
    old_pivot_y = _flat_to_numpy(self.pivot.y)
    old_pivot_z = _flat_to_numpy(self.pivot.z)
    old_pivot = vector.arr({"x": old_pivot_x, "y": old_pivot_y, "z": old_pivot_z})

    new_pivot_x = _flat_to_numpy(pivot.x)
    new_pivot_y = _flat_to_numpy(pivot.y)
    new_pivot_z = _flat_to_numpy(pivot.z)
    new_pivot = vector.arr({"x": new_pivot_x, "y": new_pivot_y, "z": new_pivot_z})

    old_error = (
        _flat_to_numpy(self.error).reshape(-1, 5, 5) if "error" in self.fields else None
    )

    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=old_error,
        old_pivot=old_pivot,
        new_pivot=new_pivot,
    )

    res_dict = {
        "dr": new_dr,
        "phi0": new_phi0,
        "kappa": kappa,
        "tanl": tanl,
        "dz": new_dz,
        "pivot": ak.zip(
            {"x": new_pivot.x, "y": new_pivot.y, "z": new_pivot.z},
            with_name="Vector3D",
        ),
    }

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

    raw_shape = _extract_index(self.dr.layout)
    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/tracks/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.
    """
    if xor(
        ("error" in self.fields),
        ("error" in other.fields),
    ):
        warnings.warn(
            "One of the helix records has an error matrix, but the other does not. "
            "Ignoring error matrix for isclose check.",
            UserWarning,
        )

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

---

helix_awk(*args, **kwargs)

helix_awk(
    helix: ak.Array,
    error: Optional[ak.Array] = 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: Optional[ak.Array] = None,
    pivot: TypeAwkPosition = (0, 0, 0)
) -> HelixAwkwardArray

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

Source code in src/pybes3/tracks/helix.py

def helix_awk(*args, **kwargs) -> HelixAwkwardArray:
    error = None
    pivot = (0, 0, 0)

    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

    error = kwargs.pop("error", error)
    pivot = kwargs.pop("pivot", pivot)

    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")

---

Old helix

Deprecated

This component is deprecated and will be removed in the future. Please use the new pybes3.tracks.helix module instead.

parse_helix(helix, library='auto')

Parse helix parameters to physical parameters.

Parameters:

Name	Type	Description	Default
helix	Union[Array, ndarray]	helix parameters, the last dimension should be 5.	required
library	Literal['ak', 'auto']	the library to use, if "auto", return a dict when input is np.ndarray, return an ak.Array when input is ak.Array. If "ak", return an ak.Array.	'auto'

Returns:

Type	Description
Union[Array, dict[str, ndarray]]	parsed physical parameters. "x", "y", "z", "r" for position, "pt", "px", "py", "pz", "p", "theta", "phi" for momentum, "charge" for charge, "r_trk" for track radius.

Source code in src/pybes3/tracks/old_helix.py

def parse_helix(
    helix: Union[ak.Array, np.ndarray],
    library: Literal["ak", "auto"] = "auto",
) -> Union[ak.Array, dict[str, np.ndarray]]:
    """
    Parse helix parameters to physical parameters.

    Parameters:
        helix: helix parameters, the last dimension should be 5.
        library: the library to use, if "auto", return a dict when input is np.ndarray, \
            return an ak.Array when input is ak.Array. If "ak", return an ak.Array.

    Returns:
        parsed physical parameters. "x", "y", "z", "r" for position, \
            "pt", "px", "py", "pz", "p", "theta", "phi" for momentum, \
            "charge" for charge, "r_trk" for track radius.
    """

    warnings.warn(
        "'parse_helix' is deprecated and will be removed in the future, "
        "use 'helix_obj' and 'helix_awk' instead.",
        DeprecationWarning,
    )

    helix0 = helix[..., 0]
    helix1 = helix[..., 1]
    helix2 = helix[..., 2]
    helix3 = helix[..., 3]
    helix4 = helix[..., 4]

    x = _helix01_to_x(helix0, helix1)
    y = _helix01_to_y(helix0, helix1)
    z = helix3
    r = np.abs(helix0)

    pt = _helix2_to_pt(helix2)
    px = _pt_helix1_to_px(pt, helix1)
    py = _pt_helix1_to_py(pt, helix1)
    pz = pt * helix4
    p = _pt_helix4_to_p(pt, helix4)
    theta = _pz_p_to_theta(pz, p)
    phi = np.arctan2(py, px)

    charge = _helix2_to_charge(helix2)

    r_trk = pt * (10 / 2.99792458)  # |pt| * [GeV/c] / 1[e] / 1[T] = |pt| * 10/3 [m]
    r_trk = r_trk * 100  # to [cm]

    res_dict = {
        "x": x,
        "y": y,
        "z": z,
        "r": r,
        "px": px,
        "py": py,
        "pz": pz,
        "pt": pt,
        "p": p,
        "theta": theta,
        "phi": phi,
        "charge": charge,
        "r_trk": r_trk,
    }

    if isinstance(helix, ak.Array) or library == "ak":
        return ak.zip(res_dict)
    else:
        return res_dict

---

_helix01_to_x(helix0, helix1)

Convert helix parameters to x location.

Parameters:

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

Returns:

Type	Description
FloatLike	x location of the helix.

Source code in src/pybes3/tracks/old_helix.py

@nb.vectorize(cache=True)
def _helix01_to_x(helix0: FloatLike, helix1: FloatLike) -> FloatLike:
    """
    Convert helix parameters to x location.

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

    Returns:
        x location of the helix.
    """
    return helix0 * np.cos(helix1)

---

_helix01_to_y(helix0, helix1)

Convert helix parameters to y location.

Parameters:

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

Returns:

Type	Description
FloatLike	y location of the helix.

Source code in src/pybes3/tracks/old_helix.py

@nb.vectorize(cache=True)
def _helix01_to_y(helix0: FloatLike, helix1: FloatLike) -> FloatLike:
    """
    Convert helix parameters to y location.

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

    Returns:
        y location of the helix.
    """
    return helix0 * np.sin(helix1)

---

_helix2_to_pt(helix2)

Convert helix parameter to pt.

Parameters:

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

Returns:

Type	Description
FloatLike	pt of the helix.

Source code in src/pybes3/tracks/old_helix.py

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

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

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

---

_helix2_to_charge(helix2)

Convert helix parameter to charge.

Parameters:

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

Returns:

Type	Description
FloatLike	charge of the helix.

Source code in src/pybes3/tracks/old_helix.py

@nb.vectorize(cache=True)
def _helix2_to_charge(
    helix2: FloatLike,
) -> FloatLike:
    """
    Convert helix parameter to charge.

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

    Returns:
        charge of the helix.
    """
    if -1e-10 < helix2 < 1e-10:
        return np.int8(0)
    return np.int8(1) if helix2 > 0 else np.int8(-1)

---

_pt_helix1_to_px(pt, helix1)

Convert pt and helix1 to px.

Parameters:

Name	Type	Description	Default
pt	FloatLike	pt of the helix.	required
helix1	FloatLike	helix[1] parameter, phi0.	required

Returns:

Type	Description
FloatLike	px of the helix.

Source code in src/pybes3/tracks/old_helix.py

@nb.vectorize(cache=True)
def _pt_helix1_to_px(pt: FloatLike, helix1: FloatLike) -> FloatLike:
    """
    Convert pt and helix1 to px.

    Parameters:
        pt: pt of the helix.
        helix1: helix[1] parameter, phi0.

    Returns:
        px of the helix.
    """
    return -pt * np.sin(helix1)

---

_pt_helix1_to_py(pt, helix1)

Convert pt and helix1 to py.

Parameters:

Name	Type	Description	Default
pt	FloatLike	pt of the helix.	required
helix1	FloatLike	helix[1] parameter, phi0.	required

Returns:

Type	Description
FloatLike	py of the helix.

Source code in src/pybes3/tracks/old_helix.py

@nb.vectorize(cache=True)
def _pt_helix1_to_py(pt: FloatLike, helix1: FloatLike) -> FloatLike:
    """
    Convert pt and helix1 to py.

    Parameters:
        pt: pt of the helix.
        helix1: helix[1] parameter, phi0.

    Returns:
        py of the helix.
    """
    return pt * np.cos(helix1)

---

_pt_helix4_to_p(pt, helix4)

Convert pt and helix4 to p.

Parameters:

Name	Type	Description	Default
pt	FloatLike	pt of the helix.	required
helix4	FloatLike	helix[4] parameter, tanl.	required

Returns:

Type	Description
FloatLike	p of the helix.

Source code in src/pybes3/tracks/old_helix.py

@nb.vectorize(cache=True)
def _pt_helix4_to_p(pt: FloatLike, helix4: FloatLike) -> FloatLike:
    """
    Convert pt and helix4 to p.

    Parameters:
        pt: pt of the helix.
        helix4: helix[4] parameter, tanl.

    Returns:
        p of the helix.
    """
    return pt * np.sqrt(1 + helix4**2)

---