From 652a4e4f2b56c763e3a6cb41672f9c747260a581 Mon Sep 17 00:00:00 2001
From: HydrogenSulfate <490868991@qq.com>
Date: Thu, 14 Sep 2023 17:22:25 +0800
Subject: [PATCH] [WIP][Feature] Add derivatives for SDF (#539)

* Add derivatives for SDF

* add sdf and its derivatives for TimeXGeometry and inital_constraint

* fix bug for ldc2d_unsteady_Re10
---
 examples/ldc/ldc2d_unsteady_Re10.py     |  2 +-
 ppsci/constraint/initial_constraint.py  |  6 +++-
 ppsci/constraint/interior_constraint.py |  4 +++
 ppsci/geometry/geometry.py              | 37 ++++++++++++++++++++++++-
 ppsci/geometry/mesh.py                  | 18 ++++++++++--
 ppsci/geometry/timedomain.py            | 27 ++++++++++++++++--
 6 files changed, 87 insertions(+), 7 deletions(-)

diff --git a/examples/ldc/ldc2d_unsteady_Re10.py b/examples/ldc/ldc2d_unsteady_Re10.py
index c8c662f15..ac22f23de 100644
--- a/examples/ldc/ldc2d_unsteady_Re10.py
+++ b/examples/ldc/ldc2d_unsteady_Re10.py
@@ -173,7 +173,7 @@
     # manually collate input data for visualization,
     # (interior+boundary) x all timestamps
     for t in range(NTIME_PDE):
-        for key in vis_points:
+        for key in geom["time_rect"].dim_keys:
             vis_points[key] = np.concatenate(
                 (
                     vis_points[key],
diff --git a/ppsci/constraint/initial_constraint.py b/ppsci/constraint/initial_constraint.py
index 351af60c7..2d8e81ce8 100644
--- a/ppsci/constraint/initial_constraint.py
+++ b/ppsci/constraint/initial_constraint.py
@@ -34,7 +34,7 @@
 
 
 class InitialConstraint(base.Constraint):
-    """Class for initial constraint.
+    """Class for initial interior constraint.
 
     Args:
         output_expr (Dict[str, Callable]): Function in dict for computing output.
@@ -53,6 +53,8 @@ class InitialConstraint(base.Constraint):
             Defaults to False.
         weight_dict (Optional[Dict[str, Callable]]): Define the weight of each
             constraint variable. Defaults to None.
+        compute_sdf_derivatives (Optional[bool]): Whether compute derivatives for SDF.
+            Defaults to False.
         name (str, optional): Name of constraint object. Defaults to "IC".
 
     Examples:
@@ -86,6 +88,7 @@ def __init__(
         criteria: Optional[Callable] = None,
         evenly: bool = False,
         weight_dict: Optional[Dict[str, Callable]] = None,
+        compute_sdf_derivatives: bool = False,
         name: str = "IC",
     ):
         self.label_dict = label_dict
@@ -107,6 +110,7 @@ def __init__(
             random,
             criteria,
             evenly,
+            compute_sdf_derivatives,
         )
         if "area" in input:
             input["area"] *= dataloader_cfg["iters_per_epoch"]
diff --git a/ppsci/constraint/interior_constraint.py b/ppsci/constraint/interior_constraint.py
index a333c82db..67559c6bc 100644
--- a/ppsci/constraint/interior_constraint.py
+++ b/ppsci/constraint/interior_constraint.py
@@ -53,6 +53,8 @@ class InteriorConstraint(base.Constraint):
             Defaults to False.
         weight_dict (Optional[Dict[str, Union[Callable, float]]]): Define the
             weight of each constraint variable. Defaults to None.
+        compute_sdf_derivatives (Optional[bool]): Whether compute derivatives for SDF.
+            Defaults to False.
         name (str, optional): Name of constraint object. Defaults to "EQ".
 
     Examples:
@@ -83,6 +85,7 @@ def __init__(
         criteria: Optional[Callable] = None,
         evenly: bool = False,
         weight_dict: Optional[Dict[str, Union[Callable, float]]] = None,
+        compute_sdf_derivatives: bool = False,
         name: str = "EQ",
     ):
         self.label_dict = label_dict
@@ -104,6 +107,7 @@ def __init__(
             random,
             criteria,
             evenly,
+            compute_sdf_derivatives,
         )
         if "area" in input:
             input["area"] *= dataloader_cfg["iters_per_epoch"]
diff --git a/ppsci/geometry/geometry.py b/ppsci/geometry/geometry.py
index 3e087c70c..ee5be07fa 100644
--- a/ppsci/geometry/geometry.py
+++ b/ppsci/geometry/geometry.py
@@ -70,6 +70,7 @@ def sample_interior(
         random="pseudo",
         criteria=None,
         evenly=False,
+        compute_sdf_derivatives: bool = False,
     ):
         """Sample random points in the geometry and return those meet criteria."""
         x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype())
@@ -106,11 +107,18 @@ def sample_interior(
         if hasattr(self, "sdf_func"):
             sdf = -self.sdf_func(x)
             sdf_dict = misc.convert_to_dict(sdf, ("sdf",))
+            sdf_derivs_dict = {}
+            if compute_sdf_derivatives:
+                sdf_derivs = -self.sdf_derivatives(x)
+                sdf_derivs_dict = misc.convert_to_dict(
+                    sdf_derivs, tuple(f"sdf__{key}" for key in self.dim_keys)
+                )
         else:
             sdf_dict = {}
+            sdf_derivs_dict = {}
         x_dict = misc.convert_to_dict(x, self.dim_keys)
 
-        return {**x_dict, **sdf_dict}
+        return {**x_dict, **sdf_dict, **sdf_derivs_dict}
 
     def sample_boundary(self, n, random="pseudo", criteria=None, evenly=False):
         """Compute the random points in the geometry and return those meet criteria."""
@@ -196,6 +204,33 @@ def periodic_point(self, x: np.ndarray, component: int):
         """Compute the periodic image of x."""
         raise NotImplementedError(f"{self}.periodic_point to be implemented")
 
+    def sdf_derivatives(self, x: np.ndarray, epsilon: float = 1e-4) -> np.ndarray:
+        """Compute derivatives of SDF function.
+
+        Args:
+            x (np.ndarray): Points for computing SDF derivatives using central
+                difference. The shape is [N, D], D is the number of dimension of
+                geometry.
+            epsilon (float): Derivative step. Defaults to 1e-4.
+
+        Returns:
+            np.ndarray: Derivatives of corresponding SDF function.
+                The shape is [N, D]. D is the number of dimension of geometry.
+        """
+        if not hasattr(self, "sdf_func"):
+            raise NotImplementedError(
+                f"{misc.typename(self)}.sdf_func should be implemented "
+                "when using 'sdf_derivatives'."
+            )
+        # Only compute sdf derivatives for those already implement `sdf_func` method.
+        sdf_derivs = np.empty_like(x)
+        for i in range(self.ndim):
+            h = np.zeros_like(x)
+            h[:, i] += epsilon / 2
+            derivs_at_i = (self.sdf_func(x + h) - self.sdf_func(x - h)) / epsilon
+            sdf_derivs[:, i : i + 1] = derivs_at_i
+        return sdf_derivs
+
     def union(self, other):
         """CSG Union."""
         from ppsci.geometry import csg
diff --git a/ppsci/geometry/mesh.py b/ppsci/geometry/mesh.py
index 0c10e47ce..cd30668a9 100644
--- a/ppsci/geometry/mesh.py
+++ b/ppsci/geometry/mesh.py
@@ -427,7 +427,14 @@ def random_points(self, n, random="pseudo", criteria=None):
         all_areas = np.full((n, 1), cuboid_volume * (_nvalid / _nsample) / n)
         return all_points, all_areas
 
-    def sample_interior(self, n, random="pseudo", criteria=None, evenly=False):
+    def sample_interior(
+        self,
+        n,
+        random="pseudo",
+        criteria=None,
+        evenly=False,
+        compute_sdf_derivatives: bool = False,
+    ):
         """Sample random points in the geometry and return those meet criteria."""
         if evenly:
             # TODO(sensen): implement uniform sample for mesh interior.
@@ -444,7 +451,14 @@ def sample_interior(self, n, random="pseudo", criteria=None, evenly=False):
         sdf = -self.sdf_func(points)
         sdf_dict = misc.convert_to_dict(sdf, ("sdf",))
 
-        return {**x_dict, **area_dict, **sdf_dict}
+        sdf_derivs_dict = {}
+        if compute_sdf_derivatives:
+            sdf_derivs = -self.sdf_derivatives(points)
+            sdf_derivs_dict = misc.convert_to_dict(
+                sdf_derivs, tuple(f"sdf__{key}" for key in self.dim_keys)
+            )
+
+        return {**x_dict, **area_dict, **sdf_dict, **sdf_derivs_dict}
 
     def union(self, other: "Mesh"):
         if not checker.dynamic_import_to_globals(["pymesh"]):
diff --git a/ppsci/geometry/timedomain.py b/ppsci/geometry/timedomain.py
index 7f9b877ab..8d9d51141 100644
--- a/ppsci/geometry/timedomain.py
+++ b/ppsci/geometry/timedomain.py
@@ -541,7 +541,12 @@ def periodic_point(self, x, component):
         return txp
 
     def sample_initial_interior(
-        self, n: int, random: str = "pseudo", criteria=None, evenly=False
+        self,
+        n: int,
+        random: str = "pseudo",
+        criteria=None,
+        evenly=False,
+        compute_sdf_derivatives: bool = False,
     ):
         """Sample random points in the time-geometry and return those meet criteria."""
         x = np.empty(shape=(n, self.ndim), dtype=paddle.get_default_dtype())
@@ -570,7 +575,25 @@ def sample_initial_interior(
                     "Sample initial interior points failed, "
                     "please check correctness of geometry and given creteria."
                 )
-        return misc.convert_to_dict(x, self.dim_keys)
+
+        # if sdf_func added, return x_dict and sdf_dict, else, only return the x_dict
+        if hasattr(self.geometry, "sdf_func"):
+            # compute sdf excluding timet t
+            sdf = -self.geometry.sdf_func(x[..., 1:])
+            sdf_dict = misc.convert_to_dict(sdf, ("sdf",))
+            sdf_derivs_dict = {}
+            if compute_sdf_derivatives:
+                # compute sdf derivatives excluding timet t
+                sdf_derivs = -self.geometry.sdf_derivatives(x[..., 1:])
+                sdf_derivs_dict = misc.convert_to_dict(
+                    sdf_derivs, tuple(f"sdf__{key}" for key in self.geometry.dim_keys)
+                )
+        else:
+            sdf_dict = {}
+            sdf_derivs_dict = {}
+        x_dict = misc.convert_to_dict(x, self.dim_keys)
+
+        return {**x_dict, **sdf_dict, **sdf_derivs_dict}
 
     def __str__(self) -> str:
         """Return the name of class"""