Heat Transfer in a Plate with a Hole¤

import os
import time
from typing import Callable

import equinox as eqx
import gmsh
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import meshio
import numpy as np
from jax import Array
from jax_autovmap import autovmap
from matplotlib.axes import Axes
from tatva import Mesh, Operator, element

jax.config.update("jax_enable_x64", True)

lx = ly = 100.0
radius = lx / 4.0
mesh = generate_refined_plate_with_hole(
    lx, ly, hole_radius=radius, mesh_size_fine=1.0, mesh_size_coarse=2.0
)

n_dofs = mesh.coords.shape[0]
plot_mesh(mesh)

k = 45e-3   # W/mm-K, thermal conductivity
rho = 7.8e-6 # kg/mm^3, density
cp = 496     # J/kg-K, specific heat capacity

T_top = 1.0
T_bottom = 0.1

op = Operator(mesh, element.Tri3())

top_nodes = jnp.where(jnp.isclose(mesh.coords[:, 1], ly, atol=1e-8))[0]
bottom_nodes = jnp.where(jnp.isclose(mesh.coords[:, 1], 0.0, atol=1e-8))[0]
fixed_dofs = jnp.concatenate([top_nodes, bottom_nodes])

@autovmap(grad_T=1)
def thermal_energy_density(grad_T: Array) -> Array:
    return 0.5 * k * jnp.dot(grad_T, grad_T)

@jax.jit
def total_thermal_energy(T_flat: Array) -> Array:
    grad_T = op.grad(T_flat)
    return op.integrate(thermal_energy_density(grad_T))

@autovmap(T=0)
def mass_energy_density(T: Array) -> Array:
    return 0.5 * rho * cp * T * T

@jax.jit
def total_mass_energy(T_flat: Array) -> Array:
    T_quad = op.eval(T_flat)
    return op.integrate(mass_energy_density(T_quad))

compute_stiffness_force = jax.jacrev(total_thermal_energy)
compute_mass_force = jax.jacrev(total_mass_energy)

@jax.jit
def K_times_x(T_prev: Array, dT: Array) -> Array:
    dT_p = dT.at[fixed_dofs].set(0.0)
    tangent = jax.jvp(compute_stiffness_force, (T_prev,), (dT_p,))[1]
    return tangent.at[fixed_dofs].set(0.0)

@jax.jit
def M_times_x(T_prev: Array, dT: Array) -> Array:
    dT_p = dT.at[fixed_dofs].set(0.0)
    tangent = jax.jvp(compute_mass_force, (T_prev,), (dT_p,))[1]
    return tangent.at[fixed_dofs].set(0.0)

@jax.jit
def A_times_x(dT: Array, T_prev: Array, dt: float) -> Array:
    return M_times_x(T_prev, dT) + dt * K_times_x(T_prev, dT)

@eqx.filter_jit
def conjugate_gradient(A: Callable, b: Array, atol=1e-8, max_iter=500):
    x = jnp.zeros_like(b)
    r = b - A(x)
    p = r
    rsold = jnp.vdot(r, r)
    iiter = 0

    def cond_fun(state):
        _, _, _, rsold, _, iiter = state
        return jnp.logical_and(jnp.sqrt(rsold) > atol, iiter < max_iter)

    def body_fun(state):
        b, p, r, rsold, x, iiter = state
        Ap = A(p)
        alpha = rsold / jnp.vdot(p, Ap)
        x = x + alpha * p
        r = r - alpha * Ap
        rsnew = jnp.vdot(r, r)
        p = r + (rsnew / rsold) * p
        return b, p, r, rsnew, x, iiter + 1

    _, _, _, _, x, iiter = jax.lax.while_loop(
        cond_fun, body_fun, (b, p, r, rsold, x, iiter)
    )
    return x, iiter

# Initial condition: T=0 everywhere except at boundaries
T_init = jnp.zeros(n_dofs)
T_init = T_init.at[top_nodes].set(T_top)
T_init = T_init.at[bottom_nodes].set(T_bottom)

total_time = 200.0
dt = 1.0
n_steps = int(total_time / dt)
save_interval = 10

f_ext = jnp.zeros(n_dofs)
T_curr = T_init
T_history = [T_curr]

print(f"Running {n_steps} steps...")
t_start = time.perf_counter()

for step in range(n_steps):
    fint = compute_stiffness_force(T_curr)
    rhs = (f_ext - fint) * dt
    rhs = rhs.at[fixed_dofs].set(0.0)

    Amat = eqx.Partial(A_times_x, T_prev=T_curr, dt=dt)
    dT, cg_iters = conjugate_gradient(Amat, rhs, atol=1e-10, max_iter=1000)

    T_curr = T_curr + dT

    if step % save_interval == 0:
        T_history.append(T_curr)
        print(f"  Step {step:4d}/{n_steps}: T_min={T_curr.min():.4f}, T_max={T_curr.max():.4f}, CG iterations={cg_iters}")

print(f"Simulation complete in {time.perf_counter() - t_start:.2f} s")

monitor_pt = np.array([0.5 * lx, 0.8 * ly])
monitor_node = int(np.argmin(np.linalg.norm(np.array(mesh.coords) - monitor_pt, axis=1)))
snapshot_times = np.array([0.0] + [(i * save_interval + 1) * dt for i in range(len(T_history)-1)])

monitor_T = np.array([float(T[monitor_node]) for T in T_history])

fig, ax = plt.subplots(figsize=(6, 3))
ax.plot(snapshot_times, monitor_T, marker='o', markersize=3)
ax.set_xlabel("Time [s]")
ax.set_ylabel("Temperature")
ax.set_title(f"Temperature evolution at {mesh.coords[monitor_node]}")
ax.grid(True)
plt.show()