From b669ef9c1759b81ec0589b342d06bb7fa170d1da Mon Sep 17 00:00:00 2001
From: emilyhunt <emily.hunt.physics@gmail.com>
Date: Fri, 26 Sep 2025 10:28:31 +0200
Subject: [PATCH] Add docs & tests for math functions.

---
 src/hr_selection_function/math.py | 44 +++++++++++++++++++--
 tests/test_math.py                | 66 +++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_math.py

diff --git a/src/hr_selection_function/math.py b/src/hr_selection_function/math.py
index 6f8898b..d820d1b 100644
--- a/src/hr_selection_function/math.py
+++ b/src/hr_selection_function/math.py
@@ -1,10 +1,28 @@
 import numpy as np
+from numpy.typing import ArrayLike
 from numba import njit
 
 
-def vectorized_multivariate_normal(values, means, covariances):
-    """This function exists here because scipy.stats.multivariate_normal can't be
-    vectorized, and is hence really really slow =(
+def vectorized_multivariate_normal(
+    values: np.ndarray, means: np.ndarray, covariances: np.ndarray
+) -> np.ndarray:
+    """Calculate the PDF of n_queries different multivariate normal distributions at 
+    n_queries different points in a fast and vectorized way. Significantly faster than
+    scipy.stats.multivariate_normal.pdf() when n_queries is large.
+
+    Parameters
+    ----------
+    values : np.ndarray
+        Values to query the PDF at. Must have shape (n_queries, n_dims).
+    means : np.ndarray
+        Means of the distributions. Must have shape (n_queries, n_dims).
+    covariances : np.ndarray
+        Covariances of the distributions. Must have shape (n_queries, n_dims, n_dims).
+
+    Returns
+    -------
+    np.ndarray
+        Array of shape (n_queries,) giving the PDF value of each distribution.
     """
     k = means.shape[1]
 
@@ -24,7 +42,25 @@ def vectorized_multivariate_normal(values, means, covariances):
 
 
 @njit
-def vectorized_1d_interpolation(values, x, y):
+def vectorized_1d_interpolation(
+    values: np.ndarray, x: np.ndarray, y: np.ndarray
+) -> np.ndarray:
+    """Vectorized, numba version of np.interp.
+
+    Parameters
+    ----------
+    values : ArrayLike
+        Values to query, of shape (n_queries,).
+    x : ArrayLike
+        x values of data, of shape (n_queries, n_data_points).
+    y : ArrayLike
+        y values of data, of shape (n_queries, n_data_points).
+
+    Returns
+    -------
+    np.ndarray
+        Interpolated values array of shape (n_queries,).
+    """
     length = values.shape[0]
     result = np.zeros(length)
     for i in range(length):
diff --git a/tests/test_math.py b/tests/test_math.py
new file mode 100644
index 0000000..9fddf73
--- /dev/null
+++ b/tests/test_math.py
@@ -0,0 +1,66 @@
+"""Tests some of the fast math functions in the library."""
+
+from hr_selection_function.math import (
+    vectorized_multivariate_normal,
+    vectorized_1d_interpolation,
+)
+from scipy.stats import multivariate_normal
+import numpy as np
+
+
+def test_vectorized_multivariate_normal():
+    rng = np.random.default_rng(42)
+    _assert_multivariate_normal_results_identical(*_create_test_data(5, 2, rng))
+    _assert_multivariate_normal_results_identical(*_create_test_data(5, 3, rng))
+    _assert_multivariate_normal_results_identical(*_create_test_data(5, 5, rng))
+    _assert_multivariate_normal_results_identical(*_create_test_data(1, 2, rng))
+    _assert_multivariate_normal_results_identical(*_create_test_data(10, 3, rng))
+
+
+def _make_valid_covariance_array(n_dims, rng):
+    covariance = rng.uniform(size=(n_dims, n_dims))
+
+    # Make it symmetric
+    covariance = np.tril(covariance) + np.triu(covariance.T, 1)
+
+    # Add offset to each term
+    return covariance + np.diag(rng.uniform(10, 100, size=n_dims))
+
+
+def _create_test_data(n_trials, n_dims, rng):
+    means = rng.uniform(-10, 10, size=(n_trials, n_dims))
+    covariances = np.asarray(
+        [_make_valid_covariance_array(n_dims, rng) for i in range(n_trials)]
+    )
+    queries = rng.uniform(-50, 50, size=(n_trials, n_dims))
+    return queries, means, covariances
+
+
+def _assert_multivariate_normal_results_identical(queries, means, covariances):
+    n_trials = means.shape[0]
+    result_expected = np.zeros(n_trials)
+    for i in range(n_trials):
+        result_expected[i] = multivariate_normal(mean=means[i], cov=covariances[i]).pdf(
+            queries[i]
+        )
+
+    result_vectorized = vectorized_multivariate_normal(queries, means, covariances)
+    np.testing.assert_allclose(result_expected, result_vectorized)
+
+
+def test_vectorized_1d_interpolation():
+    n_trials = 10
+    n_values = 50
+    x_values = np.tile(np.linspace(0, 10, num=n_values), (n_trials, 1))
+
+    offset = np.tile(np.linspace(0, 1, num=n_trials).reshape(-1, 1), (1, n_values))
+    y_values = (x_values + offset) ** 2
+    query_points = np.linspace(0, 10, num=n_trials)
+
+    result_expected = np.zeros(n_trials)
+    for i in range(n_trials):
+        result_expected[i] = np.interp(query_points[i], x_values[i], y_values[i])
+
+    result_vectorized = vectorized_1d_interpolation(query_points, x_values, y_values)
+
+    np.testing.assert_allclose(result_expected, result_vectorized)
-- 
2.51.2