#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# pylint: disable=C0103, W1401, C0302
r"""
Created on Thu Mar 7 11:17:17 2024
@author: nchevaug
2D Explicit Lagrangian Finite Volume Solver (GLACE Scheme).
Philosophy & Architecture:
This module implements the Godunov-type Lagrangian Scheme for Compressible
Euler equations (GLACE). It computes the evolution of a fluid on an unstructured
2D triangular mesh moving with the fluid velocity (Lagrangian framework).
The solver features a dual-execution architecture:
1. Pure NumPy (Fallback): Fully vectorized array operations.
2. Numba JIT (HPC Mode): Aggressively optimized loop-level parallelism
(`prange`), avoiding heap allocations and leveraging hardware registers
for the nodal Riemann solver.
"""
from abc import abstractmethod
from typing import Tuple, NamedTuple, Sequence, List, Union, assert_never, Optional, overload
#from typing import Any, Callable, Dict
#from collections import defaultdict
from llvmlite.ir import FastMathFlags
import numpy as np
import numba
#from numba import guvectorize, float32, float64, prange
from pyxfem.core.config import nbfloat
from pyxfem.core.nbkalgebra2d import Vec2, Sy22
import numba
import numba.core.types as nbct
import numba.typed as nbty
# pylint: disable=ungrouped-imports
from pyxfem.core.nbkalgebra2d import \
vec2_t, sy22_t, vec2_zero, sy22_zero,\
set_vec2,\
vec2_scal, vec2_square, vec2_norm, vec2_dot, vec2_add, vec2_add3, vec2_sub, \
sy22_add, sy22_mv, sy22_smv, sy22_r, sy22_r2, sy22_trace
from pyxfem.core.graph import T2V, extract_txy, upward_adjacencies, which_in_upward
from pyxfem.core.config import USE_NUMBA, FASTMATH, IndexType, npfloat, fzero, fone, fhalf,\
FloatArray, IndexArray
from pyxfem.core.smallsize_tensor import max_eigval22sym, gemm22, inv, det, sysmv22, sysmv22_guvec, vec_norm2
import pyxfem.core.graph as graph
[docs]
def t_tangentmap(tris_coord, out: Optional[FloatArray] = None):
''' return the "tangent map"
The F tensor between reference element and each of the elment in tris_coord
'''
shape = tris_coord.shape[:-2] +(2,2)
out = np.empty(shape, npfloat) if out is None else out
v0 = tris_coord[..., 0, :]
out[... , 0] = tris_coord[..., 1, :] - v0
out[..., 1] = tris_coord[..., 2, :] - v0
#outold = tris_coord[...,1:,:] -tris_coord[..., 0, np.newaxis, :]
#print(np.linalg.norm(out - outold.swapaxes(-1,-2)))
return out
return np.subtract(tris_coord[...,1:,:], tris_coord[..., 0, np.newaxis, :], out = out)
[docs]
class GELEES_MeshStruct(NamedTuple):
"""
A structured container for the 2D Lagrangian mesh topology, geometry,
and inverse connectivities.
This NamedTuple holds all the static and dynamically updated arrays
required for the finite volume solver, designed to be easily passed
into Numba JIT-compiled functions.
Attributes
----------
ncells : int
The total number of triangular cells in the mesh.
nverts : int
The total number of vertices (nodes) in the mesh.
t2v : IndexArray
2D array of shape `(ncells, 3)`. The cell-to-vertex connectivity matrix
mapping each triangle to its 3 defining vertex indices.
v2t_index : IndexArray
1D array of shape `(nverts + 1,)`. The index pointer array for the CSR-like
inverse connectivity (vertices-to-cells).
v2t_tris : IndexArray
1D array containing the flattened cell indices surrounding each vertex.
Sliced using `v2t_index`.
v2t_v : IndexArray
1D array containing the local vertex index (0, 1, or 2) within the
corresponding cell in `v2t_tris`. Sliced using `v2t_index`.
xy : FloatArray
2D array of shape `(nverts, 2)`. The physical (x, y) coordinates of
each vertex. This array is updated at each time step.
S : FloatArray
1D array of shape `(ncells,)`. The surface area of each triangular cell.
C : FloatArray
3D array of shape `(ncells, 3, 2)`. The corner normal vectors (C_jr)
for each vertex of each cell.
nC : FloatArray
2D array of shape `(ncells, 3)`. The Euclidean norms (magnitudes) of
the corner normal vectors.
N : FloatArray
3D array of shape `(ncells, 3, 2)`. The normalized unit direction
vectors of the corner normals.
"""
ncells: int
nverts: int
t2v: T2V
v2t_index: IndexArray
v2t_tris: IndexArray
v2t_v: IndexArray
xy: FloatArray
cell_meas: FloatArray
C: FloatArray
nC: FloatArray
N: FloatArray
[docs]
def allocate_Cjr(nc :int):
r"""
Allocates memory for the geometric properties of the mesh cells.
This function initializes empty NumPy arrays required to store the
surface areas and the corner normal vectors (and their properties)
for a 2D triangular mesh.
Parameters
----------
nc : int
The total number of triangular cells in the mesh.
Returns
-------
S : ndarray
1D array of shape `(nc,)` and type `float64`.
Allocated to store the surface area of each triangular cell.
C : ndarray
3D array of shape `(nc, 3, 2)` and type `float64`.
Allocated to store the corner normal vectors (C_jr) for the 3 vertices
of each cell in 2D space (x, y).
normC : ndarray
2D array of shape `(nc, 3)` and type `float64`.
Allocated to store the Euclidean norms (magnitudes) of the C_jr vectors.
"""
cell_measure = np.empty( (nc), dtype = npfloat)
C = np.empty( (nc, 3, 2), dtype = npfloat)
normC = np.empty( (nc, 3), dtype = npfloat)
N = np.empty( (nc, 3, 2), dtype = npfloat)
return cell_measure, C, normC, N
[docs]
def setup_MeshStruct(xy, tris):
v2t = upward_adjacencies(tris)
v2t_v = which_in_upward( v2t, tris)
ncells = tris.shape[0]
nverts = xy.shape[0]
cell_meas, C, normC, N = allocate_Cjr(ncells)
mesh = GELEES_MeshStruct(
ncells =ncells, nverts=nverts,
t2v=tris, v2t_index=v2t.index,
v2t_tris= v2t.data, v2t_v=v2t_v,
xy=xy,
cell_meas=cell_meas,
C=C, nC=normC, N=N)
if USE_NUMBA:
computeCjr_triangles_numba(mesh)
else:
computeCjr_triangles(mesh)
return mesh
[docs]
def update_MeshStruct(xy, mesh : GELEES_MeshStruct):
mesh.xy[:] = xy
computeCjr_triangles_numba(mesh)
[docs]
def computeCjr_triangles(mesh : GELEES_MeshStruct):
xy, tris = mesh.xy, mesh.t2v
txy = extract_txy(xy, tris)
F1 = t_tangentmap(txy)
computeCjr_triangles_(F1, mesh.cell_meas, mesh.C, mesh.nC, mesh.N)
[docs]
def computeCjr_triangles_(F, outMeas, outC, outnormC, outN):
det(F, outMeas)
outMeas*=0.5
F11 = F[...,0,0]
F12 = F[...,0,1]
F21 = F[...,1,0]
F22 = F[...,1,1]
#outC[:] = fhalf*np.stack([-F22+F12, F21-F11, F22, -F21, -F12, F11], axis=-1).reshape((-1, 3, 2))
np.subtract(F12, F22, outC[..., 0,0])
np.subtract(F21, F11, outC[..., 0,1])
np.copyto(outC[...,1,0], F22)
np.multiply(-fone, F21, outC[...,1,1])
np.multiply(-fone, F12, outC[...,2,0])
np.copyto(outC[...,2,1], F11)
np.multiply(fhalf, outC, outC)
vec_norm2(outC, outnormC)
np.divide(outC, outnormC[..., None], out = outN)
[docs]
@numba.njit(inline='always', fastmath = FASTMATH)
def get_vec2(index, array):
return (array[index, 0], array[index, 1])
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def get_trixy(elemid, mesh):
tri = mesh.t2v[elemid]
n0, n1, n2 = tri[0], tri[1], tri[2]
xy = mesh.xy
xy0 = get_vec2(n0, xy)
xy1 = get_vec2(n1, xy)
xy2 = get_vec2(n2, xy)
return xy0, xy1, xy2
[docs]
@numba.njit(cache = True, fastmath = FASTMATH)
def Cjr_triangle(xy0: Vec2, xy1: Vec2, xy2: Vec2):
"""
Computes the surface area and corner normal vectors (C_jr) for a single
2D triangle.
This function is optimized for Numba JIT compilation. It performs
computations using primitive types and returns tuples to ensure
zero-allocation overhead when called inside high-performance loops
(like `prange`).
Parameters
----------
xy0 : tuple of float or array-like
The (x, y) coordinates of the 0-th vertex of the triangle.
xy1 : tuple of float or array-like
The (x, y) coordinates of the 1-st vertex of the triangle.
xy2 : tuple of float or array-like
The (x, y) coordinates of the 2computeCjr_triangles_numba(mesh)
def computeCjr_triangles(F, outMeas, outC, outnormC, outN):
-nd vertex of the triangle.
Returns
-------
S : float
The surface area of the triangle.
Cj0 : tuple of float
A tuple `(x, y)` representing the corner normal vector
associated with vertex 0.
Cj1 : tuple of float
A tuple `(x, y)` representing the corner normal vector
associated with vertex 1.
Cj2 : tuple of float
A tuple `(x, y)` representing the corner normal vector
associated with vertex 2.
"""
F0 = (xy1[0] - xy0[0] , xy1[1] - xy0[1])
F1 = (xy2[0] - xy0[0] , xy2[1] - xy0[1])
S = npfloat(0.5)*F0[0]*F1[1] - npfloat(0.5)*F1[0]*F0[1]
Cj0 = (npfloat(0.5)*(F0[1]-F1[1]), npfloat(0.5)*(F1[0]-F0[0]))
Cj1 = (npfloat(0.5)*F1[1], -npfloat(0.5)*F1[0])
Cj2 = (-npfloat(0.5)*F0[1], npfloat(0.5)*F0[0])
return S, Cj0, Cj1, Cj2
[docs]
@numba.njit(parallel=True, cache = True, fastmath = FASTMATH)
def computeCjr_triangles_numba(mesh):
"""
Computes and updates the geometric properties of all triangular cells
in the mesh using Numba parallelization.
This function iterates over all cells concurrently, calculates the
surface area and the corner normal vectors (C_jr) for each triangle
based on its current vertex coordincomputeCjr_triangles_numba(mesh)
Parameters
----------
mesh : NamedTuple (e.g., GELEES_MeshStruct)
A structured object containing the mesh data. It must provide:
- `mesh.ncells` (int): The total number of cells.
- `mesh.S` (ndarray): 1D array of shape `(ncells,)` to store updated areas.
- `mesh.C` (ndarray): 3D array of shape `(ncells, 3, 2)` to store
updated corner normal vectors.
- Connectivity and coordinate data required by `get_trixy`.
Returns
-------
None
The function does not return any value. The `mesh.S` and `mesh.C`
arrays are modified in-place for maximum performance.
"""
N = mesh.ncells
outS = mesh.cell_meas
outC = mesh.C
outnC = mesh.nC
outN = mesh.N
for j in numba.prange(N): # type ignore # pylint: disable=E1133
xy0, xy1, xy2 = get_trixy(j, mesh)
Sj, Cj0, Cj1, Cj2 = Cjr_triangle( xy0, xy1, xy2)
nCj0, nCj1, nCj2 = \
vec2_norm(Cj0), vec2_norm(Cj1), vec2_norm(Cj2)
outS[j] = Sj
outnC[j, 0] = nCj0
outnC[j, 1] = nCj1
outnC[j, 2] = nCj2
for k in range(2):
outC[j, 0, k] = Cj0[k]
outC[j, 1, k] = Cj1[k]
outC[j, 2, k] = Cj2[k]
outN[j, 0, k] = Cj0[k]/nCj0
outN[j, 1, k] = Cj1[k]/nCj1
outN[j, 2, k] = Cj2[k]/nCj2
[docs]
class ConservativeLawManager():
def __init__(self, nbdata, U_shapes, U_indexes, W_shapes, W_indexes):
self.U_indexes = U_indexes
self.U_shapes = U_shapes
sizeU = U_shapes["size"]
self.W_indexes = W_indexes
self.W_shapes = W_shapes
sizeW = W_shapes["size"]
self.U = np.zeros((nbdata, sizeU), dtype=npfloat)
self.W = np.zeros((nbdata, sizeW), dtype=npfloat)
[docs]
@abstractmethod
def conservative_to_primal(self, cell_measure):
pass
[docs]
@abstractmethod
def acoustic_impedance_tensor(self, n, outZ = None):
pass
@overload
def get_primal_variable(self, names: str) -> FloatArray: ...
@overload
def get_primal_variable(self, names: Sequence[str]) -> List[FloatArray]: ...
[docs]
def get_primal_variable(self, names: Union[str, Sequence[str]]) -> FloatArray | List[FloatArray]:
""" return primal variable corresponding to it's name or a list of names; avalable variables are listed in W_indices """
def get(name:str):
s, e = self.W_indexes[name]
return self.W[:, s:e].reshape((-1,)+self.W_shapes[name])
match names :
case str() as name:
return get(name)
case [*sequ_names] :
res = []
for name in sequ_names:
res+=[get(name)]
return res
case _ as unreachable_var:
assert_never(unreachable_var)
@overload
def get_conservative_variable(self, names: str) -> FloatArray: ...
@overload
def get_conservative_variable(self, names: Sequence[str]) -> List[FloatArray]: ...
[docs]
def get_conservative_variable(self, names: Union[str, Sequence[str]]) -> FloatArray | List[FloatArray]:
""" return primal variable corresponding to it's name or a list of names; avalable variables are listed in W_indices """
def get(name:str):
s, e = self.U_indexes[name]
return self.U[:, s:e].reshape((-1,)+self.U_shapes[name])
match names :
case str() as name:
return get(name)
case [*sequ_names] :
res = []
for name in sequ_names:
res+=[get(name)]
return res
case _ as unreachable_var:
assert_never(unreachable_var)
[docs]
class GlaceExplicitLagrangianBaseSolver:
def __init__(self, conservation_law : ConservativeLawManager, mesh: GELEES_MeshStruct):
self.conservation_law = conservation_law
self.mesh = mesh
nverts = mesh.nverts
ncells = mesh.ncells
self.vA = np.zeros((nverts, 2, 2), dtype=npfloat)
self.vb = np.zeros((nverts, 2), dtype=npfloat)
self.vu = np.zeros((nverts, 2), dtype=npfloat)
self.cvp = np.zeros((ncells, 3), dtype=npfloat)
self.cvqflux = np.zeros((ncells, 3, 2), dtype=npfloat)
self.cvA = np.zeros((ncells, 3, 2, 2), dtype=npfloat)
self.cvb = np.zeros((ncells, 3, 2), dtype=npfloat)
self.cvu = np.zeros((ncells, 3, 2), dtype=npfloat)
self.qflux = np.zeros((ncells, 2), dtype=npfloat)
self.eflux = np.zeros((ncells), dtype=npfloat)
def _riemann_solver_assemble(self)-> None:
self.vA.fill(fzero)
self.vb.fill(fzero)
# order = self.mesh.t2v.ravel()
# np.add.at(self.vA, order, self.cvA.reshape((-1,2,2)))
# np.add.at(self.vb, order, self.cvb.reshape((-1,2)))
GlaceExplicitLagrangianBaseSolver._riemann_solver_assemble_numba(self.cvA, self.cvb, self.mesh, self.vA, self.vb)
self._set_boundary_conditions()
@staticmethod
@numba.njit(parallel=True, cache=True, fastmath = FASTMATH)
def _riemann_solver_assemble_numba(cvA, cvb, mesh, out_vA, out_vb)-> None:
for vid in numba.prange(mesh.nverts): # pylint: disable=E1133
s = mesh.v2t_index[vid]
e = mesh.v2t_index[vid+1]
Ar : Sy22= (out_vA[vid,0,0] ,out_vA[vid,0,1], out_vA[vid,1,1])
br : Vec2 = (out_vb[vid,0], out_vb[vid,1])
# Cell k connectd to node vid
for k in range(s, e):
tj = mesh.v2t_tris[k]
vr = mesh.v2t_v[k]
Ar = sy22_add(Ar, (cvA[tj, vr, 0,0], cvA[tj, vr, 0,1], cvA[tj, vr, 1,1]))
br = vec2_add(br, (cvb[tj, vr, 0], cvb[tj, vr, 1]))
out_vA[vid, 0, 0] = Ar[0]
out_vA[vid, 0, 1] = Ar[1]
#out_vA[vid, 1, 0] = Ar[1]
out_vA[vid, 1, 1] = Ar[2]
set_vec2(vid, out_vb, br)
def _riemann_solver(self)-> Tuple[FloatArray, FloatArray]:
#self._riemann_solver_setup_0()
#self._riemann_solver_setup_1()
self._riemann_solver_setup()
self._riemann_solver_assemble()
sysmv22(self.vA, self.vb, self.vu, parallel =True)
graph.extract_txy(self.vu, self.mesh.t2v, self.cvu)
self._riemann_solver_cvqflux()
return self.vu, self.cvqflux
def _riemann_solver_cvqflux(self):
if USE_NUMBA:
GlaceExplicitLagrangianBaseSolver._riemann_solver_cvqflux_njit(self.cvA, self.cvu, self.cvb, self.cvqflux)
else:
np.matmul(self.cvA, self.cvu[..., np.newaxis], out=self.cvqflux[..., np.newaxis])
np.subtract(self.cvb, self.cvqflux, out=self.cvqflux)
@staticmethod
@numba.njit(parallel=True, cache = True)
def _riemann_solver_cvqflux_njit(cvA, cvu, cvb, out_cvqflux):
for j in numba.prange(cvA.shape[0]): # pylint: disable=E1133
for r in numba.prange(cvA.shape[1]): # pylint: disable=E1133
Ajr : Sy22 = (cvA[j,r, 0, 0 ] , cvA[j,r, 0, 1 ], cvA[j,r, 1, 1 ])
ujr : Vec2 = (cvu[j,r, 0], cvu[j,r,1])
bjr : Vec2 = (cvb[j,r, 0], cvb[j,r,1])
tmpjr = sy22_mv(Ajr, ujr)
qjr = vec2_sub(bjr, tmpjr)
out_cvqflux[j,r, : ] = qjr
[docs]
@staticmethod
@numba.njit(parallel=True, cache=True)
def critical_timestep_pnjit(cell_measure, nC, c):
min_dt = np.inf # numba does autoamatique reduction on this variable
for j in numba.prange(len(cell_measure)): # pylint: disable=E1133
cell_measurej = cell_measure[j]
nCj0 = nC[j,0]
nCj1 = nC[j,1]
nCj2 = nC[j,2]
cj = c[j]
sum_normCj = nCj0 + nCj1 + nCj2
dt = cell_measurej / sum_normCj/ cj
# # Numba comprend automatiquement qu'il s'agit d'une réduction parallèle
min_dt = min(min_dt, dt)
return min_dt
@overload
def get_primal_variable(self, names: str) -> FloatArray: ...
@overload
def get_primal_variable(self, names: Sequence[str]) -> List[FloatArray]: ...
[docs]
def get_primal_variable(self, names: Union[str, Sequence[str]]) -> FloatArray | List[FloatArray]:
""" return primal variable corresponding to it's name or a list of names; avalable variables are listed in W_indices """
return self.conservation_law.get_primal_variable(names)
@overload
def get_conservative_variable(self, names: str) -> FloatArray: ...
@overload
def get_conservative_variable(self, names: Sequence[str]) -> List[FloatArray]: ...
[docs]
def get_conservative_variable(self, names: Union[str, Sequence[str]]) -> FloatArray | List[FloatArray]:
return self.conservation_law.get_conservative_variable(names)