Periodic Boundary Conditions

This tutorial demonstrates how to apply periodic boundary conditions in FEAX using prolongation matrices.

Problem Description

Consider the 2D Poisson equation on a unit square $\Omega = [0, 1] \times [0, 1]$:

$$-\nabla \cdot (\theta \nabla u) = f \quad \text{in } \Omega$$

with periodic BCs on left-right boundaries: $u(0, y) = u(1, y)$, and Dirichlet BCs on top-bottom: $u(x, 0) = u(x, 1) = 0$.

Mathematical Formulation

The prolongation matrix $\mathbf{P}$ relates full and reduced DOFs:

$$\mathbf{u}_{\text{full}} = \mathbf{P} \, \mathbf{u}_{\text{reduced}}$$

The system is solved in reduced space: $(\mathbf{P}^T \mathbf{K} \mathbf{P}) \, \mathbf{u}_{\text{reduced}} = \mathbf{P}^T \mathbf{f}$

Implementation

Step 1: Mesh Generation

import feax as fe
import jax.numpy as np

mesh = fe.mesh.rectangle_mesh(Nx=32, Ny=32, domain_x=1.0, domain_y=1.0)

Step 2: Problem Definition

class PoissonParametric(fe.problem.Problem):
    def get_tensor_map(self):
        def tensor_map(u_grad, theta):
            return theta * u_grad
        return tensor_map

    def get_mass_map(self):
        def mass_map(u, x, theta):
            dx, dy = x[0] - 0.5, x[1] - 0.5
            val = x[0]*np.sin(5.0*np.pi*x[1]) + np.exp(-(dx*dx + dy*dy)/0.02)
            return np.array([-val])
        return mass_map

problem = PoissonParametric(mesh=mesh, vec=1, dim=2, ele_type='QUAD4', location_fns=[])

Step 3: Periodic Boundary Conditions

import feax.flat as flat

def left_boundary(point):
    return np.isclose(point[0], 0.0, atol=1e-5)

def right_boundary(point):
    return np.isclose(point[0], 1.0, atol=1e-5)

def mapping_x(point_A):
    return np.array([point_A[0] + 1.0, point_A[1]])

periodic_pairing = flat.pbc.PeriodicPairing(
    location_master=left_boundary,
    location_slave=right_boundary,
    mapping=mapping_x,
    vec=0
)

P = flat.pbc.prolongation_matrix([periodic_pairing], mesh, vec=1)

Step 4: Dirichlet Boundary Conditions

def bottom_boundary(point):
    return np.isclose(point[1], 0.0, atol=1e-5)

def top_boundary(point):
    return np.isclose(point[1], 1.0, atol=1e-5)

bc_config = fe.DCboundary.DirichletBCConfig([
    fe.DCboundary.DirichletBCSpec(bottom_boundary, 0, 0.0),
    fe.DCboundary.DirichletBCSpec(top_boundary, 0, 0.0),
])
bc = bc_config.create_bc(problem)

Step 5: Internal Variables

theta = 1.0
theta_array = fe.TracedParams.create_uniform_volume_var(problem, theta)
traced_params = fe.TracedParams(volume_vars=(theta_array,), surface_vars=())

Step 6: Solver

solver_options = fe.KrylovSolverOptions(solver="cg", tol=1e-8)
solver = fe.create_solver(problem, bc, solver_options, linear=True, P=P)

Pass prolongation matrix P to create_solver(). With KrylovSolverOptions the reduced system $\mathbf{P}^T \mathbf{J} \mathbf{P}$ is solved matrix-free; alternatively, DirectSolverOptions or AMGSolverOptions assemble the reduced operator sparsely and factorize it (or build an AMG hierarchy from it).

Dirichlet BCs must respect the periodic pairing

A Dirichlet BC that constrains only part of a periodic equivalence class (e.g. pinning one node of a tied pair) is contradictory, and create_solver raises a ValueError at build time. Constrain interior (non-paired) nodes, or the entire class — here the top/bottom Dirichlet rows pin both partners of each left–right corner pair, which is valid. For RVE homogenization, do not pin any node: leave the DirichletBC empty (see Lattice Homogenization).

Step 7: Solve

initial_guess = np.zeros(problem.num_total_dofs_all_vars)
sol_full = solver(traced_params, initial_guess)
fe.utils.save_sol(mesh, "periodic_poisson.vtu", point_infos=[("u", sol_full.field(0))])

The solver returns a fe.Solution; sol_full.field(0) gives the (num_nodes, 1) nodal field.

3D Periodic Boundary Conditions

periodic_bc_3D builds the full set of face/edge/corner pairings automatically, but it operates on a UnitCell (which carries the geometric bounds and boundary-identification functions), not on a bare Mesh. Subclass flat.unitcell.UnitCell and implement mesh_build():

from feax.flat.pbc import periodic_bc_3D

class BoxUnitCell(flat.unitcell.UnitCell):
    def mesh_build(self, mesh_size):
        return fe.mesh.box_mesh(size=1.0, mesh_size=mesh_size, element_type='HEX8')

unitcell = BoxUnitCell(mesh_size=0.1)
mesh_3d = unitcell.mesh

pairings_3d = periodic_bc_3D(unitcell, vec=3, dim=3)
P_3d = flat.pbc.prolongation_matrix(pairings_3d, mesh_3d, vec=3)

A matching periodic_bc_2D(unitcell, vec=..., dim=2) convenience function is also available. See Lattice Homogenization for a complete RVE example.

Complete Code

import feax as fe
import feax.flat as flat
import jax.numpy as np

# Problem definition
class PoissonParametric(fe.problem.Problem):
    def get_tensor_map(self):
        def tensor_map(u_grad, theta):
            return theta * u_grad
        return tensor_map

    def get_mass_map(self):
        def mass_map(u, x, theta):
            dx = x[0] - 0.5
            dy = x[1] - 0.5
            val = x[0]*np.sin(5.0*np.pi*x[1]) + np.exp(-(dx*dx + dy*dy)/0.02)
            return np.array([-val])
        return mass_map

# Mesh
Nx, Ny = 32, 32
mesh = fe.mesh.rectangle_mesh(Nx=Nx, Ny=Ny, domain_x=1.0, domain_y=1.0)

# Create problem
problem = PoissonParametric(mesh=mesh, vec=1, dim=2, ele_type='QUAD4', location_fns=[])

# Periodic boundary conditions (left-right)
def left_boundary(point):
    return np.isclose(point[0], 0.0, atol=1e-5)

def right_boundary(point):
    return np.isclose(point[0], 1.0, atol=1e-5)

def mapping_x(point_A):
    return np.array([point_A[0] + 1.0, point_A[1]])

periodic_pairing = flat.pbc.PeriodicPairing(
    location_master=left_boundary,
    location_slave=right_boundary,
    mapping=mapping_x,
    vec=0
)

P = flat.pbc.prolongation_matrix([periodic_pairing], mesh, vec=1)

# Dirichlet boundary conditions (top-bottom = 0)
def bottom_boundary(point):
    return np.isclose(point[1], 0.0, atol=1e-5)

def top_boundary(point):
    return np.isclose(point[1], 1.0, atol=1e-5)

bc_config = fe.DCboundary.DirichletBCConfig([
    fe.DCboundary.DirichletBCSpec(bottom_boundary, 0, 0.0),
    fe.DCboundary.DirichletBCSpec(top_boundary, 0, 0.0),
])
bc = bc_config.create_bc(problem)

# Internal variables
theta = 1.0
theta_array = fe.TracedParams.create_uniform_volume_var(problem, theta)
traced_params = fe.TracedParams(volume_vars=(theta_array,), surface_vars=())

# Solver
solver_options = fe.KrylovSolverOptions(solver="cg", tol=1e-8)
solver = fe.create_solver(problem, bc, solver_options=solver_options, linear=True, P=P)

# Solve
initial_guess = np.zeros(problem.num_total_dofs_all_vars)
sol_full = solver(traced_params, initial_guess)

# Save
fe.utils.save_sol(mesh, "periodic_poisson.vtu", point_infos=[("u", sol_full.field(0))])

Vector Problems

For vector problems, apply periodicity to each component:

pairings = [
    flat.pbc.PeriodicPairing(left, right, mapping_x, vec=0),
    flat.pbc.PeriodicPairing(left, right, mapping_x, vec=1),
    flat.pbc.PeriodicPairing(left, right, mapping_x, vec=2),
]
P = flat.pbc.prolongation_matrix(pairings, mesh, vec=3)