Three-dimensional structure to maximize the fundamental frequency

Hello everyone, my code output is incorrect, and the visualizations using ParaView are even more wrong. I hope someone can help me, thank you.

import os
import time
import numpy as np

from mpi4py import MPI
from petsc4py import PETSc
from slepc4py import SLEPc

import dolfinx
from dolfinx import fem, mesh, io
from dolfinx.fem.petsc import assemble_matrix

import ufl
from ufl import inner, sym, grad, tr, Identity, dx, TestFunction, TrialFunction


from MMA import mmasub



# ============================================================
# 1. Parameters: modify here for geometry, mesh, volume fraction, BCs
# ============================================================
comm = MPI.COMM_WORLD
rank = comm.rank

# 3D design domain mesh size. Start with this coarse mesh; refine later.
nelx, nely, nelz = 60, 12, 12

# Element length. If unit is mm, rho_m=2.85e-9 corresponds to tonne/mm^3.
elem_len = 1.0
Lx, Ly, Lz = nelx * elem_len, nely * elem_len, nelz * elem_len

# Material properties: E0 / Emin / nu / rho_m from the 2D code
E0 = 78e3
Emin = 1e-6
nu = 0.30
rho_m = 2.85e-9

# Topology optimization parameters
volfrac = 0.30
q_ramp = 4.0
rmin = 2.5
tol = 0.01
maxloop = 80
objOrder = 1       # which eigenfrequency to maximize, 1 = first
nEig = 6           # number of eigenvalues to compute
save_interval = 10 # output ParaView every this many iterations; 0 = final only

# Output directory
out_dir = "paraview_3d_result"
os.makedirs(out_dir, exist_ok=True) if rank == 0 else None
comm.barrier()

r_helm = rmin / (2.0 * np.sqrt(3.0))
TWO_PI = 2.0 * np.pi


def log(msg: str):
    if rank == 0:
        print(msg, flush=True)


def ramp_ufl(r):
    """RAMP material interpolation, same as in 2D version."""
    return Emin + (1.0 - Emin) * r / (1.0 + q_ramp * (1.0 - r))


# ============================================================
# 2. Create 3D mesh and function spaces
# ============================================================
domain = mesh.create_box(
    comm,
    [np.array([0.0, 0.0, 0.0]), np.array([Lx, Ly, Lz])],
    [nelx, nely, nelz],
    cell_type=mesh.CellType.hexahedron,
)

tdim = domain.topology.dim
fdim = tdim - 1

# Displacement space: 3D vector field
V = fem.functionspace(domain, ("Lagrange", 1, (3,)))

# Design variable / density space: CG1; cell‑wise density for ParaView: DG0
V_rho = fem.functionspace(domain, ("Lagrange", 1))
D0 = fem.functionspace(domain, ("DG", 0))

n_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
n_rho = V_rho.dofmap.index_map.size_local
bs = V.dofmap.index_map_bs


# ============================================================
# 3. Boundary conditions: clamped at both ends x=0 and x=Lx
#    Modify the function below if you want a different support.
# ============================================================
def is_fixed_boundary(x):
    return np.isclose(x[0], 0.0) | np.isclose(x[0], Lx)


bc_facets = mesh.locate_entities_boundary(domain, fdim, is_fixed_boundary)
bc_dofs = fem.locate_dofs_topological(V, fdim, bc_facets)
bc = fem.dirichletbc(np.zeros(3, dtype=PETSc.ScalarType), bc_dofs, V)

constrained = np.array(bc.dof_indices()[0], dtype=np.int32)
all_dofs = np.arange(n_dofs, dtype=np.int32)
free_dofs = np.setdiff1d(all_dofs, constrained).astype(np.int32)
is_free = PETSc.IS().createGeneral(free_dofs, comm=comm)


# ============================================================
# 4. Design variables and 3D elasticity constitutive model
# ============================================================
x_design = fem.Function(V_rho, name="design")
x_design.x.array[:] = volfrac
x_design.x.scatter_forward()

xPhys = fem.Function(V_rho, name="density")
xPhys.x.array[:] = volfrac
xPhys.x.scatter_forward()

xPhys_dg0 = fem.Function(D0, name="density_cell")
u_eig = fem.Function(V, name="mode_shape")


def epsilon(u):
    return sym(grad(u))


def sigma(u, E_eff):
    """3D isotropic linear elasticity constitutive law."""
    mu_ = E_eff / (2.0 * (1.0 + nu))
    lam_ = E_eff * nu / ((1.0 + nu) * (1.0 - 2.0 * nu))
    return lam_ * tr(epsilon(u)) * Identity(3) + 2.0 * mu_ * epsilon(u)


E_eff_ufl = E0 * ramp_ufl(xPhys)
u_tr = TrialFunction(V)
v_te = TestFunction(V)

# Stiffness matrix
K_ufl = inner(sigma(u_tr, E_eff_ufl), epsilon(v_te)) * dx

# Mass matrix = structural self‑mass (no added mass)
M_ufl = rho_m * xPhys * inner(u_tr, v_te) * dx

K_form = fem.form(K_ufl)
M_form = fem.form(M_ufl)

# Residual for eigenvalue sensitivity: d/dx [u^T K u - lambda u^T M u]
lam_const = fem.Constant(domain, PETSc.ScalarType(0.0))
psi = TestFunction(V_rho)
R_eig = (
    inner(sigma(u_eig, E_eff_ufl), epsilon(u_eig)) * dx
    - lam_const * rho_m * xPhys * inner(u_eig, u_eig) * dx
)
dlam_dxPhys_form = fem.form(ufl.derivative(R_eig, xPhys, psi))

# Volume constraint
vol_form = fem.form(xPhys * dx)
dvol_dxPhys_form = fem.form(ufl.derivative(xPhys * dx, xPhys, psi))
total_vol = Lx * Ly * Lz
# ============================================================
# 5. Helmholtz density filter, same as in 2D version
# ============================================================
rho_tr = TrialFunction(V_rho)
rho_te = TestFunction(V_rho)
a_filt = (r_helm**2 * inner(grad(rho_tr), grad(rho_te)) + rho_tr * rho_te) * dx
m_filt = rho_tr * rho_te * dx

A_filt = assemble_matrix(fem.form(a_filt))
A_filt.assemble()
M_filt = assemble_matrix(fem.form(m_filt))
M_filt.assemble()

filt_solver = PETSc.KSP().create(comm)
filt_solver.setOperators(A_filt)
filt_solver.setType("preonly")
_pc = filt_solver.getPC()
_pc.setType("lu")
_pc.setFactorSolverType("mumps")
filt_solver.setFromOptions()

_filt_x, _ = A_filt.createVecs()
_filt_rhs, _ = A_filt.createVecs()
_filt_sol, _ = A_filt.createVecs()


def apply_filter(src_array):
    _filt_x.setArray(src_array)
    _filt_x.assemble()
    M_filt.mult(_filt_x, _filt_rhs)
    filt_solver.solve(_filt_rhs, _filt_sol)
    return _filt_sol.getArray().copy()


def filter_sensitivity(dobj_dxPhys_array):
    _filt_rhs.setArray(dobj_dxPhys_array)
    _filt_rhs.assemble()
    filt_solver.solve(_filt_rhs, _filt_sol)
    M_filt.mult(_filt_sol, _filt_x)
    return _filt_x.getArray().copy()


# ============================================================
# 6. Eigenvalue solver and modal mass normalization
# ============================================================
def solve_eigs(K_petsc, M_petsc, nev):
    eps = SLEPc.EPS().create(comm)
    eps.setOperators(K_petsc, M_petsc)
    eps.setProblemType(SLEPc.EPS.ProblemType.GHEP)
    eps.setDimensions(nev=nev, ncv=PETSc.DECIDE, mpd=PETSc.DECIDE)
    eps.setWhichEigenpairs(SLEPc.EPS.Which.TARGET_REAL)
    eps.setTarget(0.0)

    st = eps.getST()
    st.setType("sinvert")
    ksp = st.getKSP()
    ksp.setType("preonly")
    pc = ksp.getPC()
    pc.setType("lu")
    pc.setFactorSolverType("mumps")

    eps.setFromOptions()
    eps.solve()
    return eps


def modal_mass(M_petsc, phi_full_local):
    """Compute phi^T M phi, to avoid inconsistency due to SLEPc eigenvector scaling."""
    v, _ = M_petsc.createVecs()
    mv = v.duplicate()
    v.setArray(phi_full_local)
    v.assemble()
    M_petsc.mult(v, mv)
    den = float(np.real(v.dot(mv)))
    v.destroy()
    mv.destroy()
    return max(den, 1e-30)


def write_paraview(loop_id, freq_hz):
    """Output ParaView‑readable file."""
    xPhys_dg0.interpolate(xPhys)
    xPhys_dg0.x.scatter_forward()
    u_eig.x.scatter_forward()

    # Keep original naming: itermm
    fname = os.path.join(out_dir, f"topopt3d_itermm_{loop_id:03d}.xdmf")
    with io.XDMFFile(comm, fname, "w") as xdmf:
        xdmf.write_mesh(domain)
        xdmf.write_function(xPhys_dg0, 0.0)
        xdmf.write_function(u_eig, 0.0)

    log(f"  [ParaView] Exported {fname}, current frequency f = {freq_hz:.6g} Hz")


# ============================================================
# 7. MMA initialization
# ============================================================
xval = x_design.x.array.reshape(-1, 1).copy()
xmin_ = 0.001 * np.ones((n_rho, 1))
xmax_ = np.ones((n_rho, 1))
xold1, xold2 = xval.copy(), xval.copy()
low, upp = xmin_.copy(), xmax_.copy()

loop = 0
change = 0.2
obj0 = None

log("========== 3D topology optimization (no added point mass) ==========")
log(f"mesh: {nelx} x {nely} x {nelz}, dofs(local)={n_dofs}, rho_dofs(local)={n_rho}")
log(f"domain: Lx={Lx}, Ly={Ly}, Lz={Lz}")
log(f"boundary: fixed at x=0 and x={Lx}")


# ============================================================
# 8. Main optimization loop
# ============================================================
while change > tol and loop < maxloop:
    t0 = time.time()
    loop += 1

    # Density filtering
    xPhys.x.array[:] = apply_filter(x_design.x.array)
    xPhys.x.array[:] = np.clip(xPhys.x.array, 0.001, 1.0)
    xPhys.x.scatter_forward()

    # Assemble K/M; boundary conditions handled by sub‑matrix on free DOFs
    Kfull = assemble_matrix(K_form, bcs=[])
    Kfull.assemble()
    Mfull = assemble_matrix(M_form, bcs=[])
    Mfull.assemble()

    Kf = Kfull.createSubMatrix(is_free, is_free)
    Mf = Mfull.createSubMatrix(is_free, is_free)

    eps_solver = solve_eigs(Kf, Mf, nEig)
    nconv = eps_solver.getConverged()
    if nconv < nEig:
        log(f"  [Warning] iter {loop}: only {nconv}/{nEig} eigenpairs converged")

    eigvals = np.zeros(nEig)
    eigvecs_full = np.zeros((n_dofs, nEig))
    vr, _ = Kf.createVecs()

    n_read = min(nconv, nEig)
    for i in range(n_read):
        lam_i = eps_solver.getEigenpair(i, vr)
        eigvals[i] = float(np.real(lam_i))
        arr = np.zeros(n_dofs, dtype=PETSc.ScalarType)
        arr[free_dofs] = vr.getArray()
        eigvecs_full[:, i] = np.real(arr)

    idx = objOrder - 1
    if idx >= n_read:
        raise RuntimeError(f"Target mode objOrder={objOrder} did not converge. Refine mesh, check BCs, or increase nEig.")

    lam_obj = max(eigvals[idx], 1e-30)
    omega = np.sqrt(lam_obj)
    freq = omega / TWO_PI

    # Store current target mode shape for sensitivity and ParaView output
    u_eig.x.array[:] = eigvecs_full[:, idx]
    u_eig.x.scatter_forward()

    # Sensitivity of eigenfrequency
    lam_const.value = PETSc.ScalarType(lam_obj)
    den_modal = modal_mass(Mfull, eigvecs_full[:, idx])
    dlam_dxPhys = fem.assemble_vector(dlam_dxPhys_form).array / den_modal
    dfreq_dxPhys = (1.0 / (2.0 * TWO_PI * omega)) * dlam_dxPhys
    dfreq_dx = filter_sensitivity(dfreq_dxPhys)

    # Volume constraint and its sensitivity
    cur_vol = float(fem.assemble_scalar(vol_form))
    g_val = cur_vol / total_vol - volfrac
    dvol_dxPhys = fem.assemble_vector(dvol_dxPhys_form).array
    dg_dx = filter_sensitivity(dvol_dxPhys) / total_vol

    # Normalize objective
    if loop == 1:
        obj0 = freq
        obj_n = 1.0
    else:
        obj_n = freq / obj0

    dobj_dx = dfreq_dx / max(obj0, 1e-30)

    # MMA update: maximize=True as in the 2D code
    xval = x_design.x.array.reshape(-1, 1).copy()
    df0dx = dobj_dx.reshape(-1, 1)
    fval = np.array([[g_val]])
    dfdx = dg_dx.reshape(1, -1)

    m_con, n_var = 1, n_rho
    a0 = 1.0
    a_ = np.zeros((m_con, 1))
    c_ = 10.0 * np.ones((m_con, 1))
    d_ = np.zeros((m_con, 1))

    xmma, *_, low, upp = mmasub(
        m_con, n_var, loop, xval, xmin_, xmax_, xold1, xold2,
        np.array([[obj_n]]), df0dx, fval, dfdx, low, upp,
        a0, a_, c_, d_, maximize=True,
    )

    xnew = np.clip(xmma.ravel(), 0.001, 1.0)
    xold2 = xold1.copy()
    xold1 = xval.copy()
    change = float(np.max(np.abs(xnew - x_design.x.array)))
    x_design.x.array[:] = xnew
    x_design.x.scatter_forward()

    log(
        f" It.:{loop:3d}  Obj(f/f0):{float(obj_n):9.4f} "
        f"Vol.:{cur_vol / total_vol:7.3f}  ch.:{change:7.3f} "
        f"f[Hz]:{freq:10.4f}  ω:{omega:10.3f}  "
        f"mode#:{idx + 1}  modal_mass:{den_modal:.3e}  time:{time.time() - t0:6.2f}s"
    )

    if save_interval and (loop % save_interval == 0):
        write_paraview(loop, freq)

    # Release PETSc objects
    Kfull.destroy()
    Mfull.destroy()
    Kf.destroy()
    Mf.destroy()
    eps_solver.destroy()
    vr.destroy()


# ============================================================
# 9. Output final structure
# ============================================================
xPhys.x.array[:] = apply_filter(x_design.x.array)
xPhys.x.array[:] = np.clip(xPhys.x.array, 0.001, 1.0)
xPhys.x.scatter_forward()

# Solve once more for the final mode shape using the final density
Kfull = assemble_matrix(K_form, bcs=[])
Kfull.assemble()
Mfull = assemble_matrix(M_form, bcs=[])
Mfull.assemble()
Kf = Kfull.createSubMatrix(is_free, is_free)
Mf = Mfull.createSubMatrix(is_free, is_free)
eps_solver = solve_eigs(Kf, Mf, nEig)
nconv = eps_solver.getConverged()
vr, _ = Kf.createVecs()
idx = objOrder - 1
if nconv > idx:
    lam_obj = float(np.real(eps_solver.getEigenpair(idx, vr)))
    final_phi = np.zeros(n_dofs, dtype=PETSc.ScalarType)
    final_phi[free_dofs] = vr.getArray()
    u_eig.x.array[:] = np.real(final_phi)
    u_eig.x.scatter_forward()
    final_freq = np.sqrt(max(lam_obj, 1e-30)) / TWO_PI
else:
    final_freq = 0.0
    log("[Warning] Final mode did not converge, only output final density.")

write_paraview(999, final_freq)

Kfull.destroy()
Mfull.destroy()
Kf.destroy()
Mf.destroy()
eps_solver.destroy()
vr.destroy()
A_filt.destroy()
M_filt.destroy()
is_free.destroy()

log("========== Finished ==========")
log(f"Final file: {out_dir}/topopt3d_itermm_999.xdmf")
log("ParaView: File -> Open -> topopt3d_itermm_999.xdmf")
log("Recommended: density_cell -> Threshold (0.4~1.0); mode_shape -> Warp By Vector")