Merge branch 'astrostuff' into 'master'

Astrostuff

See merge request !84

CHANGELOG.rst

 Unreleased changes
 ------------------
 
+* A few astro helpers were added: azimuth(), zenith(), phi(), theta(), ...
+
 Version 1
 ---------
 1.0.2 / 2023-10-08

src/km3io/tools.py

@@ -640,3 +640,194 @@ class TimeConverter(object):
           DAQ/trigger hit time [ns]
         """
         return t0 - (self.__t0 - self.__t1)  # [ns]
+
+
+def angle(v1, v2, normalized=False):
+    """
+    Compute the unsigned angle between two vectors. For a stacked input, the
+    angle is computed pairwise. Inspired by the "vg" Python package.
+
+    Parameters
+    ----------
+    v1 : np.array
+        With shape `(3,)` or a `kx3` stack of vectors.
+    v2 : np.array
+        A vector or stack of vectors with the same shape as `v1`.
+    normalized : bool
+        By default, the vectors will be normalised unless `normalized` is `True`.
+
+    Returns
+    -------
+    A float or a vector of floats with the angle in radians.
+
+    """
+    dot_products = np.einsum("ij,ij->i", v1.reshape(-1, 3), v2.reshape(-1, 3))
+
+    if normalized:
+        cosines = dot_products
+    else:
+        cosines = dot_products / magnitude(v1) / magnitude(v2)
+
+    # The dot product can exceed 1 or -1 and arccos will fail unless we clip
+    angles = np.arccos(np.clip(cosines, -1.0, 1.0))
+
+    if v1.ndim == v2.ndim == 1:
+        return angles[0]
+
+    return angles
+
+
+def magnitude(v):
+    """
+    Calculates the magnitude of a vector or array of vectors.
+
+    Parameters
+    ----------
+    v : np.array
+        With shape `(3,)` or `kx3` stack of vectors.
+
+    Returns
+    -------
+    A float or a vector of floats with the magnitudes.
+    """
+    if v.ndim == 1:
+        return np.linalg.norm(v)
+    elif v.ndim == 2:
+        return np.linalg.norm(v, axis=1)
+    else:
+        ValueError("Unsupported dimensions")
+
+
+def theta(v):
+    """Neutrino direction in polar coordinates.
+
+    Parameters
+    ----------
+    v : array (x, y, z)
+
+    Notes
+    -----
+    Downgoing event: theta = 180deg
+    Horizont: 90deg
+    Upgoing: theta = 0
+
+    Angles in radians.
+
+    """
+    v = np.atleast_2d(v)
+    dir_z = v[:, 2]
+    return theta_separg(dir_z)
+
+
+def theta_separg(dir_z):
+    return np.arccos(dir_z)
+
+
+def phi(v):
+    """Neutrino direction in polar coordinates.
+
+    Parameters
+    ----------
+    v : array (x, y, z)
+
+
+    Notes
+    -----
+    ``phi``, ``theta`` is the opposite of ``zenith``, ``azimuth``.
+
+    Angles in radians.
+
+    """
+    v = np.atleast_2d(v)
+    dir_x = v[:, 0]
+    dir_y = v[:, 1]
+    return phi_separg(dir_x, dir_y)
+
+
+def phi_separg(dir_x, dir_y):
+    p = np.arctan2(dir_y, dir_x)
+    p[p < 0] += 2 * np.pi
+    return p
+
+
+def zenith(v):
+    """Return the zenith angle in radians.
+
+    Parameters
+    ----------
+    v : array (x, y, z)
+
+
+    Notes
+    -----
+    Defined as 'Angle respective to downgoing'.
+    Downgoing event: zenith = 0
+    Horizont: 90deg
+    Upgoing: zenith = 180deg
+
+    """
+    return angle_between((0, 0, -1), v)
+
+
+def azimuth(v):
+    """Return the azimuth angle in radians.
+
+    Parameters
+    ----------
+    v : array (x, y, z)
+
+
+    Notes
+    -----
+    ``phi``, ``theta`` is the opposite of ``zenith``, ``azimuth``.
+
+    This is the 'normal' azimuth definition -- beware of how you
+    define your coordinates. KM3NeT defines azimuth
+    differently than e.g. SLALIB, astropy, the AAS.org
+
+    """
+    v = np.atleast_2d(v)
+    azi = phi(v) - np.pi
+    azi[azi < 0] += 2 * np.pi
+    if len(azi) == 1:
+        return azi[0]
+    return azi
+
+
+def angle_between(v1, v2, axis=0):
+    """Returns the angle in radians between vectors 'v1' and 'v2'.
+
+    If axis=1 it evaluates the angle between arrays of vectors.
+
+    Examples
+    --------
+    >>> angle_between((1, 0, 0), (0, 1, 0))
+    1.5707963267948966
+    >>> angle_between((1, 0, 0), (1, 0, 0))
+    0.0
+    >>> angle_between((1, 0, 0), (-1, 0, 0))
+    3.141592653589793
+    >>> v1 = np.array([[1, 0, 0], [1, 0, 0]])
+    >>> v2 = np.array([[0, 1, 0], [-1, 0, 0]])
+    >>> angle_between(v1, v2, axis=1)
+    array([1.57079633, 3.14159265])
+
+    """
+    if axis == 0:
+        v1_u = unit_vector(v1)
+        v2_u = unit_vector(v2)
+        # Don't use `np.dot`, does not work with all shapes
+        return np.arccos(np.clip(np.inner(v1_u, v2_u), -1, 1))
+    elif axis == 1:
+        return angle(v1, v2)  # returns angle in deg
+    else:
+        raise ValueError("unsupported axis")
+
+
+def unit_vector(vector, **kwargs):
+    """Returns the unit vector of the vector."""
+    vector = np.array(vector)
+    out_shape = vector.shape
+    vector = np.atleast_2d(vector)
+    unit = vector / np.linalg.norm(vector, axis=1, **kwargs)[:, None]
+    return unit.reshape(out_shape)

tests/test_tools.py

@@ -5,6 +5,8 @@ import awkward as ak
 import numpy as np
 from pathlib import Path
 
+from numpy.testing import assert_almost_equal, assert_allclose
+
 from km3net_testdata import data_path
 from km3io.definitions import fitparameters as kfit
 
@@ -37,6 +39,14 @@ from km3io.tools import (
     is_3dmuon,
     is_nanobeacon,
     TimeConverter,
+    phi,
+    theta,
+    zenith,
+    azimuth,
+    angle,
+    angle_between,
+    magnitude,
+    unit_vector,
 )
 
 
@@ -747,3 +757,161 @@ class TestTimeConverter(unittest.TestCase):
             np.fabs(self.one_event.mc_t + t_MC - t_frame)
             < self.tconverter.FRAME_TIME_NS
         )
+
+
+class TestMath(unittest.TestCase):
+    def setUp(self):
+        self.v = np.array([0.26726124, 0.53452248, 0.80178373])
+        self.vecs = np.array(
+            [
+                [0.0, 0.19611614, 0.98058068],
+                [0.23570226, 0.23570226, 0.94280904],
+                [0.53452248, 0.26726124, 0.80178373],
+                [0.80178373, 0.26726124, 0.53452248],
+                [0.94280904, 0.23570226, 0.23570226],
+            ]
+        )
+
+    def test_phi(self):
+        print(phi((1, 0, 0)))
+        assert_almost_equal(0, phi((1, 0, 0)))
+        assert_almost_equal(np.pi, phi((-1, 0, 0)))
+        assert_almost_equal(np.pi / 2, phi((0, 1, 0)))
+        assert_almost_equal(np.pi / 2 * 3, phi((0, -1, 0)))
+        assert_almost_equal(np.pi / 2 * 3, phi((0, -1, 0)))
+        assert_almost_equal(0, phi((0, 0, 0)))
+        assert_almost_equal(phi(self.v), 1.10714872)
+        assert_almost_equal(
+            phi(self.vecs),
+            np.array([1.57079633, 0.78539816, 0.46364761, 0.32175055, 0.24497866]),
+        )
+
+    def test_zenith(self):
+        assert_allclose(np.pi, zenith((0, 0, 1)))
+        assert_allclose(0, zenith((0, 0, -1)))
+        assert_allclose(np.pi / 2, zenith((0, 1, 0)))
+        assert_allclose(np.pi / 2, zenith((0, -1, 0)))
+        assert_allclose(np.pi / 4 * 3, zenith((0, 1, 1)))
+        assert_allclose(np.pi / 4 * 3, zenith((0, -1, 1)))
+        assert_almost_equal(zenith(self.v), 2.5010703409103687)
+        assert_allclose(
+            zenith(self.vecs),
+            np.array([2.94419709, 2.80175574, 2.50107034, 2.13473897, 1.80873745]),
+        )
+
+    def test_azimuth(self):
+        self.assertTrue(np.allclose(np.pi, azimuth((1, 0, 0))))
+        self.assertTrue(np.allclose(0, azimuth((-1, 0, 0))))
+
+        print(azimuth((0, 1, 0)))
+        print(azimuth((0, -1, 0)))
+        print(azimuth((0, 0, 0)))
+        print(azimuth(self.v))
+        print(azimuth(self.vecs))
+        self.assertTrue(np.allclose(np.pi / 2 * 3, azimuth((0, 1, 0))))
+        self.assertTrue(np.allclose(np.pi / 2, azimuth((0, -1, 0))))
+        self.assertTrue(np.allclose(np.pi, azimuth((0, 0, 0))))
+        self.assertTrue(np.allclose(azimuth(self.v), 4.24874137138))
+        self.assertTrue(
+            np.allclose(
+                azimuth(self.vecs),
+                np.array([4.71238898, 3.92699082, 3.60524026, 3.46334321, 3.38657132]),
+            )
+        )
+
+    def test_theta(self):
+        print(theta((0, 0, -1)))
+        print(theta((0, 0, 1)))
+        print(theta((0, 1, 0)))
+        print(theta((0, -1, 0)))
+        print(theta((0, 1, 1)))
+        print(theta((0, -1, 1)))
+        print(theta(self.v))
+        print(theta(self.vecs))
+        self.assertTrue(np.allclose(0, theta((0, 0, 1))))
+        self.assertTrue(np.allclose(np.pi, theta((0, 0, -1))))
+        self.assertTrue(np.allclose(np.pi / 2, theta((0, 1, 0))))
+        self.assertTrue(np.allclose(np.pi / 2, theta((0, -1, 0))))
+        self.assertTrue(np.allclose(0, theta((0, 1, 1))))
+        self.assertTrue(np.allclose(0, theta((0, -1, 1))))
+        self.assertTrue(np.allclose(theta(self.v), 0.64052231))
+        self.assertTrue(
+            np.allclose(
+                theta(self.vecs),
+                np.array([0.19739554, 0.33983691, 0.64052231, 1.00685369, 1.3328552]),
+            )
+        )
+
+    def test_unit_vector(self):
+        v1 = (1, 0, 0)
+        v2 = (1, 1, 0)
+        v3 = (-1, 2, 0)
+        assert np.allclose(v1, unit_vector(v1))
+        assert np.allclose(np.array(v2) / np.sqrt(2), unit_vector(v2))
+        assert np.allclose(np.array(v3) / np.sqrt(5), unit_vector(v3))
+
+    def test_magnitude(self):
+        assert 1 == magnitude(np.array([1, 0, 0]))
+        assert 2 == magnitude(np.array([0, 2, 0]))
+        assert 3 == magnitude(np.array([0, 0, 3]))
+        assert np.allclose(
+            [3.74165739, 8.77496439, 13.92838828],
+            magnitude(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])),
+        )
+
+    def test_angle(self):
+        v1 = np.array([1, 0, 0])
+        v2 = np.array([0, 1, 0])
+        v3 = np.array([-1, 0, 0])
+        self.assertAlmostEqual(0, angle(v1, v1))
+        self.assertAlmostEqual(np.pi / 2, angle(v1, v2))
+        self.assertAlmostEqual(np.pi, angle(v1, v3))
+        self.assertAlmostEqual(angle(self.v, v1), 1.3002465638163236)
+        self.assertAlmostEqual(angle(self.v, v2), 1.0068536854342678)
+        self.assertAlmostEqual(angle(self.v, v3), 1.8413460897734695)
+
+        assert np.allclose(
+            [1.300246563816323, 1.0068536854342678, 1.8413460897734695],
+            angle(np.array([self.v, self.v, self.v]), np.array([v1, v2, v3])),
+        )
+
+    def test_angle_between(self):
+        v1 = (1, 0, 0)
+        v2 = (0, 1, 0)
+        v3 = (-1, 0, 0)
+        self.assertAlmostEqual(0, angle_between(v1, v1))
+        self.assertAlmostEqual(np.pi / 2, angle_between(v1, v2))
+        self.assertAlmostEqual(np.pi, angle_between(v1, v3))
+        self.assertAlmostEqual(angle_between(self.v, v1), 1.3002465638163236)
+        self.assertAlmostEqual(angle_between(self.v, v2), 1.0068536854342678)
+        self.assertAlmostEqual(angle_between(self.v, v3), 1.8413460897734695)
+
+        assert np.allclose(
+            [0.0, 0.0, 0.0]
+            - angle_between(np.array([v1, v2, v3]), np.array([v1, v2, v3]), axis=1),
+            0,
+        )
+        assert np.allclose(
+            [np.pi / 2, np.pi]
+            - angle_between(np.array([v1, v1]), np.array([v2, v3]), axis=1),
+            0,
+        )
+
+        self.assertTrue(
+            np.allclose(
+                angle_between(self.vecs, v1),
+                np.array([1.57079633, 1.3328552, 1.0068537, 0.64052231, 0.33983691]),
+            )
+        )
+        self.assertTrue(
+            np.allclose(
+                angle_between(self.vecs, v2),
+                np.array([1.37340077, 1.3328552, 1.3002466, 1.30024656, 1.3328552]),
+            )
+        )
+        self.assertTrue(
+            np.allclose(
+                angle_between(self.vecs, v3),
+                np.array([1.57079633, 1.80873745, 2.13473897, 2.50107034, 2.80175574]),
+            )
+        )