From 3cb5cb4a4d387d5b87dca93d4dff3a8170fbb325 Mon Sep 17 00:00:00 2001
From: Joe Wallwork <joe.wallwork@outlook.com>
Date: Thu, 29 Feb 2024 15:47:51 +0000
Subject: [PATCH] #58: Add displacement attribute to LaplacianSmoother

---
 demos/laplacian_smoothing.py    | 5 +++--
 movement/laplacian_smoothing.py | 4 ++++
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/demos/laplacian_smoothing.py b/demos/laplacian_smoothing.py
index b6e3bb8..0bf6257 100644
--- a/demos/laplacian_smoothing.py
+++ b/demos/laplacian_smoothing.py
@@ -85,7 +85,6 @@ def update_forcings(t):
 
 
 # We are now able to apply the mesh movement method. This works just as before. ::
-# TODO: displacement
 
 time = 0.0
 fig, axes = plt.subplots(ncols=4, nrows=3, figsize=(12, 10))
@@ -94,6 +93,8 @@ def update_forcings(t):
 
     # Move the mesh and calculate the mesh speed
     mover.move(t, update_forcings=update_forcings, fixed_boundaries=[1, 2, 3])
+    displacement = np.linalg.norm(mover.displacement)
+    print(f"time = {t:.1f} s, displacement = {displacement:.2f} m")
 
     # Plot the current mesh, adding a time label
     ax = axes[idx // 4, idx % 4]
@@ -114,6 +115,6 @@ def update_forcings(t):
 coord_data = mover.mesh.coordinates.dat.data
 linf_error = np.max(np.abs(coord_data - coord_data_init))
 print(f"l_infinity error: {linf_error:.3f} m")
-assert linf_error < 0.01
+assert np.isclose(linf_error, 0.0)
 
 # This tutorial can be downloaded as a `Python script <laplacian_smoothing.py>`__.
diff --git a/movement/laplacian_smoothing.py b/movement/laplacian_smoothing.py
index bd01b68..be774ea 100644
--- a/movement/laplacian_smoothing.py
+++ b/movement/laplacian_smoothing.py
@@ -3,6 +3,7 @@
 import ufl
 import movement.solver_parameters as solver_parameters
 from movement.mover import PrimeMover
+import numpy as np
 
 
 __all__ = ["LaplacianSmoother"]
@@ -20,6 +21,7 @@ def __init__(self, mesh, timestep, **kwargs):
         assert timestep > 0.0
         self.dt = timestep
         self.f = firedrake.Function(self.coord_space)
+        self.displacement = np.zeros(self.mesh.num_vertices())
 
     @PETSc.Log.EventDecorator()
     def _setup_solver(self, fixed_boundaries=[]):
@@ -56,5 +58,7 @@ def move(self, time, update_forcings=None, fixed_boundaries=[]):
         self._solver.solve()
 
         # Update mesh coordinates
+        coords = self._x.dat.data_with_halos.copy()
         self._x.dat.data_with_halos[:] += self.v.dat.data_with_halos * self.dt
         self.mesh.coordinates.assign(self._x)
+        self.displacement = self.mesh.coordinates.dat.data_with_halos - coords