Utility Functions¶

Functions:

`EtaPhiPtE_to_cartesian`(p4)	Transform 4-momenta from polar coordinates to Cartesian coordinates for massless particles.
`EtaPhiPtE_to_relEtaPhiPt`(p4)	Get particle features in relative polar coordinates from 4-momenta in polar coordinates.
`YPhiPtE_to_cartesian`(p4)	Transform 4-momenta from polar coordinates to Cartesian coordinates.
`cartesian_to_EtaPhiPtE`(p4)	Transform 4-momenta from Cartesian coordinates to polar coordinates for massless particles.
`cartesian_to_YPhiPtE`(p4)	Transform 4-momenta from Cartesian coordinates to polar coordinates.
`cartesian_to_relEtaPhiPt`(p4)	Get particle features in relative polar coordinates from 4-momenta in Cartesian coordinates.
`efps`(jets[, use_particle_masses, ...])	Utility for calculating EFPs for jets in JetNet format using the energyflow library.
`gen_jet_corrections`(jets[, ...])	Zero's masked particles and negative pTs.
`jet_features`(jets)	Calculates jet features by summing over particle Lorentz 4-vectors.
`relEtaPhiPt_to_EtaPhiPt`(p_polarrel, jet_features)	Get particle features in absolute polar coordinates from relative polar coordinates and jet features.
`relEtaPhiPt_to_cartesian`(p_polarrel, ...[, ...])	Get particle features in absolute Cartesian coordinates from relative polar coordinates
`to_image`(jets, im_size[, mask, maxR])	Convert jet(s) into 2D `im_size` x `im_size` or 3D `num_jets` x `im_size` x `im_size` image arrays.

jetnet.utils.EtaPhiPtE_to_cartesian(p4: np.ndarray | torch.Tensor) → np.ndarray | torch.Tensor¶

Transform 4-momenta from polar coordinates to Cartesian coordinates for massless particles.

Parameters: p4 (np.ndarray or torch.Tensor) – array of 4-momenta in polar coordinates, of shape [..., 4]. The last axis should be in order \((\eta, \phi, p_\mathrm{T}, E/c)\), where \(\eta\) is the pseudorapidity.
Returns: array of 4-momenta in polar coordinates, arranged in order \((E/c, p_x, p_y, p_z)\).
Return type: np.ndarray or torch.Tensor

jetnet.utils.EtaPhiPtE_to_relEtaPhiPt(p4: np.ndarray | torch.Tensor) → np.ndarray | torch.Tensor¶

Get particle features in relative polar coordinates from 4-momenta in polar coordinates.

Parameters: p4 (np.ndarray or torch.Tensor) – array of 4-momenta in polar coordinates, of shape [..., 4]. The last axis should be in order \((\eta, \phi, p_\mathrm{T}, E/c)\), where \(\eta\) is the pseudorapidity.
Returns: array of features in relative polar coordinates, arranged in order \((\eta^\mathrm{rel}, \phi^\mathrm{rel}, p_\mathrm{T}^\mathrm{rel})\).
Return type: np.ndarray or torch.Tensor

jetnet.utils.YPhiPtE_to_cartesian(p4: np.ndarray | torch.Tensor) → np.ndarray | torch.Tensor¶

Transform 4-momenta from polar coordinates to Cartesian coordinates.

Parameters: p4 (np.ndarray or torch.Tensor) – array of 4-momenta in Cartesian coordinates, of shape [..., 4]. The last axis should be in order \((y, \phi, E/c, p_\mathrm{T})\), where \(y\) is the rapidity.
Returns: array of 4-momenta in polar coordinates, arranged in order \((E/c, p_x, p_y, p_z)\).
Return type: np.ndarray or torch.Tensor

jetnet.utils.cartesian_to_EtaPhiPtE(p4: np.ndarray | torch.Tensor) → np.ndarray | torch.Tensor¶

Transform 4-momenta from Cartesian coordinates to polar coordinates for massless particles.

Parameters: p4 (np.ndarray or torch.Tensor) – array of 4-momenta in Cartesian coordinates, of shape [..., 4]. The last axis should be in order \((E/c, p_x, p_y, p_z)\).
Returns: array of 4-momenta in polar coordinates, arranged in order \((\eta, \phi, p_\mathrm{T}, E/c)\), where \(\eta\) is the pseudorapidity.
Return type: np.ndarray or torch.Tensor

jetnet.utils.cartesian_to_YPhiPtE(p4: np.ndarray | torch.Tensor) → np.ndarray | torch.Tensor¶

Transform 4-momenta from Cartesian coordinates to polar coordinates.

Parameters: p4 (np.ndarray or torch.Tensor) – array of 4-momenta in Cartesian coordinates, of shape [..., 4]. The last axis should be in order \((E/c, p_x, p_y, p_z)\).
Returns: array of 4-momenta in polar coordinates, arranged in order \((y, \phi, E/c, p_\mathrm{T})\), where \(y\) is the rapidity.
Return type: np.ndarray or torch.Tensor

jetnet.utils.cartesian_to_relEtaPhiPt(p4: np.ndarray | torch.Tensor) → np.ndarray | torch.Tensor¶

Get particle features in relative polar coordinates from 4-momenta in Cartesian coordinates.

Parameters: p4 (np.ndarray or torch.Tensor) – array of 4-momenta in Cartesian coordinates, of shape [..., 4]. The last axis should be in order \((E/c, p_x, p_y, p_z)\).
Returns: array of features in relative polar coordinates, arranged in order \((\eta^\mathrm{rel}, \phi^\mathrm{rel}, p_\mathrm{T}^\mathrm{rel})\).
Return type: np.ndarray or torch.Tensor

jetnet.utils.efps(jets: ndarray, use_particle_masses: bool = False, efpset_args: list | None = None, efp_jobs: int | None = None) → ndarray¶

Utility for calculating EFPs for jets in JetNet format using the energyflow library.

Parameters

jets (np.ndarray) – array of either a single or multiple jets, of shape either [num_particles, num_features] or [num_jets, num_particles, num_features], with features in order [eta, phi, pt, (optional) mass]. If no particle masses given, they are assumed to be 0.
efpset_args (List) – Args for the energyflow.efpset function to specify which EFPs to use, as defined here https://energyflow.network/docs/efp/#efpset. Defaults to the n=4, d=5, prime EFPs.
efp_jobs (int) – number of jobs to use for energyflow’s EFP batch computation. None means as many processes as there are CPUs.

Returns

1D (if inputted single jet) or 2D array of shape [num_jets, num_efps] of EFPs per jet

Return type

np.ndarray

jetnet.utils.gen_jet_corrections(jets: ArrayLike, ret_mask_separate: bool = True, zero_mask_particles: bool = True, zero_neg_pt: bool = True, pt_index: int = 2) → ArrayLike | tuple[ArrayLike, ArrayLike]¶

Zero’s masked particles and negative pTs.

Parameters

jets (ArrayLike) – jets to recorrect.
ret_mask_separate (bool, optional) – return the jet and mask separately. Defaults to True.
zero_mask_particles (bool, optional) – set features of zero-masked particles to 0. Defaults to True.
zero_neg_pt (bool, optional) – set pT to 0 for particles with negative pt. Defaults to True.
pt_index (int, optional) – index of the pT feature. Defaults to 2.

Returns

Jets of same type as input, of shape [num_jets, num_particles, num_features (including mask)] if ret_mask_separate is False, else a tuple with a tensor/array of shape [num_jets, num_particles, num_features (excluding mask)] and another binary mask tensor/array of shape [num_jets, num_particles, 1].

jetnet.utils.jet_features(jets: ndarray) → dict[str, float | numpy.ndarray]¶

Calculates jet features by summing over particle Lorentz 4-vectors.

Parameters: jets (np.ndarray) – array of either a single or multiple jets, of shape either [num_particles, num_features] or [num_jets, num_particles, num_features], with features in order [eta, phi, pt, (optional) mass]. If no particle masses given, they are assumed to be 0.
Returns: dict of float (if inputted single jet) or 1D arrays of length num_jets (if inputted multiple jets) with ‘mass’, ‘pt’, and ‘eta’ keys.
Return type: Dict[str, Union[float, np.ndarray]]

jetnet.utils.relEtaPhiPt_to_EtaPhiPt(p_polarrel: np.ndarray | torch.Tensor, jet_features: np.ndarray | torch.Tensor, jet_coord: str = 'cartesian') → np.ndarray | torch.Tensor¶

Get particle features in absolute polar coordinates from relative polar coordinates and jet features.

Parameters

p_polarrel (np.ndarray or torch.Tensor) – array of particle features in relative polar coordinates of shape [..., 3]. The last axis should be in order \((\eta^\mathrm{rel}, \phi^\mathrm{rel}, p_\mathrm{T}^\mathrm{rel})\), where \(\eta\) is the pseudorapidity.
jet_features (np.ndarray or torch.Tensor) – array of jet features in polar coordinates, of shape [..., 4]. The coordinates are specified by jet_coord.
jet_coord (str) – coordinate system of jet features. Can be either “cartesian” or “polar”. Defaults to “cartesian”. If “cartesian”, the last axis of jet_features should be in order \((E/c, p_x, p_y, p_z)\). If “polar”, the last axis of jet_features should be in order \((\eta, \phi, p_\mathrm{T}, E/c)\).

Returns

array of particle features in absolute polar coordinates, arranged in order \((\eta, \phi, p_\mathrm{T}, E/c)\).

Return type

np.ndarray or torch.Tensor

jetnet.utils.relEtaPhiPt_to_cartesian(p_polarrel: np.ndarray | torch.Tensor, jet_features: np.ndarray | torch.Tensor, jet_coord: str = 'cartesian') → np.ndarray | torch.Tensor¶

Get particle features in absolute Cartesian coordinates from relative polar coordinates: and jet features.

Parameters

p_polarrel (np.ndarray or torch.Tensor) – array of particle features in relative polar coordinates of shape [..., 3]. The last axis should be in order \((\eta^\mathrm{rel}, \phi^\mathrm{rel}, p_\mathrm{T}^\mathrm{rel})\), where \(\eta\) is the pseudorapidity.
jet_features (np.ndarray or torch.Tensor) – array of jet features in polar coordinates, of shape [..., 4]. The coordinates are specified by jet_coord.
jet_coord (str) – coordinate system of jet features. Can be either “cartesian” or “polar”. Defaults to “cartesian”. If “cartesian”, the last axis of jet_features should be in order \((E/c, p_x, p_y, p_z)\). If “polar”, the last axis of jet_features should be in order \((\eta, \phi, p_\mathrm{T}, E/c)\).

Returns

array of particle features in absolute polar coordinates, arranged in order \((E/c, p_x, p_y, p_z)\).

Return type

np.ndarray or torch.Tensor

jetnet.utils.to_image(jets: ndarray, im_size: int, mask: ndarray = None, maxR: float = 1.0) → ndarray¶

Convert jet(s) into 2D im_size x im_size or 3D num_jets x im_size x im_size image arrays.

Parameters

jets (np.ndarray) – array of jet(s) of shape [num_particles, num_features] or [num_jets, num_particles, num_features] with features in order [eta, phi, pt].
im_size (int) – number of pixels per row and column.
mask (np.ndarray) – optional binary array of masks of shape [num_particles] or [num_jets, num_particles].
maxR (float) – max radius of the jet. Defaults to 1.0.

Returns

2D or 3D array of shape [im_size, im_size] or [num_jets, im_size, im_size].

Return type

np.ndarray