File size: 2,331 Bytes
56c4b9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import numpy as np
import jax
import jax.numpy as jnp
from jax import vmap
def _compute_flux(u_padded, a, h):
"""Compute the flux terms for the Darcy equation discretization."""
# Extract interior points of a (excluding boundaries)
a_center = a[1:-1, 1:-1]
# Compute averaged coefficients at cell faces
a_east = 0.5 * (a_center + a[2:, 1:-1]) # East face
a_west = 0.5 * (a_center + a[:-2, 1:-1]) # West face
a_north = 0.5 * (a_center + a[1:-1, 2:]) # North face
a_south = 0.5 * (a_center + a[1:-1, :-2]) # South face
# X-direction fluxes
u_center_x = u_padded[1:-1, 1:-1]
flux_east = a_east * (u_padded[2:, 1:-1] - u_center_x) / h
flux_west = a_west * (u_center_x - u_padded[:-2, 1:-1]) / h
flux_x = (flux_east - flux_west) / h
# Y-direction fluxes
u_center_y = u_padded[1:-1, 1:-1]
flux_north = a_north * (u_padded[1:-1, 2:] - u_center_y) / h
flux_south = a_south * (u_center_y - u_padded[1:-1, :-2]) / h
flux_y = (flux_north - flux_south) / h
return flux_x + flux_y
@jax.jit
def solve_single(a):
"""Solve the Darcy equation for a single coefficient field (shape [N, N])."""
N = a.shape[0]
h = 1.0 / (N - 1)
rhs = jnp.ones((N-2, N-2)) # Interior RHS is 1 (original PDE scaled correctly)
u0 = jnp.zeros_like(rhs)
def matvec(u_flat):
"""Matrix-vector product for the linear operator."""
u = u_flat.reshape(N-2, N-2)
u_padded = jnp.pad(u, ((1,1), (1,1)), mode='constant') # Pad with zeros
total_flux = _compute_flux(u_padded, a, h)
return (-total_flux).ravel() # A = -div(a grad)
# Solve using Conjugate Gradient with increased maxiter for convergence
u_flat, _ = jax.scipy.sparse.linalg.cg(matvec, rhs.ravel(), x0=u0.ravel(), maxiter=2000)
u_interior = u_flat.reshape(N-2, N-2)
# Pad with boundary zeros to get full solution
return jnp.pad(u_interior, ((1,1), (1,1)), mode='constant')
def solver(a):
"""Batch solver for the Darcy equation using JAX acceleration."""
a_jax = jnp.asarray(a) # Convert input to JAX array
# Vectorize over the batch dimension using vmap
batch_solve = vmap(solve_single)
solutions = batch_solve(a_jax)
return np.asarray(solutions) # Convert back to numpy array for output
|