simsite/backend/poisson/solver.py

"""
Solveur FEniCS pour l'équation de Poisson
-Δu = f dans Ω
avec conditions de Dirichlet et/ou Neumann sur les bords
"""

import numpy as np
import time
from typing import Dict, Any, List, Tuple
import json


def solve_poisson(params: Dict[str, Any]) -> Dict[str, Any]:
    """
    Résout l'équation de Poisson avec les paramètres donnés
    
    Retourne:
        - mesh_data: données du maillage pour visualisation 3D
        - solution_data: valeurs de la solution aux noeuds
        - animation_frames: frames pour l'animation
        - statistics: min, max, temps de calcul
    """
    import dolfinx as df
    from mpi4py import MPI
    from dolfinx.fem.petsc import LinearProblem
    from petsc4py.PETSc import ScalarType
    import ufl
    
    start_time = time.time()
    
    # Extraction des paramètres
    domain = params['domain']
    mesh_params = params['mesh']
    source_params = params['source']
    dirichlet_params = params['dirichlet']
    neumann_params = params['neumann']
    
    # 1. Création du maillage
    nx, ny = mesh_params['nx'], mesh_params['ny']
    msh = df.mesh.create_rectangle(
        comm=MPI.COMM_WORLD,
        points=[[domain['x_min'], domain['y_min']], 
                [domain['x_max'], domain['y_max']]],
        n=[nx, ny],
        cell_type=df.mesh.CellType.triangle,
    )
    
    # Espace de fonctions (Lagrange P1)
    V = df.fem.functionspace(msh, ("Lagrange", 1))
    
    # 2. Extraction des données du maillage pour la visualisation 3D
    mesh_data = extract_mesh_data(msh, V)
    
    # 3. Conditions aux limites de Dirichlet
    bcs = []
    if dirichlet_params['enabled']:
        boundaries = dirichlet_params['boundaries']
        x_min, x_max = domain['x_min'], domain['x_max']
        y_min, y_max = domain['y_min'], domain['y_max']
        
        def dirichlet_boundary(x):
            on_boundary = np.zeros(x.shape[1], dtype=bool)
            if 'left' in boundaries:
                on_boundary |= np.isclose(x[0], x_min)
            if 'right' in boundaries:
                on_boundary |= np.isclose(x[0], x_max)
            if 'bottom' in boundaries:
                on_boundary |= np.isclose(x[1], y_min)
            if 'top' in boundaries:
                on_boundary |= np.isclose(x[1], y_max)
            return on_boundary
        
        facets = df.mesh.locate_entities_boundary(
            msh,
            dim=(msh.topology.dim - 1),
            marker=dirichlet_boundary,
        )
        dofs = df.fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)
        bc = df.fem.dirichletbc(value=ScalarType(dirichlet_params['value']), dofs=dofs, V=V)
        bcs.append(bc)
    
    # 4. Formulation variationnelle
    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)
    x = ufl.SpatialCoordinate(msh)
    
    # Terme source
    f = create_source_term(x, source_params)
    
    # Mesures d'intégration
    dx = ufl.Measure("dx", domain=msh)
    ds = ufl.Measure("ds", domain=msh)
    
    # Forme bilinéaire
    a = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx
    
    # Forme linéaire (terme source)
    L = ufl.inner(f, v) * dx
    
    # Condition de Neumann si activée
    if neumann_params['enabled']:
        g = create_neumann_term(x, neumann_params)
        L = L + ufl.inner(g, v) * ds
    
    # 5. Résolution
    # Essayer avec petsc_options_prefix (versions récentes de dolfinx)
    # sinon sans (versions plus anciennes)
    try:
        problem = LinearProblem(
            a, L,
            bcs=bcs,
            petsc_options_prefix="poisson",
            petsc_options={
                "ksp_type": "preonly",
                "pc_type": "lu",
            },
        )
    except TypeError:
        # Version ancienne de dolfinx sans petsc_options_prefix
        problem = LinearProblem(a, L, bcs=bcs)
    uh = problem.solve()
    
    computation_time = time.time() - start_time
    
    # 6. Extraction des données de la solution
    solution_values = uh.x.array.copy()
    
    # 7. Création des frames d'animation (évolution progressive de la solution)
    animation_frames = create_animation_frames(solution_values, num_frames=30)
    
    # 8. Préparation des résultats
    result = {
        'mesh_data': mesh_data,
        'solution_data': {
            'values': solution_values.tolist(),
            'min': float(solution_values.min()),
            'max': float(solution_values.max()),
        },
        'animation_frames': animation_frames,
        'statistics': {
            'min': float(solution_values.min()),
            'max': float(solution_values.max()),
            'mean': float(solution_values.mean()),
            'computation_time': computation_time,
            'num_dofs': len(solution_values),
            'num_elements': nx * ny * 2,  # triangles
        }
    }
    
    return result


def extract_mesh_data(msh, V) -> Dict[str, Any]:
    """Extrait les données du maillage pour la visualisation 3D"""
    
    # Coordonnées des noeuds
    geometry = msh.geometry.x
    coordinates = geometry[:, :2].tolist()  # x, y seulement
    
    # Connectivité des éléments (triangles)
    topology = msh.topology
    topology.create_connectivity(2, 0)  # faces vers sommets
    
    cells = msh.geometry.dofmap
    triangles = []
    for i in range(cells.shape[0]):
        triangles.append(cells[i].tolist())
    
    # Calcul des arêtes pour le wireframe
    edges = set()
    for tri in triangles:
        for i in range(3):
            edge = tuple(sorted([tri[i], tri[(i+1) % 3]]))
            edges.add(edge)
    
    return {
        'vertices': coordinates,
        'triangles': triangles,
        'edges': list(edges),
        'num_vertices': len(coordinates),
        'num_triangles': len(triangles),
        'bounds': {
            'x_min': float(geometry[:, 0].min()),
            'x_max': float(geometry[:, 0].max()),
            'y_min': float(geometry[:, 1].min()),
            'y_max': float(geometry[:, 1].max()),
        }
    }


def create_source_term(x, params: Dict[str, Any]):
    """Crée le terme source selon le type choisi"""
    import ufl
    
    source_type = params['type']
    amplitude = params['amplitude']
    
    if source_type == 'gaussian':
        cx, cy = params['center_x'], params['center_y']
        width = params['width']
        return amplitude * ufl.exp(-((x[0] - cx)**2 + (x[1] - cy)**2) / width)
    
    elif source_type == 'constant':
        return ufl.Constant(x.ufl_domain(), amplitude)
    
    elif source_type == 'sinusoidal':
        freq = params['frequency']
        return amplitude * ufl.sin(freq * ufl.pi * x[0]) * ufl.sin(freq * ufl.pi * x[1])
    
    else:  # custom ou default
        return amplitude * ufl.exp(-((x[0] - 0.5)**2 + (x[1] - 0.5)**2) / 0.02)


def create_neumann_term(x, params: Dict[str, Any]):
    """Crée le terme de Neumann selon le type choisi"""
    import ufl
    
    neumann_type = params['type']
    amplitude = params['amplitude']
    frequency = params['frequency']
    
    if neumann_type == 'constant':
        return ufl.Constant(x.ufl_domain(), amplitude)
    
    elif neumann_type == 'sinusoidal':
        return amplitude * ufl.sin(frequency * x[0])
    
    elif neumann_type == 'linear':
        return amplitude * x[0]
    
    else:
        return amplitude * ufl.sin(frequency * x[0])


def create_animation_frames(solution: np.ndarray, num_frames: int = 30) -> List[List[float]]:
    """
    Crée des frames d'animation montrant l'évolution progressive de la solution
    (simulation d'un processus de diffusion jusqu'à l'état stationnaire)
    """
    frames = []
    
    for i in range(num_frames):
        # Facteur de progression (0 à 1)
        t = i / (num_frames - 1)
        
        # Interpolation de 0 vers la solution finale avec une courbe ease-in-out
        ease_t = 0.5 * (1 - np.cos(np.pi * t))
        
        frame_values = (solution * ease_t).tolist()
        frames.append(frame_values)
    
    return frames


def generate_preview_mesh(params: Dict[str, Any]) -> Dict[str, Any]:
    """
    Génère uniquement les données du maillage pour un aperçu rapide
    (sans résoudre l'équation)
    """
    import dolfinx as df
    from mpi4py import MPI
    
    domain = params['domain']
    mesh_params = params['mesh']
    
    nx, ny = mesh_params['nx'], mesh_params['ny']
    msh = df.mesh.create_rectangle(
        comm=MPI.COMM_WORLD,
        points=[[domain['x_min'], domain['y_min']], 
                [domain['x_max'], domain['y_max']]],
        n=[nx, ny],
        cell_type=df.mesh.CellType.triangle,
    )
    
    V = df.fem.functionspace(msh, ("Lagrange", 1))
    
    return extract_mesh_data(msh, V)