diff --git a/graphicle/calculate.py b/graphicle/calculate.py index 0a1a1f5..d017d30 100644 --- a/graphicle/calculate.py +++ b/graphicle/calculate.py @@ -19,7 +19,6 @@ import numpy.lib.recfunctions as rfn import pyjet from pyjet import ClusterSequence, PseudoJet -from typicle import Types import graphicle as gcl @@ -32,8 +31,6 @@ "cluster_pmu", ] -_types = Types() - def azimuth_centre(pmu: gcl.MomentumArray, pt_weight: bool = True) -> float: """Calculates the central point in azimuth for a set of particles. @@ -60,6 +57,12 @@ def azimuth_centre(pmu: gcl.MomentumArray, pt_weight: bool = True) -> float: return float(np.angle(pol.sum())) +def pseudorapidity_centre(pmu: gcl.MomentumArray) -> float: + pt_norm = pmu.pt / pmu.pt.sum() + eta_wt_mid = (pmu.eta * pt_norm).sum() + return eta_wt_mid + + def combined_mass( pmu: gcl.MomentumArray | base.VoidVector, weight: base.DoubleVector | None = None, @@ -142,7 +145,7 @@ def _trace_vector( ) -> base.AnyVector: len_basis = len(basis) feat_fmt = rfn.structured_to_unstructured if is_structured else lambda x: x - color = np.zeros((len_basis, feat_dim), dtype=_types.double) + color = np.zeros((len_basis, feat_dim), dtype=" base.AnyVector: def _truthy(data: ty.Union[base.ArrayBase, base.AdjacencyBase]) -> bool: """Defines the truthy value of the graphicle data structures.""" - if len(data) == 0: - return False - return True + return not (len(data) == 0) ################################## @@ -758,8 +756,12 @@ def __eq__(self, other: base.MaskLike) -> "MaskArray": def __ne__(self, other: base.MaskLike) -> "MaskArray": return _mask_neq(self, other) - def copy(self): - return deepcopy(self) + def copy(self) -> "MaskGroup": + mask_copies = map(op.methodcaller("copy"), self._mask_arrays.values()) + return self.__class__( + cl.OrderedDict(zip(self._mask_arrays.keys(), mask_copies)), + agg_op=self._agg_op, # type: ignore + ) @property def names(self) -> ty.List[str]: diff --git a/graphicle/select.py b/graphicle/select.py index 075c5e1..c7b2b3b 100644 --- a/graphicle/select.py +++ b/graphicle/select.py @@ -31,6 +31,7 @@ "hadron_vertices", "fastjet_clusters", "leaf_masks", + "centroid_prune", ] @@ -300,7 +301,28 @@ def _partition_vertex( Parameters ---------- - mask : gcl.MaskArray + mask : MaskArray + Hard parton descendants. + pcls_in : MaskArray or ndarray[bool_] + The particles entering the hadronisation vertex. + vtx_desc : MaskArray + Particles descending from the hadronisation vertex. + final : MaskArray + Final state particles. + pmu : MomentumArray + Four momenta. + dist_strat : callable + Callable which takes two ``MomentumArray`` instances, and + returns a double array with number of rows and columns equal to + the lengths of the input momenta, respectively. Output should + represent pairwise distance between particles incident on the + hadronisation vertex, and the final state descendants. + + Returns + ------- + filtered_mask : MaskArray + Input ``MaskArray``, filtered to remove background incident on + the same hadronisation vertex. """ mask = mask.copy() parton_pmu = pmu[pcls_in] @@ -358,7 +380,12 @@ def partition_descendants( if vtx_id not in graph.edges["out"][mask.data]: continue mask.data = _partition_vertex( - mask, pcls_in, vtx_desc, graph.final, graph.pmu, dist_strat + mask, + pcls_in, + vtx_desc, + graph.final, + graph.pmu, + dist_strat, ).data return hier @@ -714,3 +741,57 @@ def any_overlap(masks: gcl.MaskGroup) -> bool: pair_checks = map(np.bitwise_and, *zip(*combos)) overlaps: bool = np.bitwise_or.reduce(tuple(pair_checks), axis=None) return overlaps + + +def centroid_prune( + pmu: gcl.MomentumArray, + radius: float, + mask: ty.Optional[gcl.MaskArray] = None, + centre: ty.Optional[ty.Tuple[float, float]] = None, +) -> gcl.MaskArray: + """For a given ``MomentumArray``, calculate the distance every + particle is from a centroid location, and return a ``MaskArray`` for + all of the particles which are within a given ``radius``. + If ``centre`` is not provided, the transverse momentum weighted + centroid will be used. + + :group: select + + .. versionadded:: 0.2.4 + + Parameters + ---------- + pmu : MomentumArray + Four-momenta for a set of particles. + radius : float + Euclidean distance in the azimuth-pseudorapidity plane from the + centroid, beyond which particles will be filtered out. + mask : MaskArray, optional + If provided, will apply the mask to the passed ``pmu``, and + output ``MaskArray`` will have the same length. + centre : tuple[float, float] + Pseudorapidity and azimuth coordinates for a user-defined + centroid. + + Returns + ------- + prune_mask : MaskArray + Mask which retains only the particles within ``radius`` of the + centroid. + """ + if mask is not None: + pmu = pmu[mask] + event_mask = np.zeros_like(mask, "