fenics_tutos/poisson_visualization.py

"""
Solution et animation de l'équation de Poisson
===============================================

Ce script résout l'équation de Poisson:
    -Δu = f  dans Ω
    u = 0    sur Γ_D (conditions de Dirichlet)
    ∂u/∂n = g sur Γ_N (conditions de Neumann)

Avec:
- Ω = [0,2] x [0,1] (domaine rectangulaire)
- f = 10*exp(-((x-0.5)² + (y-0.5)²)/0.02) (source gaussienne)
- g = sin(5x) (flux de Neumann)
"""

import dolfinx as df
from pathlib import Path
from mpi4py import MPI
from petsc4py.PETSc import ScalarType
from dolfinx.fem.petsc import LinearProblem
import numpy as np
import ufl
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.animation import FuncAnimation, PillowWriter
import warnings
warnings.filterwarnings('ignore')

# Créer le dossier de sortie
out_folder = Path("out_poisson")
out_folder.mkdir(parents=True, exist_ok=True)

print("=" * 60)
print("RÉSOLUTION DE L'ÉQUATION DE POISSON")
print("=" * 60)

# ============================================================
# 1. Création du maillage
# ============================================================
print("\n[1/5] Création du maillage...")
nx, ny = 64, 32
msh = df.mesh.create_rectangle(
    comm=MPI.COMM_WORLD,
    points=((0.0, 0.0), (2.0, 1.0)),
    n=(nx, ny),
    cell_type=df.mesh.CellType.triangle,
)
print(f"    Maillage: {nx}x{ny} éléments triangulaires")
print(f"    Domaine: [0, 2] x [0, 1]")

# Espace de fonctions
V = df.fem.functionspace(msh, ("Lagrange", 1))
print(f"    Espace d'éléments finis: Lagrange P1")
print(f"    Nombre de degrés de liberté: {V.dofmap.index_map.size_global}")

# ============================================================
# 2. Conditions aux limites
# ============================================================
print("\n[2/5] Application des conditions aux limites...")

# Conditions de Dirichlet sur x=0 et x=2
facets = df.mesh.locate_entities_boundary(
    msh,
    dim=(msh.topology.dim - 1),
    marker=lambda x: np.isclose(x[0], 0.0) | np.isclose(x[0], 2.0),
)
dofs = df.fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)
bc = df.fem.dirichletbc(value=ScalarType(0), dofs=dofs, V=V)
print("    Dirichlet: u = 0 sur x = 0 et x = 2")

# ============================================================
# 3. Formulation variationnelle
# ============================================================
print("\n[3/5] Définition de la formulation variationnelle...")

u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
x = ufl.SpatialCoordinate(msh)

# Terme source (gaussienne centrée en (0.5, 0.5))
f = 10 * ufl.exp(-((x[0] - 0.5)**2 + (x[1] - 0.5)**2) / 0.02)

# Condition de Neumann
g = ufl.sin(5 * x[0])

# Mesures d'intégration
dx = ufl.Measure("dx", domain=msh)
ds = ufl.Measure("ds", domain=msh)

# Forme bilinéaire et linéaire
a = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx
L = ufl.inner(f, v) * dx + ufl.inner(g, v) * ds

print("    Équation: -Δu = f")
print("    Source: f = 10*exp(-((x-0.5)² + (y-0.5)²)/0.02)")
print("    Neumann: g = sin(5x) sur y = 0 et y = 1")

# ============================================================
# 4. Résolution du système
# ============================================================
print("\n[4/5] Résolution du système linéaire...")

problem = LinearProblem(
    a, L,
    bcs=[bc],
    petsc_options_prefix="demo_poisson_",
    petsc_options={
        "ksp_type": "preonly",
        "pc_type": "lu",
        "ksp_error_if_not_converged": True
    },
)
uh = problem.solve()

print(f"    Solution calculée!")
print(f"    Valeur min: {uh.x.array.min():.6f}")
print(f"    Valeur max: {uh.x.array.max():.6f}")

# ============================================================
# 5. Visualisation
# ============================================================
print("\n[5/5] Génération des visualisations...")

# Extraire les coordonnées et les valeurs
coords = V.tabulate_dof_coordinates()
x_coords = coords[:, 0]
y_coords = coords[:, 1]
u_values = uh.x.array.real

# Créer une grille régulière pour l'interpolation
from scipy.interpolate import griddata

xi = np.linspace(0, 2, 200)
yi = np.linspace(0, 1, 100)
Xi, Yi = np.meshgrid(xi, yi)

# Interpoler les valeurs
Zi = griddata((x_coords, y_coords), u_values, (Xi, Yi), method='cubic')

# ============================================================
# Figure 1: Solution statique
# ============================================================
print("    Création de la figure statique...")

fig1 = plt.figure(figsize=(14, 10))

# Subplot 1: Vue 2D avec contours
ax1 = fig1.add_subplot(2, 2, 1)
contour = ax1.contourf(Xi, Yi, Zi, levels=50, cmap='viridis')
ax1.contour(Xi, Yi, Zi, levels=15, colors='white', linewidths=0.5, alpha=0.5)
plt.colorbar(contour, ax=ax1, label='u(x,y)')
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_title('Solution u(x,y) - Vue 2D avec contours')
ax1.set_aspect('equal')

# Subplot 2: Vue 3D
ax2 = fig1.add_subplot(2, 2, 2, projection='3d')
surf = ax2.plot_surface(Xi, Yi, Zi, cmap='viridis', edgecolor='none', alpha=0.9)
ax2.set_xlabel('x')
ax2.set_ylabel('y')
ax2.set_zlabel('u(x,y)')
ax2.set_title('Solution u(x,y) - Vue 3D')
ax2.view_init(elev=30, azim=45)

# Subplot 3: Coupes selon x
ax3 = fig1.add_subplot(2, 2, 3)
y_cuts = [0.1, 0.3, 0.5, 0.7, 0.9]
for y_cut in y_cuts:
    y_idx = int(y_cut * 100)
    if y_idx < Zi.shape[0]:
        ax3.plot(xi, Zi[y_idx, :], label=f'y = {y_cut}')
ax3.set_xlabel('x')
ax3.set_ylabel('u(x,y)')
ax3.set_title('Coupes de la solution selon x')
ax3.legend()
ax3.grid(True, alpha=0.3)

# Subplot 4: Coupes selon y
ax4 = fig1.add_subplot(2, 2, 4)
x_cuts = [0.3, 0.5, 0.7, 1.0, 1.5]
for x_cut in x_cuts:
    x_idx = int(x_cut * 100)
    if x_idx < Zi.shape[1]:
        ax4.plot(yi, Zi[:, x_idx], label=f'x = {x_cut}')
ax4.set_xlabel('y')
ax4.set_ylabel('u(x,y)')
ax4.set_title('Coupes de la solution selon y')
ax4.legend()
ax4.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig(out_folder / 'poisson_solution.png', dpi=150, bbox_inches='tight')
print(f"    Sauvegardé: {out_folder / 'poisson_solution.png'}")

# ============================================================
# Animation de construction de la solution
# ============================================================
print("\n    Création de l'animation de construction...")

# On va simuler la construction progressive de la solution
# en multipliant la solution finale par un facteur croissant
# (grâce à la linéarité de l'équation de Poisson)

fig_anim, axes = plt.subplots(1, 2, figsize=(14, 5))

# Nombre de frames
n_frames = 60

# Préparer les données pour l'animation
# Grâce à la linéarité: si u est solution pour (f,g), alors α*u est solution pour (α*f, α*g)
solutions = []

for i in range(n_frames + 1):
    # Facteur d'amplitude progressif (de 0 à 1)
    amplitude = i / n_frames
    
    # Simplement multiplier la solution par l'amplitude
    Zi_anim = Zi * amplitude
    solutions.append((amplitude, Zi_anim))

# Valeurs min/max pour les échelles fixes
z_min = min(s[1][~np.isnan(s[1])].min() for s in solutions)
z_max = max(s[1][~np.isnan(s[1])].max() for s in solutions)

# Configuration initiale
ax1_anim = axes[0]
ax2_anim = axes[1]

def init():
    ax1_anim.clear()
    ax2_anim.clear()
    return []

def animate(frame):
    amplitude, Zi_frame = solutions[frame]
    
    ax1_anim.clear()
    ax2_anim.clear()
    
    # Vue 2D
    contour = ax1_anim.contourf(Xi, Yi, Zi_frame, levels=50, cmap='viridis', 
                                 vmin=z_min, vmax=z_max)
    ax1_anim.contour(Xi, Yi, Zi_frame, levels=10, colors='white', linewidths=0.3, alpha=0.5)
    ax1_anim.set_xlabel('x')
    ax1_anim.set_ylabel('y')
    ax1_anim.set_title(f'Solution u(x,y) - Amplitude: {amplitude:.0%}')
    ax1_anim.set_aspect('equal')
    
    # Coupes
    y_cuts = [0.25, 0.5, 0.75]
    for y_cut in y_cuts:
        y_idx = int(y_cut * 100)
        if y_idx < Zi_frame.shape[0]:
            ax2_anim.plot(xi, Zi_frame[y_idx, :], label=f'y = {y_cut}', linewidth=2)
    
    ax2_anim.set_xlabel('x')
    ax2_anim.set_ylabel('u(x,y)')
    ax2_anim.set_title(f'Coupes de la solution - Construction: {amplitude:.0%}')
    ax2_anim.set_ylim(z_min - 0.1, z_max + 0.1)
    ax2_anim.set_xlim(0, 2)
    ax2_anim.legend(loc='upper right')
    ax2_anim.grid(True, alpha=0.3)
    
    return []

# Créer l'animation
anim = FuncAnimation(fig_anim, animate, init_func=init, 
                     frames=len(solutions), interval=100, blit=True)

# Sauvegarder en GIF
print("    Sauvegarde de l'animation (cela peut prendre un moment)...")
anim.save(out_folder / 'poisson_construction.gif', writer=PillowWriter(fps=15))
print(f"    Sauvegardé: {out_folder / 'poisson_construction.gif'}")

plt.close('all')

# ============================================================
# Animation 3D de la construction
# ============================================================
print("\n    Création de l'animation 3D...")

fig_3d = plt.figure(figsize=(10, 8))
ax_3d = fig_3d.add_subplot(111, projection='3d')

def animate_3d(frame):
    amplitude, Zi_frame = solutions[frame]
    
    ax_3d.clear()
    
    # Surface 3D
    surf = ax_3d.plot_surface(Xi, Yi, Zi_frame, cmap='viridis', 
                               edgecolor='none', alpha=0.9,
                               vmin=z_min, vmax=z_max)
    
    ax_3d.set_xlabel('x')
    ax_3d.set_ylabel('y')
    ax_3d.set_zlabel('u(x,y)')
    ax_3d.set_title(f'Équation de Poisson - Construction: {amplitude:.0%}')
    ax_3d.set_zlim(z_min - 0.1, z_max + 0.1)
    ax_3d.view_init(elev=25, azim=45 + frame * 2)  # Rotation progressive
    
    return []

anim_3d = FuncAnimation(fig_3d, animate_3d, frames=len(solutions), interval=100)
anim_3d.save(out_folder / 'poisson_construction_3d.gif', writer=PillowWriter(fps=15))
print(f"    Sauvegardé: {out_folder / 'poisson_construction_3d.gif'}")

plt.close('all')

# ============================================================
# Résumé final
# ============================================================
print("\n" + "=" * 60)
print("RÉSUMÉ")
print("=" * 60)
print(f"\nFichiers générés dans '{out_folder}':")
print(f"  1. poisson_solution.png      - Solution statique (4 vues)")
print(f"  2. poisson_construction.gif  - Animation 2D de construction")
print(f"  3. poisson_construction_3d.gif - Animation 3D de construction")
print("\nCaractéristiques de la solution:")
print(f"  - Valeur minimale: {u_values.min():.6f}")
print(f"  - Valeur maximale: {u_values.max():.6f}")
print(f"  - Valeur moyenne:  {u_values.mean():.6f}")
print("=" * 60)