Port SolCx analytic Stokes solution (uw.function.analytic.SolCx)

setup.py

@@ -243,6 +243,8 @@ def configure():
         sources=[
             "src/underworld3/function/analytic.pyx",
             "src/underworld3/function/AnalyticSolNL.c",
+            "src/underworld3/function/AnalyticSolCx.c",
+            "src/underworld3/function/solCx.c",
         ],
         extra_compile_args=extra_compile_args,
         **conf,

src/underworld3/function/AnalyticSolCx.c

@@ -0,0 +1,45 @@
+/*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*
+**                                                                                  **
+** This file forms part of the Underworld geophysics modelling application.         **
+**                                                                                  **
+** For full license and copyright information, please refer to the LICENSE.md file  **
+** located at the project root, or contact the authors.                             **
+**                                                                                  **
+**~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*/
+
+/* Thin UW3 adapters around the verbatim Velic SolCx kernel (solCx.c).
+   _Velic_solCx fills caller-provided arrays; pass NULL for components we
+   don't want. The viscosity is the analytic step (the kernel doesn't return
+   it as a field). */
+
+#include "solCx.h"
+#include "AnalyticSolCx.h"
+
+vec2 SolCx_velocity( double eta_A, double eta_B, double x_c, int n, const double x, const double z )
+{
+    double pos[2];
+    double vel[2];
+    vec2 out;
+    pos[0] = x;
+    pos[1] = z;
+    _Velic_solCx( pos, eta_A, eta_B, x_c, n, vel, NULL, NULL, NULL );
+    out.x = vel[0];
+    out.z = vel[1];
+    return out;
+}
+
+double SolCx_pressure( double eta_A, double eta_B, double x_c, int n, const double x, const double z )
+{
+    double pos[2];
+    double pressure;
+    pos[0] = x;
+    pos[1] = z;
+    _Velic_solCx( pos, eta_A, eta_B, x_c, n, NULL, &pressure, NULL, NULL );
+    return pressure;
+}
+
+double SolCx_viscosity( double eta_A, double eta_B, double x_c, int n, const double x, const double z )
+{
+    (void) n; (void) z;
+    return ( x < x_c ) ? eta_A : eta_B;
+}

src/underworld3/function/AnalyticSolCx.h

@@ -0,0 +1,31 @@
+/*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*
+**                                                                                  **
+** This file forms part of the Underworld geophysics modelling application.         **
+**                                                                                  **
+** For full license and copyright information, please refer to the LICENSE.md file  **
+** located at the project root, or contact the authors.                             **
+**                                                                                  **
+**~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*/
+
+/* UW3 wrapper for the classic Velic "SolCx" analytic Stokes solution:
+   isoviscous step in x (eta_A for x<x_c, eta_B for x>=x_c), trigonometric
+   density forcing  f = (0, +cos(pi x) sin(n pi z))  (UW3 momentum-sign
+   convention; UW2 docs quote -cos) on a unit box with free-slip walls.
+   The heavy lifting is the verbatim Velic kernel in solCx.c
+   (_Velic_solCx); these thin adapters expose it in the same vec2/scalar shape
+   as AnalyticSolNL so analytic.pyx can wrap it identically. */
+
+#ifndef __Underworld_Function_AnalyticSolCx_h__
+#define __Underworld_Function_AnalyticSolCx_h__
+
+#include "AnalyticSolNL.h"   /* for the shared vec2 typedef */
+
+/* See AnalyticSolNL.h for the rationale behind __attribute__((const)):
+   the result depends solely on the arguments, so redundant calls (e.g. the
+   x and z components evaluated separately) can be collapsed by the compiler. */
+
+vec2   SolCx_velocity(  double eta_A, double eta_B, double x_c, int n, const double x, const double z ) __attribute__((const));
+double SolCx_pressure(  double eta_A, double eta_B, double x_c, int n, const double x, const double z ) __attribute__((const));
+double SolCx_viscosity( double eta_A, double eta_B, double x_c, int n, const double x, const double z ) __attribute__((const));
+
+#endif

src/underworld3/function/analytic.pyx

@@ -16,6 +16,12 @@ cdef extern from "AnalyticSolNL.h" nogil:
     vec2   SolNL_bodyforce( double eta0, unsigned n, double r, double x, double y )
     double SolNL_viscosity( double eta0, unsigned n, double r, double x, double y )
 
+cdef extern from "AnalyticSolCx.h" nogil:
+    # vec2 already declared above (shared typedef via AnalyticSolNL.h)
+    vec2   SolCx_velocity(  double eta_A, double eta_B, double x_c, int n, double x, double z )
+    double SolCx_pressure(  double eta_A, double eta_B, double x_c, int n, double x, double z )
+    double SolCx_viscosity( double eta_A, double eta_B, double x_c, int n, double x, double z )
+
 class sympy_function_printable(sympy.Function):
     """
     This help function simply does most of the work for c-code printing.
@@ -82,4 +88,111 @@ class AnalyticSolNL_viscosity(AnalyticSolNL_base):
     _printstr = "SolNL_viscosity({})"
     def _eval_evalf(self,prec):
         from sympy import sympify
-        return sympify(SolNL_viscosity( self.args[0],self.args[1],self.args[2],self.args[3],self.args[4] ))
+        return sympify(SolNL_viscosity( self.args[0],self.args[1],self.args[2],self.args[3],self.args[4] ))
+
+
+# ----------------------------------------------------------------------------
+# SolCx: viscosity step in x (eta_A for x<x_c, eta_B for x>=x_c), trigonometric
+# density forcing  f = (0, +cos(pi x) sin(n pi z))  on a unit box with free-slip
+# walls. The +cos sign matches UW3's momentum convention (UW2 docs quote -cos).
+# args: (eta_A, eta_B, x_c, n, x, z)
+# ----------------------------------------------------------------------------
+class AnalyticSolCx_base(sympy_function_printable):
+    nargs = 6
+    _header = "AnalyticSolCx.h"
+
+class AnalyticSolCx_velocity_x(AnalyticSolCx_base):
+    _printstr = "SolCx_velocity({}).x"
+    def _eval_evalf(self,prec):
+        from sympy import sympify
+        return sympify(SolCx_velocity( self.args[0],self.args[1],self.args[2],self.args[3],self.args[4],self.args[5] ).x)
+class AnalyticSolCx_velocity_y(AnalyticSolCx_base):
+    _printstr = "SolCx_velocity({}).z"
+    def _eval_evalf(self,prec):
+        from sympy import sympify
+        return sympify(SolCx_velocity( self.args[0],self.args[1],self.args[2],self.args[3],self.args[4],self.args[5] ).z)
+class AnalyticSolCx_velocity(AnalyticSolCx_base):
+    nargs = 6
+    @classmethod
+    def eval(cls, *args ):
+        from sympy.vector import CoordSys3D
+        N = CoordSys3D("N")
+        return AnalyticSolCx_velocity_x(*args)*N.i + AnalyticSolCx_velocity_y(*args)*N.j
+
+class AnalyticSolCx_pressure(AnalyticSolCx_base):
+    _printstr = "SolCx_pressure({})"
+    def _eval_evalf(self,prec):
+        from sympy import sympify
+        return sympify(SolCx_pressure( self.args[0],self.args[1],self.args[2],self.args[3],self.args[4],self.args[5] ))
+
+class AnalyticSolCx_viscosity(AnalyticSolCx_base):
+    _printstr = "SolCx_viscosity({})"
+    def _eval_evalf(self,prec):
+        from sympy import sympify
+        return sympify(SolCx_viscosity( self.args[0],self.args[1],self.args[2],self.args[3],self.args[4],self.args[5] ))
+
+
+class SolCx:
+    r"""Velic *SolCx* analytic Stokes solution — the canonical free-slip,
+    discontinuous-viscosity benchmark: a viscosity step :math:`\eta_A` (for
+    :math:`x<x_c`) to :math:`\eta_B` (for :math:`x>x_c`), driven by the density
+    forcing :math:`f=(0,\cos(\pi x)\sin(n\pi z))` on the unit box with free-slip
+    walls. Use it to validate a UW3 Stokes solve against the exact solution::
+
+        sol = uw.function.analytic.SolCx(mesh, eta_A=1.0, eta_B=1.0e6, x_c=0.5, n=1)
+        stokes.constitutive_model.Parameters.shear_viscosity_0 = sol.fn_viscosity
+        stokes.bodyforce = sol.fn_bodyforce
+        # free-slip on all four walls + pressure nullspace, then solve and:
+        rel = sol.velocity_error(stokes.u)     # L2 error vs the analytic
+
+    Notes
+    -----
+    ``fn_velocity`` / ``fn_pressure`` wrap the compiled Velic kernel and are
+    evaluated point-wise via SymPy ``evalf`` (they are not ``lambdify``-able, so
+    use :meth:`evaluate_velocity` / :meth:`velocity_error` rather than
+    ``uw.function.evaluate``). ``fn_bodyforce`` / ``fn_viscosity`` are elementary
+    and compile through the normal JIT path when assigned to a solver.
+
+    The body-force **sign** matches UW3's momentum convention; the original
+    Underworld2 documentation quotes the opposite sign.
+    """
+    def __init__(self, mesh, eta_A=1.0, eta_B=1.0e6, x_c=0.5, n=1):
+        import sympy as _sp
+        if getattr(mesh, "dim", None) != 2:
+            raise ValueError("SolCx is a 2D analytic solution; mesh.dim must be 2.")
+        if not (float(eta_A) > 0.0 and float(eta_B) > 0.0):
+            raise ValueError("eta_A and eta_B must be positive.")
+        if not (0.0 <= float(x_c) <= 1.0):
+            raise ValueError("x_c must lie in [0, 1].")
+        if int(n) != n or int(n) < 1:
+            raise ValueError("n (z-wavenumber) must be a positive integer.")
+        self.eta_A = float(eta_A)
+        self.eta_B = float(eta_B)
+        self.x_c   = float(x_c)
+        self.n     = int(n)
+        x, y = mesh.X
+        p = (self.eta_A, self.eta_B, self.x_c, self.n)
+        self.fn_velocity  = _sp.Matrix([AnalyticSolCx_velocity_x(*p, x, y),
+                                        AnalyticSolCx_velocity_y(*p, x, y)])
+        self.fn_pressure  = AnalyticSolCx_pressure(*p, x, y)
+        # tie-break at x==x_c matches the compiled SolCx_viscosity (eta_B for x>=x_c)
+        self.fn_viscosity = _sp.Piecewise((self.eta_A, x < self.x_c), (self.eta_B, True))
+        self.fn_bodyforce = _sp.Matrix([_sp.Integer(0),
+                                        _sp.cos(_sp.pi * x) * _sp.sin(self.n * _sp.pi * y)])
+
+    def evaluate_velocity(self, coords):
+        """Exact velocity at ``coords`` (N×2 array), via the compiled kernel."""
+        import numpy as _np
+        p = (self.eta_A, self.eta_B, self.x_c, self.n)
+        out = _np.empty((len(coords), 2))
+        for i in range(len(coords)):
+            xi = float(coords[i, 0]); yi = float(coords[i, 1])
+            out[i, 0] = float(AnalyticSolCx_velocity_x(*p, xi, yi).evalf())
+            out[i, 1] = float(AnalyticSolCx_velocity_y(*p, xi, yi).evalf())
+        return out
+
+    def velocity_error(self, velocity_var):
+        """Relative L2 error of a velocity MeshVariable against the analytic."""
+        import numpy as _np
+        ua = self.evaluate_velocity(velocity_var.coords)
+        return float(_np.linalg.norm(velocity_var.data - ua) / _np.linalg.norm(ua))