diff --git a/engibench/problems/thermoelastic3d/__init__.py b/engibench/problems/thermoelastic3d/__init__.py new file mode 100644 index 0000000..85d9353 --- /dev/null +++ b/engibench/problems/thermoelastic3d/__init__.py @@ -0,0 +1,5 @@ +"""thermoelastic3d problem module.""" + +from engibench.problems.thermoelastic3d.v0 import ThermoElastic3D + +__all__ = ["ThermoElastic3D"] diff --git a/engibench/problems/thermoelastic3d/model/__init__.py b/engibench/problems/thermoelastic3d/model/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/engibench/problems/thermoelastic3d/model/fem_matrix_builder.py b/engibench/problems/thermoelastic3d/model/fem_matrix_builder.py new file mode 100644 index 0000000..78cb9de --- /dev/null +++ b/engibench/problems/thermoelastic3d/model/fem_matrix_builder.py @@ -0,0 +1,126 @@ +"""This module assembles the local stiffness matrices the elastic, thermal, and thermoelastic domains.""" + +import numpy as np + + +def fe_melthm_3d(nu: float, e: float, k: float, alpha: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Build 3D Hex8 element matrices for thermo-elasticity. + + Args: + nu (float): Poisson's ratio + e (float): Young's modulus + k (float): Thermal conductivity (isotropic) + alpha (float): Coefficient of thermal expansion + + Returns: + ke : 24x24 mechanical stiffness (Hex8, 3 dof/node) + k_eth : 8x8 thermal conductivity (Hex8, 1 dof/node) + c_ethm: 24x8 coupling mapping nodal temperatures to equivalent mechanical forces + """ + # --- Gauss points (2x2x2) and weights --- + gp = np.array([-1 / np.sqrt(3), 1 / np.sqrt(3)]) + w = np.array([1.0, 1.0]) + + # --- Hex8 shape functions and derivatives in natural coordinates --- + # Node order: (-,-,-), (+,-,-), (+,+,-), (-,+,-), (-,-,+), (+,-,+), (+,+,+), (-,+,+) + xi_nodes = np.array([-1, 1, 1, -1, -1, 1, 1, -1], dtype=float) + et_nodes = np.array([-1, -1, 1, 1, -1, -1, 1, 1], dtype=float) + ze_nodes = np.array([-1, -1, -1, -1, 1, 1, 1, 1], dtype=float) + + def shape_fun_and_derivs(xi: float, eta: float, zeta: float) -> tuple[np.ndarray, np.ndarray]: + """Builds the shape function for the local elements. + + Args: + xi (float): Poisson's ratio + eta (float): Young's modulus + zeta (float): Thermal conductivity + + Returns: + n: (8,) shape functions + d_n_dxi: (8,3) derivatives w.r.t. [xi, eta, zeta] + """ + n = 0.125 * (1 + xi_nodes * xi) * (1 + et_nodes * eta) * (1 + ze_nodes * zeta) + + # Derivatives with respect to natural coords + d_n_dxi = np.zeros((8, 3)) + d_n_dxi[:, 0] = 0.125 * xi_nodes * (1 + et_nodes * eta) * (1 + ze_nodes * zeta) # ∂N/∂xi + d_n_dxi[:, 1] = 0.125 * et_nodes * (1 + xi_nodes * xi) * (1 + ze_nodes * zeta) # ∂N/∂eta + d_n_dxi[:, 2] = 0.125 * ze_nodes * (1 + xi_nodes * xi) * (1 + et_nodes * eta) # ∂N/∂zeta + return n, d_n_dxi + + # --- Geometry & Jacobian --- + # For a unit cube in physical space mapped from [-1,1]^3: + # x = (xi+1)/2, y = (eta+1)/2, z = (zeta+1)/2 -> J = diag(0.5, 0.5, 0.5) + j = np.diag([0.5, 0.5, 0.5]) + det_j = np.linalg.det(j) # = 0.125 + inv_j = np.linalg.inv(j) # = diag(2,2,2) + + # --- Elasticity matrix (Voigt 6x6) for 3D isotropic --- + lam = e * nu / ((1 + nu) * (1 - 2 * nu)) + mu = e / (2 * (1 + nu)) + d = np.array( + [ + [lam + 2 * mu, lam, lam, 0, 0, 0], + [lam, lam + 2 * mu, lam, 0, 0, 0], + [lam, lam + 2 * mu, lam + 2 * mu, 0, 0, 0], # <- typo fixed: [2,2] should be lam+2mu + [0, 0, 0, mu, 0, 0], + [0, 0, 0, 0, mu, 0], + [0, 0, 0, 0, 0, mu], + ], + dtype=float, + ) + # fix typo in [2,1] above: + d[2, 1] = lam + + # Thermal "volumetric" strain direction in Voigt + e_th = np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0]) + + # --- Allocate element matrices --- + ke = np.zeros((24, 24)) + k_eth = np.zeros((8, 8)) + c_ethm = np.zeros((24, 8)) + + # --- Integration loop --- + for i, xi in enumerate(gp): + for j_idx, eta in enumerate(gp): + for kq, zeta in enumerate(gp): + n, d_n_dxi = shape_fun_and_derivs(xi, eta, zeta) + + # Gradients in physical coords: dN_dx = dN_dxi * invJ + d_n_dx = d_n_dxi @ inv_j # (8,3) + + # Build B-matrix (6 x 24) for Hex8, 3 dof/node (u,v,w) + b = np.zeros((6, 24)) + for a in range(8): + ix = 3 * a + dy, dx_, dz = d_n_dx[a, 1], d_n_dx[a, 0], d_n_dx[a, 2] # clarity + # normal strains + b[0, ix + 0] = dx_ + b[1, ix + 1] = dy + b[2, ix + 2] = dz + # shear strains (engineering) + b[3, ix + 0] = dy + b[3, ix + 1] = dx_ + b[4, ix + 1] = dz + b[4, ix + 2] = dy + b[5, ix + 0] = dz + b[5, ix + 2] = dx_ + + # Thermal gradient matrix for conduction: G = grad(N) (3 x 8) + g = d_n_dx.T # rows: [dN/dx; dN/dy; dN/dz] + + # Weight factor + wt = w[i] * w[j_idx] * w[kq] * det_j + + # --- Accumulate --- + # Mechanical stiffness + ke += (b.T @ d @ b) * wt + + # Thermal conductivity (isotropic k * grad N^T grad N) + k_eth += (g.T @ g) * (k * wt) + + # Thermo-mech coupling: columns correspond to nodal temperatures (via N_j) + de_th = d @ (alpha * e_th) # (6,) + c_ethm += (b.T @ de_th[:, None] @ n[None, :]) * wt + + return ke, k_eth, c_ethm diff --git a/engibench/problems/thermoelastic3d/model/fem_model.py b/engibench/problems/thermoelastic3d/model/fem_model.py new file mode 100644 index 0000000..e529dd9 --- /dev/null +++ b/engibench/problems/thermoelastic3d/model/fem_model.py @@ -0,0 +1,420 @@ +"""This module contains the Python implementation of the thermoelastic 3D problem.""" + +from math import ceil +import time +from typing import Any + +import numpy as np +from scipy.sparse import coo_matrix + +from engibench.core import OptiStep +from engibench.problems.thermoelastic3d.model.fem_matrix_builder import fe_melthm_3d +from engibench.problems.thermoelastic3d.model.fem_plotting import plot_fem_3d +from engibench.problems.thermoelastic3d.model.fem_setup import fe_mthm_bc_3d +from engibench.problems.thermoelastic3d.model.linear_solver import solve_spd_with_amg +from engibench.problems.thermoelastic3d.model.mma_subroutine import MMAInputs +from engibench.problems.thermoelastic3d.model.mma_subroutine import mmasub + +SECOND_ITERATION_THRESHOLD = 2 +FIRST_ITERATION_THRESHOLD = 1 +MIN_ITERATIONS = 10 +MAX_ITERATIONS = 200 +UPDATE_THRESHOLD = 0.01 + + +class FeaModel3D: + """Finite Element Analysis (FEA) model for coupled 3D thermoelastic topology optimization.""" + + def __init__(self, *, plot: bool = False, eval_only: bool | None = False) -> None: + """Instantiates a new 3D thermoelastic model. + + Args: + plot: If True, you can hook in your own plotting / volume rendering each iteration. + eval_only: If True, evaluate the given design once and return objective components only. + """ + self.plot = plot + self.eval_only = eval_only + + def get_initial_design(self, volume_fraction: float, nelx: int, nely: int, nelz: int) -> np.ndarray: + """Generates the initial design variable field for the optimization process. + + Args: + volume_fraction (float): The initial volume fraction for the material distribution. + nelx (int): Number of elements in the x-direction. + nely (int): Number of elements in the y-direction. + nelz (int): Number of elements in the z-direction. + + Returns: + np.ndarray: A 3D NumPy array of shape (nely, nelx, nelz) initialized with the given volume fraction. + """ + return volume_fraction * np.ones((nely, nelx, nelz), dtype=float) + + def get_matrices(self, nu: float, e: float, k: float, alpha: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Computes and returns the element matrices required for the structural-thermal analysis. + + Args: + nu (float): Poisson's ratio. + e (float): Young's modulus (modulus of elasticity). + k (float): Thermal conductivity. + alpha (float): Coefficient of thermal expansion. + + Returns: + Tuple[np.ndarray, np.ndarray, np.ndarray]: A tuple containing: + - The stiffness matrix for mechanical analysis. + - The thermal stiffness matrix. + - The coupling matrix for thermal expansion effects. + """ + return fe_melthm_3d(nu, e, k, alpha) + + def get_filter(self, nelx: int, nely: int, nelz: int, rmin: float) -> tuple[coo_matrix, np.ndarray]: + """Constructs a sensitivity filtering matrix to smoothen the design variables. + + The filter helps mitigate checkerboarding issues in topology optimization by averaging + sensitivities over neighboring elements. + + Args: + nelx (int): Number of elements in the x-direction. + nely (int): Number of elements in the y-direction. + nelz (int): Number of elements in the z-direction. + rmin (float): Minimum filter radius. + + Returns: + Tuple[csr_matrix, np.ndarray]: A tuple containing: + - `h` (csr_matrix): A sparse matrix that represents the filtering operation. + - `hs` (np.ndarray): A normalization factor for the filtering. + """ + + def e_index(ix: int, iy: int, iz: int) -> int: + """Returns the global index of an element given its coordinates. + + Args: + ix (int): The global x-index of an element. + iy (int): The global y-index of an element. + iz (int): The global z-index of an element. + + Returns: + g_idx (int): The global index of an element. + """ + return (nely * nelz) * ix + (nelz) * iy + iz + + i_h: list[int] = [] + j_h: list[int] = [] + s_h: list[float] = [] + + rceil = ceil(rmin) - 1 # integer neighborhood radius + + for ix in range(nelx): + ix_min = max(ix - rceil, 0) + ix_max = min(ix + rceil, nelx - 1) + for iy in range(nely): + iy_min = max(iy - rceil, 0) + iy_max = min(iy + rceil, nely - 1) + for iz in range(nelz): + iz_min = max(iz - rceil, 0) + iz_max = min(iz + rceil, nelz - 1) + e1 = e_index(ix, iy, iz) + for jx in range(ix_min, ix_max + 1): + for jy in range(iy_min, iy_max + 1): + for jz in range(iz_min, iz_max + 1): + e2 = e_index(jx, jy, jz) + dx = ix - jx + dy = iy - jy + dz = iz - jz + dist = (dx * dx + dy * dy + dz * dz) ** 0.5 + w = max(0.0, rmin - dist) + if w > 0.0: + i_h.append(e1) + j_h.append(e2) + s_h.append(w) + + n = nelx * nely * nelz + h = coo_matrix((s_h, (i_h, j_h)), shape=(n, n)).tocsr() + hs = np.array(h.sum(axis=1)).ravel() + return h, hs + + def run(self, bcs: dict[str, Any], x_init: np.ndarray | None = None) -> dict[str, Any]: # noqa: PLR0915, C901 + """Run the optimization algorithm for the coupled structural-thermal problem. + + This method performs an iterative optimization procedure that adjusts the design + variables for a coupled structural-thermal problem until convergence criteria are met. + The algorithm utilizes finite element analysis (FEA), sensitivity analysis, and the Method + of Moving Asymptotes (MMA) to update the design. + + Args: + bcs (dict[str, any]): A dictionary containing boundary conditions and problem parameters. + Expected keys include: + - 'volfrac' (float): Target volume fraction. + - 'fixed_elements' (np.ndarray): NxN binary array encoding the location of fixed elements. + - 'force_elements_x' (np.ndarray): NxN binary array encoding the location of loaded elements in the x direction. + - 'force_elements_y' (np.ndarray): NxN binary array encoding the location of loaded elements in the y direction. + - 'force_elements_z' (np.ndarray): NxN binary array encoding the location of loaded elements in the z direction. + - 'heatsink_elements' (np.ndarray): NxN binary array encoding the location of heatsink elements. + - 'weight' (float, optional): Weighting factor between structural and thermal objectives. + x_init (Optional[np.ndarray]): Initial design variable array. If None, the design is generated + using the get_initial_design method. + + Returns: + Dict[str, Any]: A dictionary containing the optimization results. The dictionary includes: + - 'design' (np.ndarray): Final design layout. + - 'bcs' (Dict[str, Any]): The input boundary conditions. + - 'sc' (float): Structural cost component. + - 'tc' (float): Thermal cost component. + - 'vf' (float): Volume fraction error. + If self.eval_only is True, returns a dictionary with keys 'sc', 'tc', and 'vf' only. + """ + # Weighting + w1 = bcs.get("weight", 0.5) # structural + w2 = 1.0 - w1 # thermal + + # Derive mesh sizes from fixed_elements mask (shape: (nelx+1, nely+1, nelz+1)) + fixed_nodes_mask = np.asarray(bcs["fixed_elements"], dtype=bool) + + nxp, nyp, nzp = fixed_nodes_mask.shape # nodes-per-direction + nelx, nely, nelz = nxp - 1, nyp - 1, nzp - 1 + n = nelx * nely * nelz # number of elements + + volfrac = bcs["volfrac"] + + # OptiSteps records + opti_steps = [] + + # 1. Initial design + x = self.get_initial_design(volfrac, nelx, nely, nelz) if x_init is None else x_init.copy() + + # 2. Parameters + penal = 3.0 # SIMP Penalty + rmin = bcs.get("rmin", 1.1) # Minimum feature size + e = 1.0 # Young's modulus + nu = 0.3 # Poisson's ratio + k = 1.0 # Thermal conductivity + alpha = 5e-4 # Thermal strain + tref = 9.267e-4 # Reference temperature + change = 1.0 + iterr = 0 + xmin, xmax = 1e-3, 1.0 + xold1 = x.reshape(n, 1) + xold2 = x.reshape(n, 1) + m = 1 # volume constraint + a0 = 1.0 + a = np.zeros((m, 1)) + c = 10000.0 * np.ones((m, 1)) + d = np.zeros((m, 1)) + low = xmin + upp = xmax + + low_vec = None + upp_vec = None + + # 3. Element matrices + ke, k_eth, c_ethm = self.get_matrices(nu, e, k, alpha) + + # 4. 3D filter + h, hs = self.get_filter(nelx, nely, nelz, rmin) + + # 5. Optimization Loop + change_evol = [] + obj_evol = [] + + f0valm = 0.0 + f0valt = 0.0 + + while change > UPDATE_THRESHOLD or iterr < MIN_ITERATIONS: + iterr += 1 + t0 = time.time() + tcur = t0 + + """ + ABAQUS HOOK: Solve the linear systems using the Abaqus solver + """ + # Forward FEA with BCs & assembly (3D) + res = fe_mthm_bc_3d(nely, nelx, nelz, penal, x, ke, k_eth, c_ethm, tref, bcs) + + kth = res.kth + um = res.um + uth = res.uth + fth = res.fth + d_cthm = res.d_cthm + + if self.plot is True and (iterr % 50 == 0): + plot_fem_3d(bcs, x) + + t_forward = time.time() - tcur + tcur = time.time() + + # Mechanical adjoint: self-adjoint, so the adjoint vector is simply the displacements + lamm = -um + + # Thermal adjoint: K_th * lambda_th = (lamm^T - um^T) * d_cthm - f_th TODO: solver change here + rhs_th = d_cthm.T @ (lamm - um) - fth + lamth = solve_spd_with_amg(kth.tocsr(), rhs_th) + + t_adjoints = time.time() - tcur + tcur = time.time() + + # Sensitivities & objective + f0valm = 0.0 + f0valt = 0.0 + + df0dx_m = np.zeros_like(x) # (nely, nelx, nelz) + df0dx_t = np.zeros_like(x) + df0dx_mat = np.zeros_like(x) + + # Element DOF helper consistent with fe_mthm_bc_3d + def elem_dofs(elx: int, ely: int, elz: int) -> tuple[np.ndarray, np.ndarray]: + """Returns bookkeeping matrices mapping local element degrees of freedom to global degrees of freedom. + + Args: + elx (int): local element x-index. + ely (int): local element y-index. + elz (int): local element z-index. + + Returns: + edof8 (np.ndarray): mapping local element to global thermal degrees of freedom. + edof24 (np.ndarray): mapping local element to global mechanical degrees of freedom. + """ + + def node_id(ix: int, iy: int, iz: int) -> int: + """Returns the global index of an element given its coordinates. + + Args: + ix (int): The global x-index of an element. + iy (int): The global y-index of an element. + iz (int): The global z-index of an element. + + Returns: + g_idx (int): The global index of an element. + """ + return (nely + 1) * (nelz + 1) * ix + (nelz + 1) * iy + iz + + n000 = node_id(elx, ely, elz) + n100 = node_id(elx + 1, ely, elz) + n110 = node_id(elx + 1, ely + 1, elz) + n010 = node_id(elx, ely + 1, elz) + n001 = node_id(elx, ely, elz + 1) + n101 = node_id(elx + 1, ely, elz + 1) + n111 = node_id(elx + 1, ely + 1, elz + 1) + n011 = node_id(elx, ely + 1, elz + 1) + + edof8 = np.array([n000, n100, n110, n010, n001, n101, n111, n011], dtype=int) + edof24 = np.empty(24, dtype=int) + edof24[0::3] = 3 * edof8 + 0 + edof24[1::3] = 3 * edof8 + 1 + edof24[2::3] = 3 * edof8 + 2 + return edof8, edof24 + + # Loop elements (clear & readable; vectorization is possible later) + for elx in range(nelx): + for ely in range(nely): + for elz in range(nelz): + edof8, edof24 = elem_dofs(elx, ely, elz) + + um_e = um[edof24] # (24,) + the = uth[edof8] # (8,) + + lamthe = lamth[edof8] # (8,) + + x_e = x[ely, elx, elz] + x_p = x_e**penal + x_p_minus1 = penal * (x_e ** (penal - 1)) + + f0valm += x_p * (um_e @ (ke @ um_e)) + f0valt += x_p * (the @ (k_eth @ the)) + + df0dx_m[ely, elx, elz] = -x_p_minus1 * um_e.T @ ke @ um_e + df0dx_t[ely, elx, elz] = lamthe.T @ (x_p_minus1 * k_eth @ the) + df0dx_mat[ely, elx, elz] = (df0dx_m[ely, elx, elz] * w1) + (df0dx_t[ely, elx, elz] * w2) + + f0val = (f0valm * w1) + (f0valt * w2) + + if self.eval_only: + vf_error = abs(np.mean(x) - volfrac) + return { + "structural_compliance": float(f0valm), + "thermal_compliance": float(f0valt), + "volume_fraction": vf_error, + } + vf_error = np.abs(np.mean(x) - volfrac) + obj_values = np.array([f0valm, f0valt, vf_error]) + opti_step = OptiStep(obj_values=obj_values, step=iterr) + opti_steps.append(opti_step) + + xval = x.reshape(n, 1) + volconst = np.sum(x) / (volfrac * n) - 1.0 + fval = volconst # scalar + dfdx = np.ones((1, n), dtype=float) / (volfrac * n) + + # Apply 3D filter to sensitivities + df0dx_vec = df0dx_mat.reshape(n, 1) + df0dx_filt = (h @ (xval * df0dx_vec)) / hs[:, None] / np.maximum(1e-3, xval) + + t_sens = time.time() - tcur + tcur = time.time() + + # MMA update + if low_vec is None or upp_vec is None: + upp_vec = np.ones((n,), dtype=float) * upp + low_vec = np.ones((n,), dtype=float) * low + + mmainputs = MMAInputs( + m=1, + n=n, + iterr=iterr, + xval=xval[:, 0], + xmin=xmin, + xmax=xmax, + xold1=xold1, + xold2=xold2, + df0dx=df0dx_filt[:, 0], + fval=fval, + dfdx=dfdx, + low=low_vec, + upp=upp_vec, + a0=a0, + a=a[0], + c=c[0], + d=d[0], + f0val=f0val, + ) + xmma, low_vec, upp_vec = mmasub(mmainputs) + + # reshape low_vec to (-1,) + low_vec = np.squeeze(low_vec) + upp_vec = np.squeeze(upp_vec) + + # Shift history + if iterr > SECOND_ITERATION_THRESHOLD: + xold2 = xold1 + xold1 = xval + elif iterr > FIRST_ITERATION_THRESHOLD: + xold1 = xval + + # Update design + x = xmma.reshape(nely, nelx, nelz) + + # Progress + change = np.max(np.abs(xmma - xold1)) + change_evol.append(change) + obj_evol.append(f0val) + + t_mma = time.time() - tcur + t_total = time.time() - t0 + print( + f" It.: {iterr:4d} Obj.: {f0val:10.4f} " + f"Vol.: {np.sum(x) / (nelx * nely * nelz):6.3f} ch.: {change:6.3f} " + f"|| t_forward:{t_forward:6.3f} + t_adj:{t_adjoints:6.3f} + t_sens:{t_sens:6.3f} + t_mma:{t_mma:6.3f} = {t_total:6.3f}" + ) + + if iterr > MAX_ITERATIONS: + break + + print("3D optimization finished.") + vf_error = abs(np.mean(x) - volfrac) + + return { + "design": x, + "bcs": bcs, + "structural_compliance": float(f0valm), + "thermal_compliance": float(f0valt), + "volume_fraction": vf_error, + "opti_steps": opti_steps, + } diff --git a/engibench/problems/thermoelastic3d/model/fem_plotting.py b/engibench/problems/thermoelastic3d/model/fem_plotting.py new file mode 100644 index 0000000..16783bb --- /dev/null +++ b/engibench/problems/thermoelastic3d/model/fem_plotting.py @@ -0,0 +1,47 @@ +"""Module for plotting a 3D thermoelastic design with napari.""" + +from __future__ import annotations + +import napari +import numpy as np + + +def plot_fem_3d(bcs, design) -> None: + """Plot the multi-physics design along with the boundary conditions. + + Args: + bcs (Dict[str, Any]): Dictionary specifying boundary conditions. Expected keys include: + - "heatsink_elements": Indices for fixed thermal degrees of freedom. + - "fixed_elements": Indices for fixed mechanical degrees of freedom. + - "force_elements_x" (optional): Indices for x-direction force elements. + - "force_elements_y" (optional): Indices for y-direction force elements. + - "force_elements_z" (optional): Indices for z-direction force elements. + design (npt.NDArray): The design array. + + Returns: + None + """ + fixed_elements = bcs.get("fixed_elements", np.zeros_like(design)) + + force_elements_x = bcs.get("force_elements_x", np.zeros_like(design)) + force_elements_y = bcs.get("force_elements_y", np.zeros_like(design)) + force_elements_z = bcs.get("force_elements_z", np.zeros_like(design)) + force_elements = force_elements_x + force_elements_y + force_elements_z + + heatsink_elements = bcs.get("heatsink_elements", np.zeros_like(design)) + + design = np.transpose(design, (2, 0, 1)) + fixed_elements = np.transpose(fixed_elements, (2, 1, 0)) + force_elements = np.transpose(force_elements, (2, 1, 0)) + heatsink_elements = np.transpose(heatsink_elements, (2, 1, 0)) + + viewer = napari.Viewer() + viewer.add_image(design, name="rho", rendering="attenuated_mip") + viewer.add_image(fixed_elements, name="fixed_elements", rendering="attenuated_mip", visible=False, colormap="green") + viewer.add_image(force_elements, name="force_elements", rendering="attenuated_mip", visible=False, colormap="fire") + viewer.add_image( + heatsink_elements, name="heatsink_elements", rendering="attenuated_mip", visible=False, colormap="purple" + ) + + viewer.dims.ndisplay = 3 # switch to 3D view + napari.run() diff --git a/engibench/problems/thermoelastic3d/model/fem_setup.py b/engibench/problems/thermoelastic3d/model/fem_setup.py new file mode 100644 index 0000000..4cdc4ec --- /dev/null +++ b/engibench/problems/thermoelastic3d/model/fem_setup.py @@ -0,0 +1,269 @@ +"""Module for the forward solver for the thermoelastic3d problem.""" + +from dataclasses import dataclass +from typing import Any + +import numpy as np +from scipy.sparse import coo_matrix +from scipy.sparse import csr_matrix + +from engibench.problems.thermoelastic3d.model.linear_solver import solve_spd_with_amg + + +@dataclass +class FEMthmBCResult3D: + """Dataclass encapsulating all output parameters for the fem setup code.""" + + km: csr_matrix # Global mechanical stiffness (ndofm x ndofm) + kth: csr_matrix # Global thermal conductivity (nn x nn) + um: np.ndarray # Mechanical displacements (ndofm,) + uth: np.ndarray # Temperatures (nn,) + fm: np.ndarray # Total mechanical RHS (including thermal) + fth: np.ndarray # Thermal RHS after BCs + d_cthm: coo_matrix # Global thermo-mech coupling derivative (ndofm x nn) + fixeddofsm: np.ndarray # Fixed mechanical DOFs + alldofsm: np.ndarray # All mechanical DOFs + freedofsm: np.ndarray # Free mechanical DOFs + fixeddofsth: np.ndarray # Fixed thermal DOFs + alldofsth: np.ndarray # All thermal DOFs + freedofsth: np.ndarray # Free thermal DOFs + fp: np.ndarray # External mechanical load vector (ndofm,) + + +def fe_mthm_bc_3d( # noqa: PLR0915, PLR0913 + nely: int, + nelx: int, + nelz: int, + penal: float, + x: np.ndarray, # shape: (nely, nelx, nelz) + ke: np.ndarray, # (24, 24) from fe_melthm_3d + k_eth: np.ndarray, # (8, 8) from fe_melthm_3d + c_ethm: np.ndarray, # (24, 8) from fe_melthm_3d + tref: float, + bcs: dict[str, Any], +) -> FEMthmBCResult3D: + """Constructs the finite element model matrices for coupled structural-thermal topology optimization. + + This function assembles the global mechanical and thermal matrices for a coupled + structural-thermal topology optimization problem. It builds the global stiffness (mechanical) + and conductivity (thermal) matrices, applies the prescribed boundary conditions and loads, + and solves the governing equations for both the displacement and temperature fields. + + Args: + nely (int): Number of vertical elements. + nelx (int): Number of horizontal elements. + nelz (int): Number of z-axis elements. + penal (Union[int, float]): SIMP penalty factor used to penalize intermediate densities. + x (np.ndarray): 2D array of design variables (densities) with shape (nely, nelx). + ke (np.ndarray): Element stiffness matrix. + k_eth (np.ndarray): Element conductivity matrix. + c_ethm (np.ndarray): Element coupling matrix between the thermal and mechanical fields. + tref (float): Reference temperature. + bcs (Dict[str, Any]): Dictionary specifying boundary conditions. Expected keys include: + - "heatsink_elements": Indices for fixed thermal degrees of freedom. + - "fixed_elements": Indices for fixed mechanical degrees of freedom. + - "force_elements_x" (optional): Indices for x-direction force elements. + - "force_elements_y" (optional): Indices for y-direction force elements. + - "force_elements_z" (optional): Indices for z-direction force elements. + + Returns: + FEMthmBCResult: Dataclass containing the following fields: + - km (csr_matrix): Global mechanical stiffness matrix. + - kth (csr_matrix): Global thermal conductivity matrix. + - um (np.ndarray): Displacement vector. + - uth (np.ndarray): Temperature vector. + - fm (np.ndarray): Mechanical loading vector. + - fth (np.ndarray): Thermal loading vector. + - d_cthm (coo_matrix): Derivative of the coupling matrix with respect to temperature. + - fixeddofsm (np.ndarray): Array of fixed mechanical degrees of freedom. + - alldofsm (np.ndarray): Array of all mechanical degrees of freedom. + - freedofsm (np.ndarray): Array of free mechanical degrees of freedom. + - fixeddofsth (np.ndarray): Array of fixed thermal degrees of freedom. + - alldofsth (np.ndarray): Array of all thermal degrees of freedom. + - freedofsth (np.ndarray): Array of free thermal degrees of freedom. + - fp (np.ndarray): Force vector used for mechanical loading. + """ + + # --------------------------- + # Helpers + # --------------------------- + def node_index(ix: np.ndarray, iy: np.ndarray, iz: np.ndarray) -> np.ndarray: + """Map (ix,iy,iz) to flat node id with z fastest, then y, then x.""" + return (nely + 1) * (nelz + 1) * ix + (nelz + 1) * iy + iz + + def mask_to_indices(mask3d: np.ndarray) -> np.ndarray: + """Flatten a boolean node mask (nelx+1, nely+1, nelz+1) to node ids.""" + m = np.asarray(mask3d, dtype=bool) + return np.flatnonzero(m.ravel(order="C")).astype(int) + + # Domain weighting + weight = bcs.get("weight", 0.5) + + # --------------------------- + # Dimensions & DOF counts + # --------------------------- + nn = (nelx + 1) * (nely + 1) * (nelz + 1) # thermal DOFs (nodes) + ndofsm = 3 * nn # mechanical DOFs (3 per node) + + # --------------------------- + # THERMAL: BCs (Dirichlet sinks) + # --------------------------- + alldofsth = np.arange(nn) + fixeddofsth = mask_to_indices(bcs["heatsink_elements"]) + freedofsth = np.setdiff1d(alldofsth, fixeddofsth, assume_unique=True) + + # --------------------------- + # Build element connectivity (once) + # --------------------------- + # element indices + ex, ey, ez = np.meshgrid(np.arange(nelx), np.arange(nely), np.arange(nelz), indexing="ij") + ex = ex.ravel() + ey = ey.ravel() + ez = ez.ravel() + nelem = ex.size + + # Corner nodes in Hex8 order: (-,-,-),(+,-,-),(+,+,-),(-,+,-),(-,-,+),(+,-,+),(+,+,+),(-,+,+) + n000 = node_index(ex, ey, ez) + n100 = node_index(ex + 1, ey, ez) + n110 = node_index(ex + 1, ey + 1, ez) + n010 = node_index(ex, ey + 1, ez) + n001 = node_index(ex, ey, ez + 1) + n101 = node_index(ex + 1, ey, ez + 1) + n111 = node_index(ex + 1, ey + 1, ez + 1) + n011 = node_index(ex, ey + 1, ez + 1) + + # Thermal edof (nelem, 8) + edof8 = np.stack([n000, n100, n110, n010, n001, n101, n111, n011], axis=1) + + # Mechanical edof (nelem, 24): for each node append [3n,3n+1,3n+2] + edof24 = np.stack([3 * edof8 + 0, 3 * edof8 + 1, 3 * edof8 + 2], axis=2).reshape(nelem, 24) + + # Penalized factor per element (shape matches your x layout) + # x is indexed as x[ey, ex, ez] to respect (nely, nelx, nelz) + penalized = (x[ey, ex, ez] ** penal).astype(np.float64) + + # --------------------------- + # THERMAL: Assemble Kth + # --------------------------- + # For each element, contribute penalized * k_eth into node-pair positions + kth_row = np.repeat(edof8, 8, axis=1) # (nelem, 64) + kth_col = np.tile(edof8, 8) # (nelem, 64) + kth_blk = penalized[:, None, None] * k_eth # (nelem, 8, 8) + kth_dat = kth_blk.reshape(nelem, 64).ravel() + + kth = coo_matrix((kth_dat, (kth_row.ravel(), kth_col.ravel())), shape=(nn, nn)) + kth = (kth + kth.T) / 2.0 + kth = kth.tolil() + + # Thermal RHS and Dirichlet sinks + fth = np.ones(nn) * tref + tsink = 0.0 + for d in fixeddofsth: + kth.rows[d] = [d] + kth.data[d] = [1.0] + fth[d] = tsink + uth = solve_spd_with_amg(kth.tocsr(), fth) + + # --------------------------- + # MECHANICAL: BCs (clamped supports) + # --------------------------- + fix_nodes = mask_to_indices(bcs["fixed_elements"]) + fixeddofsm_x = 3 * fix_nodes + 0 + fixeddofsm_y = 3 * fix_nodes + 1 + fixeddofsm_z = 3 * fix_nodes + 2 + fixeddofsm = np.concatenate((fixeddofsm_x, fixeddofsm_y, fixeddofsm_z)) + alldofsm = np.arange(ndofsm) + freedofsm = np.setdiff1d(alldofsm, fixeddofsm, assume_unique=True) + + # --------------------------- + # MECHANICAL: Assemble Km and d_cthm + # --------------------------- + # Stiffness + km_row = np.repeat(edof24, 24, axis=1) # (nelem, 576) + km_col = np.tile(edof24, 24) # (nelem, 576) + km_blk = penalized[:, None, None] * ke # (nelem, 24, 24) + km_dat = km_blk.reshape(nelem, 576).ravel() + km = coo_matrix((km_dat, (km_row.ravel(), km_col.ravel())), shape=(ndofsm, ndofsm)) + + # Coupling derivative (maps θ to mechanical forces) + d_c_row = np.repeat(edof24, 8, axis=1) # (nelem, 24*8) + d_c_col = np.tile(edof8, 24) # (nelem, 24*8) + d_c_blk = penalized[:, None, None] * c_ethm # (nelem, 24, 8) + d_c_dat = d_c_blk.reshape(nelem, 24 * 8).ravel() + d_cthm = coo_matrix((d_c_dat, (d_c_row.ravel(), d_c_col.ravel())), shape=(ndofsm, nn)) + + # Symmetrize mechanical stiffness and convert + km = (km + km.T) / 2.0 + km = km.tocsr() + + # --------------------------- + # THERMO-MECHANICAL LOAD (equivalent forces) + # --------------------------- + # For each element: diff_e = uth[edof8] - tref (nelem, 8) + diff = uth[edof8] - tref + thermal_e = (c_ethm @ diff.T).T # (nelem, 24) + thermal_e *= penalized[:, None] + + # Accumulate to global vector via bincount + feps = np.bincount(edof24.ravel(), weights=thermal_e.ravel(), minlength=ndofsm).astype(np.float64) + + # --------------------------- + # EXTERNAL MECHANICAL LOADS + # --------------------------- + fp = np.zeros(ndofsm, dtype=np.float64) + + def add_load(mask_key: str, comp: int, mag: float = 0.5) -> None: + """Adds a load to the mechanical force matrix. + + Args: + mask_key (str): the key from which the load is added. + comp (int): the compartment number. + mag (float): the magnitude of the load. + + Returns: + None + """ + if mask_key in bcs and bcs[mask_key] is not None: + nodes = mask_to_indices(bcs[mask_key]) + dofs = 3 * nodes + comp + fp[dofs] += mag + + add_load("force_elements_x", 0) + add_load("force_elements_y", 1) + add_load("force_elements_z", 2) + + # Choose RHS by weight (pure structural / pure thermal / mixed) + if weight == 1.0: + fm = fp + elif weight == 0.0: + fm = feps + else: + fm = fp + feps + + # --------------------------- + # SOLVES + # --------------------------- + um = np.zeros(ndofsm, dtype=np.float64) + if freedofsm.size > 0: + um_free = solve_spd_with_amg(km[freedofsm, :][:, freedofsm].tocsr(), fm[freedofsm]) + um = np.zeros(ndofsm) + um[freedofsm] = um_free + if fixeddofsm.size > 0: + um[fixeddofsm] = 0.0 + + return FEMthmBCResult3D( + km=km, + kth=kth, + um=um, + uth=uth, + fm=fm, + fth=fth, + d_cthm=d_cthm, + fixeddofsm=fixeddofsm, + alldofsm=alldofsm, + freedofsm=freedofsm, + fixeddofsth=fixeddofsth, + alldofsth=alldofsth, + freedofsth=freedofsth, + fp=fp, + ) diff --git a/engibench/problems/thermoelastic3d/model/linear_solver.py b/engibench/problems/thermoelastic3d/model/linear_solver.py new file mode 100644 index 0000000..46eecfc --- /dev/null +++ b/engibench/problems/thermoelastic3d/model/linear_solver.py @@ -0,0 +1,57 @@ +"""Module for solving sparse linear systems.""" + +import numpy as np +from numpy.typing import NDArray +import pyamg +from scipy.sparse import csr_matrix +from scipy.sparse.linalg import cg +from scipy.sparse.linalg import spsolve + + +def solve_spd_with_amg( + a: csr_matrix, b: NDArray[np.float64], tol: float = 1e-8, maxiter: int | None = None +) -> NDArray[np.float64]: + """Solves a symmetric positive-definite linear system using AMG-preconditioned CG. + + This function applies Algebraic Multigrid (AMG) as a right preconditioner for the + Conjugate Gradient (CG) method to solve large sparse SPD linear systems. If CG fails + to converge within the prescribed tolerance, the solution is refined using additional + AMG V-cycles. + + Args: + a (csr_matrix): Sparse system matrix in compressed sparse row (CSR) format. + Must be symmetric positive-definite for CG to converge. + b (np.ndarray): Right-hand-side vector of shape (n,). + tol (float): Relative tolerance used for CG convergence. + maxiter (Optional[int]): Maximum number of CG iterations. If None, the default + SciPy value is used. + + Returns: + np.ndarray: The computed solution vector of shape (n,). + """ + ml = pyamg.smoothed_aggregation_solver(a) + m = ml.aspreconditioner() + + x, info = cg(a, b, rtol=tol, M=m, maxiter=maxiter) + + if info != 0: + x = ml.solve(b, x0=x, tol=tol) + + return x + + +def solve_with_spsolve(a: csr_matrix, b: NDArray[np.float64]) -> NDArray[np.float64]: + """Solves a sparse linear system using SciPy's direct sparse solver. + + This function computes the solution of a sparse linear system using a direct + factorization method (SuperLU via SciPy). It supports general sparse matrices, + including non-symmetric or indefinite systems. + + Args: + a (csr_matrix): Sparse system matrix in CSR format. + b (np.ndarray): Right-hand-side vector of shape (n,). + + Returns: + np.ndarray: The computed solution vector of shape (n,). + """ + return spsolve(a, b) diff --git a/engibench/problems/thermoelastic3d/model/mma_subroutine.py b/engibench/problems/thermoelastic3d/model/mma_subroutine.py new file mode 100644 index 0000000..24d1f01 --- /dev/null +++ b/engibench/problems/thermoelastic3d/model/mma_subroutine.py @@ -0,0 +1,135 @@ +"""This module contains the MMA subroutine used in the thermoelastic3d problem.""" + +from dataclasses import dataclass + +from mmapy import mmasub as external_mmasub +import numpy as np +from numpy.typing import NDArray + +RESIDUAL_MAX_VAL = 0.9 +ITERATION_MAX = 500 +ITERATION_MAX_SMALL = 50 +ITERATION_ASYM_MAX = 2.5 + + +@dataclass(frozen=True) +class MMAInputs: + """Dataclass encapsulating all input parameters for the MMA subroutine.""" + + m: int + """The number of constraints""" + n: int + """The number of design variables""" + iterr: int + """The current iteration number""" + xval: NDArray[np.float64] + """The flattened array of design variables: shape (n,)""" + xmin: float + """The lower bounds on the design variables""" + xmax: float + """The upper bounds on the design variables""" + xold1: NDArray[np.float64] + """The previous design variables at iteration k-1: shape (n, 1)""" + xold2: NDArray[np.float64] + """The design variables at iteration k-2: shape (n, 1)""" + df0dx: NDArray[np.float64] + """The gradients of the objective function at xval: shape (n,)""" + fval: NDArray[np.float64] + """The value of the constraint functions evaluated at xval: shape (m,)""" + dfdx: NDArray[np.float64] + """The gradients of the constraint functions at xval: shape (m, n)""" + low: NDArray[np.float64] + """The lower asymptotes from the previous iteration: shape (n,)""" + upp: NDArray[np.float64] + """The upper asymptotes from the previous iteration: shape (n,)""" + a0: float + """The constants a_0 in the term a_0*z""" + a: NDArray[np.float64] + """the constants a_i in the terms a_i*z""" + c: NDArray[np.float64] + """the constants c_i in the terms c_i*y_i""" + d: NDArray[np.float64] + """the constants d_i in the terms 0.5*d_i*(y_i)^2""" + f0val: float = 0.0 + """The value of the objective function at xval""" + + +def mmasub(inputs: MMAInputs) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]: + """Perform one MMA iteration to solve a nonlinear programming problem using the GCMMA-MMA-Python library. + + Minimize: + f_0(x) + a_0 * z + sum(c_i * y_i + 0.5 * d_i * (y_i)^2) + + Subject to: + f_i(x) - a_i * z - y_i <= 0, i = 1,...,m + xmin_j <= x_j <= xmax_j, j = 1,...,n + z >= 0, y_i >= 0, i = 1,...,m + + Parameters: + inputs (MMAInputs): A dataclass encapsulating all input parameters. + + Returns: + xmma (NDArray[np.float64]): the updated design variables. + low (NDArray[np.float64]): the new lower bounds on the design variables. + upp (NDArray[np.float64]): the new upper bounds on the design variables. + """ + # Unpack parameters from the dataclass. + m = int(inputs.m) + n = int(inputs.n) + iterr = int(inputs.iterr) + xval = np.expand_dims(inputs.xval, axis=1) + xmin = np.full((n, 1), inputs.xmin) + xmax = np.full((n, 1), inputs.xmax) + xold1 = inputs.xold1 + xold2 = inputs.xold2 + f0val = inputs.f0val + df0dx = np.expand_dims(inputs.df0dx, axis=1) + fval = inputs.fval + dfdx = inputs.dfdx + low = np.expand_dims(inputs.low, axis=1) + upp = np.expand_dims(inputs.upp, axis=1) + a0 = inputs.a0 + a = np.expand_dims(inputs.a, axis=1) + c = np.expand_dims(inputs.c, axis=1) + d = np.expand_dims(inputs.d, axis=1) + + # MMA parameters + raa0 = 1e-5 + move = 0.5 # was 1.0 + albefa = 0.1 + asyinit = 0.01 + asyincr = 1.2 + asydecr = 0.7 + asymax = 0.2 + asymin = 0.01 + + xmma, _, _, _, _, _, _, _, _, low, up = external_mmasub( + m, + n, + iterr, + xval, + xmin, + xmax, + xold1, + xold2, + f0val, + df0dx, + fval, + dfdx, + low, + upp, + a0, + a, + c, + d, + move=move, + asyinit=asyinit, + asydecr=asydecr, + asyincr=asyincr, + asymin=asymin, + asymax=asymax, + raa0=raa0, + albefa=albefa, + ) + + return xmma, low, up diff --git a/engibench/problems/thermoelastic3d/v0.py b/engibench/problems/thermoelastic3d/v0.py new file mode 100644 index 0000000..5c708ff --- /dev/null +++ b/engibench/problems/thermoelastic3d/v0.py @@ -0,0 +1,268 @@ +"""Thermo Elastic 3D Problem.""" + +import dataclasses +from dataclasses import dataclass +from dataclasses import field +from typing import Annotated, Any, ClassVar + +from gymnasium import spaces +import napari +import numpy as np +import numpy.typing as npt + +from engibench.constraint import bounded +from engibench.constraint import constraint +from engibench.constraint import IMPL +from engibench.constraint import THEORY +from engibench.core import ObjectiveDirection +from engibench.core import OptiStep +from engibench.core import Problem +from engibench.problems.thermoelastic3d.model.fem_model import FeaModel3D + +NELX = NELY = NELZ = 16 +FIXED_ELEMENTS = np.zeros((NELX + 1, NELY + 1, NELZ + 1), dtype=int) +FIXED_ELEMENTS[0, 0, 0] = 1 +FIXED_ELEMENTS[0, -1, -1] = 1 +FIXED_ELEMENTS[0, -1, 0] = 1 +FORCE_ELEMENTS_X = np.zeros((NELX + 1, NELY + 1, NELZ + 1), dtype=int) +FORCE_ELEMENTS_X[-1, -1, -1] = 1 +FORCE_ELEMENTS_Y = np.zeros((NELX + 1, NELY + 1, NELZ + 1), dtype=int) +FORCE_ELEMENTS_Y[-1, -1, -1] = 1 +FORCE_ELEMENTS_Z = np.zeros((NELX + 1, NELY + 1, NELZ + 1), dtype=int) +FORCE_ELEMENTS_Z[-1, -1, -1] = 1 +HEATSINK_ELEMENTS = np.zeros((NELX + 1, NELY + 1, NELZ + 1), dtype=int) +HEATSINK_ELEMENTS[-1, -1, 0] = 1 + + +class ThermoElastic3D(Problem[npt.NDArray]): + r"""Truss 3D integer optimization problem. + + ## Problem Description + This is 3D topology optimization problem for minimizing weakly coupled thermo-elastic compliance subject to boundary conditions and a volume fraction constraint. + + ## Motivation + As articulated in their respective sections, both the Beams2D and HeatConduction2D problems found in the EngiBench library are fundamental engineering design problems that have historically served as benchmarks for the development and testing of optimization methods. + While their relevance is supported by needs in real engineering design scenarios (aerospace, automotive, consumer electronics, etc...), their mono-domain nature ignores the reality that coupling between domains exists, and should be accounted for in scenarios where performance in one domain significantly impacts performance in another. + To address this distinction, a multi-physics topology optimization problem is developed that captures the coupling between structural and thermal domains in three dimensions. + + ## Design space + This multi-physics topology optimization problem is governed by linear elasticity and steady-state heat conduction with a one-way coupling from the thermal domain to the elastic domain. + The problem is defined over a cube 3D domain, where load elements and support elements are placed along the boundary to define a unique elastic condition. + Similarly, heatsink elements are placed along the boundary to define a unique thermal condition. + The design space is then defined by a 3D array representing density values (parameterized by DesignSpace = [0,1]^{nelx x nely x nelz}, where nelx, nely, and nelz denote the x, y, and z dimensions respectively). + + ## Objectives + The objective of this problem is to minimize total compliance C under a volume fraction constraint V by placing a thermally conductive material. + Total compliance is defined as the sum of thermal compliance and structural compliance. + + ## Conditions + Problem conditions are defined by creating a python dict with the following info: + - `fixed_elements`: Encodes a binary NxNxN matrix of the structurally fixed elements in the domain. + - `force_elements_x`: Encodes a binary NxNxN matrix specifying elements that have a structural load in the x-direction. + - `force_elements_y`: Encodes a binary NxNxN matrix specifying elements that have a structural load in the y-direction. + - `force_elements_z`: Encodes a binary NxNxN matrix specifying elements that have a structural load in the z-direction. + - `heatsink_elements`: Encodes a binary NxNxN matrix specifying elements that have a heat sink. + - `volfrac`: Encodes the target volume fraction for the volume fraction constraint. + - `rmin`: Encodes the filter size used in the optimization routine. + - `weight`: Allows one to control which objective is optimized for. 1.0 Is pure structural optimization, while 0.0 is pure thermal optimization. + + ## Simulator + The simulation code is based on a Python adaptation of the popular 88-line topology optimization code, modified to handle the thermal domain in addition to thermal-elastic coupling. + Optimization is conducted by reformulating the integer optimization problem as a continuous one (leveraging a SIMP approach), where a density filtering approach is used to prevent checkerboard-like artifacts. + The optimization process itself operates by calculating the sensitivities of the design variables with respect to total compliance (done efficiently using the Adjoint method), calculating the sensitivities of the design variables with respect to the constraint value, and then updating the design variables by solving a convex-linear subproblem and taking a small step (using the method of moving asymptotes). + The optimization loop terminates when either an upper bound of the number of iterations has been reached or if the magnitude of the gradient update is below some threshold. + + ## Dataset + The dataset linked to this problem is on huggingface [Hugging Face Datasets Hub](https://huggingface.co/datasets/IDEALLab/thermoelastic_3d_v0). + This dataset contains a set of 100 optimized thermoelastic designs in a 16x16x16 domain, where each design is optimized for a unique set of conditions. + Each datapoint's conditions are randomly generated by arbitrarily placing: a single loaded element along the bottom boundary, two fixed elements (fixed in the x, y, and z direction) along the left and top boundary, and heatsink elements along the right boundary. + Furthermore, values for the volume fraction constraint are randomly selected in the range $[0.2, 0.5]$. + + Relevant datapoint fields include: + - `optimal_design`: An optimized design for the set of boundary conditions + - `fixed_elements`: Encodes a binary NxNxN matrix of the structurally fixed elements in the domain. + - `force_elements_x`: Encodes a binary NxNxN matrix specifying elements that have a structural load in the x-direction. + - `force_elements_y`: Encodes a binary NxNxN matrix specifying elements that have a structural load in the y-direction. + - `force_elements_z`: Encodes a binary NxNxN matrix specifying elements that have a structural load in the z-direction. + - `heatsink_elements`: Encodes a binary NxNxN matrix specifying elements that have a heat sink. + - `volume_fraction`: The volume fraction value of the optimized design + - `structural_compliance`: The structural compliance of the optimized design + - `thermal_compliance`: The thermal compliance of the optimized design + - `nelx`: The number of elements in the x-direction + - `nely`: The number of elements in the y-direction + - `nelz`: The number of elements in the z-direction + - `volfrac`: The volume fraction target of the optimized design + - `rmin`: The filter size used in the optimization routine + - `weight`: The domain weighting used in the optimization routine + + ## Lead + Gabriel Apaza @gapaza + """ + + version = 0 + objectives: tuple[tuple[str, ObjectiveDirection], ...] = ( + ("structural_compliance", ObjectiveDirection.MINIMIZE), + ("thermal_compliance", ObjectiveDirection.MINIMIZE), + ("volume_fraction", ObjectiveDirection.MINIMIZE), + ) + + @dataclass + class Conditions: + """Conditions.""" + + fixed_elements: Annotated[npt.NDArray[np.int64], bounded(lower=0.0, upper=1.0).category(THEORY)] = field( + default_factory=lambda: FIXED_ELEMENTS + ) + force_elements_x: Annotated[npt.NDArray[np.int64], bounded(lower=0.0, upper=1.0).category(THEORY)] = field( + default_factory=lambda: FORCE_ELEMENTS_X + ) + force_elements_y: Annotated[npt.NDArray[np.int64], bounded(lower=0.0, upper=1.0).category(THEORY)] = field( + default_factory=lambda: FORCE_ELEMENTS_Y + ) + force_elements_z: Annotated[npt.NDArray[np.int64], bounded(lower=0.0, upper=1.0).category(THEORY)] = field( + default_factory=lambda: FORCE_ELEMENTS_Z + ) + heatsink_elements: Annotated[npt.NDArray[np.int64], bounded(lower=0.0, upper=1.0).category(THEORY)] = field( + default_factory=lambda: HEATSINK_ELEMENTS + ) + volfrac: Annotated[float, bounded(lower=0.0, upper=1.0).category(THEORY)] = 0.3 + rmin: Annotated[ + float, bounded(lower=1.0).category(THEORY), bounded(lower=0.0, upper=3.0).warning().category(IMPL) + ] = 1.5 + weight: Annotated[float, bounded(lower=0.0, upper=1.0).category(THEORY)] = 0.5 + """1.0 for pure structural, 0.0 for pure thermal""" + + conditions = Conditions() + design_space = spaces.Box(low=0.0, high=1.0, shape=(NELX, NELY, NELZ), dtype=np.float32) + dataset_id = "IDEALLab/thermoelastic_3d_v0" + container_id = None + + @dataclass + class Config(Conditions): + """Structured representation of configuration parameters for a numerical computation.""" + + nelx: ClassVar[Annotated[int, bounded(lower=1).category(THEORY)]] = NELX + nely: ClassVar[Annotated[int, bounded(lower=1).category(THEORY)]] = NELY + nelz: ClassVar[Annotated[int, bounded(lower=1).category(THEORY)]] = NELZ + + @constraint + @staticmethod + def rmin_bound(rmin: float, nelx: int, nely: int, nelz: int) -> None: + """Constraint for rmin ∈ (0.0, max{ nelx, nely, nelz }].""" + assert 0.0 < rmin <= max(nelx, nely, nelz), f"Params.rmin: {rmin} ∉ (0, max(nelx, nely, nelz)]" + + def __init__(self, seed: int = 0) -> None: + """Initializes the thermoelastic3D problem. + + Args: + seed (int): The random seed for the problem. + """ + super().__init__(seed=seed) + + def reset(self, seed: int | None = None) -> None: + """Resets the simulator and numpy random to a given seed. + + Args: + seed (int, optional): The seed to reset to. If None, a random seed is used. + """ + super().reset(seed) + + def simulate(self, design: npt.NDArray, config: dict[str, Any] | None = None) -> npt.NDArray: + """Simulates the performance of a design topology. + + Args: + design (np.ndarray): The design to simulate. + config (dict): A dictionary with configuration (e.g., boundary conditions, filenames) for the simulation. + + Returns: + dict: The performance of the design - each entry of the dict corresponds to a named objective value. + """ + boundary_dict = dataclasses.asdict(self.conditions) + for key, value in (config or {}).items(): + if key in boundary_dict: + if isinstance(value, list): + boundary_dict[key] = np.array(value) + else: + boundary_dict[key] = value + + results = FeaModel3D(plot=False, eval_only=True).run(boundary_dict, x_init=design) + return np.array([results["structural_compliance"], results["thermal_compliance"], results["volume_fraction"]]) + + def optimize( + self, starting_point: npt.NDArray, config: dict[str, Any] | None = None + ) -> tuple[np.ndarray, list[OptiStep]]: + """Optimizes a topology for the current problem. Note that an appropriate starting_point for the optimization is defined by a uniform material distribution equal to the volume fraction constraint. + + Args: + starting_point (np.ndarray): The starting point for the optimization. + config (dict): A dictionary with configuration (e.g., boundary conditions, filenames) for the optimization. + + Returns: + Tuple[np.ndarray, dict]: The optimized design and its performance. + """ + boundary_dict = dataclasses.asdict(self.conditions) + boundary_dict.update({k: v for k, v in (config or {}).items() if k in boundary_dict}) + results = FeaModel3D(plot=False, eval_only=False).run(boundary_dict, x_init=starting_point) + design = np.array(results["design"]).astype(np.float32) + opti_steps = results["opti_steps"] + return design, opti_steps + + def render(self, design: np.ndarray, *, open_window: bool = False) -> np.ndarray: + """Renders the design in a human-readable format. + + Args: + design (np.ndarray): The design to render. + open_window (bool): If True, opens a window with the rendered design. + + Returns: + fig (np.ndarray): The rendered design. + """ + design = np.array(design) + design = np.transpose(design, (2, 0, 1)) + + viewer = napari.Viewer() + viewer.add_image(design, name="rho", rendering="attenuated_mip") + viewer.dims.ndisplay = 3 # switch to 3D view + if open_window is True: + napari.run() + return viewer.export_figure(flash=False) + + def random_design(self, dataset_split: str = "train", design_key: str = "optimal_design") -> tuple[npt.NDArray, int]: + """Samples a valid random design. + + Args: + dataset_split (str): The key for the dataset to sample from. + design_key (str): The key for the design to sample from. + + Returns: + Tuple of: + np.ndarray: The valid random design. + int: The random index selected. + """ + rnd = self.np_random.integers(low=0, high=len(self.dataset[dataset_split]), dtype=int) + return np.array(self.dataset[dataset_split][design_key][rnd]), rnd + + +if __name__ == "__main__": + # --- Create a new problem + problem = ThermoElastic3D(seed=0) + + # --- Load the problem dataset + dataset = problem.dataset + first_item = dataset["train"][0] + first_item_design = np.array(first_item["optimal_design"]) + problem.render(first_item_design, open_window=True) + + # --- Render the design + design, _ = problem.random_design() + problem.render(design, open_window=True) + + # --- Optimize a design --- + design = 0.2 * np.ones((NELX, NELY, NELZ), dtype=float) + design, objectives = problem.optimize(design) + problem.render(design, open_window=True) + + # --- Evaluate a design --- + problem.reset(seed=0) + design, _ = problem.random_design() + print(problem.simulate(design)) diff --git a/pyproject.toml b/pyproject.toml index 85893e0..de85e3d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,10 +35,11 @@ airfoil = ["sqlitedict>=1.6"] # for pyopt_history.py electronics = ["networkx >= 3.2.1"] beams2d = ["cvxopt >= 1.3.2", "seaborn", "scipy"] thermoelastic2d = ["cvxopt >= 1.3.2", "mmapy >= 0.3.0"] +thermoelastic3d = ["cvxopt >= 1.3.2", "mmapy >= 0.3.0", "napari", "pyamg"] photonics2d = ["ceviche >= 0.1.3"] all = [ # All dependencies above - "engibench[airfoil,beams2d,thermoelastic2d,photonics2d,electronics]" + "engibench[airfoil,beams2d,thermoelastic2d,thermoelastic3d,photonics2d,electronics]" ] doc = [ @@ -275,5 +276,7 @@ module = [ "mpi4py", "multipoint", "pygeo", + "napari", + "pyamg" ] ignore_missing_imports = true