fix examples (#15)

* Fix compliance calculation

* Remove topopt helpers, this is not in scope for this package

* Fix examples

examples/compliance_2d.py

@@ -3,76 +3,90 @@
 import autograd.numpy as np
 import matplotlib.pyplot as plt
 import nlopt
+import scipy.ndimage
 from autograd import value_and_grad
+from autograd.extend import defvjp, primitive
 
 from tofea.fea2d import FEA2D_K
-from tofea.topopt_helpers import simp_parametrization
 
-max_its = 100
-volfrac = 0.5
-sigma = 0.5
-shape = (160, 80)
-nelx, nely = shape
-emin, emax = 1e-6, 1
+gaussian_filter = primitive(scipy.ndimage.gaussian_filter)
+defvjp(
+    gaussian_filter,
+    lambda ans, x, *args, **kwargs: lambda g: gaussian_filter(g, *args, **kwargs),  # noqa: ARG005
+)
 
-dofs = np.arange(2 * (nelx + 1) * (nely + 1)).reshape(nelx + 1, nely + 1, 2)
-fixed = np.zeros_like(dofs, dtype=bool)
-load = np.zeros_like(dofs)
 
-fixed[0, :, :] = 1
-load[-1, -1, 1] = 1
+def simp_parametrization(shape, sigma, vmin, vmax, penalty=3.0):
+    def _parametrization(x):
+        x = np.reshape(x, shape)
+        x = gaussian_filter(x, sigma)
+        x = vmin + (vmax - vmin) * x**penalty
+        return x
 
-fem = FEA2D_K(fixed)
-parametrization = simp_parametrization(shape, sigma, emin, emax)
-x0 = np.full(shape, volfrac)
+    return _parametrization
 
 
-def objective(x):
-    x = parametrization(x)
-    xp = np.pad(x, (0, 1), mode="constant", constant_values=emin)
-    xp = (xp - emin) / (emax - emin)
-    body = 1e-3 * np.stack([np.zeros_like(xp), xp], axis=-1)
-    x = fem.compliance(x, load + body)
-    return np.sum(x)
+def main():
+    max_its = 50
+    volfrac = 0.3
+    sigma = 1.0
+    shape = (120, 60)
+    nelx, nely = shape
+    emin, emax = 1e-6, 1
 
+    dofs = np.arange(2 * (nelx + 1) * (nely + 1)).reshape(nelx + 1, nely + 1, 2)
+    fixed = np.zeros_like(dofs, dtype=bool)
+    load = np.zeros_like(dofs)
 
-def volume(x):
-    x = parametrization(x)
-    x = (x - emin) / (emax - emin)
-    return np.mean(x)
+    fixed[0, :, :] = 1
+    load[-1, -1, 1] = 1
 
+    fem = FEA2D_K(fixed)
+    parametrization = simp_parametrization(shape, sigma, emin, emax)
+    x0 = np.full(shape, volfrac)
 
-plt.ion()
-fig, ax = plt.subplots(1, 1)
-im = ax.imshow(parametrization(x0).T, cmap="gray_r", vmin=emin, vmax=emax)
-fig.tight_layout()
+    plt.ion()
+    fig, ax = plt.subplots(1, 1, tight_layout=True)
+    im = ax.imshow(parametrization(x0).T, cmap="gray_r", vmin=emin, vmax=emax)
 
+    @value_and_grad
+    def objective(x):
+        x = parametrization(x)
+        d = fem.displacement(x, load)
+        c = fem.compliance(x, d)
+        return c
 
-def volume_constraint(x, gd):
-    v, g = value_and_grad(volume)(x)
-    if gd.size > 0:
-        gd[:] = g
-    return v - volfrac
+    @value_and_grad
+    def volume(x):
+        return np.mean(x)
 
+    def volume_constraint(x, gd):
+        v, g = volume(x)
+        if gd.size > 0:
+            gd[:] = g.ravel()
+        return v - volfrac
 
-def nlopt_obj(x, gd):
-    c, dc = value_and_grad(objective)(x)
+    def nlopt_obj(x, gd):
+        v, g = objective(x)
 
-    if gd.size > 0:
-        gd[:] = dc.ravel()
+        if gd.size > 0:
+            gd[:] = g.ravel()
 
-    im.set_data(parametrization(x).T)
-    plt.pause(0.01)
+        im.set_data(parametrization(x).T)
+        plt.pause(0.01)
 
-    return c
+        return v
 
+    opt = nlopt.opt(nlopt.LD_MMA, x0.size)
+    opt.add_inequality_constraint(volume_constraint)
+    opt.set_min_objective(nlopt_obj)
+    opt.set_lower_bounds(0)
+    opt.set_upper_bounds(1)
+    opt.set_maxeval(max_its)
+    opt.optimize(x0.ravel())
 
-opt = nlopt.opt(nlopt.LD_CCSAQ, x0.size)
-opt.add_inequality_constraint(volume_constraint, 1e-3)
-opt.set_min_objective(nlopt_obj)
-opt.set_lower_bounds(0)
-opt.set_upper_bounds(1)
-opt.set_maxeval(max_its)
-opt.optimize(x0.ravel())
+    plt.show(block=True)
 
-plt.show(block=True)
+
+if __name__ == "__main__":
+    main()

examples/heat_2d.py

@@ -3,72 +3,89 @@
 import autograd.numpy as np
 import matplotlib.pyplot as plt
 import nlopt
+import scipy.ndimage
 from autograd import value_and_grad
+from autograd.extend import defvjp, primitive
 
 from tofea.fea2d import FEA2D_T
-from tofea.topopt_helpers import simp_parametrization
 
-max_its = 100
-volfrac = 0.5
-sigma = 0.5
-shape = (100, 100)
-nelx, nely = shape
-cmin, cmax = 1.0, 1e6
+gaussian_filter = primitive(scipy.ndimage.gaussian_filter)
+defvjp(
+    gaussian_filter,
+    lambda ans, x, *args, **kwargs: lambda g: gaussian_filter(g, *args, **kwargs),  # noqa: ARG005
+)
 
-fixed = np.zeros((nelx + 1, nely + 1), dtype="?")
-load = np.zeros_like(fixed)
 
-fixed[40:60, -1] = 1
-load[:, :-10] = 1
+def simp_parametrization(shape, sigma, vmin, vmax, penalty=3.0):
+    def _parametrization(x):
+        x = np.reshape(x, shape)
+        x = gaussian_filter(x, sigma)
+        x = vmin + (vmax - vmin) * x**penalty
+        return x
 
-fem = FEA2D_T(fixed)
-parametrization = simp_parametrization(shape, sigma, cmin, cmax)
-x0 = np.full(shape, volfrac)
+    return _parametrization
 
 
-def objective(x):
-    x = parametrization(x)
-    x = fem.temperature(x, load)
-    return np.sum(x)
+def main():
+    max_its = 100
+    volfrac = 0.5
+    sigma = 1.0
+    shape = (100, 100)
+    nelx, nely = shape
+    cmin, cmax = 1e-4, 1
 
+    fixed = np.zeros((nelx + 1, nely + 1), dtype="?")
+    load = np.zeros_like(fixed)
 
-def volume(x):
-    x = parametrization(x)
-    x = (x - cmin) / (cmax - cmin)
-    return np.mean(x)
+    fixed[shape[0] // 2 - 5 : shape[0] // 2 + 5, -1] = 1
+    load[:, 0] = 1
+    load[(0, -1), :] = 1
 
+    fem = FEA2D_T(fixed)
+    parametrization = simp_parametrization(shape, sigma, cmin, cmax)
+    x0 = np.full(shape, volfrac)
 
-plt.ion()
-fig, ax = plt.subplots(1, 1)
-im = ax.imshow(parametrization(x0).T, cmap="gray_r", vmin=cmin, vmax=cmax)
-fig.tight_layout()
+    plt.ion()
+    fig, ax = plt.subplots(1, 1, tight_layout=True)
+    im = ax.imshow(parametrization(x0).T, cmap="gray_r", vmin=cmin, vmax=cmax)
 
+    @value_and_grad
+    def objective(x):
+        x = parametrization(x)
+        t = fem.temperature(x, load)
+        return np.mean(t)
 
-def volume_constraint(x, gd):
-    v, g = value_and_grad(volume)(x)
-    if gd.size > 0:
-        gd[:] = g
-    return v - volfrac
+    @value_and_grad
+    def volume(x):
+        return np.mean(x)
 
+    def volume_constraint(x, gd):
+        v, g = volume(x)
+        if gd.size > 0:
+            gd[:] = g.ravel()
+        return v - volfrac
 
-def nlopt_obj(x, gd):
-    c, dc = value_and_grad(objective)(x)
+    def nlopt_obj(x, gd):
+        v, g = objective(x)
 
-    if gd.size > 0:
-        gd[:] = dc.ravel()
+        if gd.size > 0:
+            gd[:] = g.ravel()
 
-    im.set_data(parametrization(x).T)
-    plt.pause(0.01)
+        im.set_data(parametrization(x).T)
+        plt.pause(0.01)
 
-    return c
+        return v
 
+    opt = nlopt.opt(nlopt.LD_MMA, x0.size)
+    opt.add_inequality_constraint(volume_constraint)
+    opt.set_min_objective(nlopt_obj)
+    opt.set_lower_bounds(0)
+    opt.set_upper_bounds(1)
+    opt.set_maxeval(max_its)
+    opt.optimize(x0.ravel())
 
-opt = nlopt.opt(nlopt.LD_CCSAQ, x0.size)
-opt.add_inequality_constraint(volume_constraint, 1e-3)
-opt.set_min_objective(nlopt_obj)
-opt.set_lower_bounds(0)
-opt.set_upper_bounds(1)
-opt.set_maxeval(max_its)
-opt.optimize(x0.ravel())
+    plt.show(block=True)
 
-plt.show(block=True)
+
+if __name__ == "__main__":
+    main()

tofea/fea2d.py

@@ -74,13 +74,12 @@ def dofmap(self) -> NDArray[np.uint32]:
     def displacement(self, x: NDArray, b: NDArray) -> NDArray:
         return self.solve(x, b)
 
-    def compliance(self, x: NDArray, b: NDArray) -> NDArray:
-        displacement = self.displacement(x, b)
+    def compliance(self, x: NDArray, displacement: NDArray) -> NDArray:
         dofmap = np.reshape(self.e2sdofmap.T, (-1, *self.out_shape))
         c = anp.einsum(
             "ixy,ij,jxy->xy", displacement[dofmap], self.element, displacement[dofmap]
         )
-        return c
+        return anp.sum(x * c)
 
 
 class FEA2D_T(FEA2D):