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