tessl/pypi-warp-lang

A Python framework for high-performance simulation and graphics programming that JIT compiles Python functions to efficient GPU/CPU kernel code.

Overview

Eval results

Files

Finite Element Method (FEM)

Name: tessl/pypi-warp-lang
Author: tessl

Warp's finite element framework provides comprehensive tools for solving partial differential equations using the finite element method. It includes geometry definitions, function spaces, quadrature, field operations, and integration capabilities for numerical simulation of physical phenomena.

Capabilities

Geometry Types

Geometric domains for finite element discretization.

class Grid2D:
    """Regular 2D grid geometry."""
    
    def __init__(self, res_x: int, res_y: int):
        """Create 2D grid with specified resolution."""

class Grid3D:
    """Regular 3D grid geometry."""
    
    def __init__(self, res_x: int, res_y: int, res_z: int):
        """Create 3D grid with specified resolution."""

class Tetmesh:
    """Tetrahedral mesh geometry."""
    
    def __init__(self, vertices: array, indices: array):
        """
        Create tetrahedral mesh from vertices and element connectivity.
        
        Args:
            vertices: Array of vertex positions
            indices: Array of tetrahedron vertex indices
        """

class Hexmesh:
    """Hexahedral mesh geometry."""
    
    def __init__(self, vertices: array, indices: array):
        """Create hexahedral mesh from vertices and connectivity."""

class Trimesh2D:
    """2D triangular mesh geometry."""
    
    def __init__(self, vertices: array, indices: array):
        """Create 2D triangular mesh."""

class Trimesh3D:  
    """3D triangular surface mesh geometry."""
    
    def __init__(self, vertices: array, indices: array):
        """Create 3D triangular surface mesh."""

class Quadmesh2D:
    """2D quadrilateral mesh geometry."""
    
    def __init__(self, vertices: array, indices: array):
        """Create 2D quadrilateral mesh."""

class Quadmesh3D:
    """3D quadrilateral surface mesh geometry."""
    
    def __init__(self, vertices: array, indices: array):
        """Create 3D quadrilateral surface mesh."""

class Nanogrid:
    """Dense voxel grid for level set methods."""
    
    def __init__(self, data: array, transform: transform = None):
        """Create nanogrid from voxel data."""

class AdaptiveNanogrid:
    """Adaptive sparse voxel grid."""
    
    def __init__(self, field: Field, tolerance: float = 1e-6):
        """Create adaptive nanogrid from field with specified tolerance."""

Function Spaces

Mathematical spaces defining basis functions and degrees of freedom.

class FunctionSpace:
    """Function space defining basis functions over geometry."""
    
    @property
    def dimension(self) -> int:
        """Number of degrees of freedom."""
    
    @property
    def geometry(self) -> Geometry:
        """Underlying geometry."""

class BasisSpace:
    """Basis function space on geometry elements."""

class PointBasisSpace:
    """Point-based basis space for meshfree methods."""

def make_polynomial_space(geometry: Geometry, 
                         degree: int, 
                         dtype: type = float) -> FunctionSpace:
    """
    Create polynomial function space of specified degree.
    
    Args:
        geometry: Geometric domain
        degree: Polynomial degree (1=linear, 2=quadratic, etc.)
        dtype: Scalar type for degrees of freedom
        
    Returns:
        Function space with polynomial basis functions
    """

def make_polynomial_basis_space(geometry: Geometry, 
                               degree: int,
                               family: str = "lagrange") -> BasisSpace:
    """Create polynomial basis space with specified family."""

def make_collocated_function_space(geometry: Geometry, 
                                  shape: tuple, 
                                  degree: int = 1) -> FunctionSpace:
    """Create collocated vector/tensor function space."""

def make_covariant_function_space(geometry: Geometry, 
                                 degree: int = 1) -> FunctionSpace:
    """Create covariant (edge-based) function space."""

def make_contravariant_function_space(geometry: Geometry,
                                     degree: int = 1) -> FunctionSpace:
    """Create contravariant (face-based) function space."""

Fields

Field representations and operations over function spaces.

class DiscreteField:
    """Field represented by discrete degrees of freedom."""
    
    def __init__(self, space: FunctionSpace, dof_values: array = None):
        """
        Create discrete field on function space.
        
        Args:
            space: Function space defining basis
            dof_values: Degree of freedom values (allocated if None)
        """

class ImplicitField:
    """Field defined by implicit function."""
    
    def __init__(self, func: Callable, space: FunctionSpace = None):
        """Create field from implicit function."""

class UniformField:
    """Field with constant value everywhere."""
    
    def __init__(self, value, space: FunctionSpace = None):
        """Create uniform field with constant value."""

class NonconformingField:
    """Field with non-conforming discretization."""

def make_discrete_field(space: FunctionSpace, 
                       dof_values: array = None) -> DiscreteField:
    """Create discrete field on function space."""

def make_test(space: FunctionSpace) -> DiscreteField:
    """Create test function for weak forms."""

def make_trial(space: FunctionSpace) -> DiscreteField:
    """Create trial function for weak forms."""

def make_restriction(field: Field, domain: Domain) -> Field:
    """Restrict field to subdomain."""

Quadrature

Numerical integration schemes for finite element computations.

class Quadrature:
    """Base class for quadrature rules."""

class RegularQuadrature:
    """Regular quadrature on structured grids."""
    
    def __init__(self, geometry: Geometry, order: int):
        """
        Create regular quadrature rule.
        
        Args:
            geometry: Geometric domain
            order: Integration order (accuracy)
        """

class NodalQuadrature:
    """Quadrature using mesh nodes as integration points."""
    
    def __init__(self, geometry: Geometry):
        """Create nodal quadrature rule."""

class PicQuadrature:
    """Particle-in-cell quadrature for particle methods."""
    
    def __init__(self, positions: array, weights: array = None):
        """
        Create PIC quadrature from particle data.
        
        Args:
            positions: Particle positions
            weights: Particle weights (uniform if None)
        """

class ExplicitQuadrature:
    """Explicit quadrature with specified points and weights."""
    
    def __init__(self, points: array, weights: array):
        """Create quadrature from explicit points and weights."""

Domain and Topology

Domain definitions and topological operations.

class GeometryDomain:
    """Geometric domain for finite element computations."""
    
    def __init__(self, geometry: Geometry):
        """Create domain from geometry."""

class Cells:
    """Domain representing geometry cells/elements."""

class Sides:
    """Domain representing element faces/edges."""

class BoundarySides:
    """Domain representing boundary faces."""

class FrontierSides:
    """Domain representing interfaces between regions."""

class SpacePartition:
    """Partition of function space for parallel computation."""

class SpaceRestriction:
    """Restriction of function space to subdomain."""

class SpaceTopology:
    """Topological information for function space."""

def make_space_partition(space: FunctionSpace, 
                        partitions: int) -> SpacePartition:
    """Partition function space for parallel computation."""

def make_space_restriction(space: FunctionSpace, 
                          domain: Domain) -> SpaceRestriction:
    """Create space restriction to subdomain."""

Integration and Interpolation

Core operations for finite element computations.

def integrate(integrand: Callable, 
             domain: Domain, 
             quadrature: Quadrature = None,
             output: Field = None) -> Field:
    """
    Integrate function over domain using specified quadrature.
    
    Args:
        integrand: Function to integrate
        domain: Integration domain
        quadrature: Quadrature rule (default if None)
        output: Output field (allocated if None)
        
    Returns:
        Integrated field values
    """

def interpolate(field: Field, 
               space: FunctionSpace, 
               domain: Domain = None) -> DiscreteField:
    """
    Interpolate field onto function space.
    
    Args:
        field: Input field to interpolate
        space: Target function space
        domain: Interpolation domain (geometry if None)
        
    Returns:
        Discrete field on target space
    """

Differential Operators

Operators for computing derivatives and differential operations.

def grad(field: Field) -> Field:
    """Compute gradient of scalar field."""

def div(field: Field) -> Field:
    """Compute divergence of vector field."""

def curl(field: Field) -> Field:
    """Compute curl of vector field."""

def D(field: Field, direction: int) -> Field:
    """Partial derivative in specified direction."""

def grad_outer(field: Field) -> Field:
    """Outer product of gradient."""

def div_outer(field: Field) -> Field:
    """Outer divergence operation."""

def deformation_gradient(field: Field) -> Field:
    """Compute deformation gradient tensor."""

def grad_average(field: Field) -> Field:
    """Average gradient across elements."""

def grad_jump(field: Field) -> Field:
    """Jump in gradient across element boundaries."""

Field Operations

Operations for manipulating and querying fields.

def inner(a: Field, b: Field) -> Field:
    """Inner product of two fields."""

def outer(a: Field, b: Field) -> Field:
    """Outer product of two fields."""

def average(field: Field) -> Field:
    """Average field values across elements."""

def jump(field: Field) -> Field:
    """Jump in field across element boundaries."""

def lookup(field: Field, coordinates: array) -> array:
    """Evaluate field at specified coordinates."""

def partition_lookup(field: Field, partition: SpacePartition) -> Field:
    """Lookup field values in partitioned space."""

def at_node(field: Field, node_index: int):
    """Get field value at specific node."""

def measure(domain: Domain) -> Field:
    """Compute measure (area/volume) of domain."""

def measure_ratio(domain: Domain, reference: Domain) -> Field:
    """Ratio of domain measure to reference."""

def normal(domain: Domain) -> Field:
    """Outward normal vector on boundary."""

def position(domain: Domain) -> Field:
    """Position coordinates on domain."""

Element Operations

Operations for working with individual elements.

def cells(geometry: Geometry) -> Cells:
    """Get cells/elements of geometry."""

def element_coordinates(element: ElementIndex) -> Coords:
    """Get coordinates of element."""

def element_closest_point(element: ElementIndex, point: vec3) -> vec3:
    """Find closest point on element to query point."""

def node_count(element: ElementIndex) -> int:
    """Number of nodes in element."""

def node_index(element: ElementIndex, local_index: int) -> int:
    """Global node index from element and local index."""

def degree(space: FunctionSpace) -> int:
    """Polynomial degree of function space."""

def to_cell_side(side: int) -> int:
    """Convert side index to cell-relative index."""

def to_inner_cell(side: int) -> int:
    """Get inner cell adjacent to side."""

def to_outer_cell(side: int) -> int:
    """Get outer cell adjacent to side."""

Caching and Memory Management

Utilities for managing temporary storage and caching.

class TemporaryStore:
    """Storage manager for temporary arrays."""
    
    def __init__(self, device: Device = None):
        """Create temporary storage on device."""

def borrow_temporary(shape: tuple, 
                    dtype: type, 
                    device: Device = None) -> array:
    """Borrow temporary array from pool."""

def borrow_temporary_like(arr: array, 
                         dtype: type = None, 
                         device: Device = None) -> array:
    """Borrow temporary array with same shape as existing array."""

def set_default_temporary_store(store: TemporaryStore) -> None:
    """Set default temporary storage."""

Adaptivity

Adaptive mesh refinement and field operations.

def adaptive_nanogrid_from_field(field: Field, 
                                tolerance: float = 1e-6) -> AdaptiveNanogrid:
    """Create adaptive nanogrid from field with error tolerance."""

def adaptive_nanogrid_from_hierarchy(hierarchy, 
                                    level: int = -1) -> AdaptiveNanogrid:
    """Create adaptive nanogrid from mesh hierarchy."""

Usage Examples

Basic FEM Workflow

import warp as wp
import warp.fem as fem

# Create geometry
grid = fem.Grid2D(res_x=32, res_y=32)

# Create function space
space = fem.make_polynomial_space(grid, degree=1)

# Create fields
u = fem.make_trial(space)  # Solution field
v = fem.make_test(space)   # Test function

# Define weak form for Poisson equation: ∫∇u·∇v dx = ∫f·v dx
@wp.kernel
def poisson_integrand(
    coords: fem.Coords,
    u: fem.Field,
    v: fem.Field,
    f: fem.Field
):
    return fem.inner(fem.grad(u), fem.grad(v)) - f * v

# Integrate over domain
cells_domain = fem.cells(grid)
weak_form = fem.integrate(poisson_integrand, cells_domain)

Working with Different Geometries

# Tetrahedral mesh
vertices = wp.array([[0,0,0], [1,0,0], [0,1,0], [0,0,1]], dtype=wp.vec3)
tets = wp.array([[0,1,2,3]], dtype=wp.int32)
tet_mesh = fem.Tetmesh(vertices, tets)

# Function space on tetrahedral mesh
tet_space = fem.make_polynomial_space(tet_mesh, degree=2)

# Vector field
vector_space = fem.make_collocated_function_space(tet_mesh, shape=(3,), degree=1)
displacement = fem.make_discrete_field(vector_space)

Quadrature and Integration

# Custom quadrature
quadrature = fem.RegularQuadrature(grid, order=2)

# Integrate with specific quadrature
result = fem.integrate(
    my_integrand_function,
    fem.cells(grid),
    quadrature=quadrature
)

# PIC quadrature for particle methods
particle_pos = wp.array([[0.1, 0.2], [0.8, 0.9]], dtype=wp.vec2)
particle_weights = wp.ones(2, dtype=float)
pic_quad = fem.PicQuadrature(particle_pos, particle_weights)

Types

# Core FEM types
Geometry = typing.Union[Grid2D, Grid3D, Tetmesh, Hexmesh, Trimesh2D, Trimesh3D, Quadmesh2D, Quadmesh3D]
Field = typing.Union[DiscreteField, ImplicitField, UniformField, NonconformingField]
Domain = typing.Union[GeometryDomain, Cells, Sides, BoundarySides, FrontierSides]

# Index types
ElementIndex = int
QuadraturePointIndex = int
NULL_ELEMENT_INDEX = -1
NULL_QP_INDEX = -1

# Coordinate types
Coords = vec3  # Element coordinates

# Sample type for field evaluation
class Sample:
    element_index: ElementIndex
    qp_index: QuadraturePointIndex
    coords: Coords

def make_free_sample(coords: Coords) -> Sample:
    """Create sample at free coordinates."""

# DOF mapper for complex field structures
class DofMapper:
    """Maps between different DOF organizations."""

class SymmetricTensorMapper:
    """Maps symmetric tensor components to DOFs."""
    
class SkewSymmetricTensorMapper:
    """Maps skew-symmetric tensor components to DOFs."""

# Polynomial utilities
class Polynomial:
    """Polynomial representation for basis functions."""
    
    def __init__(self, coefficients: array, degree: int):
        """Create polynomial from coefficients."""

Install with Tessl CLI