Feat: Implement leading order derivatives of arbitrary order

photosurfactant/first_order.py

@@ -2,39 +2,12 @@
 
 from .parameters import Parameters
 from .leading_order import LeadingOrder
+from .utils import to_arr, hyperbolic, polyder, Y
 from enum import Enum
 from functools import wraps
 from typing import Callable
 import numpy as np
 
-Y = np.poly1d([1, 0])  # Polynomial for differentiation
-
-
-def hyperbolic(n):
-    """Return the n-th derivative of cosh."""
-    return np.sinh if n % 2 else np.cosh
-
-
-def polyder(p: np.poly1d, n: int):
-    """Wrap np.polyder and np.polyint."""
-    return p.deriv(n) if n > 0 else p.integ(-n)
-
-
-def to_arr(vals, unknowns):
-    """Convert a dictionary of values to an array."""
-    arr = np.zeros(len(unknowns), dtype=complex)
-    for key in vals.keys():
-        if key not in unknowns:
-            raise ValueError(f"Unknown key: {key}")
-
-    for i, key in enumerate(unknowns):
-        try:
-            arr[i] = vals[key]
-        except KeyError:
-            pass
-
-    return arr
-
 
 class Symbols(Enum):  # TODO: This is unnecessary
     """Symbols used in the first order solution."""
@@ -509,7 +482,8 @@ def kinetic_flux(self, k):
             @ (
                 (
                     self.first._c(k, 1)
-                    + Variables.S[np.newaxis, :] * self.leading.d_c(1)[:, np.newaxis]
+                    + Variables.S[np.newaxis, :]
+                    * self.leading.c(1, y_order=1)[:, np.newaxis]
                 )
                 * (1 - self.leading.gamma_tr - self.leading.gamma_ci)
                 - self.leading.c(1)[:, np.newaxis]
@@ -543,7 +517,7 @@ def mass_balance(self, k):
         """Mass balance boundary condition."""
         cond = (self.params.k_tr * self.params.chi_tr) * (
             self.first._c(k, 1, y_order=1)
-            + Variables.S[np.newaxis, :] * self.leading.d2_c(1)[:, np.newaxis]
+            + Variables.S[np.newaxis, :] * self.leading.c(1, y_order=2)[:, np.newaxis]
         ) + self.params.P @ np.array([Variables.J_tr, Variables.J_ci])
 
         if k == 0:

photosurfactant/leading_order.py

@@ -1,6 +1,7 @@
 """Leading order solution to the photosurfactant model."""
 
 from .parameters import Parameters
+from .utils import hyperbolic, polyder, Y
 import numpy as np
 
 
@@ -143,58 +144,35 @@ def _set_root_auto(self):
         else:
             raise ValueError("No valid solution found.")
 
-    def c_ci(self, y):
-        """Concentration of cis surfactant at leading order."""
-        params = self.params
-        return self.A_0 - self.B_0 * np.cosh(y * np.sqrt(params.zeta))
-
-    def c_tr(self, y):
+    def c_tr(self, y, y_order=0):
         """Concentration of trans surfactant at leading order."""
-        return (self.params.alpha + 1) * self.A_0 - self.params.eta * self.c_ci(y)
-
-    def c(self, y):
-        """Concentration of surfactant at leading order."""
-        return np.array([self.c_tr(y), self.c_ci(y)])
-
-    def d_c_ci(self, y):
-        """First derivative of the cis concentration at leading order."""
         params = self.params
-        return -self.B_0 * np.sqrt(params.zeta) * np.sinh(y * np.sqrt(params.zeta))
-
-    def d_c_tr(self, y):
-        """First derivative of the trans concentration at leading order."""
-        return -self.params.eta * self.d_c_ci(y)
-
-    def d_c(self, y):
-        """First derivative of the surfactant concentration at leading order."""
-        return np.array([self.d_c_tr(y), self.d_c_ci(y)])
+        return params.alpha * self.A_0 * polyder(Y**0, y_order)(
+            y
+        ) + params.eta * self.B_0 * np.sqrt(params.zeta) ** y_order * hyperbolic(
+            y_order
+        )(
+            y * np.sqrt(params.zeta)
+        )
 
-    def d2_c_ci(self, y):
-        """Second derivative of the cis concentration at leading order."""
+    def c_ci(self, y, y_order=0):
+        """Concentration of cis surfactant at leading order."""
         params = self.params
-        return -self.B_0 * params.zeta * np.cosh(y * np.sqrt(params.zeta))
+        return self.A_0 * polyder(Y**0, y_order)(y) - self.B_0 * np.sqrt(
+            params.zeta
+        ) ** y_order * hyperbolic(y_order)(y * np.sqrt(params.zeta))
 
-    def d2_c_tr(self, y):
-        """Second derivative of the trans concentration at leading order."""
-        return -self.params.eta * self.d2_c_ci(y)
+    def c(self, y, y_order=0):
+        """Concentration of surfactant at leading order."""
+        return np.array([self.c_tr(y, y_order), self.c_ci(y, y_order)])
 
-    def d2_c(self, y):
-        """Second derivative of the surfactant concentration at leading order."""
-        return np.array([self.d2_c_tr(y), self.d2_c_ci(y)])
+    def i_c_tr(self, y, y_s=0):
+        """Integral of the trans concentration at leading order."""
+        return self.c_tr(y, y_order=-1) - self.c_tr(y_s, y_order=-1)
 
     def i_c_ci(self, y, y_s=0):
         """Integral of the cis concentration at leading order."""
-        params = self.params
-        return (
-            self.A_0 * y
-            - self.B_0 / np.sqrt(params.zeta) * np.sinh(y * np.sqrt(params.zeta))
-        ) - (self.i_c_ci(y_s) if y_s else 0)
-
-    def i_c_tr(self, y, y_s=0):
-        """Integral of the trans concentration at leading order."""
-        return (self.params.alpha + 1) * self.A_0 * (
-            y - y_s
-        ) - self.params.eta * self.i_c_ci(y, y_s)
+        return self.c_ci(y, y_order=-1) - self.c_ci(y_s, y_order=-1)
 
     def i_c(self, y, y_s):
         """Integral of the surfactant concentration at leading order."""

photosurfactant/utils.py

@@ -1,6 +1,36 @@
 """Utility functions for the photosurfactant model."""
 
 from argparse import ArgumentParser
+import numpy as np
+
+
+Y = np.poly1d([1, 0])  # Polynomial for differentiation
+
+
+def hyperbolic(n):
+    """Return the n-th derivative of cosh."""
+    return np.sinh if n % 2 else np.cosh
+
+
+def polyder(p: np.poly1d, n: int):
+    """Wrap np.polyder and np.polyint."""
+    return p.deriv(n) if n > 0 else p.integ(-n)
+
+
+def to_arr(vals, unknowns):
+    """Convert a dictionary of values to an array."""
+    arr = np.zeros(len(unknowns), dtype=complex)
+    for key in vals.keys():
+        if key not in unknowns:
+            raise ValueError(f"Unknown key: {key}")
+
+    for i, key in enumerate(unknowns):
+        try:
+            arr[i] = vals[key]
+        except KeyError:
+            pass
+
+    return arr
 
 
 def parameter_parser(parser: ArgumentParser):

test/test_first_fourier.py

@@ -134,7 +134,7 @@ def test_p_2():
                     1.0j
                     * k
                     * first._psi(k, y)[np.newaxis, :]
-                    * leading.d_c_ci(y)
+                    * leading.c_ci(y, y_order=1)
                     * params.Pen_ci
                     / (alpha + eta)
                     * np.array([eta**2 - eta, eta**2 + alpha])[:, np.newaxis]
@@ -183,7 +183,7 @@ def test_p():
                     1.0j
                     * k
                     * first._psi(k, y)[np.newaxis, :]
-                    * leading.d_c_ci(y)
+                    * leading.c_ci(y, y_order=1)
                     * params.Pen_ci
                     / (alpha + eta)
                     * np.array([eta**2 - eta, eta**2 + alpha])[:, np.newaxis]
@@ -222,7 +222,7 @@ def test_bulk_concentration():
                     1.0j
                     * k
                     * first._psi(k, y)
-                    * (params.P @ leading.d_c(y)[:, np.newaxis])
+                    * (params.P @ leading.c(y, y_order=1)[:, np.newaxis])
                 )
                 for k in wavenumbers
             ]

test/test_first_order.py

@@ -135,7 +135,7 @@ def test_bulk_concentrations(func):
 
     eq_tr = np.array(
         [
-            leading.d_c_tr(y) * first.v(xx, y)
+            leading.c_tr(y, y_order=1) * first.v(xx, y)
             - 1
             / params.Pen_tr
             * (first.c_tr(xx, y, x_order=2) + first.c_tr(xx, y, y_order=2))
@@ -146,7 +146,7 @@ def test_bulk_concentrations(func):
     )
     eq_ci = np.array(
         [
-            leading.d_c_ci(y) * first.v(xx, y)
+            leading.c_ci(y, y_order=1) * first.v(xx, y)
             - 1
             / params.Pen_ci
             * (first.c_ci(xx, y, x_order=2) + first.c_ci(xx, y, y_order=2))
@@ -221,14 +221,14 @@ def test_kinetic_flux(func):
 
     eq_tr = params.Bit_tr * (
         params.k_tr
-        * (first.c_tr(xx, 1) + first.S(xx) * leading.d_c_tr(1))
+        * (first.c_tr(xx, 1) + first.S(xx) * leading.c_tr(1, y_order=1))
         * (1 - leading.gamma_tr - leading.gamma_ci)
         - params.k_tr * leading.c_tr(1) * (first.gamma_tr(xx) + first.gamma_ci(xx))
         - first.gamma_tr(xx)
     ) - first.J_tr(xx)
     eq_ci = params.Bit_ci * (
         params.k_ci
-        * (first.c_ci(xx, 1) + first.S(xx) * leading.d_c_ci(1))
+        * (first.c_ci(xx, 1) + first.S(xx) * leading.c_ci(1, y_order=1))
         * (1 - leading.gamma_tr - leading.gamma_ci)
         - params.k_ci * leading.c_ci(1) * (first.gamma_tr(xx) + first.gamma_ci(xx))
         - first.gamma_ci(xx)
@@ -317,10 +317,10 @@ def test_mass_balance(func):
     xx = np.linspace(-params.L, params.L, 100)
 
     eq_tr = params.k_tr * params.chi_tr / params.Pen_tr * (
-        first.c_tr(xx, 1, y_order=1) + first.S(xx) * leading.d2_c_tr(1)
+        first.c_tr(xx, 1, y_order=1) + first.S(xx) * leading.c_tr(1, y_order=2)
     ) + first.J_tr(xx)
     eq_ci = params.k_ci * params.chi_ci / params.Pen_ci * (
-        first.c_ci(xx, 1, y_order=1) + first.S(xx) * leading.d2_c_ci(1)
+        first.c_ci(xx, 1, y_order=1) + first.S(xx) * leading.c_ci(1, y_order=2)
     ) + first.J_ci(xx)
 
     assert np.allclose(eq_tr, 0)

test/test_leading_order.py

@@ -13,12 +13,12 @@ def test_bulk_concentrations():
     yy = np.linspace(0, 1, 100)
 
     eq_tr = (
-        (1 / params.Pen_tr) * leading.d2_c_tr(yy)
+        (1 / params.Pen_tr) * leading.c_tr(yy, y_order=2)
         - params.Dam_tr * leading.c_tr(yy)
         + params.Dam_ci * leading.c_ci(yy)
     )
     eq_ci = (
-        (1 / params.Pen_ci) * leading.d2_c_ci(yy)
+        (1 / params.Pen_ci) * leading.c_ci(yy, y_order=2)
         + params.Dam_tr * leading.c_tr(yy)
         - params.Dam_ci * leading.c_ci(yy)
     )
@@ -63,10 +63,12 @@ def test_mass_balance():
     leading = LeadingOrder(params)
 
     eq_tr = (
-        params.k_tr * params.chi_tr / params.Pen_tr * leading.d_c_tr(1) + leading.J_tr
+        params.k_tr * params.chi_tr / params.Pen_tr * leading.c_tr(1, y_order=1)
+        + leading.J_tr
     )
     eq_ci = (
-        params.k_ci * params.chi_ci / params.Pen_ci * leading.d_c_ci(1) + leading.J_ci
+        params.k_ci * params.chi_ci / params.Pen_ci * leading.c_ci(1, y_order=1)
+        + leading.J_ci
     )
 
     assert np.allclose(eq_tr, 0)
@@ -92,4 +94,4 @@ def test_no_flux():
     params = Parameters()
     leading = LeadingOrder(params)
 
-    assert np.allclose(leading.d_c(0), 0)
+    assert np.allclose(leading.c(0, y_order=1), 0)