PoissonFEM/postprocessing.py at main · matteo-cadoni/PoissonFEM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
from matplotlib.tri import CubicTriInterpolator, Triangulation

import triangle as tr


def postprocessing(geometry, triangulation, nodes, solution, degree):
    """
    Plot the domain, mesh, and finite element solution.

    Parameters
    ----------
    geometry : dict
        Geometry dictionary passed to triangle.
    triangulation : dict
        Output of tr.triangulate(...).
    nodes : ndarray
        Array of node coordinates, shape (n_nodes, 2).
    solution : ndarray
        Solution values at the nodes.
    degree : str
        Element degree ('linear', optionally others later).
    """
    # -----------------------------------------------------------------
    # Plot geometry and triangulation
    # -----------------------------------------------------------------
    tr.compare(plt, geometry, triangulation)
    plt.show()

    # If you want to save the mesh figure, do it before show() or keep a handle.
    # plt.savefig("lamado.png", dpi=300, bbox_inches="tight")

    # -----------------------------------------------------------------
    # Plot solution
    # -----------------------------------------------------------------
    if degree == "linear":
        n_nodes = nodes.shape[0]

        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")

        surf = ax.plot_trisurf(
            nodes[:n_nodes, 0],
            nodes[:n_nodes, 1],
            triangulation["triangles"],
            solution,
            cmap=cm.coolwarm,
            linewidth=0.1,
            edgecolors="k",
            antialiased=True,
        )

        fig.colorbar(surf, shrink=0.5, aspect=5)
        ax.view_init(elev=12, azim=20)

        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_zlabel("u")

        plt.show()


def plot_vector_field(nodes, triangulation, solution, levels=13):
    """
    Optional helper: plot contours and normalized electric field direction.
    """
    triang = Triangulation(nodes[:, 0], nodes[:, 1], triangulation["triangles"])
    interpolator = CubicTriInterpolator(triang, solution)

    # Gradient at mesh nodes
    Ex, Ey = interpolator.gradient(nodes[:, 0], nodes[:, 1])
    E_norm = np.sqrt(Ex**2 + Ey**2)

    # Avoid division by zero
    E_norm[E_norm == 0] = 1.0

    fig, ax = plt.subplots()
    ax.set_aspect("equal")
    ax.use_sticky_edges = False
    ax.margins(0.07)

    ax.tricontour(
        nodes[:, 0], nodes[:, 1], triangulation["triangles"], solution, levels
    )

    ax.quiver(
        nodes[:, 0],
        nodes[:, 1],
        -Ex / E_norm,
        -Ey / E_norm,
        units="xy",
        scale=10.0,
        zorder=3,
        color="blue",
        width=0.007,
        headwidth=3.0,
        headlength=4.0,
    )

    ax.set_xlabel("x")
    ax.set_ylabel("y")
    plt.show()


def plot_contours(nodes, triangulation, solution, levels=13):
    """
    Optional helper: plot contour lines of the solution.
    """
    plt.figure()
    plt.tricontour(
        nodes[:, 0], nodes[:, 1], triangulation["triangles"], solution, levels
    )
    plt.xlabel("x")
    plt.ylabel("y")
    plt.gca().set_aspect("equal")
    plt.show()