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]))

deepxde/pinnx/__init__.py

@@ -0,0 +1,39 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+__all__ = [
+    "callbacks",
+    "fnspace",
+    "geometry",
+    "grad",
+    "icbc",
+    "metrics",
+    "nn",
+    "problem",
+    "utils",
+    "Trainer",
+]
+
+from . import callbacks
+from . import fnspace
+from . import geometry
+from . import grad
+from . import icbc
+from . import metrics
+from . import nn
+from . import problem
+from . import utils
+from ._trainer import Trainer