1import enum
  2from enum import Enum
  3
  4import ngsolve
  5from ngsolve import (
  6    BBND,
  7    CF,
  8    H1,
  9    L2,
 10    QUAD,
 11    TRIG,
 12    BilinearForm,
 13    Discontinuous,
 14    IfPos,
 15    LinearForm,
 16    Normalize,
 17    dx,
 18    specialcf,
 19)
 20from ngsolve.solve_implementation import GridFunction
 21from ngstrefftz import (
 22    EmbeddedTrefftzFES,
 23    TrefftzEmbedding,
 24)
 25from xfem import (
 26    IF,
 27    NEG,
 28    POS,
 29    CutInfo,
 30    dCut,
 31)
 32
 33from ngstSpaceKit.diffops import grad, hesse
 34from ngstSpaceKit.mesh_properties import (
 35    throw_on_wrong_mesh_dimension,
 36    throw_on_wrong_mesh_eltype,
 37)
 38
 39
 40def ImmersedP1FE(
 41    mesh: ngsolve.comp.Mesh,
 42    lsetp1: ngsolve.CoefficientFunction,
 43    beta_neg: float,
 44    beta_pos: float,
 45    dirichlet: str = "",
 46    dgjumps: bool = False,
 47    check_mesh: bool = True,
 48    stats: dict | None = None,
 49) -> EmbeddedTrefftzFES:
 50    r"""
 51    `check_mesh`: test, if the `mesh` is compatible with this space
 52
 53    `stats`: use the `stats` flag of the `TrefftzEmbeddin` method
 54
 55    This Immersed P1 space is tailored towards solving the following interface problem.
 56
 57    Let $\Omega$ be a domain, which is decomposed by a cut $\Gamma$ into $\Omega = \Omega^- \cup \Gamma \cup \Omega^+$.
 58    Let $\beta$ be a piecewise constant coefficient
 59    \begin{align}
 60    \beta(x) &:=
 61        \begin{cases}
 62            \beta^-, &\text{if } x \in \Omega^- \\\\
 63            \beta^+, &\text{if } x \in \Omega^+ \\\\
 64        \end{cases}, \\\\
 65    \beta^-, \beta^+ &> 0.
 66    \end{align}
 67
 68    Then, find $u$, s.t.
 69
 70    \begin{align}
 71        -\operatorname{div} (\beta \nabla u) &= f \text{ in } \Omega^- \cup \Omega^+, \\\\
 72                                        〚u〛 &= 0 \text{ on } \Gamma, \\\\
 73	        〚\beta \nabla u \cdot n_\Gamma〛 &= 0 \text{ on } \Gamma, \\\\
 74	                                       u &= 0 \text{ on } \partial \Gamma. \\\\
 75    \end{align}
 76
 77    In particular, the functions in this space fulfil the property
 78    \begin{align}
 79                                〚u〛 &= 0 \text{ on } \Gamma, \\\\
 80	〚\beta \nabla u \cdot n_\Gamma〛 &= 0 \text{ on } \Gamma, \\\\
 81    \end{align}
 82    as well as being continuous at mesh vertices.
 83    The space consists of piecewise linear functions.
 84
 85    Actually, the returned space consists of vectorial functions of order $1$,
 86    that have to be interpreted in the following way: the first component represents the
 87    piecewise linear function in $\Omega^-$, the second component in $\Omega^+$.
 88    Formally, for $v = \begin{pmatrix} v^- \\\\ v^+ \end{pmatrix}$ we define the piecewise linear function
 89    \begin{align}
 90        \hat{v}(x) &:=
 91            \begin{cases}
 92                v^-(x), &\text{if } x \in \Omega^- \\\\
 93                v^+(x), &\text{if } x \in \Omega^+ \\\\
 94            \end{cases}.
 95    \end{align}
 96
 97    `lsetp1`: The levelset function $p$ used to describe the cut as $\Gamma := \\{x \in \Omega \mid p(x) = 0 \\}$,
 98        as well as $\Omega^- := \\{x \in \Omega \mid p(x) < 0 \\}$ and $\Omega^- := \\{x \in \Omega \mid p(x) > 0 \\}$.
 99        $p$ needs to be affine linear on each element. E.g. set `lsetp1` as a `Gridfunction(H1(mesh, order=1))`.
100
101    `beta_neg`: diffusion coefficient for $\Omega^-$. Should be a positive number.
102
103    `beta_pos`: diffusion coefficient for $\Omega^+$. Should be a positive number.
104
105    # Conforming Trefftz Formulation
106    - $\mathbb{V}_h := [\mathbb{P}^{1, \text{disc}}(\mathcal{T}_h)]^2$
107    - $\mathbb{Q}_h := \mathbb{V}_h$
108    - $\mathbb{Z}_h := \mathbb{P}^{1}(\mathcal{T}_h)$
109    - \begin{align}
110      \mathcal{C}_K(v_h, z_h) &:=
111          \sum_{p \text{ is vertex}} \hat{v}_h(p) z_h(p) \\\\
112      \mathcal{D}_K(y_h, z_h) &:=
113          \sum_{p \text{ is vertex}} y_h(p) z_h(p) \\\\
114      \end{align}
115    - \begin{align}
116      (\mathcal{L}_K v_h, q_h) :=
117        \int_\Gamma 〚\hat{v}_h〛 〚\hat{q}_h〛 \;dS + \int_\Gamma〚\beta \nabla \hat{v}_h \cdot n_\Gamma〛 〚\beta \nabla \hat{q}_h \cdot n_\Gamma〛\;dS
118      \end{align}
119    """
120
121    if check_mesh:
122        throw_on_wrong_mesh_dimension(mesh, 2)
123        throw_on_wrong_mesh_eltype(mesh, TRIG)
124
125    ci = CutInfo(mesh, lsetp1)
126    dS = dCut(lsetp1, IF, definedonelements=ci.GetElementsOfType(IF), order=2)
127    n = Normalize(grad(lsetp1, mesh.dim))
128    h = specialcf.mesh_size
129
130    fes = L2(mesh, order=1, dgjumps=dgjumps) * L2(
131        mesh, order=1, dgjumps=dgjumps
132    )
133    conformity_space = H1(mesh, order=1, dirichlet=dirichlet)
134
135    (u_neg, u_pos), (v_neg, v_pos) = fes.TnT()
136
137    top = u_neg * v_neg * dx(definedonelements=ci.GetElementsOfType(POS))
138    top += u_pos * v_pos * dx(definedonelements=ci.GetElementsOfType(NEG))
139    top += 1 / h * (u_pos - u_neg) * (v_pos - v_neg) * dS
140    top += (
141        h
142        * ((beta_pos * grad(u_pos) - beta_neg * grad(u_neg)) * n)
143        * ((beta_pos * grad(v_pos) - beta_neg * grad(v_neg)) * n)
144        * dS
145    )
146
147    uc, vc = conformity_space.TnT()
148
149    cop = IfPos(lsetp1, u_pos, u_neg) * vc * dx(element_vb=BBND)
150
151    crhs = uc * vc * dx(element_vb=BBND)
152
153    emb = TrefftzEmbedding(
154        top=top,
155        trhs=None,
156        cop=cop,
157        crhs=crhs,
158        ndof_trefftz=0,
159        stats=stats,
160    )
161    imp1fe = EmbeddedTrefftzFES(emb)
162    return imp1fe
163
164
165class ImmersedQ1Impl(Enum):
166    """
167    Represents a Conforming Trefftz implementation
168    for the `ImmersedQ1FE` space.
169    """
170
171    Canonical = enum.auto()
172    """
173    Formulation is described by <https://doi.org/10.1002/num.20318>.
174    """
175
176    NonConforming = enum.auto()
177    """
178    Stable formulation, that is not continuous acrosss the cut interface.
179    """
180
181    Overloaded = enum.auto()
182    """
183    Technically overconstrains the Trefftz condition,
184    but it may still work.
185    """
186
187
188def ImmersedQ1FE(
189    mesh: ngsolve.comp.Mesh,
190    lsetq1: ngsolve.GridFunction,
191    beta_neg: float,
192    beta_pos: float,
193    dirichlet: str = "",
194    dgjumps: bool = False,
195    stats: dict | None = None,
196    impl: ImmersedQ1Impl = ImmersedQ1Impl.NonConforming,
197) -> EmbeddedTrefftzFES:
198    r"""
199    This is the version of `ImmersedP1FE` for quadrilateral meshes.
200    Refer to the documentation of `ImmersedP1FE` for most details.
201
202    `lsetq1`: The levelset function, as a `GridFunction` over the `H1` space with `order=1`.
203        In geleral, the cut may not be piecewise linear, as the `H1` space of `order=1` contains bilinear functions on quads.
204        You can use `straighten_levelset` inorder to produce a levelset function
205        with a piecewise linear cut.
206
207    `impl`: declare what Trefftz implementation you want to use
208
209    # Canonical Conforming Trefftz Formulation
210    This implements the formulation of <https://doi.org/10.1002/num.20318>.
211    - $\mathbb{V}_h := [\mathbb{Q}^{1, \text{disc}}(\mathcal{T}_h)]^2$
212    - $\mathbb{Q}_h := \mathbb{Q}^{1, \text{disc}}(\mathcal{T}_h) \times \mathbb{Q}^{0, \text{disc}}(\mathcal{T}_h)$
213    - $\mathbb{Z}_h := \mathbb{Q}^{1}(\mathcal{T}_h)$
214    - \begin{align}
215      \mathcal{C}_K(v_h, z_h) &:=
216          \sum_{p \text{ is vertex}} \hat{v}_h(p) z_h(p) \\\\
217      \mathcal{D}_K(y_h, z_h) &:=
218          \sum_{p \text{ is vertex}} y_h(p) z_h(p) \\\\
219      \end{align}
220    - \begin{align}
221      (\mathcal{L}_K (v_h, (q_{h, 1}, q_{h, 0})) :=
222        \int_\Gamma 〚\hat{v}_h〛q_{h,1} \;dS + \int_\Gamma〚\beta \nabla \hat{v}_h \cdot n_\Gamma〛 q_{h,0} \;dS
223      \end{align}
224
225    # Non-Conforming Conforming Trefftz Formulation
226    - $\mathbb{V}_h := [\mathbb{Q}^{1, \text{disc}}(\mathcal{T}_h)]^2$
227    - $\mathbb{Q}_h := [\mathbb{Q}^{0, \text{disc}}(\mathcal{T}_h)]^4$
228    - $\mathbb{Z}_h := \mathbb{Q}^{1}(\mathcal{T}_h)$
229    - \begin{align}
230    \mathcal{C}_K(v_h, z_h) &:=
231        \sum_{p \text{ is vertex}} \hat{v}_h(p) z_h(p) \\\\
232    \mathcal{D}_K(y_h, z_h) &:=
233        \sum_{p \text{ is vertex}} y_h(p) z_h(p) \\\\
234    \end{align}
235    - \begin{align}
236    (\mathcal{L}_K v_h, (q_h, q_\tau, q_n, q_{\tau n})) &:=
237        \int_\Gamma \frac{1}{h}〚\hat{v}_h〛 q_h \;dS + \int_\Gamma〚\nabla \hat{v}_h \cdot \tau_\Gamma〛 q_\tau \;dS \\\\
238        &+ \int_\Gamma〚\beta \nabla \hat{v}_h \cdot n_\Gamma〛 q_n \;dS + \int_\Gamma h〚\beta n_\Gamma^T \mathbf{H}_{\hat{v}_h} \tau_\Gamma〛 q_{\tau n} \;dS
239    \end{align}
240
241    # Overloaded Conforming Trefftz Formulation
242    - $\mathbb{V}_h := [\mathbb{Q}^{1, \text{disc}}(\mathcal{T}_h)]^2$
243    - $\mathbb{Q}_h := \mathbb{V}_h$
244    - $\mathbb{Z}_h := \mathbb{Q}^{1}(\mathcal{T}_h)$
245    - \begin{align}
246    \mathcal{C}_K(v_h, z_h) &:=
247        \sum_{p \text{ is vertex}} \hat{v}_h(p) z_h(p) \\\\
248    \mathcal{D}_K(y_h, z_h) &:=
249        \sum_{p \text{ is vertex}} y_h(p) z_h(p) \\\\
250    \end{align}
251    - \begin{align}
252    (\mathcal{L}_K v_h, q_h) :=
253        \int_\Gamma 〚\hat{v}_h〛 〚\hat{q}_h〛 \;dS + \int_\Gamma〚\beta \nabla \hat{v}_h \cdot n_\Gamma〛 〚\beta \nabla \hat{q}_h \cdot n_\Gamma〛\;dS
254    \end{align}
255    """
256
257    throw_on_wrong_mesh_dimension(mesh, 2)
258    throw_on_wrong_mesh_eltype(mesh, QUAD)
259
260    if not (
261        (
262            isinstance(lsetq1.space, H1)
263            or isinstance(lsetq1.space, Discontinuous)
264        )
265        and lsetq1.space.globalorder == 1
266    ):
267        raise ValueError(
268            f"lsetq1 must be a GridFunction on an H1(order=1) or Discontinuous(H1(order=1)) space. You got: {lsetq1.space}"
269        )
270
271    ci = CutInfo(mesh, lsetq1)
272    dS = dCut(lsetq1, IF, definedonelements=ci.GetElementsOfType(IF), order=4)
273    n = Normalize(grad(lsetq1, mesh.dim))
274    h = specialcf.mesh_size
275
276    fes = L2(mesh, order=1, dgjumps=dgjumps) * L2(
277        mesh, order=1, dgjumps=dgjumps
278    )
279
280    conformity_space = H1(mesh, order=1, dirichlet=dirichlet)
281
282    match impl:
283        case ImmersedQ1Impl.Overloaded:
284            (u_neg, u_pos), (v_neg, v_pos) = fes.TnT()
285
286            top = (
287                u_neg * v_neg * dx(definedonelements=ci.GetElementsOfType(POS))
288            )
289            top += (
290                u_pos * v_pos * dx(definedonelements=ci.GetElementsOfType(NEG))
291            )
292            top += 1 * (u_pos - u_neg) * (v_pos - v_neg) * dS
293            top += (
294                h
295                * ((beta_pos * grad(u_pos) - beta_neg * grad(u_neg)) * n)
296                * ((beta_pos * grad(v_pos) - beta_neg * grad(v_neg)) * n)
297                * dS
298            )
299        case ImmersedQ1Impl.Canonical:  # old
300            fes_test = L2(mesh, order=1) * L2(mesh, order=0)
301
302            (u_neg, u_pos) = fes.TrialFunction()
303            v_bilin, v_const = fes_test.TestFunction()
304
305            top = (
306                u_neg
307                * v_bilin
308                * dx(definedonelements=ci.GetElementsOfType(POS))
309            )
310            top += (
311                u_pos
312                * v_bilin
313                * dx(definedonelements=ci.GetElementsOfType(NEG))
314            )
315            top += 1 / h**2 * (u_pos - u_neg) * v_bilin * dS
316            top += (
317                h
318                * ((beta_pos * grad(u_pos) - beta_neg * grad(u_neg)) * n)
319                * v_const
320                * dS
321            )
322        case ImmersedQ1Impl.NonConforming:  # new version
323            fes_test = (
324                L2(mesh, order=1)
325                * L2(mesh, order=0)
326                * L2(mesh, order=0)
327                * L2(mesh, order=0)
328                * L2(mesh, order=0)
329            )
330
331            (u_neg, u_pos) = fes.TrialFunction()
332            v_bilin, v1, v2, v3, v4 = fes_test.TestFunction()
333
334            top = (
335                u_neg
336                * v_bilin
337                * dx(definedonelements=ci.GetElementsOfType(POS))
338            )
339            top += (
340                u_pos
341                * v_bilin
342                * dx(definedonelements=ci.GetElementsOfType(NEG))
343            )
344            t = CF((n[1], -n[0]))
345            top += 1 / h * (u_pos - u_neg) * v1 * dS
346            top += (grad(u_pos) - grad(u_neg) | t) * v2 * dS
347            top += (
348                (beta_pos * grad(u_pos) - beta_neg * grad(u_neg) | n) * v3 * dS
349            )
350            top += (
351                h
352                * (
353                    beta_pos * (hesse(u_pos) * t | n)
354                    - beta_neg * (hesse(u_neg) * t | n)
355                )
356                * v4
357                * dS
358            )
359
360    uc, vc = conformity_space.TnT()
361
362    cop = IfPos(lsetq1, u_pos, u_neg) * vc * dx(element_vb=BBND)
363    crhs = uc * vc * dx(element_vb=BBND)
364
365    emb = TrefftzEmbedding(
366        top=top,
367        trhs=None,
368        cop=cop,
369        crhs=crhs,
370        ndof_trefftz=0,
371        stats=stats,
372    )
373    imq1fe = EmbeddedTrefftzFES(emb)
374    return imq1fe
375
376
377def straighten_levelset(lsetq1: GridFunction) -> GridFunction:
378    """
379    Produces a new levelset function with an element-wise
380    straight cut.
381
382    This is interesting for straightening of levelset functions
383    on quad meshes.
384    """
385    eps = 1e-9
386    fes = L2(lsetq1.space.mesh, order=1)
387    u, v = fes.TnT()
388    op = (hesse(u) | hesse(v)) * dx + u * v * dCut(
389        lsetq1, IF, element_boundary=True
390    )
391    emb = TrefftzEmbedding(op, eps=eps)
392
393    etfes = EmbeddedTrefftzFES(emb)
394    u, v = etfes.TnT()
395    a = BilinearForm(u * v * dx).Assemble()
396    f = LinearForm(lsetq1 * v * dx).Assemble()
397    inv = a.mat.Inverse(inverse="sparsecholesky")
398    lsetp1_straight_trefftz = GridFunction(etfes)
399    lsetp1_straight_trefftz.vec.data = inv * f.vec
400
401    lsetp1_straight = GridFunction(fes)
402    lsetp1_straight.vec.data = emb.Embed(lsetp1_straight_trefftz.vec)
403
404    lsetp1_straight_final = GridFunction(
405        Discontinuous(H1(lsetq1.space.mesh, order=1))
406    )
407    lsetp1_straight_final.Set(lsetp1_straight)
408    return lsetp1_straight_final