diff --git a/.pylintrc b/.pylintrc index b8491cb8..12c0e82c 100644 --- a/.pylintrc +++ b/.pylintrc @@ -7,3 +7,6 @@ init-import=yes [MESSAGES CONTROL] disable=C0116, RST203, RST301 + +[MASTER] +ignore=conf.py diff --git a/doc/conf.py b/doc/conf.py index 9c098231..af38d6d3 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -59,7 +59,7 @@ 'members': True, 'undoc-members': True, 'show-inheritance': True, - 'special-members': '__init__, __call__', + 'special-members': '__init__, __call__, __eq__', } html_copy_source = False # do not copy rst files html_show_copyright = False diff --git a/tensorwaves/physics/helicityformalism/kinematics.py b/tensorwaves/physics/helicityformalism/kinematics.py index e80ae258..fe18c2c1 100644 --- a/tensorwaves/physics/helicityformalism/kinematics.py +++ b/tensorwaves/physics/helicityformalism/kinematics.py @@ -1,16 +1,262 @@ +r"""Kinematic based calculations for the helicity formalism. -class Kinematics(): - def __init__(self): - pass +It's responsibilities are defined by the interface +:class:`.interfaces.Kinematics`. - def register_invariant_mass(self): - pass +Here, the main responsibility is the conversion of general kinematic +information of a reaction to helicity formalism specific quantities - def register_helicity_angles(self): - pass +:math:`(s, \theta, \phi)` - def register_subsystem(self): - pass +The basic building blocks are the :class:`~HelicityKinematics` and +:class:`~SubSystem`. +""" - def convert(events): - pass \ No newline at end of file +import logging +from typing import Union + +import amplitf.kinematics as tfa_kin + +import numpy as np + +from tensorwaves.interfaces import Kinematics + + +class SubSystem(): + """Represents a part of a decay chain. + + A SubSystem resembles a decaying state and its ingoing and outgoing state. + It is uniquely defined by + + * :attr:`final_states` + * :attr:`recoil_state` + * :attr:`parent_recoil_state` + """ + + def __init__(self, final_states, recoil_state, parent_recoil_state): + """Fully initializes the :class:`~SubSystem`. + + Args: + final_states: `tuple` of `tuple` s containing unique ids. + Represents the final state content of the decay products. + recoil_state: `tuple` of unique ids representing the recoil partner + of the decaying state. + parent_recoil_state: `tuple` of unique ids representing the recoil + partner of the parent state. + """ + self._final_states = tuple(tuple(x) for x in final_states) + self._recoil_state = tuple(recoil_state) + self._parent_recoil_state = tuple(parent_recoil_state) + + @property + def final_states(self): + """Get final state content of the decay products.""" + return self._final_states + + @property + def recoil_state(self): + """Get final state content of the recoil partner.""" + return self._recoil_state + + @property + def parent_recoil_state(self): + """Get final state content of the recoil partner of the parent.""" + return self._parent_recoil_state + + def __eq__(self, other): + """Equal testing operator.""" + if self._final_states != other._final_states: + return False + if self._recoil_state != other._recoil_state: + return False + if self._parent_recoil_state != other._parent_recoil_state: + return False + return True + + def __hash__(self): + """Hash function to use SubSystem as key.""" + return hash( + (self._final_states, + self._recoil_state, + self._parent_recoil_state) + ) + + +class HelicityKinematics(Kinematics): + """Kinematics of the helicity formalism. + + General usage is + + 1. Register kinematic variables via the three methods + (:meth:`register_invariant_mass`, :meth:`register_helicity_angles`, + :meth:`register_subsystem`) first. + 2. Then convert events to these kinematic variables. + + For additional functionality check :meth:`phase_space_volume` and + :meth:`is_within_phase_space`. + """ + + def __init__(self, fs_id_event_pos_mapping=None): + """Initialize the a blank HelicityKinematics. + + Args: + fs_id_event_pos_mapping: Optional mapping between particle unique + ids and the position in the event array. + """ + self._registered_inv_masses = dict() + self._registered_subsystems = dict() + self._fs_id_event_pos_mapping = fs_id_event_pos_mapping + + @property + def phase_space_volume(self): + """Get volume of the defined phase space. + + Return: + `float` + """ + return 1.0 + + def is_within_phase_space(self, events): + """Check whether events lie within the phase space definition.""" + raise NotImplementedError("Currently not implemented.") + + def register_invariant_mass(self, final_state: Union[tuple, list]): + """Register an invariant mass :math:`s`. + + Args: + final_state: collection of particle unique id's + + Return: + A `str` key representing the invariant mass. It can be used to + retrieve this invariant mass from the dataset returned by + :meth:`~convert`. + """ + logging.debug("registering inv mass in kinematics") + final_state = tuple(final_state) + if final_state not in self._registered_inv_masses: + label = 'mSq' + for particle_uid in final_state: + label += '_' + str(particle_uid) + + self._registered_inv_masses[final_state] = label + return self._registered_inv_masses[final_state] + + def register_helicity_angles(self, subsystem: SubSystem): + r"""Register helicity angles :math:`(\theta, \phi)` of a `SubSystem`. + + Args: + subsystem: SubSystem to which the registered angles correspond. + + Return: + A pair of `str` keys representing the angles. They can be used to + retrieve the angles from the dataset returned by :meth:`~convert`. + """ + logging.debug("registering helicity angles in kinematics") + if subsystem not in self._registered_subsystems: + suffix = '' + for final_state in subsystem.final_states: + suffix += '+' + for particle_uid in final_state: + suffix += str(particle_uid) + '_' + suffix = suffix[:-1] + if subsystem.recoil_state: + suffix += '_vs_' + for particle_uid in subsystem.recoil_state: + suffix += str(particle_uid) + '_' + suffix = suffix[:-1] + + self._registered_subsystems[subsystem] = ( + 'theta' + suffix, 'phi' + suffix) + return self._registered_subsystems[subsystem] + + def register_subsystem(self, subsystem: SubSystem): + r"""Register all kinematic variables of the :class:`~SubSystem`. + + Args: + subsystem: SubSystem to which the registered kinematic variables + correspond. + + Return: + A tuple of `str` keys representing the :math:`(s, \theta, \phi)`. + They can be used to retrieve the kinematic data from the dataset + returned by :meth:`~convert`. + """ + state_fs = [] + for fs_uid in subsystem.final_states: + state_fs += fs_uid + invmass_name = self.register_invariant_mass(list(set(state_fs))) + angle_names = self.register_helicity_angles(subsystem) + + return (invmass_name,) + angle_names + + def _convert_ids_to_indices(self, ids: tuple): + """Convert unique ids to event indices. + + Uses the :attr:`_fs_id_event_pos_mapping`. + """ + if self._fs_id_event_pos_mapping: + return [self._fs_id_event_pos_mapping[i] for i in ids] + + return ids + + def convert(self, events): + r"""Convert events to the registered kinematics variables. + + Args: + events: A three dimensional numpy array of the shape + :math:`(n_{\mathrm{part}}, n_{\mathrm{evts}}, 4)`. + + * :math:`n_{\mathrm{part}}` is the number of particles + * :math:`n_{\mathrm{evts}}` is the number of events + + The third dimension correspond to the four momentum info + :math:`(p_x, p_y, p_z, E)`. + + Return: + A `dict` containing the registered kinematic variables as keys + and their corresponding values. This is also known as a dataset. + """ + logging.info('converting %s events', len(events[0])) + + dataset = {} + + for four_momenta_ids, inv_mass_name \ + in self._registered_inv_masses.items(): + four_momenta = np.sum( + events[self._convert_ids_to_indices(four_momenta_ids), :], + axis=0) + + dataset[inv_mass_name] = tfa_kin.mass_squared( + np.array(four_momenta)) + + for subsys, angle_names in self._registered_subsystems.items(): + topology = [ + np.sum( + events[self._convert_ids_to_indices(x), :], + axis=0) + for x in subsys.final_states + ] + if subsys.recoil_state: + topology = [ + topology, + np.sum( + events[self._convert_ids_to_indices( + subsys.recoil_state), :], + axis=0), + ] + if subsys.parent_recoil_state: + topology = [ + topology, + np.sum( + events[self._convert_ids_to_indices( + subsys.parent_recoil_state), :], + axis=0), + ] + + values = tfa_kin.nested_helicity_angles(topology) + + # the last two angles is always what we are interested + dataset[angle_names[0]] = values[-2].numpy() + dataset[angle_names[1]] = values[-1].numpy() + + return dataset diff --git a/tests/physics/helicityformalism/test_helicityangles.py b/tests/physics/helicityformalism/test_helicityangles.py new file mode 100644 index 00000000..aa919831 --- /dev/null +++ b/tests/physics/helicityformalism/test_helicityangles.py @@ -0,0 +1,154 @@ +"""Tests regarding the helicity kinematics.""" + +import numpy as np + +import pytest + +from tensorwaves.physics.helicityformalism.kinematics import ( + HelicityKinematics, SubSystem) + +# Initial State: J/Psi +# Final State: pi0 gamma pi0 pi0 +TEST_DATA = { + "events": + { + 0: [ + (0.514208, -0.184219, 1.23296, 1.35527), + (0.0727385, -0.0528868, 0.826163, 0.841933), + (-0.162529, 0.29976, -0.411133, 0.550927), + (0.0486171, 0.151922, 0.370309, 0.425195), + (-0.0555915, -0.100214, -0.0597338, 0.186869), + (0.238921, 0.266712, -1.20442, 1.26375), + (0.450724, -0.439515, -0.360076, 0.737698), + (0.552298, 0.440006, 0.644927, 0.965809), + (-0.248155, -0.158587, -0.229673, 0.397113), + (1.33491, 0.358535, 0.0457548, 1.38955) + ], + 1: [ + (-0.305812, 0.284, -0.630057, 0.755744), + (0.784483, 0.614347, -0.255334, 1.02861), + (-0.20767, 0.272796, 0.0990739, 0.356875), + (0.404557, 0.510467, -0.276426, 0.70757), + (0.47713, 0.284575, -0.775431, 0.953902), + (-0.204775, -0.0197981, 0.0799868, 0.220732), + (0.00590727, 0.709346, -0.190877, 0.734602), + (0.329157, -0.431973, 0.272873, 0.607787), + (-0.201436, -0.534829, 0.256253, 0.626325), + (-0.196357, 0.00211926, -0.33282, 0.386432) + ], + 2: [ + (-0.061663, -0.0211864, 0.144596, 0.208274), + (-0.243319, -0.283044, -0.234866, 0.461193), + (0.82872, -0.0465425, -0.599834, 1.03294), + (0.263003, -0.089236, 0.686187, 0.752466), + (0.656892, -0.107848, 0.309898, 0.746588), + (0.521569, -0.0448683, 0.43283, 0.692537), + (-0.517582, -0.676002, -0.0734335, 0.865147), + (-0.975278, -0.0207817, -0.934467, 1.35759), + (-0.41665, 0.237646, 0.691269, 0.852141), + (-0.464203, -0.358114, 0.13307, 0.616162) + ], + 3: [ + (-0.146733, -0.0785946, -0.747499, 0.777613), + (-0.613903, -0.278416, -0.335962, 0.765168), + (-0.458522, -0.526014, 0.911894, 1.15616), + (-0.716177, -0.573154, -0.780069, 1.21167), + (-1.07843, -0.0765127, 0.525267, 1.20954), + (-0.555715, -0.202046, 0.691605, 0.919879), + (0.0609506, 0.406171, 0.624387, 0.759452), + (0.0938229, 0.012748, 0.0166676, 0.165716), + (0.866241, 0.455769, -0.717849, 1.22132), + (-0.674348, -0.0025409, 0.153994, 0.704759) + ] + }, + "angles": + { + (((1, 2, 3), (0,)), (), ()): [ + (-0.914298, 2.79758), + (-0.994127, 2.51292), + (0.769715, -1.07396), + (-0.918418, -1.88051), + (0.462214, 1.06433), + (0.958535, -2.30129), + (0.496489, 2.36878), + (-0.674376, -2.46888), + (0.614968, 0.568649), + (-0.0330843, -2.8792) + ], + (((2, 3), (1,)), (0,), ()): [ + (-0.772533, 1.04362), + (0.163659, 1.87349), + (0.556365, 0.160733), + (0.133251, -2.81088), + (-0.0264361, 2.84379), + (0.227188, 2.29128), + (-0.166924, 2.24539), + (0.652761, -1.20272), + (0.443122, 0.615838), + (0.503577, 2.98067) + ], + (((2,), (3,)), (1,), (0,)): [ + (0.460324, -2.77203), + (-0.410464, 1.45339), + (0.248566, -2.51096), + (-0.301959, 2.71085), + (-0.522502, -1.12706), + (0.787267, -3.01323), + (0.488066, 2.07305), + (0.954167, 0.502648), + (-0.553114, -1.23689), + (0.00256349, 1.7605) + ] + } +} + + +@pytest.mark.parametrize( + "test_events, expected_angles", + [ + ( + TEST_DATA["events"], + TEST_DATA["angles"] + ) + ] +) +def test_helicity_angles_correctness(test_events, expected_angles): + """Test the correctness of the helicity theta and phi angles.""" + subsys_angle_names = {} + kin = HelicityKinematics() + for subsys in expected_angles.keys(): + temp_names = kin.register_subsystem(SubSystem(*subsys)) + subsys_angle_names.update( + {subsys: [temp_names[1], temp_names[2]]} + ) + + data = np.array(tuple(np.array(v) for v in test_events.values())) + kinematic_vars = kin.convert(data) + + assert len(kinematic_vars) == 3 * len(expected_angles.keys()) + number_of_events = len(data[0]) + for subsys, angle_names in subsys_angle_names.items(): + for name in angle_names: + assert len(kinematic_vars[name]) == number_of_events + + expected_values = np.array(np.array(expected_angles[subsys]).T) + # test cos(theta) + np.testing.assert_array_almost_equal( + np.cos(kinematic_vars[angle_names[0]]), + expected_values[0], + 1e-6 + ) + # test phi + if subsys == (((2,), (3,)), (1,), (0,)): + for kin_var, expected in zip(kinematic_vars[angle_names[1]], + expected_values[1]): + assert ( + round(kin_var, 4) == round(expected - np.pi, 4) + or round(kin_var, 4) == round(expected + np.pi, 4) + ) + else: + np.testing.assert_array_almost_equal( + kinematic_vars[angle_names[1]], + expected_values[1], + 1e-6 + ) diff --git a/tox.ini b/tox.ini index 8a571f8b..51359ec5 100644 --- a/tox.ini +++ b/tox.ini @@ -21,6 +21,7 @@ filename = ./tests/*.py exclude = __pycache__ + doc/conf.py ignore = D102 # method docstring D107 # init docstring @@ -30,6 +31,8 @@ rst-roles = class, func, ref, + attr, + meth, rst-directives = envvar, exception,