optim working

Author: andrinr

.gitignore

@@ -141,3 +141,5 @@ dmypy.json
 # Ignore UV dev environments
 uv.lock
 uv.lock.bak
+
+tmp_img/

examples/ansys/demo_2.ipynb

@@ -106,13 +106,20 @@ def get_tensor_map(self) -> Callable:
         """
 
         def stress(u_grad, theta):
-            E = 70e3
+            Emax = 70.0e3
+            Emin = 1e-3 * Emax
+            penal = 3.0
+
+            E = Emin + (Emax - Emin) * theta[0] ** penal
+
             nu = 0.3
             mu = E / (2.0 * (1.0 + nu))
             lmbda = E * nu / ((1 + nu) * (1 - 2 * nu))
 
             epsilon = 0.5 * (u_grad + u_grad.T)
             sigma = lmbda * jnp.trace(epsilon) * jnp.eye(self.dim) + 2 * mu * epsilon
+
+            sigma = lmbda * jnp.trace(epsilon) * jnp.eye(self.dim) + 2.0 * mu * epsilon
             return sigma
 
         return stress
@@ -123,13 +130,7 @@ def get_surface_maps(self) -> list[Callable]:
         Returns:
             List of van Neumann boundary condition value functions.
         """
-
-        def surface_map(u, x):
-            return jnp.array([0.0, 0.0, 100.0])
-
-        return [surface_map]
-
-        # return self.van_neumann_value_fns
+        return self.van_neumann_value_fns
 
     def set_params(self, params: jnp.ndarray) -> None:
         """Set density parameters for topology optimization.
@@ -280,6 +281,41 @@ def apply_fn(inputs: dict) -> dict:
     Returns:
         Dictionary containing the compliance of the structure.
     """
+    from typing import TypeVar
+
+    T = TypeVar("T")
+
+    def stop_grads_int(x: T) -> T:
+        """Stops gradient computation.
+
+        We cannot use jax.lax.stop_gradient directly because Tesseract meshes are
+        nested dictionaries with arrays and integers, and jax.lax.stop_gradient
+        does not support integers.
+
+        Args:
+                x: Input value.
+
+        Returns:
+                Value with stopped gradients.
+        """
+
+        def stop(x):
+            return jax._src.ad_util.stop_gradient_p.bind(x)
+
+        return jax.tree_util.tree_map(stop, x)
+
+    # stop grads on all inputs except rho
+
+    # problem, fwd_pred = setup(
+    #     pts=stop_grads_int(inputs["hex_mesh"]["points"][: inputs["hex_mesh"]["n_points"]]),
+    #     cells=stop_grads_int(inputs["hex_mesh"]["faces"][: inputs["hex_mesh"]["n_faces"]]),
+    #     dirichlet_mask=stop_grads_int(inputs["dirichlet_mask"]),
+    #     dirichlet_values=stop_grads_int(inputs["dirichlet_values"]),
+    #     van_neumann_mask=stop_grads_int(inputs["van_neumann_mask"]),
+    #     van_neumann_values=stop_grads_int(inputs["van_neumann_values"]),
+    # )
+
+    # no stop grads
     problem, fwd_pred = setup(
         pts=inputs["hex_mesh"]["points"][: inputs["hex_mesh"]["n_points"]],
         cells=inputs["hex_mesh"]["faces"][: inputs["hex_mesh"]["n_faces"]],
@@ -289,7 +325,7 @@ def apply_fn(inputs: dict) -> dict:
         van_neumann_values=inputs["van_neumann_values"],
     )
 
-    rho = inputs["rho"][: inputs["hex_mesh"]["n_points"]]
+    rho = inputs["rho"][: inputs["hex_mesh"]["n_faces"]]
 
     sol_list = fwd_pred(rho)
     compliance = problem.compute_compliance(sol_list[0])

examples/ansys/fem_shapeopt_mesh.vtk

@@ -2,8 +2,8 @@
 
 images = []
 
-for i in range(60):
-    filename = f"img/mesh_optim_{i:03d}.png"
+for i in range(30):
+    filename = f"tmp_img/mesh_optim_{i:03d}.png"
     images.append(imageio.imread(filename))
     print(f"Added {filename} to gif.")
 imageio.mimsave("mesh_optim.gif", images, fps=10)

examples/ansys/fem_tess/tesseract_api.py

@@ -313,7 +313,9 @@ def vector_jacobian_product(
         epsilon=inputs.epsilon,
     )
     if inputs.normalize_jacobian:
-        n_elements = inputs.Nx * inputs.Ny * inputs.Nz
+        n_elements = (
+            inputs.grid_elements[0] * inputs.grid_elements[1] * inputs.grid_elements[2]
+        )
         jac = jac / n_elements
     # Reduce the cotangent vector to the shape of the Jacobian, to compute VJP by hand
     vjp = np.einsum("klmn,lmn->k", jac, cotangent_vector["sdf"]).astype(np.float32)

examples/ansys/gf.py

@@ -30,17 +30,8 @@ class InputSchema(BaseModel):
             "Sizing field values defined on a regular grid for mesh adaptation."
         )
     )
-    Lx: float = Field(
-        default=60.0,
-        description=("Length of the domain in the x direction. "),
-    )
-    Ly: float = Field(
-        default=30.0,
-        description=("Length of the domain in the y direction. "),
-    )
-    Lz: float = Field(
-        default=30.0,
-        description=("Length of the domain in the z direction. "),
+    domain_size: tuple[float, float, float] = Field(
+        description=("Size of the domain in x, y, z directions.")
     )
 
     max_points: int = Field(
@@ -326,10 +317,13 @@ def apply(inputs: InputSchema) -> OutputSchema:
     Returns:
         OutputSchema, outputs of the function.
     """
+    Lx = inputs.domain_size[0]
+    Ly = inputs.domain_size[1]
+    Lz = inputs.domain_size[2]
     pts, cells = generate_mesh(
-        Lx=inputs.Lx,
-        Ly=inputs.Ly,
-        Lz=inputs.Lz,
+        Lx=Lx,
+        Ly=Ly,
+        Lz=Lz,
         sizing_field=inputs.sizing_field,
         max_levels=inputs.max_subdivision_levels,
     )
@@ -339,9 +333,9 @@ def apply(inputs: InputSchema) -> OutputSchema:
     cells_padded = jnp.zeros((inputs.max_cells, 8), dtype=cells.dtype)
     cells_padded = cells_padded.at[: cells.shape[0], :].set(cells)
 
-    xs = jnp.linspace(-inputs.Lx / 2, inputs.Lx / 2, inputs.field_values.shape[0])
-    ys = jnp.linspace(-inputs.Ly / 2, inputs.Ly / 2, inputs.field_values.shape[1])
-    zs = jnp.linspace(-inputs.Lz / 2, inputs.Lz / 2, inputs.field_values.shape[2])
+    xs = jnp.linspace(-Lx / 2, Lx / 2, inputs.field_values.shape[0])
+    ys = jnp.linspace(-Ly / 2, Ly / 2, inputs.field_values.shape[1])
+    zs = jnp.linspace(-Lz / 2, Lz / 2, inputs.field_values.shape[2])
 
     interpolator = RegularGridInterpolator(
         (xs, ys, zs),
@@ -395,19 +389,23 @@ def vector_jacobian_product(
     assert vjp_inputs == {"field_values"}
     assert vjp_outputs == {"mesh_cell_values"}
 
+    Lx = inputs.domain_size[0]
+    Ly = inputs.domain_size[1]
+    Lz = inputs.domain_size[2]
+
     pts, cells = generate_mesh(
-        Lx=inputs.Lx,
-        Ly=inputs.Ly,
-        Lz=inputs.Lz,
+        Lx=Lx,
+        Ly=Ly,
+        Lz=Lz,
         sizing_field=inputs.sizing_field,
         max_levels=inputs.max_subdivision_levels,
     )
 
     cell_centers = jnp.mean(pts[cells], axis=1)
 
-    xs = jnp.linspace(-inputs.Lx / 2, inputs.Lx / 2, inputs.field_values.shape[0])
-    ys = jnp.linspace(-inputs.Ly / 2, inputs.Ly / 2, inputs.field_values.shape[1])
-    zs = jnp.linspace(-inputs.Lz / 2, inputs.Lz / 2, inputs.field_values.shape[2])
+    xs = jnp.linspace(-Lx / 2, Lx / 2, inputs.field_values.shape[0])
+    ys = jnp.linspace(-Ly / 2, Ly / 2, inputs.field_values.shape[1])
+    zs = jnp.linspace(-Lz / 2, Lz / 2, inputs.field_values.shape[2])
     xs, ys, zs = jnp.meshgrid(xs, ys, zs, indexing="ij")
 
     field_cotangent_vector = griddata(