Note
Go to the end to download the full example code.
Shape optimization: minimize surface area under volume constraint#
Use graphlow to deform a volumetric mesh with a small MLP and gradient descent: minimize surface area while keeping total volume fixed and penalizing large deformations. The result approaches a sphere (fixed volume, minimal area).
This example walks through: loading a mesh, building a deformation model, defining a differentiable cost, and running the optimization loop. Helper functions are defined at module level with clear roles.
Imports and constants#
import itertools
import pathlib
from collections.abc import Callable
import numpy as np
import phlower_tensor as pt
import pyvista as pv
import torch
import graphlow
from graphlow.utils.topology_helper import TopologyDim
# Optimization weights and hyperparameters
WEIGHT_VOLUME_CONSTRAINT = 10.0
WEIGHT_NORM_CONSTRAINT = 1.0
MLP_HIDDEN = 64
DEFORMATION_FACTOR = 1.0
LEARNING_RATE = 1e-2
# Reject step if any cell volume < this * (total / n_cells)
TINY_CELL_RATIO = 1e-3
Step 1: Deformation model (MLP)#
The mesh is deformed by adding a learned delta to each point: delta = MLP(points). We use a small 2-layer MLP with tanh and optional bias.
def make_deformation_mlp(
hidden_size: int,
use_bias: bool = False,
) -> tuple[
Callable[[pt.PhlowerTensor, float], pt.PhlowerTensor],
list[torch.nn.Parameter],
]:
"""
Create the deformation MLP and its trainable parameters.
Returns
-------
deform : collections.abc.Callable
Function that maps points to deformed points.
params : list[torch.nn.Parameter]
Trainable parameters for the optimizer.
"""
scale = hidden_size**0.5
w1 = torch.nn.Parameter(torch.randn(3, hidden_size) / scale)
w2 = torch.nn.Parameter(torch.randn(hidden_size, 3) / scale)
if use_bias:
b1 = torch.nn.Parameter(torch.randn(hidden_size) / hidden_size)
b2 = torch.nn.Parameter(torch.randn(3) / 3)
params = [w1, b1, w2, b2]
else:
b1 = 0.0
b2 = 0.0
params = [w1, w2]
activation = torch.nn.Identity()
def deform(
points: pt.PhlowerTensor, factor: float = 1.0
) -> pt.PhlowerTensor:
hidden = torch.tanh(points.to_tensor() @ w1 + b1)
delta = pt.phlower_tensor(
activation(hidden @ w2 + b2),
dimension=points.dimension,
)
return points + delta * factor
return deform, params
Step 2: Cost function#
Differentiable cost: normalized surface area + volume penalty + deformation
regularizer. Returns None for cost/area/volume when any cell volume
becomes too small (step rejected).
def compute_cost(
mesh: graphlow.TensorMesh[pt.PhlowerTensor],
surface: graphlow.TensorMesh[pt.PhlowerTensor],
initial_points: pt.PhlowerTensor,
initial_total_volume: pt.PhlowerTensor,
surface_initial_total_area: pt.PhlowerTensor,
deformed_points: pt.PhlowerTensor,
weight_volume: float,
weight_norm: float,
tiny_cell_ratio: float,
) -> tuple[
pt.PhlowerTensor | None,
pt.PhlowerTensor | None,
pt.PhlowerTensor | None,
]:
"""
Evaluate cost and optional area/volume for the current deformation.
Updates mesh and surface points in place from deformed_points, then
computes cell volumes and surface face areas. If any cell volume is
below a small fraction of the mean, returns ``(None, None, None)`` to
signal step rejection.
Returns
-------
cost : pt.PhlowerTensor or None
Scalar objective value, or None when the step is rejected.
total_area : pt.PhlowerTensor or None
Total surface area, or None when the step is rejected.
total_volume : pt.PhlowerTensor or None
Total volume, or None when the step is rejected.
"""
mesh.points = deformed_points
volumes = mesh.geometry.cell_volumes()
total_volume = torch.sum(volumes)
surface.points = surface.gather_parent_point_data(deformed_points)
areas = surface.geometry.face_areas()
total_area = torch.sum(areas)
deformation = deformed_points - initial_points
min_vol = initial_total_volume.to_tensor() / mesh.n_cells * tiny_cell_ratio
if torch.any(volumes.to_tensor() < min_vol):
return None, None, None
cost_area = total_area / surface_initial_total_area
volume_penalty = (
(total_volume - initial_total_volume) / initial_total_volume
) ** 2
scale_vol = initial_total_volume ** (2 / 3)
norm_penalty = torch.exp(torch.mean(deformation * deformation / scale_vol))
cost = (
cost_area + weight_volume * volume_penalty + weight_norm * norm_penalty
)
return cost, total_area, total_volume
Step 3: Build example mesh (hex grid)#
Create a small hexahedral grid. You can replace with your own mesh
(e.g. graphlow.read() or graphlow.from_pyvista()).
def generate_hex_grid(ni: int, nj: int, nk: int) -> pv.UnstructuredGrid:
"""
Build a hexahedral grid for the example.
Returns
-------
pv.UnstructuredGrid
Hexahedral grid.
"""
n_cells = (ni - 1) * (nj - 1) * (nk - 1)
x = np.arange(ni, dtype=np.float32)
y = np.arange(nj, dtype=np.float32)
z = np.arange(nk, dtype=np.float32)
X, Y, Z = np.meshgrid(x, y, z, indexing="ij")
points = np.array([X.ravel(), Y.ravel(), Z.ravel()]).T
indices = np.arange(ni * nj * nk).reshape(ni, nj, nk)
hex_lists = []
for k in range(nk - 1):
for j in range(nj - 1):
for i in range(ni - 1):
v0 = indices[i, j, k]
v1 = indices[i + 1, j, k]
v2 = indices[i + 1, j + 1, k]
v3 = indices[i, j + 1, k]
v4 = indices[i, j, k + 1]
v5 = indices[i + 1, j, k + 1]
v6 = indices[i + 1, j + 1, k + 1]
v7 = indices[i, j + 1, k + 1]
hex_lists.append([8, v0, v1, v2, v3, v4, v5, v6, v7])
flat = list(itertools.chain.from_iterable(hex_lists))
celltypes = np.full(n_cells, pv.CellType.HEXAHEDRON)
return pv.UnstructuredGrid(flat, celltypes, points)
Step 4: Visualization#
Color by vertex index (turbo cmap) so that vertex movement is easy to see.
def create_plotter(grid: pv.UnstructuredGrid) -> pv.Plotter:
"""
Build a plotter for GIF animation.
Returns
-------
pv.Plotter
Plotter configured for frame capture.
"""
plotter = pv.Plotter(window_size=[800, 600])
init_grid = grid.copy()
grid["index"] = np.arange(grid.n_points)
plotter.add_mesh(init_grid, show_edges=True, color="white", opacity=0.1)
plotter.add_mesh(
grid,
scalars="index",
show_edges=True,
lighting=False,
cmap="turbo",
opacity=0.8,
)
plotter.open_gif("shape_optimization.gif")
plotter.show_bounds(mesh=init_grid, location="outer")
plotter.camera_position = "iso"
return plotter
Step 5: Run optimization#
Load (or create) mesh, center it, enable gradients, build MLP and cost, then run gradient descent. Optionally save meshes and report final mean radius vs. sphere radius.
def main(
mesh_path: pathlib.Path | None = None,
n_steps: int = 2000,
use_bias: bool = False,
save_dir: pathlib.Path | None = None,
log_interval: int = 100,
) -> None:
"""
Run shape optimization: minimize area under fixed volume.
If mesh_path is None, use a small generated hex grid. Otherwise load
the mesh from mesh_path. When save_dir is set, write mesh.VTU every
log_interval steps.
"""
# ----- Load or create mesh -----
if mesh_path is not None and mesh_path.exists():
pv_mesh: pv.DataSet = pv.read(mesh_path)
pv_mesh: pv.UnstructuredGrid = pv_mesh.cast_to_unstructured_grid()
else:
pv_mesh = generate_hex_grid(10, 10, 10)
pv_mesh.points = pv_mesh.points - np.mean(
pv_mesh.points, axis=0, keepdims=True
)
mesh = graphlow.from_pyvista(pv_mesh, "phlower")
mesh.requires_grad(True)
if mesh.topology.mesh_dim() < TopologyDim.VOLUME:
raise RuntimeError(
"Volume mesh required (face_registry / cell_volumes)"
)
initial_volumes = mesh.geometry.cell_volumes().clone()
initial_total_volume = torch.sum(initial_volumes).detach().clone()
initial_points = mesh.points.clone()
surface = mesh.extract_surface()
surface_initial_areas = surface.geometry.face_areas().detach().clone()
surface_initial_total_area = torch.sum(surface_initial_areas)
# ----- Deformation model and optimizer -----
deform_fn, params = make_deformation_mlp(MLP_HIDDEN, use_bias=use_bias)
optimizer = torch.optim.Adam(params, lr=LEARNING_RATE)
deform_factor = [DEFORMATION_FACTOR]
def cost_fn(
deformed_points: pt.PhlowerTensor,
) -> tuple[
pt.PhlowerTensor | None,
pt.PhlowerTensor | None,
pt.PhlowerTensor | None,
]:
return compute_cost(
mesh,
surface,
initial_points,
initial_total_volume,
surface_initial_total_area,
deformed_points,
WEIGHT_VOLUME_CONSTRAINT,
WEIGHT_NORM_CONSTRAINT,
TINY_CELL_RATIO,
)
if save_dir is not None:
save_dir = pathlib.Path(save_dir)
save_dir.mkdir(parents=True, exist_ok=True)
plotter = create_plotter(mesh.pvmesh)
plotter.write_frame()
print("Step | area | vol ratio | cost")
print("------|----------|-----------|----------")
for step in range(1, n_steps + 1):
optimizer.zero_grad()
deformed_points = deform_fn(initial_points, deform_factor[0])
cost, area, volume = cost_fn(deformed_points)
if cost is None:
deform_factor[0] *= 0.9
continue
if step % log_interval == 0:
vol_ratio = volume / initial_total_volume
print(
f"{step:5d} | {area.to_tensor():.3e} | "
f"{vol_ratio.to_tensor():.3e} | {cost.to_tensor():.3e}"
)
mesh.pvmesh.points = (
mesh.backend.to_torch(mesh.points).detach().numpy()
)
plotter.write_frame()
if save_dir is not None:
mesh.pvmesh.save(str(save_dir / f"mesh.{step:06d}.vtu"))
cost.backward()
optimizer.step()
plotter.close()
# ----- Final shape: mean radius vs. sphere radius -----
final_deformed = deform_fn(initial_points, deform_factor[0])
surface_deformed_points = surface.gather_parent_point_data(final_deformed)
actual_mean_radius = (
torch.mean(
torch.linalg.norm(surface_deformed_points.to_tensor(), dim=1)
)
.detach()
.numpy()
)
v0 = initial_total_volume.to_tensor().numpy()
desired_radius = (v0 * 3 / 4 / np.pi) ** (1 / 3)
print(f"\nInitial total volume: {v0:.5f}")
print(f"Desired sphere radius (same volume): {desired_radius:.5f}")
print(f"Actual mean boundary radius: {actual_mean_radius:.5f}")
Step 6: Run the example#
if __name__ == "__main__":
main(n_steps=2000, log_interval=100)
Step | area | vol ratio | cost
------|----------|-----------|----------
100 | 4.210e+02 | 9.710e-01 | 1.876e+00
200 | 3.963e+02 | 9.667e-01 | 1.830e+00
300 | 3.929e+02 | 9.679e-01 | 1.822e+00
400 | 3.939e+02 | 9.787e-01 | 1.818e+00
500 | 3.930e+02 | 9.794e-01 | 1.816e+00
600 | 3.853e+02 | 9.541e-01 | 1.817e+00
700 | 3.906e+02 | 9.764e-01 | 1.813e+00
800 | 3.916e+02 | 9.819e-01 | 1.813e+00
900 | 3.772e+02 | 9.299e-01 | 1.829e+00
1000 | 3.882e+02 | 9.721e-01 | 1.810e+00
1100 | 3.876e+02 | 9.712e-01 | 1.810e+00
1200 | 3.818e+02 | 9.503e-01 | 1.814e+00
1300 | 3.875e+02 | 9.725e-01 | 1.809e+00
1400 | 3.833e+02 | 9.572e-01 | 1.811e+00
1500 | 3.873e+02 | 9.727e-01 | 1.808e+00
1600 | 3.683e+02 | 9.024e-01 | 1.857e+00
1700 | 3.872e+02 | 9.729e-01 | 1.808e+00
1800 | 3.870e+02 | 9.726e-01 | 1.808e+00
1900 | 3.863e+02 | 9.701e-01 | 1.808e+00
2000 | 3.869e+02 | 9.726e-01 | 1.808e+00
Initial total volume: 729.00000
Desired sphere radius (same volume): 5.58315
Actual mean boundary radius: 5.56019
Total running time of the script: (0 minutes 14.586 seconds)