simulation.py

from firedrake import *
import numpy as np
import pyvista as pv
from mpi4py import MPI

def setup_firedrake():
    mesh = UnitIntervalMesh(100)
    # The learning scripts work on a 1D mixed space here, with c and mu stored
    # as the two components of the same Firedrake Function.
    V = FunctionSpace(mesh, "Lagrange", 1)
    W = V * V

    # Deterministic oscillatory initial condition; using the DOF index keeps it
    # aligned with the saved 1D target files from ch_fh.py.
    rng = np.random.default_rng(12)
    u_ic = Function(W, name="Initial_condition")
    num_dofs = u_ic.sub(0).dat.data.shape[0]
    ic = np.zeros((num_dofs, 2))
    ic[:, 0] = [0.5 + 0.2 * sin(pi*i/4) for i in range(num_dofs)]
    ic[:, 1] = 0
    u_ic.sub(0).dat.data[:] = ic[:, 0]
    u_ic.sub(1).dat.data[:] = ic[:, 1]
    return mesh, V, W, u_ic

def solve_one_step(u_old, dfdc_f, u, c, mu, c_test, mu_test, dt, M, lmbda):
    u_ = Function(u.function_space(), name="Solution_Old")
    u_.assign(u_old)
    c_ = u_.sub(0)
    mu_ = u_.sub(1)

    # Semi-implicit CH step:
    #   c - c_old = (dt/2) M Laplacian(mu + mu_old)
    #   mu = learned df/dc - lambda^2 Laplacian(c)
    F0 = (inner(c, c_test) - inner(c_, c_test)) * dx + (dt/2) * M * dot(grad(mu + mu_), grad(c_test)) * dx
    F1 = inner(mu, mu_test) * dx - inner(dfdc_f, mu_test) * dx - lmbda**2 * dot(grad(c), grad(mu_test)) * dx
    F = F0 + F1

    # Solve nonlinear/linear system for u (holds global Function `u`)
    solve(F == 0, u, solver_parameters={
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "mumps"
    })
    return u

def load_target_data(num_timesteps, V, comm=None, rank=None):
    print("Loading target from PVD (pyvista)...")
    c_target_list = []

    for i in range(num_timesteps):
        # The historical master branch reads its generated VTU sequence from a
        # fixed path, assuming the Firedrake mesh and VTU point ordering match.
        reader = pv.get_reader(f"/home/rnp/firedrake/ch_learn/ch_fh/ch_fh_{i}.vtu")
        data = reader.read()
        arr_global = data.point_data["Volume Fraction"].astype(np.float64)

        f = Function(V, name=f"target_{i}")
        f.dat.data[:] = arr_global
        c_target_list.append(f)
    return c_target_list, None, None