vectorized version of symmetric implicit closure

fiberoripy/closures.py

@@ -3,7 +3,6 @@
 from itertools import permutations
 
 import numpy as np
-from scipy.optimize import root_scalar
 
 
 def compute_closure(a, closure="IBOF"):
@@ -679,51 +678,61 @@ def orthotropic_fitted_closures(a, closure="ORF"):
     return A
 
 
-def _loss_siq(s, evs):
-    """Compute loss for SIQ.
+def symmetric_quadratic_closure(a):
+    """Generate a symmetric quadratic closure.
     Parameters
     ----------
-    s : float
-        trace of the unknown second-order tensor in SIQ.
-    evs : 3x1 numpy array
-        eigenvalues of the input 2nd-order FOT.
+    a : (Mx)3x3 numpy array
+        (Array of) Second order fiber orientation tensor.
+
     Returns
     -------
-    s_root : float
-        approximated root of the loss function.
+    A : (Mx)3x3x3x3 numpy array
+        (Array of) Fourth order fiber orientation tensor.
+    References
+    ----------
+    .. [1] Karl, Tobias and Gatti, Davide and Frohnapfel, Bettina and Böhlke, Thomas,
+       'Asymptotic fiber orientation states of the quadratically
+       closed Folgar--Tucker equation and a subsequent closure
+       improvement',
+       Journal of Rheology 65(5) : 999-1022,
+       https://doi.org/10.1122/8.0000245
+    Notes
+    ----------
+        In general, the SQC does contract to its second-order input.
     """
-    d = len(evs)
-    s_root = (d + 4.0) * s - np.sum(np.sqrt(1.5 * evs + s**2.0))
+    assert_fot_properties(a)
 
-    return s_root
+    a4 = (
+        1.0
+        / 3.0
+        * (
+            np.einsum("...ij, ...kl -> ...ijkl", a, a)
+            + np.einsum("...ik, ...lj -> ...ijkl", a, a)
+            + np.einsum("...il, ...kj -> ...ijkl", a, a)
+        )
+    )
 
+    a4 /= 1.0 / 3.0 * (1.0 + 2.0 * np.einsum("...ij, ...ij -> ...", a, a))
 
-def _grad_loss_siq(s, evs):
-    """Compute gradient of the loss function for SIQ.
-    Parameters
-    ----------
-    s : float
-        trace of the unknown second-order tensor in SIQ.
-    evs : 3x1 numpy array
-        eigenvalues of the input 2nd-order FOT.
-    Returns
-    -------
-    k : float
-        gradient of the loss function at s.
-    """
-    k = 4.0 + np.sum(1.0 - s / np.sqrt(1.5 * evs + s**2.0))
-    return k
+    return a4
 
 
-def symmetric_implicit_closure(a):
+def symmetric_implicit_closure(a, eps_newton=1.0e-12, n_iter_newton=25):
     """Generate SIQ closure.
     Parameters
     ----------
-    a : 3x3 numpy array
-        Second order fiber orientation tensor.
+    a : ...x3x3 numpy array
+        (Array of) second order fiber orientation tensor.
+    eps_newton : float
+        convergence criterion for newton algorithm
+        optional: default 1.0e-12
+    n_iter_newton : int
+        number of maximum iterations in newton algorithm
+        optional: default 25
     Returns
     -------
-    3x3x3x3 numpy array
+    ...x3x3x3x3 numpy array
         Fourth order fiber orientation tensor.
     References
     ----------
@@ -732,65 +741,42 @@ def symmetric_implicit_closure(a):
        tensors', Journal of Non-Newtonian Fluid Mechanics 318 : 105049,
        https://doi.org/10.1016/j.jnnfm.2023.105049
     """
-    assert_fot_properties(a)
 
-    if not a.ndim == 2:
-        raise ValueError("Currently SIQ does not support broadcasting")
+    assert_fot_properties(a)
 
     evs, r = np.linalg.eigh(a)
 
-    s_root = root_scalar(
-        f=_loss_siq, fprime=_grad_loss_siq, args=(evs,), x0=1.0
-    ).root
+    d = evs.shape[-1]
+    s = np.ones(a.shape[:-2])
+    err, it = 1.0e12, 0
 
-    mus = np.sqrt(1.5 * evs + s_root**2) - s_root
-    b = r @ np.diag(mus) @ r.T
+    while err > eps_newton and it < n_iter_newton:
 
-    return (
-        1.0
-        / 3.0
-        * (
-            np.einsum("...ij, ...kl -> ...ijkl", b, b)
-            + np.einsum("...ik, ...lj -> ...ijkl", b, b)
-            + np.einsum("...il, ...kj -> ...ijkl", b, b)
+        f = (d + 4.0) * s - np.sum(
+            np.sqrt(1.5 * evs + s[..., None] ** 2.0), axis=-1
+        )
+        f_prime = 4.0 + np.sum(
+            1.0 - s[..., None] / np.sqrt(1.5 * evs + s[..., None] ** 2.0),
+            axis=-1,
         )
-    )
 
+        s -= f / f_prime
 
-def symmetric_quadratic_closure(a):
-    """Generate a symmetric quadratic closure.
-    Parameters
-    ----------
-    a : (Mx)3x3 numpy array
-        (Array of) Second order fiber orientation tensor.
-    Returns
-    -------
-    A : (Mx)3x3x3x3 numpy array
-        (Array of) Fourth order fiber orientation tensor.
-    References
-    ----------
-    .. [1] Karl, Tobias and Gatti, Davide and Frohnapfel, Bettina and Böhlke, Thomas,
-       'Asymptotic fiber orientation states of the quadratically
-       closed Folgar--Tucker equation and a subsequent closure
-       improvement',
-       Journal of Rheology 65(5) : 999-1022,
-       https://doi.org/10.1122/8.0000245
-    Notes
-    ----------
-        In general, the SQC does not contract to its second-order input.
-    """
-    assert_fot_properties(a)
+        err = np.linalg.norm(f)
+        it += 1
 
-    a4 = (
+    if err > eps_newton:
+        raise ValueError("Newton algorithm did not converge.")
+
+    mus = np.sqrt(1.5 * evs + s[..., None] ** 2) - s[..., None]
+    b = np.einsum("...ij, ...j, ...kj -> ...ik", r, mus, r)
+
+    return (
         1.0
         / 3.0
         * (
-            np.einsum("...ij, ...kl -> ...ijkl", a, a)
-            + np.einsum("...ik, ...lj -> ...ijkl", a, a)
-            + np.einsum("...il, ...kj -> ...ijkl", a, a)
+            np.einsum("...ij, ...kl -> ...ijkl", b, b)
+            + np.einsum("...ik, ...lj -> ...ijkl", b, b)
+            + np.einsum("...il, ...kj -> ...ijkl", b, b)
         )
     )
-
-    a4 /= 1.0 / 3.0 * (1.0 + 2.0 * np.einsum("...ij, ...ij -> ...", a, a))
-
-    return a4