Add pinnx submodule

deepxde/geometry/geometry.py

@@ -4,8 +4,184 @@
 import numpy as np
 
 
-class Geometry(abc.ABC):
+class AbstractGeometry(abc.ABC):
+    def __init__(self, dim: int):
+        assert isinstance(dim, int), "dim must be an integer"
+        self.dim = dim
+        self.idstr = type(self).__name__
+
+    @abc.abstractmethod
+    def inside(self, x) -> np.ndarray[bool]:
+        """
+        Check if x is inside the geometry (including the boundary).
+
+        Args:
+            x: A 2D array of shape (n, dim), where `n` is the number of points and
+               `dim` is the dimension of the geometry.
+
+        Returns:
+            A boolean array of shape (n,) where each element is True if the point is inside the geometry.
+        """
+
+    @abc.abstractmethod
+    def on_boundary(self, x) -> np.ndarray[bool]:
+        """
+        Check if x is on the geometry boundary.
+
+        Args:
+            x: A 2D array of shape (n, dim), where `n` is the number of points and
+               `dim` is the dimension of the geometry.
+
+        Returns:
+            A boolean array of shape (n,) where each element is True if the point is on the boundary.
+        """
+
+    def distance2boundary(self, x, dirn):
+        """
+        Compute the distance to the boundary.
+
+        Args:
+            x: A 2D array of shape (n, dim), where `n` is the number of points and
+                `dim` is the dimension of the geometry.
+            dirn: A 2D array of shape (n, dim), where `n` is the number of points and
+                `dim` is the dimension of the geometry. The direction of the distance
+                computation. If `dirn` is not provided, the distance is computed in the
+                normal direction.
+        """
+        raise NotImplementedError("{}.distance2boundary to be implemented".format(self.idstr))
+
+    def mindist2boundary(self, x):
+        """
+        Compute the minimum distance to the boundary.
+
+        Args:
+            x: A 2D array of shape (n, dim), where `n` is the number of points and
+                `dim` is the dimension of the geometry.
+        """
+        raise NotImplementedError("{}.mindist2boundary to be implemented".format(self.idstr))
+
+    def boundary_constraint_factor(
+        self,
+        x,
+        smoothness: Literal["C0", "C0+", "Cinf"] = "C0+"
+    ):
+        """
+        Compute the hard constraint factor at x for the boundary.
+
+        This function is used for the hard-constraint methods in Physics-Informed Neural Networks (PINNs).
+        The hard constraint factor satisfies the following properties:
+
+        - The function is zero on the boundary and positive elsewhere.
+        - The function is at least continuous.
+
+        In the ansatz `boundary_constraint_factor(x) * NN(x) + boundary_condition(x)`, when `x` is on the boundary,
+        `boundary_constraint_factor(x)` will be zero, making the ansatz be the boundary condition, which in
+        turn makes the boundary condition a "hard constraint".
+
+        Args:
+            x: A 2D array of shape (n, dim), where `n` is the number of points and
+                `dim` is the dimension of the geometry. Note that `x` should be a tensor type
+                of backend (e.g., `tf.Tensor` or `torch.Tensor`), not a numpy array.
+            smoothness (string, optional): A string to specify the smoothness of the distance function,
+                e.g., "C0", "C0+", "Cinf". "C0" is the least smooth, "Cinf" is the most smooth.
+                Default is "C0+".
+
+                - C0
+                The distance function is continuous but may not be non-differentiable.
+                But the set of non-differentiable points should have measure zero,
+                which makes the probability of the collocation point falling in this set be zero.
+
+                - C0+
+                The distance function is continuous and differentiable almost everywhere. The
+                non-differentiable points can only appear on boundaries. If the points in `x` are
+                all inside or outside the geometry, the distance function is smooth.
+
+                - Cinf
+                The distance function is continuous and differentiable at any order on any
+                points. This option may result in a polynomial of HIGH order.
+
+        Returns:
+            A tensor of a type determined by the backend, which will have a shape of (n, 1).
+            Each element in the tensor corresponds to the computed distance value for the respective point in `x`.
+        """
+        raise NotImplementedError("{}.boundary_constraint_factor to be implemented".format(self.idstr))
+
+    def boundary_normal(self, x):
+        """
+        Compute the unit normal at x for Neumann or Robin boundary conditions.
+
+        Args:
+            x: A 2D array of shape (n, dim), where `n` is the number of points and
+               `dim` is the dimension of the geometry.
+        """
+        raise NotImplementedError("{}.boundary_normal to be implemented".format(self.idstr))
+
+    @abc.abstractmethod
+    def uniform_points(self, n, boundary: bool = True) -> np.ndarray:
+        """
+        Compute the equispaced point locations in the geometry.
+
+        Args:
+            n: The number of points.
+            boundary: If True, include the boundary points.
+        """
+
+    @abc.abstractmethod
+    def random_points(self, n, random: str = "pseudo") -> np.ndarray:
+        """
+        Compute the random point locations in the geometry.
+
+        Args:
+            n: The number of points.
+            random: The random distribution. One of the following: "pseudo" (pseudorandom),
+                "LHS" (Latin hypercube sampling), "Halton" (Halton sequence),
+                "Hammersley" (Hammersley sequence), or "Sobol" (Sobol sequence
+        """
+
+    @abc.abstractmethod
+    def uniform_boundary_points(self, n) -> np.ndarray:
+        """
+        Compute the equispaced point locations on the boundary.
+
+        Args:
+            n: The number of points.
+        """
+
+    @abc.abstractmethod
+    def random_boundary_points(self, n, random: str = "pseudo") -> np.ndarray:
+        """Compute the random point locations on the boundary."""
+
+    def periodic_point(self, x, component):
+        """
+        Compute the periodic image of x for periodic boundary condition.
+
+        Args:
+            x: A 2D array of shape (n, dim), where `n` is the number of points and
+                `dim` is the dimension of the geometry.
+            component: The component of the periodic direction.
+        """
+        raise NotImplementedError("{}.periodic_point to be implemented".format(self.idstr))
+
+    def background_points(self, x, dirn, dist2npt, shift):
+        """
+        Compute the background points for the collocation points.
+
+        Args:
+            x: A 2D array of shape (n, dim), where `n` is the number of points and
+                `dim` is the dimension of the geometry.
+            dirn: The direction of the background points. One of the following: -1 (left),
+                or 1 (right), or 0 (both direction).
+            dist2npt: A function which converts distance to the number of extra
+                points (not including x).
+            shift: The number of shift.
+        """
+        raise NotImplementedError("{}.background_points to be implemented".format(self.idstr))
+
+
+
+class Geometry(AbstractGeometry):
     def __init__(self, dim, bbox, diam):
+        super().__init__(dim)
         self.dim = dim
         self.bbox = bbox
         self.diam = min(diam, np.linalg.norm(bbox[1] - bbox[0]))