#!/usr/bin/env python3
# -*- coding: utf-8 -*
# pylint: disable=invalid-name
r"""
Created march 2026
@author: nchevaug
2D Algebraic Kernel for High-Performance Numba JIT Compilation.
Philosophy & Architecture:
This module provides fundamental 2D vector and symmetric tensor operations
specifically tailored for Numba's `nopython` mode (`@njit`).
To achieve maximum performance and avoid the massive overhead of heap
allocations inside hot loops (e.g., cell or vertex iterations), this module
STRICTLY avoids NumPy arrays for small local variables. Instead, it relies
entirely on Python `Tuple` of floats.
Why Tuples?
When compiled with Numba and LLVM, homogenous tuples of primitive types
(like `Tuple[float, float]`) are not treated as standard Python objects.
They are lowered directly into C-style structs passed by value, which the
LLVM backend maps directly into the processor's hardware registers.
This guarantees zero-allocation overhead, completely bypasses the Python
garbage collector, and allows aggressive compiler optimizations like loop
unrolling and SIMD vectorization (enhanced by `fastmath=True`).
Storage Convention:
- Vectors (Vec2) are stored as 2-tuples: `(x, y)`
- Symmetric 2x2 Matrices (Sy22) are stored as 3-tuples: `(A00, A01, A11)`.
This exploits the symmetry ($A_{01} = A_{10}$) to save register space
and reduce floating-point operations.
"""
from typing import Tuple
from pyxfem.core.config import USE_NUMBA
# ==========================================
# STATIC ANALYSIS TYPE ALIASES (Type Hinting)
# ==========================================
# These types guide the IDE (VS Code, Pylance, MyPy) for autocompletion
# and static type checking. They are ignored by Numba at runtime.
# Vec2 = Tuple[float, float]
# Sy22 = Tuple[float, float, float]
#
if USE_NUMBA:
from math import sqrt
import numba
import numba.core.types as nct
# pylint: disable=ungrouped-imports
from pyxfem.core.config import FASTMATH, nbfloat, fzero,npfloat
# ==========================================
# EXPLICIT NUMBA SIGNATURE TYPES
# ==========================================
sy22_t = nct.UniTuple(nbfloat, 3)
vec2_t = nct.UniTuple(nbfloat, 2)
# Vec2 = vec2_t
# Sy22 = sy22_t
# Vec2 = Tuple[nbfloat, nbfloat]
# Sy22 = Tuple[nbfloat, nbfloat, nbfloat]
Vec2 = Tuple[npfloat, npfloat]
Sy22 = Tuple[npfloat, npfloat, npfloat]
Ge22 = Tuple[npfloat, npfloat, npfloat, npfloat]
# ==========================================
# constants
# ==========================================
vec2_zero: Vec2 = (fzero, fzero)
sy22_zero: Sy22 = (fzero, fzero, fzero)
# ==========================================
# 2D VECTOR OPERATIONS
# ==========================================
[docs]
@numba.njit(inline='always', fastmath = FASTMATH)
def set_vec2(index, array, v):
array[index, 0], array[index, 1] = v[0], v[1]
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def vec2_add(a: Vec2, b: Vec2) -> Vec2:
r"""
Adds two 2D vectors.
Math: $a + b$
"""
return (a[0] + b[0], a[1] + b[1])
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def vec2_add3(a: Vec2, b: Vec2, c: Vec2) -> Vec2:
r"""
Adds three 2D vectors.
Math: $a + b + c$
"""
return (a[0] + b[0]+c[0], a[1] + b[1]+c[1])
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def vec2_scal(alpha: npfloat, x: Vec2) -> Vec2:
r"""
Multiplies a 2D vector by a scalar.
Math: $\alpha x$
"""
return (alpha*x[0], alpha*x[1])
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def vec2_sub(a: Vec2, b: Vec2) -> Vec2:
r"""
Subtracts one 2D vector from another.
Math: $a - b$
"""
return (a[0] - b[0], a[1] - b[1])
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def vec2_norm(a: Vec2) -> npfloat:
r"""
Computes the Euclidean (L2) norm of a 2D vector.
Math: $\|a\| = \sqrt{a_0^2 + a_1^2}$
"""
return npfloat(sqrt((a[0]**2 +a[1]**2)))
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def vec2_dot(a: Vec2, b: Vec2) -> npfloat:
r"""
Computes the dot product of two 2D vectors.
Math: $a \cdot b = a_0 b_0 + a_1 b_1$
"""
return a[0]*b[0] +a[1]*b[1]
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def vec2_square(a: Vec2) -> npfloat:
r"""
Computes the squared norm of a 2D vector.
Math: $\|a\|^2 = a_0^2 + a_1^2$
"""
return a[0]**2 +a[1]**2
# ==========================================
# 2X2 SYMMETRIC TENSOR OPERATIONS
# ==========================================
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def sy22_mv(A: Sy22, x: Vec2) -> Vec2:
r"""
Performs a matrix-vector multiplication for a symmetric 2x2 matrix.
Math: $A x$
Note: The matrix is implicitly reconstructed from its components $(A_{00}, A_{01}, A_{11})$.
"""
A00, A01, A11 = A[0], A[1], A[2]
x0, x1 = x[0], x[1]
return (A00*x0 + A01*x1, A01*x0+ A11*x1)
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def sy22_add(A: Sy22, B: Sy22) -> Sy22:
r"""
Adds two symmetric 2x2 matrices.
Math: $A + B$
"""
return (A[0]+B[0], A[1]+B[1], A[2]+B[2])
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def sy22_r(alpha: npfloat, x: Vec2) -> Sy22:
r"""
Computes the symmetric rank-1 update (dyadic product).
Math: $\alpha (x \otimes x)$
"""
A00 = alpha * x[0] * x[0]
A01 = alpha * x[0] * x[1]
A11 = alpha * x[1] * x[1]
return (A00, A01, A11)
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def sy22_r2(alpha: npfloat, x: Vec2, y: Vec2) -> Sy22:
r"""
Computes the scalar-penalized symmetric rank-2 update.
Math: $\alpha (x \otimes y + y \otimes x)$
"""
A00 = alpha*2*x[0] * y[0]
A01 = alpha*(x[0] * y[1] + y[0] * x[1])
A11 = alpha*2*(x[1] * y[1])
return (A00, A01, A11)
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def sy22_trace(A: Sy22) -> npfloat:
r"""
Computes the trace of a symmetric 2x2 matrix.
Math: $\text{Tr}(A) = A_{00} + A_{11}$
"""
return A[0] + A[2]
[docs]
@numba.njit(cache=True, fastmath = FASTMATH)
def sy22_smv(A: Sy22, b: Vec2) -> Vec2:
r"""
Solves the linear system $A x = b$ for a symmetric 2x2 matrix
using explicit Cramer's rule inversion.
Math: $x = A^{-1} b$
"""
A00, A01, A11 = A[0], A[1], A[2]
b0, b1 = b[0], b[1]
detA = A00*A11 - A01*A01
x0 = (A11*b0 - A01*b1)/detA
x1 = (A00*b1 - A01*b0)/detA
return (x0, x1)
# # The two following function where not used yet :need to b checked
# # for efficiency
# @numba.njit(cache=True)
# def gemm22_numba(alpha, A, B, beta, C):
# C[0,0] = beta*C[0,0] + alpha*(A[0,0]*B[0,0] + A[0,1]*B[1,0])
# C[0,1] = beta*C[0,1] + alpha*(A[0,0]*B[0,1] + A[0,1]*B[1,1])
# C[1,0] = beta*C[1,0] + alpha*(A[1,0]*B[0,0] + A[1,1]*B[1,0])
# C[1,1] = beta*C[1,1] + alpha*(A[1,0]*B[0,1] + A[1,1]*B[1,1])
# @numba.njit(cache=True)
# def gemv22_numba(A,x,y, beta =fzero):
# x0, x1= x[0], x[1]
# y[0] = beta*y[0] + A[0,0]*x0 + A[0,1]*x1
# y[1] = beta*y[1] + A[1,0]*x0 + A[1,1]*x1