adtlert 0.1.0__py3-none-any.whl

adtlert/__init__.py

@@ -0,0 +1,43 @@
+"""Top-level package for differentiable time-lapse ERT on Torch."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from adtlert.fem import P1ElementData, build_p1_element_data, p1_shape_functions
+from adtlert.forward import (
+    ERT2p5DApparentResistivityFunction,
+    ERTForward2p5D,
+    ERTForwardModeling,
+    ForwardResponse,
+    apparent_resistivity_autograd,
+)
+from adtlert.inversion import (
+    ERTInversion,
+    InversionConfig,
+    TimeLapseERTInversion,
+    WindowedTimeLapseERTInversion,
+)
+from adtlert.mesh import Mesh
+from adtlert.survey import Survey
+
+try:
+    __version__ = version("adtlert")
+except PackageNotFoundError:  # pragma: no cover - source tree without installation
+    __version__ = "0+unknown"
+
+__all__ = [
+    "ERTForward2p5D",
+    "ERT2p5DApparentResistivityFunction",
+    "ERTForwardModeling",
+    "ERTInversion",
+    "ForwardResponse",
+    "InversionConfig",
+    "Mesh",
+    "P1ElementData",
+    "Survey",
+    "TimeLapseERTInversion",
+    "WindowedTimeLapseERTInversion",
+    "apparent_resistivity_autograd",
+    "build_p1_element_data",
+    "p1_shape_functions",
+    "__version__",
+]

adtlert/fem/__init__.py

@@ -0,0 +1,72 @@
+"""Finite-element basis and quadrature utilities."""
+
+from adtlert.fem.assembly import (
+    COORouting,
+    assemble_global_bcoo,
+    assemble_helmholtz_operator,
+    assemble_local_mass,
+    assemble_local_mass_p2,
+    assemble_local_stiffness,
+    assemble_local_stiffness_p2,
+    build_coo_routing,
+    build_coo_routing_from_connectivity,
+)
+from adtlert.fem.boundary import (
+    BoundaryRouting,
+    assemble_boundary_bcoo,
+    assemble_local_boundary_mass,
+    assemble_local_boundary_mass_p2,
+    build_boundary_routing,
+    build_boundary_routing_from_connectivity,
+    robin_boundary_coefficients,
+)
+from adtlert.fem.p1 import (
+    P1ElementData,
+    build_p1_element_data,
+    p1_shape_functions,
+    reference_shape_gradients,
+    triangle_quadrature,
+)
+from adtlert.fem.p2 import P2ElementData, build_p2_element_data, p2_shape_functions, reference_shape_gradients_p2
+from adtlert.fem.tet_p1 import (
+    TetrahedronP1Data,
+    TetrahedronP2Data,
+    build_tetrahedron_p1_data,
+    build_tetrahedron_p2_data,
+    p2_tetrahedron_shape_values,
+    p2_triangle_mass_template,
+)
+
+__all__ = [
+    "BoundaryRouting",
+    "COORouting",
+    "assemble_boundary_bcoo",
+    "P1ElementData",
+    "assemble_global_bcoo",
+    "assemble_helmholtz_operator",
+    "assemble_local_boundary_mass",
+    "assemble_local_boundary_mass_p2",
+    "assemble_local_mass",
+    "assemble_local_mass_p2",
+    "assemble_local_stiffness",
+    "assemble_local_stiffness_p2",
+    "build_boundary_routing",
+    "build_boundary_routing_from_connectivity",
+    "build_p1_element_data",
+    "build_p2_element_data",
+    "build_coo_routing",
+    "build_coo_routing_from_connectivity",
+    "p1_shape_functions",
+    "p2_shape_functions",
+    "reference_shape_gradients",
+    "reference_shape_gradients_p2",
+    "robin_boundary_coefficients",
+    "triangle_quadrature",
+    "P2ElementData",
+    "TetrahedronP1Data",
+    "TetrahedronP2Data",
+    "build_tetrahedron_p1_data",
+    "build_tetrahedron_p2_data",
+    "p2_tetrahedron_shape_values",
+    "p2_triangle_mass_template",
+]

adtlert/fem/assembly.py

@@ -0,0 +1,220 @@
+"""Tensorized assembly utilities for finite-element sparse operators."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from adtlert.utils.torch_runtime import Array
+from adtlert.utils.torch_runtime import torch_np
+from adtlert.utils.torch_runtime import BCOO
+import numpy as np
+
+from adtlert.fem.p1 import P1ElementData
+from adtlert.fem.p2 import P2ElementData
+from adtlert.mesh import Mesh
+from adtlert.utils.dtypes import FLOAT_DTYPE, INT_DTYPE, NP_FLOAT_DTYPE
+
+
+@dataclass(frozen=True)
+class COORouting:
+    """Reusable COO routing indices for sparse element assembly."""
+
+    indices: Array
+    shape: tuple[int, int]
+
+
+def _expand_coefficient(coefficients: Array | float, element_data: P1ElementData) -> Array:
+    coefficient_array = torch_np.asarray(coefficients, dtype=FLOAT_DTYPE)
+    cell_count = element_data.cell_areas.shape[0]
+    quadrature_count = element_data.quadrature_weights.shape[0]
+
+    if coefficient_array.ndim == 0:
+        return torch_np.broadcast_to(coefficient_array, (cell_count, quadrature_count))
+    if coefficient_array.shape == (cell_count,):
+        return torch_np.broadcast_to(coefficient_array[:, None], (cell_count, quadrature_count))
+    if coefficient_array.shape == (cell_count, quadrature_count):
+        return coefficient_array
+
+    raise ValueError(
+        "coefficients must be scalar, shape (num_cells,), or shape (num_cells, num_quadrature)"
+    )
+
+
+def _static_scalar_value(value: Array | float) -> float | None:
+    try:
+        value_array = np.asarray(value)
+    except Exception:
+        return None
+    if value_array.ndim != 0:
+        return None
+    return float(value_array)
+
+
+def _expand_coefficient_generic(
+    coefficients: Array | float,
+    *,
+    cell_count: int,
+    quadrature_count: int,
+) -> Array:
+    coefficient_array = torch_np.asarray(coefficients, dtype=FLOAT_DTYPE)
+
+    if coefficient_array.ndim == 0:
+        return torch_np.broadcast_to(coefficient_array, (cell_count, quadrature_count))
+    if coefficient_array.shape == (cell_count,):
+        return torch_np.broadcast_to(coefficient_array[:, None], (cell_count, quadrature_count))
+    if coefficient_array.shape == (cell_count, quadrature_count):
+        return coefficient_array
+
+    raise ValueError(
+        "coefficients must be scalar, shape (num_cells,), or shape (num_cells, num_quadrature)"
+    )
+
+
+def build_coo_routing(mesh: Mesh) -> COORouting:
+    """Precompute the COO routing pattern for cellwise assembly."""
+
+    cells = np.asarray(mesh.cells, dtype=np.int32)
+    local_dof = int(cells.shape[1])
+    row_indices = np.broadcast_to(cells[:, :, None], (mesh.cell_count, local_dof, local_dof)).reshape(-1)
+    col_indices = np.broadcast_to(cells[:, None, :], (mesh.cell_count, local_dof, local_dof)).reshape(-1)
+    indices = np.stack((row_indices, col_indices), axis=1).astype(np.int32)
+    return COORouting(indices=torch_np.asarray(indices, dtype=INT_DTYPE), shape=(mesh.node_count, mesh.node_count))
+
+
+def build_coo_routing_from_connectivity(connectivity: Array, node_count: int) -> COORouting:
+    """Precompute COO routing for arbitrary local connectivity."""
+
+    connectivity_np = np.asarray(connectivity, dtype=np.int32)
+    local_dof = int(connectivity_np.shape[1])
+    row_indices = np.broadcast_to(
+        connectivity_np[:, :, None],
+        (connectivity_np.shape[0], local_dof, local_dof),
+    ).reshape(-1)
+    col_indices = np.broadcast_to(
+        connectivity_np[:, None, :],
+        (connectivity_np.shape[0], local_dof, local_dof),
+    ).reshape(-1)
+    indices = np.stack((row_indices, col_indices), axis=1).astype(np.int32)
+    return COORouting(indices=torch_np.asarray(indices, dtype=INT_DTYPE), shape=(node_count, node_count))
+
+
+def assemble_local_stiffness(element_data: P1ElementData, conductivity: Array | float) -> Array:
+    """Assemble elementwise diffusion matrices with batched einsums."""
+
+    static_conductivity = _static_scalar_value(conductivity)
+    if static_conductivity is not None:
+        dtype = NP_FLOAT_DTYPE
+        gradients = np.asarray(element_data.gradients, dtype=dtype)
+        weights = (
+            2.0
+            * np.asarray(element_data.cell_areas, dtype=dtype)[:, None]
+            * np.asarray(element_data.quadrature_weights, dtype=dtype)[None, :]
+        )
+        local = np.einsum("eqid,eqjd,eq->eij", gradients, gradients, static_conductivity * weights)
+        return torch_np.asarray(local, dtype=FLOAT_DTYPE)
+
+    conductivity_values = _expand_coefficient(conductivity, element_data)
+    weights = 2.0 * element_data.cell_areas[:, None] * element_data.quadrature_weights[None, :]
+    return torch_np.einsum(
+        "eqid,eqjd,eq->eij",
+        element_data.gradients,
+        element_data.gradients,
+        conductivity_values * weights,
+    )
+
+
+def assemble_local_mass(
+    element_data: P1ElementData, coefficients: Array | float = 1.0
+) -> Array:
+    """Assemble elementwise mass matrices with optional scalar coefficients."""
+
+    static_coefficient = _static_scalar_value(coefficients)
+    if static_coefficient is not None:
+        dtype = NP_FLOAT_DTYPE
+        shape_values = np.asarray(element_data.shape_values, dtype=dtype)
+        weights = (
+            2.0
+            * np.asarray(element_data.cell_areas, dtype=dtype)[:, None]
+            * np.asarray(element_data.quadrature_weights, dtype=dtype)[None, :]
+        )
+        local = np.einsum("qi,qj,eq->eij", shape_values, shape_values, static_coefficient * weights)
+        return torch_np.asarray(local, dtype=FLOAT_DTYPE)
+
+    coefficient_values = _expand_coefficient(coefficients, element_data)
+    weights = 2.0 * element_data.cell_areas[:, None] * element_data.quadrature_weights[None, :]
+    return torch_np.einsum(
+        "qi,qj,eq->eij",
+        element_data.shape_values,
+        element_data.shape_values,
+        coefficient_values * weights,
+    )
+
+
+def assemble_local_stiffness_p2(
+    element_data: P2ElementData,
+    conductivity: Array | float,
+) -> Array:
+    """Assemble quadratic elementwise diffusion matrices."""
+
+    conductivity_values = _expand_coefficient_generic(
+        conductivity,
+        cell_count=element_data.cell_areas.shape[0],
+        quadrature_count=element_data.quadrature_weights.shape[0],
+    )
+    weights = 2.0 * element_data.cell_areas[:, None] * element_data.quadrature_weights[None, :]
+    return torch_np.einsum(
+        "eqid,eqjd,eq->eij",
+        element_data.gradients,
+        element_data.gradients,
+        conductivity_values * weights,
+    )
+
+
+def assemble_local_mass_p2(
+    element_data: P2ElementData,
+    coefficients: Array | float = 1.0,
+) -> Array:
+    """Assemble quadratic elementwise mass matrices."""
+
+    coefficient_values = _expand_coefficient_generic(
+        coefficients,
+        cell_count=element_data.cell_areas.shape[0],
+        quadrature_count=element_data.quadrature_weights.shape[0],
+    )
+    weights = 2.0 * element_data.cell_areas[:, None] * element_data.quadrature_weights[None, :]
+    return torch_np.einsum(
+        "qi,qj,eq->eij",
+        element_data.shape_values,
+        element_data.shape_values,
+        coefficient_values * weights,
+    )
+
+
+def assemble_global_bcoo(local_matrices: Array, routing: COORouting) -> BCOO:
+    """Reduce local element matrices into a global sparse operator."""
+
+    sparse_matrix = BCOO(
+        (local_matrices.reshape(-1), routing.indices),
+        shape=routing.shape,
+        unique_indices=False,
+    )
+    return sparse_matrix.sum_duplicates()
+
+
+def assemble_helmholtz_operator(
+    element_data: P1ElementData,
+    routing: COORouting,
+    conductivity: Array | float,
+    wavenumber_sq: Array | float,
+) -> BCOO:
+    """Assemble a sparse Helmholtz-type operator K + k^2 M."""
+
+    conductivity_values = torch_np.asarray(conductivity, dtype=FLOAT_DTYPE)
+    local_operator = assemble_local_stiffness(
+        element_data=element_data,
+        conductivity=conductivity,
+    ) + assemble_local_mass(
+        element_data=element_data,
+        coefficients=conductivity_values * wavenumber_sq,
+    )
+    return assemble_global_bcoo(local_operator, routing)

adtlert/fem/boundary.py

@@ -0,0 +1,176 @@
+"""Boundary assembly utilities for 2.5D ERT."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import numpy as np
+from scipy.special import k0 as besselk0
+from scipy.special import k1 as besselk1
+
+from adtlert.utils.torch_runtime import Array
+from adtlert.utils.torch_runtime import torch_np
+from adtlert.utils.torch_runtime import BCOO
+
+from adtlert.mesh import Mesh
+from adtlert.utils.dtypes import FLOAT_DTYPE, INT_DTYPE, NP_FLOAT_DTYPE
+
+
+@dataclass(frozen=True)
+class BoundaryRouting:
+    """Reusable COO routing indices for boundary-edge assembly."""
+
+    indices: Array
+    shape: tuple[int, int]
+
+
+def build_boundary_routing(mesh: Mesh) -> BoundaryRouting:
+    """Precompute the COO routing pattern for boundary edge assembly."""
+
+    return build_boundary_routing_from_connectivity(mesh.boundary_edges, mesh.node_count)
+
+
+def build_boundary_routing_from_connectivity(boundary_connectivity: Array, node_count: int) -> BoundaryRouting:
+    """Precompute COO routing for arbitrary boundary connectivity."""
+
+    connectivity_np = np.asarray(boundary_connectivity, dtype=np.int32)
+    edge_count = connectivity_np.shape[0]
+    local_dof = int(connectivity_np.shape[1])
+    row_indices = np.broadcast_to(connectivity_np[:, :, None], (edge_count, local_dof, local_dof)).reshape(-1)
+    col_indices = np.broadcast_to(connectivity_np[:, None, :], (edge_count, local_dof, local_dof)).reshape(-1)
+    indices = np.stack((row_indices, col_indices), axis=1).astype(np.int32)
+    return BoundaryRouting(indices=torch_np.asarray(indices, dtype=INT_DTYPE), shape=(node_count, node_count))
+
+
+def _static_scalar_value(value: Array | float) -> float | None:
+    try:
+        value_array = np.asarray(value)
+    except Exception:
+        return None
+    if value_array.ndim != 0:
+        return None
+    return float(value_array)
+
+
+def assemble_local_boundary_mass(mesh: Mesh, coefficients: Array | float) -> Array:
+    """Assemble midpoint-based Robin edge matrices."""
+
+    static_coefficient = _static_scalar_value(coefficients)
+    if static_coefficient is not None:
+        dtype = NP_FLOAT_DTYPE
+        reference = np.asarray([[2.0, 1.0], [1.0, 2.0]], dtype=dtype) / 6.0
+        lengths = np.asarray(mesh.boundary_edge_lengths, dtype=dtype)
+        local = static_coefficient * lengths[:, None, None] * reference[None, :, :]
+        return torch_np.asarray(local, dtype=FLOAT_DTYPE)
+
+    coefficient_array = torch_np.asarray(coefficients, dtype=FLOAT_DTYPE)
+    if coefficient_array.ndim == 0:
+        coefficient_array = torch_np.broadcast_to(coefficient_array, (mesh.boundary_edges.shape[0],))
+    elif coefficient_array.shape != (mesh.boundary_edges.shape[0],):
+        raise ValueError("boundary coefficients must be scalar or shape (num_boundary_edges,)")
+
+    reference = torch_np.asarray([[2.0, 1.0], [1.0, 2.0]], dtype=FLOAT_DTYPE) / 6.0
+    return coefficient_array[:, None, None] * mesh.boundary_edge_lengths[:, None, None] * reference
+
+
+def assemble_local_boundary_mass_p2(lengths: Array, coefficients: Array | float) -> Array:
+    """Assemble quadratic Robin edge matrices."""
+
+    static_coefficient = _static_scalar_value(coefficients)
+    if static_coefficient is not None:
+        dtype = NP_FLOAT_DTYPE
+        length_array_np = np.asarray(lengths, dtype=dtype)
+        reference = np.asarray(
+            [
+                [4.0, 2.0, -1.0],
+                [2.0, 16.0, 2.0],
+                [-1.0, 2.0, 4.0],
+            ],
+            dtype=dtype,
+        ) / 30.0
+        local = static_coefficient * length_array_np[:, None, None] * reference[None, :, :]
+        return torch_np.asarray(local, dtype=FLOAT_DTYPE)
+
+    length_array = torch_np.asarray(lengths, dtype=FLOAT_DTYPE)
+    coefficient_array = torch_np.asarray(coefficients, dtype=FLOAT_DTYPE)
+    if coefficient_array.ndim == 0:
+        coefficient_array = torch_np.broadcast_to(coefficient_array, (length_array.shape[0],))
+    elif coefficient_array.shape != (length_array.shape[0],):
+        raise ValueError("boundary coefficients must be scalar or shape (num_boundary_edges,)")
+
+    reference = torch_np.asarray(
+        [
+            [4.0, 2.0, -1.0],
+            [2.0, 16.0, 2.0],
+            [-1.0, 2.0, 4.0],
+        ],
+        dtype=FLOAT_DTYPE,
+    ) / 30.0
+    return coefficient_array[:, None, None] * length_array[:, None, None] * reference
+
+
+def assemble_boundary_bcoo(local_matrices: Array, routing: BoundaryRouting) -> BCOO:
+    """Reduce local boundary edge matrices into a global sparse operator."""
+
+    sparse_matrix = BCOO(
+        (local_matrices.reshape(-1), routing.indices),
+        shape=routing.shape,
+        unique_indices=False,
+    )
+    return sparse_matrix.sum_duplicates()
+
+
+def _expand_cell_coefficient(coefficients: Array | float, mesh: Mesh) -> Array:
+    coefficient_array = torch_np.asarray(coefficients, dtype=FLOAT_DTYPE)
+    if coefficient_array.ndim == 0:
+        return torch_np.broadcast_to(coefficient_array, (mesh.cell_count,))
+    if coefficient_array.shape == (mesh.cell_count,):
+        return coefficient_array
+    raise ValueError("cell coefficients must be scalar or shape (num_cells,)")
+
+
+def robin_boundary_coefficients(
+    mesh: Mesh,
+    conductivity: Array | float,
+    source_center: Array,
+    wavenumber: float,
+) -> Array:
+    """Compute 2.5D Robin coefficients for the mixed boundary condition."""
+
+    static_conductivity = _static_scalar_value(conductivity)
+    if static_conductivity is None:
+        conductivity_values = _expand_cell_coefficient(conductivity, mesh)
+        boundary_sigma = conductivity_values[mesh.boundary_edge_cells]
+    else:
+        boundary_sigma = np.full(mesh.boundary_edges.shape[0], static_conductivity, dtype=NP_FLOAT_DTYPE)
+
+    centers = np.asarray(mesh.boundary_edge_centers, dtype=float)
+    normals = np.asarray(mesh.boundary_edge_normals, dtype=float)
+    source = np.asarray(source_center, dtype=float)
+    mirrored_source = np.asarray([source[0], -source[1]], dtype=float)
+    r1 = centers - source
+    r2 = centers - mirrored_source
+    r1_abs = np.linalg.norm(r1, axis=1)
+    r2_abs = np.linalg.norm(r2, axis=1)
+    top_boundary = np.asarray(mesh.surface_edge_mask, dtype=bool)
+
+    geometry = np.zeros(centers.shape[0], dtype=float)
+    valid = (r1_abs > 1e-12) & (r2_abs > 1e-12) & (~top_boundary)
+    if np.any(valid):
+        denominator = besselk0(r1_abs[valid] * wavenumber) + besselk0(r2_abs[valid] * wavenumber)
+        stable = denominator > 1e-12
+        valid_indices = np.flatnonzero(valid)
+        if np.any(stable):
+            stable_indices = valid_indices[stable]
+            numerator = (
+                np.sum(r1[stable_indices] * normals[stable_indices], axis=1)
+                / r1_abs[stable_indices]
+                * besselk1(r1_abs[stable_indices] * wavenumber)
+                + np.sum(r2[stable_indices] * normals[stable_indices], axis=1)
+                / r2_abs[stable_indices]
+                * besselk1(r2_abs[stable_indices] * wavenumber)
+            )
+            geometry[stable_indices] = wavenumber * numerator / denominator[stable]
+
+    if static_conductivity is not None:
+        return torch_np.asarray(boundary_sigma * geometry, dtype=FLOAT_DTYPE)
+    return boundary_sigma * torch_np.asarray(geometry, dtype=FLOAT_DTYPE)

adtlert/fem/p1.py

@@ -0,0 +1,134 @@
+"""P1 triangle basis functions and batched quadrature data."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from adtlert.utils.torch_runtime import Array
+from adtlert.utils.torch_runtime import torch_np
+import numpy as np
+
+from adtlert.mesh import Mesh
+from adtlert.utils.dtypes import FLOAT_DTYPE, NP_FLOAT_DTYPE
+
+
+@dataclass(frozen=True)
+class TriangleQuadrature:
+    """Reference triangle quadrature points and weights."""
+
+    points: Array
+    weights: Array
+
+
+@dataclass(frozen=True)
+class P1ElementData:
+    """Batched P1 element tensors for tensorized assembly."""
+
+    quadrature_points: Array
+    quadrature_weights: Array
+    shape_values: Array
+    reference_gradients: Array
+    gradients: Array
+    cell_quadrature_points: Array
+    cell_areas: Array
+
+
+def triangle_quadrature(order: int = 2) -> TriangleQuadrature:
+    """Return a low-order Gauss rule on the reference triangle."""
+
+    if order == 1:
+        points = torch_np.asarray([[1.0 / 3.0, 1.0 / 3.0]], dtype=FLOAT_DTYPE)
+        weights = torch_np.asarray([0.5], dtype=FLOAT_DTYPE)
+        return TriangleQuadrature(points=points, weights=weights)
+
+    if order == 2:
+        points = torch_np.asarray(
+            [
+                [1.0 / 6.0, 1.0 / 6.0],
+                [2.0 / 3.0, 1.0 / 6.0],
+                [1.0 / 6.0, 2.0 / 3.0],
+            ],
+            dtype=FLOAT_DTYPE,
+        )
+        weights = torch_np.asarray([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0], dtype=FLOAT_DTYPE)
+        return TriangleQuadrature(points=points, weights=weights)
+
+    raise ValueError(f"unsupported triangle quadrature order: {order}")
+
+
+def p1_shape_functions(points: Array) -> Array:
+    """Evaluate P1 shape functions at reference coordinates."""
+
+    point_array = torch_np.asarray(points, dtype=FLOAT_DTYPE)
+    xi = point_array[..., 0]
+    eta = point_array[..., 1]
+    return torch_np.stack((1.0 - xi - eta, xi, eta), axis=-1)
+
+
+def reference_shape_gradients() -> Array:
+    """Return the constant reference gradients of the P1 basis."""
+
+    return torch_np.asarray(
+        [
+            [-1.0, -1.0],
+            [1.0, 0.0],
+            [0.0, 1.0],
+        ],
+        dtype=FLOAT_DTYPE,
+    )
+
+
+def build_p1_element_data(mesh: Mesh, quadrature_order: int = 2) -> P1ElementData:
+    """Prepare batched element tensors for later FEM assembly."""
+
+    dtype = NP_FLOAT_DTYPE
+    if quadrature_order == 1:
+        quadrature_points = np.asarray([[1.0 / 3.0, 1.0 / 3.0]], dtype=dtype)
+        quadrature_weights = np.asarray([0.5], dtype=dtype)
+    elif quadrature_order == 2:
+        quadrature_points = np.asarray(
+            [
+                [1.0 / 6.0, 1.0 / 6.0],
+                [2.0 / 3.0, 1.0 / 6.0],
+                [1.0 / 6.0, 2.0 / 3.0],
+            ],
+            dtype=dtype,
+        )
+        quadrature_weights = np.asarray([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0], dtype=dtype)
+    else:
+        raise ValueError(f"unsupported triangle quadrature order: {quadrature_order}")
+
+    xi = quadrature_points[..., 0]
+    eta = quadrature_points[..., 1]
+    shape_values = np.stack((1.0 - xi - eta, xi, eta), axis=-1).astype(dtype)
+    ref_gradients = np.asarray(
+        [
+            [-1.0, -1.0],
+            [1.0, 0.0],
+            [0.0, 1.0],
+        ],
+        dtype=dtype,
+    )
+
+    cell_nodes = np.asarray(mesh.nodes, dtype=dtype)[np.asarray(mesh.cells, dtype=np.int32)]
+    edge_1 = cell_nodes[:, 1] - cell_nodes[:, 0]
+    edge_2 = cell_nodes[:, 2] - cell_nodes[:, 0]
+    jacobians = np.stack((edge_1, edge_2), axis=-1)
+    inverse_jacobian_t = np.linalg.inv(jacobians).transpose((0, 2, 1))
+    cell_gradients = np.einsum("eij,nj->eni", inverse_jacobian_t, ref_gradients)
+    gradients = np.broadcast_to(
+        cell_gradients[:, None, :, :],
+        (mesh.cell_count, quadrature_points.shape[0], 3, 2),
+    )
+    cell_quadrature_points = np.einsum("qi,eid->eqd", shape_values, cell_nodes)
+    cell_areas = 0.5 * np.abs(np.linalg.det(jacobians))
+
+    return P1ElementData(
+        quadrature_points=torch_np.asarray(quadrature_points, dtype=FLOAT_DTYPE),
+        quadrature_weights=torch_np.asarray(quadrature_weights, dtype=FLOAT_DTYPE),
+        shape_values=torch_np.asarray(shape_values, dtype=FLOAT_DTYPE),
+        reference_gradients=torch_np.asarray(ref_gradients, dtype=FLOAT_DTYPE),
+        gradients=torch_np.asarray(gradients, dtype=FLOAT_DTYPE),
+        cell_quadrature_points=torch_np.asarray(cell_quadrature_points, dtype=FLOAT_DTYPE),
+        cell_areas=torch_np.asarray(cell_areas, dtype=FLOAT_DTYPE),
+    )