Refactor convolve_moments and convolve_moments_self functions for improved clarity and consistency; enhance tests for edge cases.

Author: HugoMVale

src/polykin/distributions/misc.py

@@ -8,7 +8,8 @@
 import numpy as np
 from numpy.polynomial.laguerre import Laguerre
 
-from polykin.utils.types import FloatArray, FloatArrayLike, FloatOrArrayLike
+from polykin.utils.types import (FloatArray, FloatArrayLike, FloatVector,
+                                 FloatVectorLike)
 
 __all__ = [
     'convolve_moments',
@@ -18,119 +19,130 @@
 ]
 
 
-def convolve_moments(q0: float,
-                     q1: float,
-                     q2: float,
-                     r0: float,
-                     r1: float,
-                     r2: float
-                     ) -> tuple[float, float, float]:
-    r"""Compute the first three moments of the convolution of two distributions.
+def convolve_moments(
+    q: FloatVectorLike,
+    r: FloatVectorLike,
+) -> FloatVector:
+    r"""Compute the moments of the convolution of two distributions.
 
-    If $P = Q * R$ is the convolution of $Q$ and $R$, defined as:
+    If $P = Q * R$ is the convolution of distributions $Q$ and $R$, defined as:
 
-    $$ P_n = \sum_{i=0}^{n} Q_{n-i}R_{i} $$
+    $$ P_n = \sum_{m=0}^{n} Q_{n-m}R_{m} $$
 
-    then the first three moments of $P$ are related to the moments of $Q$ and
-    $R$ by:
+    and $p_i$, $q_i$ and $r_i$ denote the $i$-th moments of $P$, $Q$ and $R$,
+    respectively,
 
     \begin{aligned}
-    p_0 &= q_0 r_0 \\
-    p_1 &= q_1 r_0 + q_0 r_1 \\
-    p_2 &= q_2 r_0 + 2 q_1 r_1 + q_0 r_2
+        p_i & = \sum_{n=0}^{\infty} n^i P_n \\
+        q_i & = \sum_{n=0}^{\infty} n^i Q_n \\
+        r_i & = \sum_{n=0}^{\infty} n^i R_n        
     \end{aligned}
 
-    where $p_i$, $q_i$ and $r_i$ denote the $i$-th moments of $P$, $Q$ and $R$,
-    respectively.    
+    then the moments of $P$ are related to the moments of $Q$ and $R$ by:
+
+    $$ p_i = \sum_{j=0}^{i} \binom{i}{j} q_j r_{i-j} $$
 
     Parameters
     ----------
-    q0 : float
-        0-th moment of $Q$.
-    q1 : float
-        1-st moment of $Q$.
-    q2 : float
-        2-nd moment of $Q$.
-    r0 : float
-        0-th moment of $R$.
-    r1 : float
-        1-st moment of $R$.
-    r2 : float
-        2-nd moment of $R$.
+    q : FloatVectorLike (N)
+        Moments of $Q$, denoted $(q_0, q_1, \ldots)$.
+    r : FloatVectorLike (N)
+        Moments of $R$, denoted $(r_0, r_1, \ldots)$.
 
     Returns
     -------
-    tuple[float, float, float]
-        0-th, 1-st and 2-nd moments of $P=Q*R$.
+    FloatVector (N)
+        Moments of $P=Q*R$, denoted $(p_0, p_1, \ldots)$.
 
     Examples
     --------
     >>> from polykin.distributions import convolve_moments
-    >>> convolve_moments(1., 1e2, 2e4, 1., 50., 5e4)
-    (1.0, 150.0, 80000.0)
+    >>> convolve_moments([1.0, 1e2, 2e4], [1.0, 50.0, 5e4])
+    array([1.0e+00, 1.5e+02, 8.0e+04])
     """
-    p0 = q0*r0
-    p1 = q1*r0 + q0*r1
-    p2 = q2*r0 + 2*q1*r1 + q0*r2
-    return p0, p1, p2
+    if len(q) != len(r):
+        raise ValueError("`q` and `r` must have the same length.")
+
+    p = np.zeros(len(q))
+    for i in range(len(q)):
+        for j in range(i+1):
+            p[i] += comb(i, j)*q[j]*r[i-j]
+
+    return p
 
 
-def convolve_moments_self(q0: float,
-                          q1: float,
-                          q2: float,
-                          order: int = 1
-                          ) -> tuple[float, float, float]:
-    r"""Compute the first three moments of the k-th order convolution of a
-    distribution with itself.
+def convolve_moments_self(
+    q: FloatVectorLike,
+    order: int
+) -> FloatVector:
+    r"""Compute the moments of the k-th order convolution of a distribution
+    with itself.
 
-    If $P^k$ is the $k$-th order convolution of $Q$ with itself, defined as:
+    If $P^{(k)}$ is the $k$-th order convolution of $Q$ with itself, defined as:
 
     \begin{aligned}
-    P^1_n &= Q*Q = \sum_{i=0}^{n} Q_{n-i} Q_{i} \\
-    P^2_n &= (Q*Q)*Q = \sum_{i=0}^{n} Q_{n-i} P^1_{i} \\
-    P^3_n &= ((Q*Q)*Q)*Q = \sum_{i=0}^{n} Q_{n-i} P^2_{i} \\
+    P^{(1)}_n &= Q*Q = \sum_{m=0}^{n} Q_{n-m} Q_{m} \\
+    P^{(2)}_n &= (Q*Q)*Q = \sum_{m=0}^{n} Q_{n-m} P^{(1)}_{m} \\
+    P^{(3)}_n &= ((Q*Q)*Q)*Q = \sum_{m=0}^{n} Q_{n-m} P^{(2)}_{m} \\
     ...
     \end{aligned}
 
-    then the first three moments of $P^k$ are related to the moments of $Q$ by:
+    then the moments of $P^{(k)}$ are related to the moments of $Q$ by:
 
     \begin{aligned}
     p_0 &= q_0^{k+1}  \\
     p_1 &= (k+1) q_0^k q_1 \\
-    p_2 &= (k+1) q_0^{k-1} (k q_1^2 +q_0 q_2)
+    p_2 &= (k+1) q_0^{k-1} (k q_1^2 +q_0 q_2) \\
+    \ldots
     \end{aligned}
 
-    where $p_i$ and $q_i$ denote the $i$-th moments of $P^k$ and $Q$,
+    where $p_i$ and $q_i$ denote the $i$-th moments of $P^{(k)}$ and $Q$,
     respectively.    
 
     Parameters
     ----------
-    q0 : float
-        0-th moment of $Q$.
-    q1 : float
-        1-st moment of $Q$.
-    q2 : float
-        2-nd moment of $Q$.
+    q : FloatVectorLike (N)
+        Moments of $Q$, denoted $(q_0, q_1, \ldots)$.
+    order : int
+        Order of the convolution.
 
     Returns
     -------
-    tuple[float, float, float]
-        0-th, 1-st and 2-nd moments of $P^k=(Q*Q)*...$.
+    FloatVector (N)
+        Moments of $P^{(k)}=(Q*Q)*...$, denoted $(p_0, p_1, \ldots)$.
 
     Examples
     --------
     >>> from polykin.distributions import convolve_moments_self
-    >>> convolve_moments_self(1., 1e2, 2e4, 2)
-    (1.0, 300.0, 120000.0)
+    >>> convolve_moments_self([1e0, 1e2, 2e4], 2)
+    array([1.0e+00, 3.0e+02, 1.2e+05])
     """
-    p0 = q0**(order+1)
-    p1 = (order+1) * q0**order * q1
-    p2 = (order+1) * q0**(order-1) * (order*q1**2 + q0*q2)
-    return p0, p1, p2
+
+    q = np.asarray(q, dtype=float)
+
+    if order == 0:
+        return q.copy()
+
+    N = len(q)
+    if N <= 3:
+        # Use closed-form expressions for the first three moments
+        p = np.empty(N)
+        if N > 0:
+            p[0] = q[0]**(order + 1)
+        if N > 1:
+            p[1] = (order + 1) * q[0]**order * q[1]
+        if N > 2:
+            p[2] = (order + 1) * q[0]**(order - 1) * (order*q[1]**2 + q[0]*q[2])
+    else:
+        p = q.copy()
+        for _ in range(order):
+            p = convolve_moments(q, p)
+
+    return p
 
 
 def convert_polymer_standards(
-    M1: FloatOrArrayLike,
+    M1: float | FloatArrayLike,
     K1: float,
     K2: float,
     a1: float,
@@ -159,7 +171,7 @@ def convert_polymer_standards(
 
     Parameters
     ----------
-    M1 : FloatOrArrayLike
+    M1 : float | FloatArrayLike
         Molar mass in standard 1.
     K1 : float
         Mark-Houwink coefficient for standard 1.
@@ -177,30 +189,30 @@ def convert_polymer_standards(
 
     Examples
     --------
-    A sample of PMMA was mesured to have a molar mass of 100 kg/mol in PS
+    A sample of PMMA was mesured to have a molar mass of 1e5 g/mol in PS
     equivalent weight. What is the sample molar mass in actual PMMA weight?
 
     >>> from polykin.distributions import convert_polymer_standards
     >>> a1 = 0.77      # PS in THF
     >>> K1 = 6.82e-3   # PS in THF
     >>> a2 = 0.69      # PMMA in THF
     >>> K2 = 1.28e-2   # PMMA in THF
-    >>> M2 = convert_polymer_standards(100, K1, K2, a1, a2) 
-    >>> print(f"{M2:.2f} kg/mol")
-    85.68 kg/mol
+    >>> M2 = convert_polymer_standards(1e5, K1, K2, a1, a2) 
+    >>> print(f"{M2:.2e} g/mol")
+    1.19e+05 g/mol
     """
-    M1 = np.asarray(M1, dtype=np.float64)
+    M1 = np.asarray(M1, dtype=float)
     return (K1/K2)**(1/(1 + a2)) * M1**((1 + a1)/(1 + a2))
 
 
 def reconstruct_Laguerre(
-    moments: FloatArrayLike,
+    moments: FloatVectorLike,
 ) -> Callable[[FloatArrayLike], FloatArray]:
-    r"""Reconstruct a differential number distribution from its first `k` 
+    r"""Reconstruct a differential number distribution from a finite set of 
     moments using a Laguerre-series approximation.
 
-    According to Bamford and Tompa, a number distribution can be expressed as
-    an (infinite) expansion in Laguerre polynomials:
+    According to Bamford and Tompa, a number distribution $P(n)$ can be expressed
+    as an (infinite) expansion in Laguerre polynomials:
 
     $$ P(n) = \frac{e^{-\rho}}{(DP_n)^2} 
               \sum_{m=0}^{\infty} \gamma_m L_m(\rho) $$
@@ -231,8 +243,8 @@ def reconstruct_Laguerre(
 
     Parameters
     ----------
-    moments : FloatArrayLike
-        First `k` raw moments of the number distribution (`λ₀`, `λ₁`, ..., `λ_k`).
+    moments : FloatVectorLike
+        Moments of $P$, denoted $(\lambda_0, \lambda_1, \ldots)$.
 
     Returns
     -------

tests/distributions/test_misc.py

@@ -3,12 +3,46 @@
 # Copyright Hugo Vale 2025
 
 import numpy as np
+import pytest
 from numpy import allclose, isclose
 
 from polykin.distributions import (Flory, convert_polymer_standards,
+                                   convolve_moments, convolve_moments_self,
                                    reconstruct_Laguerre)
 
 
+def test_convolve_moments():
+    q0 = 1.0
+    q1 = q0 * 100.0
+    q2 = q1 * 200.0
+    p0, p1, p2 = convolve_moments([q0, q1, q2], [q0, q1, q2])
+    assert isclose(p0*p2/p1**2, 1.5)
+    # different lengths
+    with pytest.raises(ValueError):
+        convolve_moments([q0, q1], [q0, q1, q2])
+
+
+def test_convolve_moments_self():
+
+    def convolve_moments_self_iter(q, order):
+        r = q
+        for _ in range(order):
+            r = convolve_moments(q, r)
+        return r
+
+    for order in range(1, 10):
+        q0 = 1.0
+        q1 = q0 * 100.0
+        q2 = q1 * 200.0
+        s = convolve_moments_self_iter([q0, q1, q2], order)
+        p = convolve_moments_self([q0, q1, q2], order)
+        assert allclose(s, p)
+
+    pA = convolve_moments_self([q0, q1, q2], order=4)
+    pB = convolve_moments_self([q0, q1, q2, q2*100], order=4)
+    assert allclose(pA, pB[:3])
+
+
 def test_convert_polymer_standards():
     a1 = 0.77      # PS in THF
     K1 = 6.82e-3   # PS in THF