From c0969c8e289af956ae427fcf2af4f9083bce58fd Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Fri, 31 May 2024 11:57:32 +0100 Subject: [PATCH 01/12] cleared deprecated transforms, added momentum split #175 --- graphicle/transform.py | 315 +++++++++++++++++++++++++---------------- 1 file changed, 193 insertions(+), 122 deletions(-) diff --git a/graphicle/transform.py b/graphicle/transform.py index c13b2e6..70e6eda 100644 --- a/graphicle/transform.py +++ b/graphicle/transform.py @@ -4,158 +4,229 @@ Utilities for manipulating the graph structure of particle data. -.. deprecated:: 0.3.1 - Module is out of date, and will be removed in 0.4.0. """ -import networkx as _nx +import cmath +import operator as op +import typing as ty + import numpy as np -from deprecation import deprecated -from typicle import Types -from typicle.convert import cast_array +import numpy.typing as npt import graphicle as gcl -from . import base +_complex_unpack = op.attrgetter("real", "imag") + + +class SphericalAngle(ty.NamedTuple): + """Pair of inclination and azimuthal angles, respectively.""" + + theta: float + phi: float + + +class SphericalAxis(ty.NamedTuple): + x: float + y: float + z: float + -__all__ = ["particle_as_node", "centre_angle", "centre_pseudorapidity"] +def _angles_to_axis(angles: SphericalAngle) -> SphericalAxis: + axis = gcl.calculate._angles_to_axis(np.array([angles.phi, angles.theta])) + return SphericalAxis(*axis.tolist()) -_types = Types() +def _axis_to_angles(axis: SphericalAxis) -> SphericalAngle: + phi_polar = complex(axis.x, axis.y) + pt, phi = cmath.polar(phi_polar) + theta_polar = complex(pt, axis.z) + theta = cmath.phase(theta_polar) + return SphericalAngle(theta=theta, phi=phi) -@deprecated( - deprecated_in="0.3.1", - removed_in="0.4.0", - details="See ``networkx.line_graph()`` for potential replacement.", -) -def particle_as_node(adj_list: gcl.AdjacencyList) -> gcl.AdjacencyList: - """Converts an ``AdjacencyList`` in which the particles are - represented as edges, to one in which the particles are the nodes. - The order of the nodes in the resulting ``AdjacencyList`` retains - the same particle ordering of the initial edge list. - :group: transform +def _momentum_to_numpy(momenta: gcl.MomentumArray) -> gcl.base.DoubleVector: + return momenta.data.view(np.float64).reshape(-1, 4) - .. versionadded:: 0.1.0 + +def _cos_sin(angle: float) -> ty.Tuple[float, float]: + """Returns a tuple containing the sine and cosine of an angle.""" + return _complex_unpack(cmath.rect(1.0, angle)) + + +def soft_hard_axis(momenta: gcl.MomentumArray) -> SphericalAngle: + """Calculates the axis defined by the plane swept out between the + hardest and softest particles in ``momenta``. Parameters ---------- - adj_list : AdjacencyList - The edge-as-particle representation. + momenta : MomentumArray + Point cloud of particle four-momenta. Returns ------- - node_adj : AdjacencyList - The node-as-particle representation. + SphericalAngle + Normal axis to the soft-hard plane, defined in terms of + inclination and azimuthal angles. + """ + data = momenta.data.view(np.float64).reshape(-1, 4) + softest = data[momenta.pt.argmin()] + hardest = data[momenta.pt.argmax()] + axis = np.cross(softest[:3], hardest[:3]) + return _axis_to_angles(SphericalAxis(*axis.tolist())) - Examples - -------- - >>> from graphicle import transform - >>> # restructuring existing graph: - >>> graph.adj = transform.particle_as_node(graph.adj) + +def rotation_matrix( + angle: float, axis: SphericalAngle +) -> npt.NDArray[np.float64]: + """Computes the matrix operator to rotate a 3D vector with respect + to an arbitrary ``axis`` by a given ``angle``. + + Parameters + ---------- + angle : float + Desired angular displacement after matrix multiplication. + axis : SphericalAngle + Inclination and azimuthal angles, defining the axis about which + the rotation is to be performed. + + Returns + ------- + ndarray[float64] + A 3x3 matrix which, when acting to the left of a 3D vector, will + rotate it about the provided ``axis`` by the given ``angle``. + + Notes + ----- + This is a matrix implementation of Rodrigues' rotation formula [1]_. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula """ - # create the networkx edge graph - nx_edge_graph = _nx.MultiDiGraph() - graph_dicts = adj_list.to_dicts() - nx_edge_graph.add_edges_from(graph_dicts["edges"]) - # transform into node graph (with edge tuples rep'ing nodes) - nx_node_graph = _nx.line_graph(G=nx_edge_graph, create_using=_nx.DiGraph) - # create a node index for each particle - edges = adj_list.edges - num_pcls = len(edges) - node_idxs = np.empty(num_pcls, dtype=_types.int) - # translate nodes represented with edge tuples into node indices - edge_node_type = _types.edge.copy() - edge_node_type.append(("key", _types.int)) # if > 1 pcl between vtxs - edges_as_nodes = cast_array( - np.array(tuple(nx_node_graph.nodes)), edge_node_type - ) - - def check_sign(x): - """Returns -1 if x <= 0, and +1 if x > 0.""" - sign = np.sign(x) - sign = sign + int(not sign) # if sign = 0 => sign = +1 - return sign - - # node labels set as particle indices in original edge array - for i, node_triplet in enumerate(edges_as_nodes): - key = node_triplet["key"] - node = node_triplet[["src", "dst"]] - sign = -1 * check_sign(node["src"] * node["dst"]) - node_idxs[i] = sign * (np.where(edges == node)[0][key] + 1) - nx_node_graph = _nx.relabel_nodes( - nx_node_graph, - {n: idx for n, idx in zip(nx_node_graph, node_idxs)}, - ) - return gcl.AdjacencyList(np.array(nx_node_graph.edges)) - - -@deprecated( - deprecated_in="0.3.1", - removed_in="0.4.0", - details="Use ``calculate.resultant_coords()`` and " - "``data.MomentumArray.shift_phi()`` instead.", -) -def centre_angle( - angle: base.DoubleVector, pt: base.DoubleVector -) -> base.DoubleVector: - """Shifts angles so transverse momentum weighted centroid is at - ``0``. - - :group: transform - - .. versionadded:: 0.1.0 + axis_3d = _angles_to_axis(axis) + skew_sym = np.zeros((3, 3), dtype=np.float64) + upper_idxs = np.triu_indices(3, 1) + skew_sym[upper_idxs] = -axis_3d.z, axis_3d.y, -axis_3d.x + skew_sym[np.tril_indices(3, -1)] = -skew_sym[upper_idxs] + cos_alpha, sin_alpha = _cos_sin(angle) + rot = sin_alpha * skew_sym + (1.0 - cos_alpha) * (skew_sym @ skew_sym) + np.fill_diagonal(rot, rot.diagonal() + 1.0) + return rot + + +def split_momentum( + momentum: gcl.MomentumArray, + z: float, + angle: float, + axis: ty.Union[ty.Tuple[float, float], SphericalAngle], +) -> gcl.MomentumArray: + """Splits the momentum of the given particle into two momenta. + Energy and 3-momentum is conserved. Hardness and collinearity of the + split are determined by ``z`` and ``angle``. Parameters ---------- - angle : array - Angular displacements. - pt : array - Transverse momenta. + momentum : MomentumArray + Four-momentum prior to splitting. Must contain only one element. + z : float + Energy fraction retained by the first child after the split. + Must be in range ``0.0 < z <= 0.5``. + angle : float + Angular displacement of the first child after the split. + axis : SphericalAngle or tuple[float, float], optional + The theta and phi values of the axis vector about which to + rotate the momentum. If ``None``, will choose the axis normal to + the plane swept out by the hardest and softest momentum + constituents. Returns ------- - centred_angle : array - Shifted angular displacements, with centroid at 0. + MomentumArray + Four-momenta of two particles produced by splitting. + + See Also + -------- + soft_hard_axis : Axis of plane swept by softest and hardest momenta. + rotation_matrix : Matrix to rotate 3-vectors about a given axis. """ - # convert angles into complex polar positions - pos = np.exp(1.0j * angle) - # obtain weighted sum positions ie. un-normalised midpoint - pos_wt_mid = (pos * pt).sum() - # convert to U(1) rotation operator e^(-i delta x) - rot_op = (pos_wt_mid / np.abs(pos_wt_mid)).conjugate() - # rotate positions so midpoint is at 0 - pos_centred = rot_op * pos - return np.angle(pos_centred) # type: ignore - - -@deprecated( - deprecated_in="0.3.1", - removed_in="0.4.0", - details="Use ``calculate.resultant_coords()`` and " - "``data.MomentumArray.shift_eta()`` instead.", -) -def centre_pseudorapidity( - eta: base.DoubleVector, pt: base.DoubleVector -) -> base.DoubleVector: - """Shifts pseudorapidities so pt weighted midpoint is at ``0``. - - :group: transform - - .. versionadded:: 0.1.0 + if len(momentum) > 1: + raise ValueError("momentum must have only one element.") + if not isinstance(axis, SphericalAngle): + axis = SphericalAngle(*axis) + parent = _momentum_to_numpy(momentum) + children = np.tile(parent, (2, 1)) + children[0, :] *= z + children[0, :3] @= rotation_matrix(angle, axis).T + children[1, :] -= children[0, :] + return gcl.MomentumArray(children) + + +def split_hardest( + momenta: gcl.MomentumArray, + z: float, + angle: float, + axis: ty.Optional[ty.Union[ty.Tuple[float, float], SphericalAngle]] = None, +) -> gcl.MomentumArray: + """Splits the momentum of the hardest particle into two momenta. + Energy and 3-momentum is conserved over the whole MomentumArray. + Hardness and collinearity of the split are determined by function + parameters. Parameters ---------- - eta : ndarray[float64] - Values of pseudorapidity for the particle set. - pt : ndarray[float64] - Values of transverse momenta for the particle set. + momenta : MomentumArray + Set of four-momenta, representing the point cloud of particles, + prior to splitting. + z : float + Energy fraction retained by the first child after the split. + Must be in range ``0.0 < z <= 0.5``. + angle : float + Angular displacement of the first child after the split. + axis : SphericalAngle or tuple[float, float], optional + The theta and phi values of the axis vector about which to + rotate the momenta. If ``None``, will choose the axis normal to + the plane swept out by the hardest and softest momentum + constituents. Returns ------- - eta_centred : ndarray[float64] - Pseudorapidity values relative to the centre of transverse - momentum. + MomentumArray + Set of four-momenta after splitting, such that the length is + increased by one. The highest transverse momentum element has + been removed from the set, and replaced with two momenta + elements. The first and second children of the split are the + penultimate and final elements of the MomentumArray, + respectively. + + See Also + -------- + soft_hard_axis : Axis of plane swept by softest and hardest momenta. + rotation_matrix : Matrix to rotate 3-vectors about a given axis. + + Notes + ----- + This function is intended to check the IRC safety of our GNN jet + clustering algorithms. It is implemented from the description given + in a jet tagging paper [1]_, which defined the IRC safe message + passing procedure used in this work. + + References + ---------- + .. [1] https://doi.org/10.48550/arXiv.2109.14636 """ - pt_norm = pt / pt.sum() - eta_wt_mid = (eta * pt_norm).sum() - return eta - eta_wt_mid # type: ignore + if not (0.0 < z <= 0.5): + raise ValueError("z must be in range (0, 0.5]") + if axis is None: + if len(momenta) < 2: + raise ValueError( + "If axis is not provided, there must be at least two elements " + "in momenta." + ) + axis = soft_hard_axis(momenta) + data = _momentum_to_numpy(momenta) + hard_idx = momenta.pt.argmax() + out = np.empty((len(momenta) + 1, 4), dtype=np.float64) + out[:hard_idx] = data[:hard_idx] + out[hard_idx:-2] = data[(hard_idx + 1) :] + children = split_momentum(momenta[hard_idx], z, angle, axis) + out[-2:, ...] = _momentum_to_numpy(children)[...] + return gcl.MomentumArray(out) From db9fd197846a044303a6c6000d5abba28ee73810 Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Fri, 31 May 2024 12:19:05 +0100 Subject: [PATCH 02/12] added libtool to micromamba env --- .github/workflows/tests.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 53df731..794ca85 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -19,6 +19,7 @@ jobs: create-args: >- python=${{ matrix.pyver }} pytest>=6.2.5 + libtool init-shell: bash cache-environment: true post-cleanup: 'all' From e5a436495be6731054806b551f55e39af81fcc8d Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Fri, 31 May 2024 12:46:29 +0100 Subject: [PATCH 03/12] classmethod samples spherical uniform momenta #178 --- graphicle/data.py | 76 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/graphicle/data.py b/graphicle/data.py index 2cbc334..fdf0a9c 100644 --- a/graphicle/data.py +++ b/graphicle/data.py @@ -1333,6 +1333,82 @@ def __attrs_post_init__(self): self.dtype = np.dtype(list(zip("xyze", it.repeat(" "MomentumArray": + """Returns a ``MomentumArray`` whose elements are sampled from + uniform distributions of energy and 3-momentum. + + Parameters + ---------- + size : int + Number of elements. + max_energy : float + Upper bound for the energy of the momenta. + massless : float + Probability for any given momentum element to be massless. + Default is ``0.0``. + seed : int, optional + Random number generator seed. Use the same seed for + consistent results. + + Returns + ------- + MomentumArray + Set of momenta, sampled uniformly in the energy component, + and with uniform probability density over the surface of a + 3-sphere for the spatial components. + + Raises + ------ + ValueError + If massless is not within the interval [0.0, 1.0]. + + Notes + ----- + Some 'massless' particles may have small, but finite masses. + This is due to numerical errors associated with floating point + calculations. + """ + if not (0.0 <= massless <= 1.0): + raise ValueError( + "Probability of massless particles must be in the " + "interval [0, 1]." + ) + rng = np.random.default_rng(seed=seed) + energy = rng.uniform(0.0, max_energy, size) + p_mag = energy + if massless == 0.0: + mass = rng.uniform(0.0, energy) + p_mag = calculate._root_diff_two_squares(energy, mass) + elif massless < 1.0: + p_mag = p_mag.copy() + mask = rng.random(size) < massless + masked_e = energy[~mask] + mass = rng.uniform(0.0, masked_e) + p_mag[~mask] = calculate._root_diff_two_squares(masked_e, mass) + theta_polar = p_mag * np.exp( + 1.0j * np.arccos(1.0 - 2.0 * rng.random(size, dtype=np.float64)) + ) + phi_polar = theta_polar.real * np.exp( + 1.0j * rng.uniform(-np.pi, np.pi, size) + ) + return cls( + np.hstack( + [ + phi_polar.real.reshape(-1, 1), + phi_polar.imag.reshape(-1, 1), + theta_polar.imag.reshape(-1, 1), + energy.reshape(-1, 1), + ] + ) + ) + @property def __array_interface__(self) -> ty.Dict[str, ty.Any]: return self._data.__array_interface__ From 3ff4bfa3c71feeee1afbce3dfe3c0a3f40441102 Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Fri, 31 May 2024 15:48:33 +0100 Subject: [PATCH 04/12] test on macos-13 to avoid arm-only --- .github/workflows/tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 53df731..9ea75e4 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -8,7 +8,7 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - os: [ubuntu-latest, macos-latest] + os: [ubuntu-latest, macos-13] pyver: ["3.8", "3.9", "3.10", "3.11"] steps: - uses: actions/checkout@v3 From 0859e21ed23840202ec1f3c1e0beca6af6d90ef3 Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Fri, 31 May 2024 15:57:34 +0100 Subject: [PATCH 05/12] update micromamba version --- .github/workflows/tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 9ea75e4..8541e4f 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -14,7 +14,7 @@ jobs: - uses: actions/checkout@v3 - uses: mamba-org/setup-micromamba@v1 with: - micromamba-version: '1.3.1-0' + micromamba-version: '1.5.6-0' environment-name: test-env create-args: >- python=${{ matrix.pyver }} From bf124c06ac41e84db7d76124247ba862858df61a Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Fri, 31 May 2024 16:09:46 +0100 Subject: [PATCH 06/12] added docstrings to transform routines #175 --- graphicle/transform.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/graphicle/transform.py b/graphicle/transform.py index 70e6eda..e5847db 100644 --- a/graphicle/transform.py +++ b/graphicle/transform.py @@ -18,13 +18,25 @@ class SphericalAngle(ty.NamedTuple): - """Pair of inclination and azimuthal angles, respectively.""" + """Pair of inclination and azimuthal angles, respectively. + + :group: transform + + .. versionadded:: 0.4.0 + """ theta: float phi: float class SphericalAxis(ty.NamedTuple): + """Axis vector in 3D cartesian coordinates. + + :group: transform + + .. versionadded:: 0.4.0 + """ + x: float y: float z: float @@ -56,6 +68,10 @@ def soft_hard_axis(momenta: gcl.MomentumArray) -> SphericalAngle: """Calculates the axis defined by the plane swept out between the hardest and softest particles in ``momenta``. + :group: transform + + .. versionadded:: 0.4.0 + Parameters ---------- momenta : MomentumArray @@ -80,6 +96,10 @@ def rotation_matrix( """Computes the matrix operator to rotate a 3D vector with respect to an arbitrary ``axis`` by a given ``angle``. + :group: transform + + .. versionadded:: 0.4.0 + Parameters ---------- angle : float @@ -123,6 +143,10 @@ def split_momentum( Energy and 3-momentum is conserved. Hardness and collinearity of the split are determined by ``z`` and ``angle``. + :group: transform + + .. versionadded:: 0.4.0 + Parameters ---------- momentum : MomentumArray @@ -171,6 +195,10 @@ def split_hardest( Hardness and collinearity of the split are determined by function parameters. + :group: transform + + .. versionadded:: 0.4.0 + Parameters ---------- momenta : MomentumArray From fa420107efe1f96a99cbaf9727b372517e39665b Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Fri, 31 May 2024 16:11:49 +0100 Subject: [PATCH 07/12] module level docstring updated --- graphicle/transform.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/graphicle/transform.py b/graphicle/transform.py index e5847db..1b72fbf 100644 --- a/graphicle/transform.py +++ b/graphicle/transform.py @@ -2,8 +2,7 @@ ``graphicle.transform`` ======================= -Utilities for manipulating the graph structure of particle data. - +Utilities for mutating particle data. """ import cmath import operator as op From e34b6dd545c77cfa5cf8a85f3e776a3105b29648 Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Fri, 31 May 2024 16:19:40 +0100 Subject: [PATCH 08/12] removed deprecated packages for 0.4.0 release --- graphicle/calculate.py | 236 ----------------------------------------- 1 file changed, 236 deletions(-) diff --git a/graphicle/calculate.py b/graphicle/calculate.py index f32ff96..d13c1d1 100644 --- a/graphicle/calculate.py +++ b/graphicle/calculate.py @@ -13,14 +13,11 @@ import operator as op import typing as ty import warnings -from functools import lru_cache, partial -import networkx as nx import numba as nb import numpy as np import numpy.lib.recfunctions as rfn import scipy.optimize as spo -from deprecation import deprecated from . import base @@ -28,13 +25,9 @@ from graphicle.data import * __all__ = [ - "azimuth_centre", - "pseudorapidity_centre", - "rapidity_centre", "weighted_centroid", "resultant_coords", "combined_mass", - "flow_trace", "aggregate_momenta", "cluster_coeff_distbn", "thrust", @@ -49,96 +42,6 @@ ) -@deprecated( - deprecated_in="0.3.1", - removed_in="0.4.0", - details="Use ``weighted_centroid()`` or ``resultant_coords()`` instead.", -) -def azimuth_centre(pmu: "MomentumArray", pt_weight: bool = True) -> float: - """Calculates the central point in azimuth for a set of particles. - - :group: calculate - - .. versionadded:: 0.1.7 - - .. versionchanged:: 0.3.1 - The ``pt_weight`` parameter implementation has been corrected. - Previous versions resulted in :math:`p_T`-weighting when - ``False``, and :math:`p_T^2`-weighting when ``True``. - - Parameters - ---------- - pmu : MomentumArray - Four-momenta of the particles. - pt_weight : bool - If ``True``, will weight the contributions of each particle by - transverse momentum. Similar to finding a centre-of-mass, but - for transverse momentum. Default is ``True``. - - Returns - ------- - float - The centre of the particle set in the azimuth dimension. - """ - pol = pmu._xy_pol - if not pt_weight: - pol = pol / pmu.pt - return np.angle(pol.sum()).item() - - -@deprecated( - deprecated_in="0.3.1", - removed_in="0.4.0", - details="Use ``weighted_centroid()`` or ``resultant_coords()`` instead.", -) -def pseudorapidity_centre(pmu: "MomentumArray") -> float: - """Calculates the central point in pseudorapidity for a set of - particles. - - :group: calculate - - .. versionadded:: 0.1.7 - - Parameters - ---------- - pmu : MomentumArray - Four-momenta of the particles. - - Returns - ------- - float - The :math:`p_T` weighted centre of the particle set in the - pseudorapidity dimension. - """ - return ((pmu.eta * pmu.pt).sum() / pmu.pt.sum()).item() - - -@deprecated( - deprecated_in="0.3.1", - removed_in="0.4.0", - details="Use ``weighted_centroid()`` or ``resultant_coords()`` instead.", -) -def rapidity_centre(pmu: "MomentumArray") -> float: - """Calculates the central point in rapidity for a set of particles. - - :group: calculate - - .. versionadded:: 0.2.11 - - Parameters - ---------- - pmu : MomentumArray - Four-momenta of the particles. - - Returns - ------- - float - The :math:`p_T` weighted centre of the particle set in the - rapidity dimension. - """ - return (pmu.rapidity * pmu.pt).sum() / pmu.pt.sum() - - def weighted_centroid( pmu: "MomentumArray", pseudo: bool = True ) -> ty.Tuple[float, float]: @@ -286,145 +189,6 @@ def combined_mass( return mass -def _diffuse(colors: ty.List[base.AnyVector], feats: ty.List[base.AnyVector]): - color_shape = colors[0].shape - av_color = np.zeros((color_shape[0], color_shape[1]), dtype=" 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=" ty.Dict[str, base.DoubleVector]: - """Performs flow tracing from specified particles in an event, back - to the hard partons. - - :group: calculate - - .. versionadded:: 0.1.0 - - Parameters - ---------- - graph : Graphicle - Full particle event, containing hard partons, showering and - hadronisation. - mask : MaskArray or MaskGroup or ndarray[bool_] - Boolean mask identifying which particles should have their - ancestry traced. - prop : ArrayBase or ndarray[any] - Property to trace back, *eg.* 4-momentum, charge, *etc*. Must be - the same shape as arrays stored in graph. Can be structured, - unstructured, or a graphicle array, though unstructured arrays - must be 1D. - exclusive : bool - If ``True``, double counting from descendant particles in the - hard event will be switched off. *eg.* for event ``t > b W+``, - descendants of ``b`` will show no contribution from ``t``, as - ``b`` is a subset of ``t``. Default is ``False``. - target : set[int], optional - Highlights specific partons in the hard event to decompose - properties with respect to. If unset, will use all partons in - hard event, except for incoming partons. - - Returns - ------- - trace_array : dict[str, ndarray] - Keys are parton names, arrays represent the contributions of - hard partons traced down to the properties of the selected - subset of particles specified by mask. - """ - if isinstance(prop, base.ArrayBase): - prop = prop.data - # encoding graph features onto NetworkX - nx_graph = nx.DiGraph() - graph_dict = graph.adj.to_dicts(edge_data={"feat": prop}) - example_feat = graph_dict["edges"][0][2]["feat"] - try: - feat_dim = len(example_feat) - dtype = example_feat.dtype - except TypeError: - feat_dim = 1 - dtype = np.dtype(type(example_feat)) - is_structured = dtype.names is not None - nx_graph.add_edges_from(graph_dict["edges"]) - # identify the hard ancestors to which we trace - hard_mask = graph.hard_mask.copy() - del hard_mask["incoming"] - hard_graph = graph[hard_mask.bitwise_or()] - if target: # restrict hard partons to user specified pdgs - target_mask = hard_graph.pdg.mask( - list(target), blacklist=False, sign_sensitive=True - ) - hard_graph = hard_graph[target_mask] - names, vtxs = tuple(hard_graph.pdg.name), tuple(hard_graph.edges["dst"]) - # out vertices of user specified particles - focus_pcls = graph.edges[mask]["dst"] - trc = np.array( - [ - _trace_vector( - nx_graph, pcl, vtxs, feat_dim, is_structured, exclusive - ) - for pcl in focus_pcls - ] - ) - _trace_vector.cache_clear() - traces = dict() - array_fmt: ty.Callable[[base.AnyVector], base.AnyVector] = ( - partial(rfn.unstructured_to_structured, dtype=dtype) # type: ignore - if is_structured - else lambda x: x.squeeze() - ) - - for i, name in enumerate(names): - traces[name] = array_fmt(trc[:, i, :]) - return traces - - @nb.njit("float64[:](float64[:], float64[:], float64)", cache=True) def _rapidity( energy: base.DoubleVector, z: base.DoubleVector, zero_tol: float From 96669810911c5175522e6a945abcb2674cb940b3 Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Fri, 31 May 2024 16:22:48 +0100 Subject: [PATCH 09/12] automated tests for momentum transforms #178 --- tests/test_transform.py | 78 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 tests/test_transform.py diff --git a/tests/test_transform.py b/tests/test_transform.py new file mode 100644 index 0000000..1424e10 --- /dev/null +++ b/tests/test_transform.py @@ -0,0 +1,78 @@ +""" +``test_transform`` +================== + +Unit tests for the data structures, probing their attributes and +methods. +""" +import math + +import numpy as np +import pytest + +import graphicle as gcl + +DoubleVector = gcl.base.DoubleVector + + +def unit_vec_about_axis(vec: DoubleVector, axis: DoubleVector) -> DoubleVector: + """Normalises ``vec``, such that its closest distance from ``axis`` + is unity. This enables us to consider it as a point along a 2D unit + circle about ``axis``, albeit embedded in 3D space. + """ + if not math.isclose(np.linalg.norm(axis), 1.0): + raise ValueError("axis must be unit length.") + mag = np.linalg.norm(vec) + axis_angle = np.arccos(np.dot(vec, axis) / mag) + radius = mag * np.sin(axis_angle) + return vec / radius + + +def angular_distance_about_axis( + vec_1: DoubleVector, vec_2: DoubleVector, axis: DoubleVector +) -> float: + """Computes the angle subtended by the chord connecting projections + of ``vec_1`` and ``vec_2`` in the plane of the unit circle about + ``axis``. + + Will always give a positive number, as direction is not considered. + """ + chord = unit_vec_about_axis(vec_2, axis) - unit_vec_about_axis(vec_1, axis) + chord_length = np.linalg.norm(chord) + return abs(2.0 * np.arcsin(0.5 * chord_length).item()) + + +def test_momentum_split(): + rng = np.random.default_rng() + for _ in range(100): + energy_fraction = rng.uniform(1.0e-4, 0.5) + angles = rng.uniform(-1.0, 1.0, 3) + angles[0] = np.arccos(angles[0]) + angles[1:] *= np.pi + rotation_angle = angles[2].item() + rotation_axis = gcl.transform.SphericalAngle(*angles[:2].tolist()) + parent = gcl.MomentumArray.from_spherical_uniform(1, 100.0) + children = gcl.transform.split_momentum( + parent, energy_fraction, rotation_angle, rotation_axis + ) + child_sum = np.sum(children, axis=0) + assert np.allclose(parent, child_sum) + first_child = children[0].copy() + actual_angle = angular_distance_about_axis( + parent._data[0, :3], + first_child._data[0, :3], + np.array(gcl.transform._angles_to_axis(rotation_axis)), + ) + assert math.isclose(abs(rotation_angle), actual_angle) + # invert the split of the first child: + first_child_inv = children[0].copy() + first_child_inv /= energy_fraction + first_child_inv._data[0, :3] @= gcl.transform.rotation_matrix( + -rotation_angle, + rotation_axis, + ).T + assert np.allclose(parent, first_child_inv) + + +if __name__ == "__main__": + pytest.main() From 7976a698689a9f44be59e06b7024ae12d6569e4c Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Sat, 1 Jun 2024 10:47:49 +0100 Subject: [PATCH 10/12] backwards compat for inplace matmul #178 --- graphicle/transform.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/graphicle/transform.py b/graphicle/transform.py index 1b72fbf..b812f31 100644 --- a/graphicle/transform.py +++ b/graphicle/transform.py @@ -178,7 +178,11 @@ def split_momentum( parent = _momentum_to_numpy(momentum) children = np.tile(parent, (2, 1)) children[0, :] *= z - children[0, :3] @= rotation_matrix(angle, axis).T + # backwards compatibility, switched inplace syntax for explicit ufunc + # children[0, :3] @= rotation_matrix(angle, axis).T + np.matmul( + children[0, :3], rotation_matrix(angle, axis).T, out=children[0, :3] + ) children[1, :] -= children[0, :] return gcl.MomentumArray(children) From bf1e6584cc81973f38ef8217cd55f43edac092b3 Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Sat, 1 Jun 2024 10:55:46 +0100 Subject: [PATCH 11/12] matmul inplace backwards compatibility #178 --- tests/test_transform.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/test_transform.py b/tests/test_transform.py index 1424e10..7136aad 100644 --- a/tests/test_transform.py +++ b/tests/test_transform.py @@ -67,10 +67,14 @@ def test_momentum_split(): # invert the split of the first child: first_child_inv = children[0].copy() first_child_inv /= energy_fraction - first_child_inv._data[0, :3] @= gcl.transform.rotation_matrix( - -rotation_angle, - rotation_axis, - ).T + np.matmul( + first_child_inv._data[0, :3], + gcl.transform.rotation_matrix( + -rotation_angle, + rotation_axis, + ).T, + out=first_child_inv._data[0, :3], + ) assert np.allclose(parent, first_child_inv) From dc0df526da8c2193c14ec259c7753b4790e9ac52 Mon Sep 17 00:00:00 2001 From: Jacan Chaplais Date: Sat, 1 Jun 2024 11:04:34 +0100 Subject: [PATCH 12/12] versionadded directive #178 --- graphicle/data.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/graphicle/data.py b/graphicle/data.py index fdf0a9c..eb7399f 100644 --- a/graphicle/data.py +++ b/graphicle/data.py @@ -1344,6 +1344,8 @@ def from_spherical_uniform( """Returns a ``MomentumArray`` whose elements are sampled from uniform distributions of energy and 3-momentum. + .. versionadded:: 0.4.0 + Parameters ---------- size : int