From c0999ec376ee8ff5c972fbfc7e254ce268661022 Mon Sep 17 00:00:00 2001
From: mingrui <mingrui.zhang@taichi.graphics>
Date: Thu, 29 Feb 2024 03:10:24 +0800
Subject: [PATCH] Add laplacian smooth for monitor function

---
 movement/__init__.py         |  1 +
 movement/monge_ampere.py     | 17 +++++++++++++++
 movement/monitor_smoother.py | 41 ++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+)
 create mode 100644 movement/monitor_smoother.py

diff --git a/movement/__init__.py b/movement/__init__.py
index 083d787..bf98365 100644
--- a/movement/__init__.py
+++ b/movement/__init__.py
@@ -3,3 +3,4 @@
 from movement.laplacian import *  # noqa
 from movement.monge_ampere import *  # noqa
 from movement.spring import *  # noqa
+from movement.monitor_smoother import *  # noqa
diff --git a/movement/monge_ampere.py b/movement/monge_ampere.py
index af98cd6..773385a 100644
--- a/movement/monge_ampere.py
+++ b/movement/monge_ampere.py
@@ -5,6 +5,7 @@
 import numpy as np
 import movement.solver_parameters as solver_parameters
 from movement.mover import PrimeMover
+from movement.monitor_smoother import laplacian_smoothing
 
 
 __all__ = [
@@ -74,6 +75,9 @@ def __init__(self, mesh, monitor_function, **kwargs):
         self.rtol = kwargs.pop("rtol", 1.0e-08)
         self.dtol = kwargs.pop("dtol", 2.0)
         self.fix_boundary_nodes = kwargs.pop("fix_boundary_nodes", False)
+        self.smooth_monitor_values = kwargs.pop("smooth_monitor_values", False)
+        self.monitor_smoother = kwargs.pop("monitor_smoother", "laplacian")
+        self.laplacian_smoother_num = kwargs.pop("laplacian_smoother_num", 40)
         super().__init__(mesh, monitor_function=monitor_function)
 
         # Create function spaces
@@ -365,6 +369,19 @@ def move(self):
 
             # Update monitor function
             self.monitor.interpolate(self.monitor_function(self.mesh))
+
+            # Perform monitor function values smooth
+            if self.smooth_monitor_values:
+                if self.monitor_smoother == "laplacian":
+                    self.monitor = laplacian_smoothing(
+                        mesh=self.mesh,
+                        monitor_values=self.monitor,
+                        smooth_num=self.laplacian_smoother_num,
+                    )
+                else:
+                    raise NotImplementedError(
+                        f"Monitor smoother {self.monitor_smoother} not implemented."
+                    )
             firedrake.assemble(self.L_P0, tensor=self.volume)
             self.volume.interpolate(self.volume / self.original_volume)
             self.mesh.coordinates.assign(self.xi)
diff --git a/movement/monitor_smoother.py b/movement/monitor_smoother.py
new file mode 100644
index 0000000..f8475bc
--- /dev/null
+++ b/movement/monitor_smoother.py
@@ -0,0 +1,41 @@
+import firedrake as fd
+
+__all__ = ["laplacian_smoothing"]
+
+
+def laplacian_smoothing(mesh, monitor_values, smooth_num=40):
+    """
+    Laplacian smoothing for monitor function values
+
+    Solve the equation below:
+    m_{smooth} = dt * K * \nabla^2 m_{smooth} + m
+
+    K = N * dx^2 / (4 * dt)
+    N (smooth_num): the number of number of applications of a (1, −2, 1) filter
+
+    Reference:
+    Mariana C A Clare, Joseph Gregory Wallwork , Stephan C Kramer, Hilary Weller, Colin J Cotter, Matthew Piggott, On the use of mesh movement methods to help overcome the multi-scale
+    challenges associated with hydro-morphodynamic modelling
+    https://eartharxiv.org/repository/view/1751/
+
+
+    """
+    V = fd.FunctionSpace(mesh, "CG", 1)
+    u = fd.TrialFunction(V)
+    v = fd.TestFunction(V)
+
+    f = monitor_values
+    dx = mesh.cell_sizes.dat.data[:].mean()
+    K = smooth_num * dx**2 / 4
+    RHS = f * v * fd.dx(domain=mesh)
+    LHS = (K * fd.dot(fd.grad(v), fd.grad(u)) + v * u) * fd.dx(domain=mesh)
+    bc = fd.DirichletBC(V, f, "on_boundary")
+
+    monitor_smoothed = fd.Function(V)
+    fd.solve(
+        LHS == RHS,
+        monitor_smoothed,
+        solver_parameters={"ksp_type": "cg", "pc_type": "none"},
+        bcs=bc,
+    )
+    return monitor_values.project(monitor_smoothed)