From 7d825287664272557aafdb4f7da5d13d5c3b99f6 Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Mon, 21 Aug 2023 19:35:39 -0400
Subject: [PATCH 01/12] add typing and type annotations for ICBC

---
 deepxde/icbc/boundary_conditions.py | 110 +++++++++++++++++++---------
 deepxde/types.py                    |   8 ++
 2 files changed, 82 insertions(+), 36 deletions(-)
 create mode 100644 deepxde/types.py

diff --git a/deepxde/icbc/boundary_conditions.py b/deepxde/icbc/boundary_conditions.py
index e1f863a08..31bb03ac7 100644
--- a/deepxde/icbc/boundary_conditions.py
+++ b/deepxde/icbc/boundary_conditions.py
@@ -14,8 +14,10 @@
 import numbers
 from abc import ABC, abstractmethod
 from functools import wraps
+from typing import Any, Callable, List, Optional, overload, Union
 
 import numpy as np
+from numpy.typing import NDArray, ArrayLike
 
 from .. import backend as bkd
 from .. import config
@@ -23,6 +25,8 @@
 from .. import gradients as grad
 from .. import utils
 from ..backend import backend_name
+from ..geometry import Geometry
+from ..types import _Tensor, _TensorOrTensors
 
 
 class BC(ABC):
@@ -30,11 +34,10 @@ class BC(ABC):
 
     Args:
         geom: A ``deepxde.geometry.Geometry`` instance.
-        on_boundary: A function: (x, Geometry.on_boundary(x)) -> True/False.
-        component: The output component satisfying this BC.
+        on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
+        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
     """
-
-    def __init__(self, geom, on_boundary, component):
+    def __init__(self, geom: Geometry, on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int]):
         self.geom = geom
         self.on_boundary = lambda x, on: np.array(
             [on_boundary(x[i], on[i]) for i in range(len(x))]
@@ -45,32 +48,39 @@ def __init__(self, geom, on_boundary, component):
             utils.return_tensor(self.geom.boundary_normal)
         )
 
-    def filter(self, X):
+    def filter(self, X: NDArray[Any]):
         return X[self.on_boundary(X, self.geom.on_boundary(X))]
 
-    def collocation_points(self, X):
+    def collocation_points(self, X: NDArray[Any]):
         return self.filter(X)
 
-    def normal_derivative(self, X, inputs, outputs, beg, end):
+    def normal_derivative(self, X: NDArray[Any], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int):
         dydx = grad.jacobian(outputs, inputs, i=self.component, j=None)[beg:end]
         n = self.boundary_normal(X, beg, end, None)
         return bkd.sum(dydx * n, 1, keepdims=True)
 
     @abstractmethod
-    def error(self, X, inputs, outputs, beg, end, aux_var=None):
+    def error(self, X: NDArray[Any], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int, aux_var: Union[NDArray[np.float_], None] = None):
         """Returns the loss."""
         # aux_var is used in PI-DeepONet, where aux_var is the input function evaluated
         # at x.
 
 
 class DirichletBC(BC):
-    """Dirichlet boundary conditions: y(x) = func(x)."""
-
-    def __init__(self, geom, func, on_boundary, component=0):
+    """Dirichlet boundary conditions: `y(x) = func(x)`.
+    
+    Args:
+        geom: A ``deepxde.geometry.Geometry`` instance.
+        func: A function: `x` -> `y`.
+        on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
+        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
+    """
+    
+    def __init__(self, geom: Geometry, func: Callable[[NDArray[np.float_]], NDArray[np.float_]], on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int] = 0):
         super().__init__(geom, on_boundary, component)
         self.func = npfunc_range_autocache(utils.return_tensor(func))
 
-    def error(self, X, inputs, outputs, beg, end, aux_var=None):
+    def error(self, X, inputs, outputs, beg, end, aux_var = None):
         values = self.func(X, beg, end, aux_var)
         if bkd.ndim(values) == 2 and bkd.shape(values)[1] != 1:
             raise RuntimeError(
@@ -81,9 +91,16 @@ def error(self, X, inputs, outputs, beg, end, aux_var=None):
 
 
 class NeumannBC(BC):
-    """Neumann boundary conditions: dy/dn(x) = func(x)."""
+    """Neumann boundary conditions: `dy/dn(x) = func(x)`.
+    
+    Args:
+        geom: A ``deepxde.geometry.Geometry`` instance.
+        func: A function: `x` -> `dy/dn`.
+        on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
+        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
+    """
 
-    def __init__(self, geom, func, on_boundary, component=0):
+    def __init__(self, geom: Geometry, func: Callable[[NDArray[np.float_]], NDArray[np.float_]], on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int] = 0):
         super().__init__(geom, on_boundary, component)
         self.func = npfunc_range_autocache(utils.return_tensor(func))
 
@@ -93,9 +110,16 @@ def error(self, X, inputs, outputs, beg, end, aux_var=None):
 
 
 class RobinBC(BC):
-    """Robin boundary conditions: dy/dn(x) = func(x, y)."""
+    """Robin boundary conditions: `dy/dn(x) = func(x, y)`.
+    
+    Args:
+        geom: A ``deepxde.geometry.Geometry`` instance.
+        func: A function: `(x, y)` -> `dy/dn`.
+        on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
+        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
+    """
 
-    def __init__(self, geom, func, on_boundary, component=0):
+    def __init__(self, geom: Geometry, func: Callable[[NDArray[np.float_]], NDArray[np.float_]], on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int] = 0):
         super().__init__(geom, on_boundary, component)
         self.func = func
 
@@ -106,9 +130,17 @@ def error(self, X, inputs, outputs, beg, end, aux_var=None):
 
 
 class PeriodicBC(BC):
-    """Periodic boundary conditions on component_x."""
+    """Periodic boundary conditions on component_x.
+    
+    Args:
+        geom: A ``deepxde.geometry.Geometry`` instance.
+        component_x: The component of the input satisfying this BC.
+        on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
+        derivative_order: The derivative order of the output satisfying this BC.
+        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
+    """
 
-    def __init__(self, geom, component_x, on_boundary, derivative_order=0, component=0):
+    def __init__(self, geom: Geometry, component_x: int, on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], derivative_order: int = 0, component: Union[List[int], int] = 0):
         super().__init__(geom, on_boundary, component)
         self.component_x = component_x
         self.derivative_order = derivative_order
@@ -122,7 +154,7 @@ def collocation_points(self, X):
         X2 = self.geom.periodic_point(X1, self.component_x)
         return np.vstack((X1, X2))
 
-    def error(self, X, inputs, outputs, beg, end, aux_var=None):
+    def error(self, X, inputs, outputs, beg, end, aux_var = None):
         mid = beg + (end - beg) // 2
         if self.derivative_order == 0:
             yleft = outputs[beg:mid, self.component : self.component + 1]
@@ -135,16 +167,16 @@ def error(self, X, inputs, outputs, beg, end, aux_var=None):
 
 
 class OperatorBC(BC):
-    """General operator boundary conditions: func(inputs, outputs, X) = 0.
+    """General operator boundary conditions: `func(inputs, outputs, X) = 0`.
 
     Args:
-        geom: ``Geometry``.
-        func: A function takes arguments (`inputs`, `outputs`, `X`)
-            and outputs a tensor of size `N x 1`, where `N` is the length of `inputs`.
-            `inputs` and `outputs` are the network input and output tensors,
-            respectively; `X` are the NumPy array of the `inputs`.
+        geom: A ``deepxde.geometry.Geometry`` instance.
+        func: A function takes arguments `(inputs, outputs, X)` 
+        and outputs a tensor of size `N x 1`, where `N` is the length of `inputs`.
+        `inputs` and `outputs` are the network input and output tensors,
+        respectively; `X` are the NumPy array of the `inputs`.
         on_boundary: (x, Geometry.on_boundary(x)) -> True/False.
-
+        
     Warning:
         If you use `X` in `func`, then do not set ``num_test`` when you define
         ``dde.data.PDE`` or ``dde.data.TimePDE``, otherwise DeepXDE would throw an
@@ -153,7 +185,7 @@ class OperatorBC(BC):
         which cannot be fixed in an easy way for all backends.
     """
 
-    def __init__(self, geom, func, on_boundary):
+    def __init__(self, geom: Geometry, func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor], on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]]):
         super().__init__(geom, on_boundary, 0)
         self.func = func
 
@@ -161,7 +193,7 @@ def error(self, X, inputs, outputs, beg, end, aux_var=None):
         return self.func(inputs, outputs, X)[beg:end]
 
 
-class PointSetBC:
+class PointSetBC(BC):
     """Dirichlet boundary condition for a set of points.
 
     Compare the output (that associates with `points`) with `values` (target data).
@@ -181,7 +213,7 @@ class PointSetBC:
         shuffle: Randomize the order on each pass through the data when batching.
     """
 
-    def __init__(self, points, values, component=0, batch_size=None, shuffle=True):
+    def __init__(self, points: ArrayLike, values: ArrayLike, component: Union[List[int], int] = 0, batch_size: Union[int, None] = None, shuffle: bool = True):
         self.points = np.array(points, dtype=config.real(np))
         self.values = bkd.as_tensor(values, dtype=config.real(bkd.lib))
         self.component = component
@@ -233,7 +265,7 @@ def error(self, X, inputs, outputs, beg, end, aux_var=None):
         return outputs[beg:end, self.component] - self.values
 
 
-class PointSetOperatorBC:
+class PointSetOperatorBC(BC):
     """General operator boundary conditions for a set of points.
 
     Compare the function output, func, (that associates with `points`)
@@ -249,7 +281,7 @@ class PointSetOperatorBC:
             tensors, respectively; `X` are the NumPy array of the `inputs`.
     """
 
-    def __init__(self, points, values, func):
+    def __init__(self, points: ArrayLike, values: ArrayLike, func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor]):
         self.points = np.array(points, dtype=config.real(np))
         if not isinstance(values, numbers.Number) and values.shape[1] != 1:
             raise RuntimeError("PointSetOperatorBC should output 1D values")
@@ -262,6 +294,11 @@ def collocation_points(self, X):
     def error(self, X, inputs, outputs, beg, end, aux_var=None):
         return self.func(inputs, outputs, X)[beg:end] - self.values
 
+@overload
+def npfunc_range_autocache(func: Callable[[NDArray[np.float_]], NDArray[np.float_]]) -> NDArray[np.float_]: ...
+
+@overload
+def npfunc_range_autocache(func: Callable[[NDArray[np.float_], NDArray[np.float_]], Optional[NDArray[np.float_]]]) -> NDArray[np.float_]: ...
 
 def npfunc_range_autocache(func):
     """Call a NumPy function on a range of the input ndarray.
@@ -288,25 +325,26 @@ def npfunc_range_autocache(func):
     # key. However, IC/BC is only for dde.data.PDE, where the ndarray is fixed. So we
     # can simply use id of X as the key, as what we do for gradients.
 
-    cache = {}
+    cache = {}    
 
     @wraps(func)
-    def wrapper_nocache(X, beg, end, _):
+    def wrapper_nocache(X: NDArray[np.float_], beg: int, end: int, _) -> NDArray[np.float_]:
         return func(X[beg:end])
 
     @wraps(func)
-    def wrapper_nocache_auxiliary(X, beg, end, aux_var):
+    def wrapper_nocache_auxiliary(X: NDArray[np.float_], beg: int, end: int, aux_var: NDArray[np.float_]) -> NDArray[np.float_]:
+        aux_var: callable
         return func(X[beg:end], aux_var[beg:end])
 
     @wraps(func)
-    def wrapper_cache(X, beg, end, _):
+    def wrapper_cache(X: NDArray[np.float_], beg: int, end: int, _) -> NDArray[np.float_]:
         key = (id(X), beg, end)
         if key not in cache:
             cache[key] = func(X[beg:end])
         return cache[key]
 
     @wraps(func)
-    def wrapper_cache_auxiliary(X, beg, end, aux_var):
+    def wrapper_cache_auxiliary(X: NDArray[np.float_], beg: int, end: int, aux_var: NDArray[np.float_]) -> NDArray[np.float_]:
         # Even if X is the same one, aux_var could be different
         key = (id(X), beg, end)
         if key not in cache:
diff --git a/deepxde/types.py b/deepxde/types.py
new file mode 100644
index 000000000..646402e0e
--- /dev/null
+++ b/deepxde/types.py
@@ -0,0 +1,8 @@
+import numpy as np
+from collections.abc import Sequence
+from typing import Callable, Dict, List, Optional, Tuple, Union, TypeVar
+from numpy.typing import NDArray, ArrayLike
+
+# Tensor from any backend
+_Tensor = TypeVar("_Tensor")
+_TensorOrTensors = Union[_Tensor, Sequence[_Tensor]]
\ No newline at end of file

From 25cd00473d8655e93f1fadc373e3081826c693e0 Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Mon, 21 Aug 2023 19:44:38 -0400
Subject: [PATCH 02/12] update

---
 deepxde/icbc/boundary_conditions.py |  8 ++++----
 deepxde/icbc/initial_conditions.py  | 16 ++++++++++------
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/deepxde/icbc/boundary_conditions.py b/deepxde/icbc/boundary_conditions.py
index 31bb03ac7..da8556131 100644
--- a/deepxde/icbc/boundary_conditions.py
+++ b/deepxde/icbc/boundary_conditions.py
@@ -48,19 +48,19 @@ def __init__(self, geom: Geometry, on_boundary: Callable[[NDArray[Any], NDArray[
             utils.return_tensor(self.geom.boundary_normal)
         )
 
-    def filter(self, X: NDArray[Any]):
+    def filter(self, X: NDArray[Any]) -> NDArray[np.bool_]:
         return X[self.on_boundary(X, self.geom.on_boundary(X))]
 
-    def collocation_points(self, X: NDArray[Any]):
+    def collocation_points(self, X: NDArray[Any]) -> NDArray[Any]:
         return self.filter(X)
 
-    def normal_derivative(self, X: NDArray[Any], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int):
+    def normal_derivative(self, X: NDArray[Any], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int) -> _Tensor:
         dydx = grad.jacobian(outputs, inputs, i=self.component, j=None)[beg:end]
         n = self.boundary_normal(X, beg, end, None)
         return bkd.sum(dydx * n, 1, keepdims=True)
 
     @abstractmethod
-    def error(self, X: NDArray[Any], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int, aux_var: Union[NDArray[np.float_], None] = None):
+    def error(self, X: NDArray[Any], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int, aux_var: Union[NDArray[np.float_], None] = None) -> _Tensor:
         """Returns the loss."""
         # aux_var is used in PI-DeepONet, where aux_var is the input function evaluated
         # at x.
diff --git a/deepxde/icbc/initial_conditions.py b/deepxde/icbc/initial_conditions.py
index a7ca57f2f..d8a94cca6 100644
--- a/deepxde/icbc/initial_conditions.py
+++ b/deepxde/icbc/initial_conditions.py
@@ -2,17 +2,21 @@
 
 __all__ = ["IC"]
 
+from typing import Any, Callable, List, Optional, overload, Union
+
 import numpy as np
+from numpy.typing import NDArray, ArrayLike
 
 from .boundary_conditions import npfunc_range_autocache
 from .. import backend as bkd
 from .. import utils
+from ..geometry import Geometry
+from ..types import _Tensor, _TensorOrTensors
 
-
-class IC:
+class IC():
     """Initial conditions: y([x, t0]) = func([x, t0])."""
 
-    def __init__(self, geom, func, on_initial, component=0):
+    def __init__(self, geom: Geometry, func: Callable[[NDArray[np.float_]], NDArray[np.float_]], on_initial: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int] = 0):
         self.geom = geom
         self.func = npfunc_range_autocache(utils.return_tensor(func))
         self.on_initial = lambda x, on: np.array(
@@ -20,13 +24,13 @@ def __init__(self, geom, func, on_initial, component=0):
         )
         self.component = component
 
-    def filter(self, X):
+    def filter(self, X: NDArray[np.float_]) -> NDArray[np.bool_]:
         return X[self.on_initial(X, self.geom.on_initial(X))]
 
-    def collocation_points(self, X):
+    def collocation_points(self, X: NDArray[np.float_]) -> NDArray[np.float_]:
         return self.filter(X)
 
-    def error(self, X, inputs, outputs, beg, end, aux_var=None):
+    def error(self, X: NDArray[np.float_], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int, aux_var: Union[NDArray[np.float_], None] = None) -> _Tensor:
         values = self.func(X, beg, end, aux_var)
         if bkd.ndim(values) == 2 and bkd.shape(values)[1] != 1:
             raise RuntimeError(

From 68a85fa2484d5b78124e1f45f2cf12c15aa58641 Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Mon, 21 Aug 2023 19:45:15 -0400
Subject: [PATCH 03/12] format

---
 deepxde/icbc/boundary_conditions.py | 134 ++++++++++++++++++++++------
 deepxde/icbc/initial_conditions.py  |  21 ++++-
 2 files changed, 126 insertions(+), 29 deletions(-)

diff --git a/deepxde/icbc/boundary_conditions.py b/deepxde/icbc/boundary_conditions.py
index da8556131..78fdc5efd 100644
--- a/deepxde/icbc/boundary_conditions.py
+++ b/deepxde/icbc/boundary_conditions.py
@@ -37,7 +37,13 @@ class BC(ABC):
         on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
         component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
     """
-    def __init__(self, geom: Geometry, on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int]):
+
+    def __init__(
+        self,
+        geom: Geometry,
+        on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
+        component: Union[List[int], int],
+    ):
         self.geom = geom
         self.on_boundary = lambda x, on: np.array(
             [on_boundary(x[i], on[i]) for i in range(len(x))]
@@ -54,13 +60,28 @@ def filter(self, X: NDArray[Any]) -> NDArray[np.bool_]:
     def collocation_points(self, X: NDArray[Any]) -> NDArray[Any]:
         return self.filter(X)
 
-    def normal_derivative(self, X: NDArray[Any], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int) -> _Tensor:
+    def normal_derivative(
+        self,
+        X: NDArray[Any],
+        inputs: _TensorOrTensors,
+        outputs: _Tensor,
+        beg: int,
+        end: int,
+    ) -> _Tensor:
         dydx = grad.jacobian(outputs, inputs, i=self.component, j=None)[beg:end]
         n = self.boundary_normal(X, beg, end, None)
         return bkd.sum(dydx * n, 1, keepdims=True)
 
     @abstractmethod
-    def error(self, X: NDArray[Any], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int, aux_var: Union[NDArray[np.float_], None] = None) -> _Tensor:
+    def error(
+        self,
+        X: NDArray[Any],
+        inputs: _TensorOrTensors,
+        outputs: _Tensor,
+        beg: int,
+        end: int,
+        aux_var: Union[NDArray[np.float_], None] = None,
+    ) -> _Tensor:
         """Returns the loss."""
         # aux_var is used in PI-DeepONet, where aux_var is the input function evaluated
         # at x.
@@ -68,19 +89,25 @@ def error(self, X: NDArray[Any], inputs: _TensorOrTensors, outputs: _Tensor, beg
 
 class DirichletBC(BC):
     """Dirichlet boundary conditions: `y(x) = func(x)`.
-    
+
     Args:
         geom: A ``deepxde.geometry.Geometry`` instance.
         func: A function: `x` -> `y`.
         on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
         component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
     """
-    
-    def __init__(self, geom: Geometry, func: Callable[[NDArray[np.float_]], NDArray[np.float_]], on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int] = 0):
+
+    def __init__(
+        self,
+        geom: Geometry,
+        func: Callable[[NDArray[np.float_]], NDArray[np.float_]],
+        on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
+        component: Union[List[int], int] = 0,
+    ):
         super().__init__(geom, on_boundary, component)
         self.func = npfunc_range_autocache(utils.return_tensor(func))
 
-    def error(self, X, inputs, outputs, beg, end, aux_var = None):
+    def error(self, X, inputs, outputs, beg, end, aux_var=None):
         values = self.func(X, beg, end, aux_var)
         if bkd.ndim(values) == 2 and bkd.shape(values)[1] != 1:
             raise RuntimeError(
@@ -92,7 +119,7 @@ def error(self, X, inputs, outputs, beg, end, aux_var = None):
 
 class NeumannBC(BC):
     """Neumann boundary conditions: `dy/dn(x) = func(x)`.
-    
+
     Args:
         geom: A ``deepxde.geometry.Geometry`` instance.
         func: A function: `x` -> `dy/dn`.
@@ -100,7 +127,13 @@ class NeumannBC(BC):
         component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
     """
 
-    def __init__(self, geom: Geometry, func: Callable[[NDArray[np.float_]], NDArray[np.float_]], on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int] = 0):
+    def __init__(
+        self,
+        geom: Geometry,
+        func: Callable[[NDArray[np.float_]], NDArray[np.float_]],
+        on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
+        component: Union[List[int], int] = 0,
+    ):
         super().__init__(geom, on_boundary, component)
         self.func = npfunc_range_autocache(utils.return_tensor(func))
 
@@ -111,7 +144,7 @@ def error(self, X, inputs, outputs, beg, end, aux_var=None):
 
 class RobinBC(BC):
     """Robin boundary conditions: `dy/dn(x) = func(x, y)`.
-    
+
     Args:
         geom: A ``deepxde.geometry.Geometry`` instance.
         func: A function: `(x, y)` -> `dy/dn`.
@@ -119,7 +152,13 @@ class RobinBC(BC):
         component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
     """
 
-    def __init__(self, geom: Geometry, func: Callable[[NDArray[np.float_]], NDArray[np.float_]], on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int] = 0):
+    def __init__(
+        self,
+        geom: Geometry,
+        func: Callable[[NDArray[np.float_]], NDArray[np.float_]],
+        on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
+        component: Union[List[int], int] = 0,
+    ):
         super().__init__(geom, on_boundary, component)
         self.func = func
 
@@ -131,7 +170,7 @@ def error(self, X, inputs, outputs, beg, end, aux_var=None):
 
 class PeriodicBC(BC):
     """Periodic boundary conditions on component_x.
-    
+
     Args:
         geom: A ``deepxde.geometry.Geometry`` instance.
         component_x: The component of the input satisfying this BC.
@@ -140,7 +179,14 @@ class PeriodicBC(BC):
         component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
     """
 
-    def __init__(self, geom: Geometry, component_x: int, on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], derivative_order: int = 0, component: Union[List[int], int] = 0):
+    def __init__(
+        self,
+        geom: Geometry,
+        component_x: int,
+        on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
+        derivative_order: int = 0,
+        component: Union[List[int], int] = 0,
+    ):
         super().__init__(geom, on_boundary, component)
         self.component_x = component_x
         self.derivative_order = derivative_order
@@ -154,7 +200,7 @@ def collocation_points(self, X):
         X2 = self.geom.periodic_point(X1, self.component_x)
         return np.vstack((X1, X2))
 
-    def error(self, X, inputs, outputs, beg, end, aux_var = None):
+    def error(self, X, inputs, outputs, beg, end, aux_var=None):
         mid = beg + (end - beg) // 2
         if self.derivative_order == 0:
             yleft = outputs[beg:mid, self.component : self.component + 1]
@@ -171,12 +217,12 @@ class OperatorBC(BC):
 
     Args:
         geom: A ``deepxde.geometry.Geometry`` instance.
-        func: A function takes arguments `(inputs, outputs, X)` 
+        func: A function takes arguments `(inputs, outputs, X)`
         and outputs a tensor of size `N x 1`, where `N` is the length of `inputs`.
         `inputs` and `outputs` are the network input and output tensors,
         respectively; `X` are the NumPy array of the `inputs`.
         on_boundary: (x, Geometry.on_boundary(x)) -> True/False.
-        
+
     Warning:
         If you use `X` in `func`, then do not set ``num_test`` when you define
         ``dde.data.PDE`` or ``dde.data.TimePDE``, otherwise DeepXDE would throw an
@@ -185,7 +231,12 @@ class OperatorBC(BC):
         which cannot be fixed in an easy way for all backends.
     """
 
-    def __init__(self, geom: Geometry, func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor], on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]]):
+    def __init__(
+        self,
+        geom: Geometry,
+        func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor],
+        on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
+    ):
         super().__init__(geom, on_boundary, 0)
         self.func = func
 
@@ -213,7 +264,14 @@ class PointSetBC(BC):
         shuffle: Randomize the order on each pass through the data when batching.
     """
 
-    def __init__(self, points: ArrayLike, values: ArrayLike, component: Union[List[int], int] = 0, batch_size: Union[int, None] = None, shuffle: bool = True):
+    def __init__(
+        self,
+        points: ArrayLike,
+        values: ArrayLike,
+        component: Union[List[int], int] = 0,
+        batch_size: Union[int, None] = None,
+        shuffle: bool = True,
+    ):
         self.points = np.array(points, dtype=config.real(np))
         self.values = bkd.as_tensor(values, dtype=config.real(bkd.lib))
         self.component = component
@@ -281,7 +339,12 @@ class PointSetOperatorBC(BC):
             tensors, respectively; `X` are the NumPy array of the `inputs`.
     """
 
-    def __init__(self, points: ArrayLike, values: ArrayLike, func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor]):
+    def __init__(
+        self,
+        points: ArrayLike,
+        values: ArrayLike,
+        func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor],
+    ):
         self.points = np.array(points, dtype=config.real(np))
         if not isinstance(values, numbers.Number) and values.shape[1] != 1:
             raise RuntimeError("PointSetOperatorBC should output 1D values")
@@ -294,11 +357,22 @@ def collocation_points(self, X):
     def error(self, X, inputs, outputs, beg, end, aux_var=None):
         return self.func(inputs, outputs, X)[beg:end] - self.values
 
+
 @overload
-def npfunc_range_autocache(func: Callable[[NDArray[np.float_]], NDArray[np.float_]]) -> NDArray[np.float_]: ...
+def npfunc_range_autocache(
+    func: Callable[[NDArray[np.float_]], NDArray[np.float_]]
+) -> NDArray[np.float_]:
+    ...
+
 
 @overload
-def npfunc_range_autocache(func: Callable[[NDArray[np.float_], NDArray[np.float_]], Optional[NDArray[np.float_]]]) -> NDArray[np.float_]: ...
+def npfunc_range_autocache(
+    func: Callable[
+        [NDArray[np.float_], NDArray[np.float_]], Optional[NDArray[np.float_]]
+    ]
+) -> NDArray[np.float_]:
+    ...
+
 
 def npfunc_range_autocache(func):
     """Call a NumPy function on a range of the input ndarray.
@@ -325,26 +399,34 @@ def npfunc_range_autocache(func):
     # key. However, IC/BC is only for dde.data.PDE, where the ndarray is fixed. So we
     # can simply use id of X as the key, as what we do for gradients.
 
-    cache = {}    
+    cache = {}
 
     @wraps(func)
-    def wrapper_nocache(X: NDArray[np.float_], beg: int, end: int, _) -> NDArray[np.float_]:
+    def wrapper_nocache(
+        X: NDArray[np.float_], beg: int, end: int, _
+    ) -> NDArray[np.float_]:
         return func(X[beg:end])
 
     @wraps(func)
-    def wrapper_nocache_auxiliary(X: NDArray[np.float_], beg: int, end: int, aux_var: NDArray[np.float_]) -> NDArray[np.float_]:
+    def wrapper_nocache_auxiliary(
+        X: NDArray[np.float_], beg: int, end: int, aux_var: NDArray[np.float_]
+    ) -> NDArray[np.float_]:
         aux_var: callable
         return func(X[beg:end], aux_var[beg:end])
 
     @wraps(func)
-    def wrapper_cache(X: NDArray[np.float_], beg: int, end: int, _) -> NDArray[np.float_]:
+    def wrapper_cache(
+        X: NDArray[np.float_], beg: int, end: int, _
+    ) -> NDArray[np.float_]:
         key = (id(X), beg, end)
         if key not in cache:
             cache[key] = func(X[beg:end])
         return cache[key]
 
     @wraps(func)
-    def wrapper_cache_auxiliary(X: NDArray[np.float_], beg: int, end: int, aux_var: NDArray[np.float_]) -> NDArray[np.float_]:
+    def wrapper_cache_auxiliary(
+        X: NDArray[np.float_], beg: int, end: int, aux_var: NDArray[np.float_]
+    ) -> NDArray[np.float_]:
         # Even if X is the same one, aux_var could be different
         key = (id(X), beg, end)
         if key not in cache:
diff --git a/deepxde/icbc/initial_conditions.py b/deepxde/icbc/initial_conditions.py
index d8a94cca6..a38a89880 100644
--- a/deepxde/icbc/initial_conditions.py
+++ b/deepxde/icbc/initial_conditions.py
@@ -13,10 +13,17 @@
 from ..geometry import Geometry
 from ..types import _Tensor, _TensorOrTensors
 
-class IC():
+
+class IC:
     """Initial conditions: y([x, t0]) = func([x, t0])."""
 
-    def __init__(self, geom: Geometry, func: Callable[[NDArray[np.float_]], NDArray[np.float_]], on_initial: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]], component: Union[List[int], int] = 0):
+    def __init__(
+        self,
+        geom: Geometry,
+        func: Callable[[NDArray[np.float_]], NDArray[np.float_]],
+        on_initial: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
+        component: Union[List[int], int] = 0,
+    ):
         self.geom = geom
         self.func = npfunc_range_autocache(utils.return_tensor(func))
         self.on_initial = lambda x, on: np.array(
@@ -30,7 +37,15 @@ def filter(self, X: NDArray[np.float_]) -> NDArray[np.bool_]:
     def collocation_points(self, X: NDArray[np.float_]) -> NDArray[np.float_]:
         return self.filter(X)
 
-    def error(self, X: NDArray[np.float_], inputs: _TensorOrTensors, outputs: _Tensor, beg: int, end: int, aux_var: Union[NDArray[np.float_], None] = None) -> _Tensor:
+    def error(
+        self,
+        X: NDArray[np.float_],
+        inputs: _TensorOrTensors,
+        outputs: _Tensor,
+        beg: int,
+        end: int,
+        aux_var: Union[NDArray[np.float_], None] = None,
+    ) -> _Tensor:
         values = self.func(X, beg, end, aux_var)
         if bkd.ndim(values) == 2 and bkd.shape(values)[1] != 1:
             raise RuntimeError(

From 74cee5bc2d2bd54585a6b9b44d5a24f9c0f2d549 Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Mon, 21 Aug 2023 22:12:36 -0400
Subject: [PATCH 04/12] update

---
 deepxde/icbc/boundary_conditions.py | 12 ++++++------
 deepxde/icbc/initial_conditions.py  |  4 ++--
 deepxde/types.py                    |  3 +--
 3 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/deepxde/icbc/boundary_conditions.py b/deepxde/icbc/boundary_conditions.py
index 78fdc5efd..7e8ec398a 100644
--- a/deepxde/icbc/boundary_conditions.py
+++ b/deepxde/icbc/boundary_conditions.py
@@ -14,7 +14,7 @@
 import numbers
 from abc import ABC, abstractmethod
 from functools import wraps
-from typing import Any, Callable, List, Optional, overload, Union
+from typing import Any, Callable, List, Optional, overload, Tuple, Union
 
 import numpy as np
 from numpy.typing import NDArray, ArrayLike
@@ -26,7 +26,7 @@
 from .. import utils
 from ..backend import backend_name
 from ..geometry import Geometry
-from ..types import _Tensor, _TensorOrTensors
+from ..types import _Tensor
 
 
 class BC(ABC):
@@ -63,7 +63,7 @@ def collocation_points(self, X: NDArray[Any]) -> NDArray[Any]:
     def normal_derivative(
         self,
         X: NDArray[Any],
-        inputs: _TensorOrTensors,
+        inputs: Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor],
         outputs: _Tensor,
         beg: int,
         end: int,
@@ -76,7 +76,7 @@ def normal_derivative(
     def error(
         self,
         X: NDArray[Any],
-        inputs: _TensorOrTensors,
+        inputs: Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor],
         outputs: _Tensor,
         beg: int,
         end: int,
@@ -234,7 +234,7 @@ class OperatorBC(BC):
     def __init__(
         self,
         geom: Geometry,
-        func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor],
+        func: Callable[[Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor], _Tensor, NDArray[np.float_]], _Tensor],
         on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
     ):
         super().__init__(geom, on_boundary, 0)
@@ -343,7 +343,7 @@ def __init__(
         self,
         points: ArrayLike,
         values: ArrayLike,
-        func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor],
+        func: Callable[[Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor], _Tensor, NDArray[np.float_]], _Tensor],
     ):
         self.points = np.array(points, dtype=config.real(np))
         if not isinstance(values, numbers.Number) and values.shape[1] != 1:
diff --git a/deepxde/icbc/initial_conditions.py b/deepxde/icbc/initial_conditions.py
index a38a89880..d0e6167f8 100644
--- a/deepxde/icbc/initial_conditions.py
+++ b/deepxde/icbc/initial_conditions.py
@@ -11,7 +11,7 @@
 from .. import backend as bkd
 from .. import utils
 from ..geometry import Geometry
-from ..types import _Tensor, _TensorOrTensors
+from ..types import _Tensor
 
 
 class IC:
@@ -40,7 +40,7 @@ def collocation_points(self, X: NDArray[np.float_]) -> NDArray[np.float_]:
     def error(
         self,
         X: NDArray[np.float_],
-        inputs: _TensorOrTensors,
+        inputs: Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor],
         outputs: _Tensor,
         beg: int,
         end: int,
diff --git a/deepxde/types.py b/deepxde/types.py
index 646402e0e..1b5f2d110 100644
--- a/deepxde/types.py
+++ b/deepxde/types.py
@@ -4,5 +4,4 @@
 from numpy.typing import NDArray, ArrayLike
 
 # Tensor from any backend
-_Tensor = TypeVar("_Tensor")
-_TensorOrTensors = Union[_Tensor, Sequence[_Tensor]]
\ No newline at end of file
+_Tensor = TypeVar("_Tensor")
\ No newline at end of file

From 45e8ac9bf7dcfa4231b26216b77427abd08dd161 Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Mon, 21 Aug 2023 22:16:09 -0400
Subject: [PATCH 05/12] fix sequence

---
 deepxde/icbc/boundary_conditions.py | 10 +++++-----
 deepxde/icbc/initial_conditions.py  |  4 ++--
 deepxde/types.py                    |  6 +++---
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/deepxde/icbc/boundary_conditions.py b/deepxde/icbc/boundary_conditions.py
index 7e8ec398a..4d0d9dde7 100644
--- a/deepxde/icbc/boundary_conditions.py
+++ b/deepxde/icbc/boundary_conditions.py
@@ -26,7 +26,7 @@
 from .. import utils
 from ..backend import backend_name
 from ..geometry import Geometry
-from ..types import _Tensor
+from ..types import _Tensor, _TensorOrTensors
 
 
 class BC(ABC):
@@ -63,7 +63,7 @@ def collocation_points(self, X: NDArray[Any]) -> NDArray[Any]:
     def normal_derivative(
         self,
         X: NDArray[Any],
-        inputs: Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor],
+        inputs: _TensorOrTensors,
         outputs: _Tensor,
         beg: int,
         end: int,
@@ -76,7 +76,7 @@ def normal_derivative(
     def error(
         self,
         X: NDArray[Any],
-        inputs: Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor],
+        inputs: _TensorOrTensors,
         outputs: _Tensor,
         beg: int,
         end: int,
@@ -234,7 +234,7 @@ class OperatorBC(BC):
     def __init__(
         self,
         geom: Geometry,
-        func: Callable[[Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor], _Tensor, NDArray[np.float_]], _Tensor],
+        func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor],
         on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
     ):
         super().__init__(geom, on_boundary, 0)
@@ -343,7 +343,7 @@ def __init__(
         self,
         points: ArrayLike,
         values: ArrayLike,
-        func: Callable[[Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor], _Tensor, NDArray[np.float_]], _Tensor],
+        func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor],
     ):
         self.points = np.array(points, dtype=config.real(np))
         if not isinstance(values, numbers.Number) and values.shape[1] != 1:
diff --git a/deepxde/icbc/initial_conditions.py b/deepxde/icbc/initial_conditions.py
index d0e6167f8..a38a89880 100644
--- a/deepxde/icbc/initial_conditions.py
+++ b/deepxde/icbc/initial_conditions.py
@@ -11,7 +11,7 @@
 from .. import backend as bkd
 from .. import utils
 from ..geometry import Geometry
-from ..types import _Tensor
+from ..types import _Tensor, _TensorOrTensors
 
 
 class IC:
@@ -40,7 +40,7 @@ def collocation_points(self, X: NDArray[np.float_]) -> NDArray[np.float_]:
     def error(
         self,
         X: NDArray[np.float_],
-        inputs: Union[Tuple[_Tensor, ...], List[_Tensor], _Tensor],
+        inputs: _TensorOrTensors,
         outputs: _Tensor,
         beg: int,
         end: int,
diff --git a/deepxde/types.py b/deepxde/types.py
index 1b5f2d110..853400b19 100644
--- a/deepxde/types.py
+++ b/deepxde/types.py
@@ -1,7 +1,7 @@
 import numpy as np
-from collections.abc import Sequence
-from typing import Callable, Dict, List, Optional, Tuple, Union, TypeVar
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union, TypeVar
 from numpy.typing import NDArray, ArrayLike
 
 # Tensor from any backend
-_Tensor = TypeVar("_Tensor")
\ No newline at end of file
+_Tensor = TypeVar("_Tensor")
+_TensorOrTensors = Union[_Tensor, Sequence[_Tensor]]
\ No newline at end of file

From 9f4c722f061d29a0ea5833a168145d0c1b2226ea Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Mon, 21 Aug 2023 22:18:55 -0400
Subject: [PATCH 06/12] rename as tensor

---
 deepxde/icbc/boundary_conditions.py | 18 +++++++++---------
 deepxde/icbc/initial_conditions.py  |  8 ++++----
 deepxde/types.py                    |  4 ++--
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/deepxde/icbc/boundary_conditions.py b/deepxde/icbc/boundary_conditions.py
index 4d0d9dde7..48c680622 100644
--- a/deepxde/icbc/boundary_conditions.py
+++ b/deepxde/icbc/boundary_conditions.py
@@ -26,7 +26,7 @@
 from .. import utils
 from ..backend import backend_name
 from ..geometry import Geometry
-from ..types import _Tensor, _TensorOrTensors
+from ..types import Tensor, TensorOrTensors
 
 
 class BC(ABC):
@@ -63,11 +63,11 @@ def collocation_points(self, X: NDArray[Any]) -> NDArray[Any]:
     def normal_derivative(
         self,
         X: NDArray[Any],
-        inputs: _TensorOrTensors,
-        outputs: _Tensor,
+        inputs: TensorOrTensors,
+        outputs: Tensor,
         beg: int,
         end: int,
-    ) -> _Tensor:
+    ) -> Tensor:
         dydx = grad.jacobian(outputs, inputs, i=self.component, j=None)[beg:end]
         n = self.boundary_normal(X, beg, end, None)
         return bkd.sum(dydx * n, 1, keepdims=True)
@@ -76,12 +76,12 @@ def normal_derivative(
     def error(
         self,
         X: NDArray[Any],
-        inputs: _TensorOrTensors,
-        outputs: _Tensor,
+        inputs: TensorOrTensors,
+        outputs: Tensor,
         beg: int,
         end: int,
         aux_var: Union[NDArray[np.float_], None] = None,
-    ) -> _Tensor:
+    ) -> Tensor:
         """Returns the loss."""
         # aux_var is used in PI-DeepONet, where aux_var is the input function evaluated
         # at x.
@@ -234,7 +234,7 @@ class OperatorBC(BC):
     def __init__(
         self,
         geom: Geometry,
-        func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor],
+        func: Callable[[TensorOrTensors, Tensor, NDArray[np.float_]], Tensor],
         on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
     ):
         super().__init__(geom, on_boundary, 0)
@@ -343,7 +343,7 @@ def __init__(
         self,
         points: ArrayLike,
         values: ArrayLike,
-        func: Callable[[_TensorOrTensors, _Tensor, NDArray[np.float_]], _Tensor],
+        func: Callable[[TensorOrTensors, Tensor, NDArray[np.float_]], Tensor],
     ):
         self.points = np.array(points, dtype=config.real(np))
         if not isinstance(values, numbers.Number) and values.shape[1] != 1:
diff --git a/deepxde/icbc/initial_conditions.py b/deepxde/icbc/initial_conditions.py
index a38a89880..db047c30d 100644
--- a/deepxde/icbc/initial_conditions.py
+++ b/deepxde/icbc/initial_conditions.py
@@ -11,7 +11,7 @@
 from .. import backend as bkd
 from .. import utils
 from ..geometry import Geometry
-from ..types import _Tensor, _TensorOrTensors
+from ..types import Tensor, TensorOrTensors
 
 
 class IC:
@@ -40,12 +40,12 @@ def collocation_points(self, X: NDArray[np.float_]) -> NDArray[np.float_]:
     def error(
         self,
         X: NDArray[np.float_],
-        inputs: _TensorOrTensors,
-        outputs: _Tensor,
+        inputs: TensorOrTensors,
+        outputs: Tensor,
         beg: int,
         end: int,
         aux_var: Union[NDArray[np.float_], None] = None,
-    ) -> _Tensor:
+    ) -> Tensor:
         values = self.func(X, beg, end, aux_var)
         if bkd.ndim(values) == 2 and bkd.shape(values)[1] != 1:
             raise RuntimeError(
diff --git a/deepxde/types.py b/deepxde/types.py
index 853400b19..0addf34e0 100644
--- a/deepxde/types.py
+++ b/deepxde/types.py
@@ -3,5 +3,5 @@
 from numpy.typing import NDArray, ArrayLike
 
 # Tensor from any backend
-_Tensor = TypeVar("_Tensor")
-_TensorOrTensors = Union[_Tensor, Sequence[_Tensor]]
\ No newline at end of file
+Tensor = TypeVar("Tensor")
+TensorOrTensors = Union[Tensor, Sequence[Tensor]]
\ No newline at end of file

From a491700e99cde833c85eef764fdeebf65c6288d9 Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Mon, 21 Aug 2023 22:22:46 -0400
Subject: [PATCH 07/12] format fix

---
 deepxde/icbc/boundary_conditions.py | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/deepxde/icbc/boundary_conditions.py b/deepxde/icbc/boundary_conditions.py
index 48c680622..fb777ac7f 100644
--- a/deepxde/icbc/boundary_conditions.py
+++ b/deepxde/icbc/boundary_conditions.py
@@ -35,7 +35,8 @@ class BC(ABC):
     Args:
         geom: A ``deepxde.geometry.Geometry`` instance.
         on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
-        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
+        component: The output component satisfying this BC, should be provided 
+            if ``BC.error`` involves derivatives and the output has multiple components.
     """
 
     def __init__(
@@ -94,7 +95,8 @@ class DirichletBC(BC):
         geom: A ``deepxde.geometry.Geometry`` instance.
         func: A function: `x` -> `y`.
         on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
-        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
+        component: The output component satisfying this BC, should be provided 
+            if ``BC.error`` involves derivatives and the output has multiple components.
     """
 
     def __init__(
@@ -124,7 +126,8 @@ class NeumannBC(BC):
         geom: A ``deepxde.geometry.Geometry`` instance.
         func: A function: `x` -> `dy/dn`.
         on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
-        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
+        component: The output component satisfying this BC, should be provided 
+            if ``BC.error`` involves derivatives and the output has multiple components.
     """
 
     def __init__(
@@ -149,7 +152,8 @@ class RobinBC(BC):
         geom: A ``deepxde.geometry.Geometry`` instance.
         func: A function: `(x, y)` -> `dy/dn`.
         on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
-        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
+        component: The output component satisfying this BC, should be provided 
+            if ``BC.error`` involves derivatives and the output has multiple components.
     """
 
     def __init__(
@@ -176,7 +180,8 @@ class PeriodicBC(BC):
         component_x: The component of the input satisfying this BC.
         on_boundary: A function: `(x, Geometry.on_boundary(x))` -> True/False.
         derivative_order: The derivative order of the output satisfying this BC.
-        component: The output component satisfying this BC, should be provided if ``BC.error`` involves derivatives and the output has multiple components.
+        component: The output component satisfying this BC, should be provided 
+            if ``BC.error`` involves derivatives and the output has multiple components.
     """
 
     def __init__(
@@ -218,9 +223,9 @@ class OperatorBC(BC):
     Args:
         geom: A ``deepxde.geometry.Geometry`` instance.
         func: A function takes arguments `(inputs, outputs, X)`
-        and outputs a tensor of size `N x 1`, where `N` is the length of `inputs`.
-        `inputs` and `outputs` are the network input and output tensors,
-        respectively; `X` are the NumPy array of the `inputs`.
+            and outputs a tensor of size `N x 1`, where `N` is the length of `inputs`.
+            `inputs` and `outputs` are the network input and output tensors,
+            respectively; `X` are the NumPy array of the `inputs`.
         on_boundary: (x, Geometry.on_boundary(x)) -> True/False.
 
     Warning:
@@ -255,7 +260,7 @@ class PointSetBC(BC):
         points: An array of points where the corresponding target values are known and
             used for training.
         values: A scalar or a 2D-array of values that gives the exact solution of the problem.
-        component: Integer or a list of integers. The output components satisfying this BC.
+            omponent: Integer or a list of integers. The output components satisfying this BC.
             List of integers only supported for the backend PyTorch.
         batch_size: The number of points per minibatch, or `None` to return all points.
             This is only supported for the backend PyTorch and PaddlePaddle.

From 3eacc22741b670bca6d560f0a1457328f105a5e7 Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Tue, 22 Aug 2023 22:11:15 -0400
Subject: [PATCH 08/12] update annotation representation

---
 deepxde/icbc/boundary_conditions.py | 24 +++++++++++-------------
 deepxde/icbc/initial_conditions.py  |  7 ++++---
 deepxde/types.py                    |  7 +++----
 3 files changed, 18 insertions(+), 20 deletions(-)

diff --git a/deepxde/icbc/boundary_conditions.py b/deepxde/icbc/boundary_conditions.py
index fb777ac7f..01bd7a26d 100644
--- a/deepxde/icbc/boundary_conditions.py
+++ b/deepxde/icbc/boundary_conditions.py
@@ -1,3 +1,4 @@
+from __future__ import annotations
 """Boundary conditions."""
 
 __all__ = [
@@ -14,7 +15,7 @@
 import numbers
 from abc import ABC, abstractmethod
 from functools import wraps
-from typing import Any, Callable, List, Optional, overload, Tuple, Union
+from typing import Any, Callable, overload
 
 import numpy as np
 from numpy.typing import NDArray, ArrayLike
@@ -43,7 +44,7 @@ def __init__(
         self,
         geom: Geometry,
         on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
-        component: Union[List[int], int],
+        component: list[int] | int,
     ):
         self.geom = geom
         self.on_boundary = lambda x, on: np.array(
@@ -81,7 +82,7 @@ def error(
         outputs: Tensor,
         beg: int,
         end: int,
-        aux_var: Union[NDArray[np.float_], None] = None,
+        aux_var: NDArray[np.float_] | None = None,
     ) -> Tensor:
         """Returns the loss."""
         # aux_var is used in PI-DeepONet, where aux_var is the input function evaluated
@@ -104,7 +105,7 @@ def __init__(
         geom: Geometry,
         func: Callable[[NDArray[np.float_]], NDArray[np.float_]],
         on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
-        component: Union[List[int], int] = 0,
+        component: list[int] | int = 0,
     ):
         super().__init__(geom, on_boundary, component)
         self.func = npfunc_range_autocache(utils.return_tensor(func))
@@ -135,7 +136,7 @@ def __init__(
         geom: Geometry,
         func: Callable[[NDArray[np.float_]], NDArray[np.float_]],
         on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
-        component: Union[List[int], int] = 0,
+        component: list[int] | int = 0,
     ):
         super().__init__(geom, on_boundary, component)
         self.func = npfunc_range_autocache(utils.return_tensor(func))
@@ -161,7 +162,7 @@ def __init__(
         geom: Geometry,
         func: Callable[[NDArray[np.float_]], NDArray[np.float_]],
         on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
-        component: Union[List[int], int] = 0,
+        component: list[int] | int = 0,
     ):
         super().__init__(geom, on_boundary, component)
         self.func = func
@@ -190,7 +191,7 @@ def __init__(
         component_x: int,
         on_boundary: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
         derivative_order: int = 0,
-        component: Union[List[int], int] = 0,
+        component: list[int] | int = 0,
     ):
         super().__init__(geom, on_boundary, component)
         self.component_x = component_x
@@ -273,8 +274,8 @@ def __init__(
         self,
         points: ArrayLike,
         values: ArrayLike,
-        component: Union[List[int], int] = 0,
-        batch_size: Union[int, None] = None,
+        component: list[int] | int = 0,
+        batch_size: int | None = None,
         shuffle: bool = True,
     ):
         self.points = np.array(points, dtype=config.real(np))
@@ -372,9 +373,7 @@ def npfunc_range_autocache(
 
 @overload
 def npfunc_range_autocache(
-    func: Callable[
-        [NDArray[np.float_], NDArray[np.float_]], Optional[NDArray[np.float_]]
-    ]
+    func: Callable[[NDArray[np.float_], NDArray[np.float_]], NDArray[np.float_]]
 ) -> NDArray[np.float_]:
     ...
 
@@ -416,7 +415,6 @@ def wrapper_nocache(
     def wrapper_nocache_auxiliary(
         X: NDArray[np.float_], beg: int, end: int, aux_var: NDArray[np.float_]
     ) -> NDArray[np.float_]:
-        aux_var: callable
         return func(X[beg:end], aux_var[beg:end])
 
     @wraps(func)
diff --git a/deepxde/icbc/initial_conditions.py b/deepxde/icbc/initial_conditions.py
index db047c30d..cc4e0bd6e 100644
--- a/deepxde/icbc/initial_conditions.py
+++ b/deepxde/icbc/initial_conditions.py
@@ -1,8 +1,9 @@
+from __future__ import annotations
 """Initial conditions."""
 
 __all__ = ["IC"]
 
-from typing import Any, Callable, List, Optional, overload, Union
+from typing import Any, Callable
 
 import numpy as np
 from numpy.typing import NDArray, ArrayLike
@@ -22,7 +23,7 @@ def __init__(
         geom: Geometry,
         func: Callable[[NDArray[np.float_]], NDArray[np.float_]],
         on_initial: Callable[[NDArray[Any], NDArray[Any]], NDArray[np.bool_]],
-        component: Union[List[int], int] = 0,
+        component: list[int] | int = 0,
     ):
         self.geom = geom
         self.func = npfunc_range_autocache(utils.return_tensor(func))
@@ -44,7 +45,7 @@ def error(
         outputs: Tensor,
         beg: int,
         end: int,
-        aux_var: Union[NDArray[np.float_], None] = None,
+        aux_var: NDArray[np.float_] | None = None,
     ) -> Tensor:
         values = self.func(X, beg, end, aux_var)
         if bkd.ndim(values) == 2 and bkd.shape(values)[1] != 1:
diff --git a/deepxde/types.py b/deepxde/types.py
index 0addf34e0..8931c142c 100644
--- a/deepxde/types.py
+++ b/deepxde/types.py
@@ -1,7 +1,6 @@
-import numpy as np
-from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union, TypeVar
-from numpy.typing import NDArray, ArrayLike
+from __future__ import annotations
+from typing import Sequence, TypeVar
 
 # Tensor from any backend
 Tensor = TypeVar("Tensor")
-TensorOrTensors = Union[Tensor, Sequence[Tensor]]
\ No newline at end of file
+TensorOrTensors = Tensor | Sequence[Tensor]
\ No newline at end of file

From f8c2d99930ea248caddeffac4fbc68ed03c86b91 Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Tue, 22 Aug 2023 22:46:10 -0400
Subject: [PATCH 09/12] a = b | c is only valid in py3.10

---
 deepxde/types.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/deepxde/types.py b/deepxde/types.py
index 8931c142c..8cff698e8 100644
--- a/deepxde/types.py
+++ b/deepxde/types.py
@@ -1,6 +1,8 @@
 from __future__ import annotations
-from typing import Sequence, TypeVar
+from typing import Sequence, TypeVar, Union
 
 # Tensor from any backend
 Tensor = TypeVar("Tensor")
-TensorOrTensors = Tensor | Sequence[Tensor]
\ No newline at end of file
+TensorOrTensors = Union[Tensor, Sequence[Tensor]]
+SparseTensor = TypeVar("SparseTensor")
+dtype = TypeVar("dtype")
\ No newline at end of file

From 19fdaed9c3d4d81dffba41e071f6aa19ca23b2c4 Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Mon, 4 Sep 2023 14:55:00 -0400
Subject: [PATCH 10/12] typing for geometry

---
 deepxde/backend/backend.py      | 109 ++++++++++++++++++--------------
 deepxde/geometry/csg.py         |  15 ++---
 deepxde/geometry/geometry.py    |  47 ++++++++------
 deepxde/geometry/geometry_1d.py |   9 +--
 deepxde/geometry/geometry_2d.py |  58 +++++++++--------
 deepxde/geometry/geometry_3d.py |   9 ++-
 deepxde/geometry/geometry_nd.py |  54 +++++++++-------
 deepxde/geometry/pointcloud.py  |   5 +-
 deepxde/geometry/sampler.py     |   8 ++-
 deepxde/geometry/timedomain.py  |  11 ++--
 deepxde/utils/external.py       |   4 +-
 11 files changed, 185 insertions(+), 144 deletions(-)

diff --git a/deepxde/backend/backend.py b/deepxde/backend/backend.py
index b744f9174..866522949 100644
--- a/deepxde/backend/backend.py
+++ b/deepxde/backend/backend.py
@@ -1,3 +1,9 @@
+from __future__ import annotations
+
+from numbers import Number
+from typing import Sequence, overload
+from ..types import Tensor, dtype, SparseTensor, TensorOrTensors
+from numpy.typing import NDArray, ArrayLike
 """This file defines the unified tensor framework interface required by DeepXDE.
 
 The principles of this interface:
@@ -26,7 +32,7 @@
 # Tensor, data type and context interfaces
 
 
-def data_type_dict():
+def data_type_dict() -> dict[str, object]:
     """Returns a dictionary from data type string to the data type.
 
     The dictionary should include at least:
@@ -58,7 +64,7 @@ def data_type_dict():
     """
 
 
-def is_gpu_available():
+def is_gpu_available() -> bool:
     """Returns a bool indicating if GPU is currently available.
 
     Returns:
@@ -66,11 +72,11 @@ def is_gpu_available():
     """
 
 
-def is_tensor(obj):
+def is_tensor(obj: object) -> bool:
     """Returns True if `obj` is a backend-native type tensor."""
 
 
-def shape(input_tensor):
+def shape(input_tensor: Tensor) -> Sequence[int]:
     """Return the shape of the tensor.
 
     Args:
@@ -81,7 +87,7 @@ def shape(input_tensor):
     """
 
 
-def size(input_tensor):
+def size(input_tensor: Tensor) -> int:
     """Return the total number of elements in the input tensor.
 
     Args:
@@ -92,7 +98,7 @@ def size(input_tensor):
     """
 
 
-def ndim(input_tensor):
+def ndim(input_tensor: Tensor) -> int:
     """Returns the number of dimensions of the tensor.
 
     Args:
@@ -103,7 +109,7 @@ def ndim(input_tensor):
     """
 
 
-def transpose(tensor, axes=None):
+def transpose(tensor: Tensor, axes: Sequence[int] | int | None = None) -> Tensor:
     """Reverse or permute the axes of a tensor; returns the modified array.
 
     For a tensor with two axes, transpose gives the matrix transpose.
@@ -117,7 +123,7 @@ def transpose(tensor, axes=None):
     """
 
 
-def reshape(tensor, shape):
+def reshape(tensor: Tensor, shape: Sequence[int]) -> Tensor:
     """Gives a new shape to a tensor without changing its data.
 
     Args:
@@ -130,7 +136,7 @@ def reshape(tensor, shape):
     """
 
 
-def Variable(initial_value, dtype=None):
+def Variable(initial_value: Number, dtype: dtype = None) -> Tensor:
     """Return a trainable variable.
 
     Args:
@@ -140,7 +146,7 @@ def Variable(initial_value, dtype=None):
     """
 
 
-def as_tensor(data, dtype=None):
+def as_tensor(data: ArrayLike, dtype: dtype = None) -> Tensor:
     """Convert the data to a Tensor.
 
     If the data is already a tensor and has the same dtype, directly return.
@@ -155,7 +161,7 @@ def as_tensor(data, dtype=None):
     """
 
 
-def sparse_tensor(indices, values, shape):
+def sparse_tensor(indices: Sequence[Sequence[Number, Number]], values: Tensor, shape: Sequence[int]) -> SparseTensor:
     """Construct a sparse tensor based on given indices, values and shape.
 
     Args:
@@ -170,7 +176,7 @@ def sparse_tensor(indices, values, shape):
     """
 
 
-def from_numpy(np_array):
+def from_numpy(np_array: NDArray) -> Tensor:
     """Create a tensor that shares the underlying numpy array memory, if possible.
 
     Args:
@@ -181,7 +187,7 @@ def from_numpy(np_array):
     """
 
 
-def to_numpy(input_tensor):
+def to_numpy(input_tensor: Tensor) -> NDArray:
     """Create a numpy ndarray that shares the same underlying storage, if possible.
 
     Args:
@@ -192,7 +198,7 @@ def to_numpy(input_tensor):
     """
 
 
-def concat(values, axis):
+def concat(values: TensorOrTensors, axis: int) -> Tensor:
     """Returns the concatenation of the input tensors along the given dim.
 
     Args:
@@ -204,7 +210,7 @@ def concat(values, axis):
     """
 
 
-def stack(values, axis):
+def stack(values: TensorOrTensors, axis: int) -> Tensor:
     """Returns the stack of the input tensors along the given dim.
 
     Args:
@@ -216,7 +222,7 @@ def stack(values, axis):
     """
 
 
-def expand_dims(tensor, axis):
+def expand_dims(tensor: Tensor, axis: int) -> Tensor:
     """Expand dim for tensor along given axis.
 
     Args:
@@ -228,7 +234,7 @@ def expand_dims(tensor, axis):
     """
 
 
-def reverse(tensor, axis):
+def reverse(tensor: Tensor, axis: int) -> Tensor:
     """Reverse the order of elements along the given axis.
 
     Args:
@@ -240,7 +246,7 @@ def reverse(tensor, axis):
     """
 
 
-def roll(tensor, shift, axis):
+def roll(tensor: Tensor, shift: int | Sequence[int], axis: int | Sequence[int]) -> Tensor:
     """Roll the tensor along the given axis (axes).
 
     Args:
@@ -261,67 +267,67 @@ def roll(tensor, shift, axis):
 # implementation in each framework.
 
 
-def lgamma(x):
+def lgamma(x: Tensor) -> Tensor:
     """Computes the natural logarithm of the absolute value of the gamma function of x
     element-wise.
     """
 
 
-def elu(x):
+def elu(x: Tensor) -> Tensor:
     """Computes the exponential linear function."""
 
 
-def relu(x):
+def relu(x: Tensor) -> Tensor:
     """Applies the rectified linear unit activation function."""
 
 
-def gelu(x):
+def gelu(x: Tensor) -> Tensor:
     """Computes Gaussian Error Linear Unit function."""
 
 
-def selu(x):
+def selu(x: Tensor) -> Tensor:
     """Computes scaled exponential linear."""
 
 
-def sigmoid(x):
+def sigmoid(x: Tensor) -> Tensor:
     """Computes sigmoid of x element-wise."""
 
 
-def silu(x):
+def silu(x: Tensor) -> Tensor:
     """Sigmoid Linear Unit (SiLU) function, also known as the swish function.
     silu(x) = x * sigmoid(x).
     """
 
 
-def sin(x):
+def sin(x: Tensor) -> Tensor:
     """Computes sine of x element-wise."""
 
 
-def cos(x):
+def cos(x: Tensor) -> Tensor:
     """Computes cosine of x element-wise."""
 
 
-def exp(x):
+def exp(x: Tensor) -> Tensor:
     """Computes exponential of x element-wise."""
 
 
-def square(x):
+def square(x: Tensor) -> Tensor:
     """Returns the square of the elements of input."""
 
 
-def abs(x):
+def abs(x: Tensor) -> Tensor:
     """Computes the absolute value element-wise."""
 
 
-def minimum(x, y):
+def minimum(x: Tensor, y: Tensor) -> Tensor:
     """Returns the minimum of x and y (i.e. x < y ? x : y) element-wise."""
 
 
-def tanh(x):
+def tanh(x: Tensor) -> Tensor:
     """Computes hyperbolic tangent of x element-wise."""
 
 
-def pow(x, y):
+def pow(x: Tensor, y: Number | Tensor) -> Tensor:
     """Computes the power of one value to another: x ^ y."""
 
 
@@ -332,15 +338,15 @@ def pow(x, y):
 # implementation in each framework.
 
 
-def mean(input_tensor, dim, keepdims=False):
+def mean(input_tensor: Tensor, dim: int | Sequence[int], keepdims: bool = False) -> Tensor:
     """Returns the mean value of the input tensor in the given dimension dim."""
 
 
-def reduce_mean(input_tensor):
+def reduce_mean(input_tensor: Tensor) -> Tensor:
     """Returns the mean value of all elements in the input tensor."""
 
 
-def sum(input_tensor, dim, keepdims=False):
+def sum(input_tensor: Tensor, dim: int | Sequence[int], keepdims: Tensor = False):
     """Returns the sum of the input tensor along the given dim.
 
     Args:
@@ -353,7 +359,7 @@ def sum(input_tensor, dim, keepdims=False):
     """
 
 
-def reduce_sum(input_tensor):
+def reduce_sum(input_tensor: Tensor) -> Tensor:
     """Returns the sum of all elements in the input tensor.
 
     Args:
@@ -364,7 +370,7 @@ def reduce_sum(input_tensor):
     """
 
 
-def prod(input_tensor, dim, keepdims=False):
+def prod(input_tensor: Tensor, dim: int | Sequence[int], keepdims: bool = False) -> Tensor:
     """Returns the product of the input tensor along the given dim.
 
     Args:
@@ -377,7 +383,7 @@ def prod(input_tensor, dim, keepdims=False):
     """
 
 
-def reduce_prod(input_tensor):
+def reduce_prod(input_tensor: Tensor) -> Tensor:
     """Returns the product of all elements in the input tensor.
 
     Args:
@@ -388,7 +394,7 @@ def reduce_prod(input_tensor):
     """
 
 
-def min(input_tensor, dim, keepdims=False):
+def min(input_tensor: Tensor, dim: int | Sequence[int], keepdims: bool = False) -> Tensor:
     """Returns the minimum of the input tensor along the given dim.
 
     Args:
@@ -401,7 +407,7 @@ def min(input_tensor, dim, keepdims=False):
     """
 
 
-def reduce_min(input_tensor):
+def reduce_min(input_tensor: Tensor) -> Tensor:
     """Returns the minimum of all elements in the input tensor.
 
     Args:
@@ -412,7 +418,7 @@ def reduce_min(input_tensor):
     """
 
 
-def max(input_tensor, dim, keepdims=False):
+def max(input_tensor: Tensor, dim: int | Sequence[int], keepdims: bool = False) -> Tensor:
     """Returns the maximum of the input tensor along the given dim.
 
     Args:
@@ -425,7 +431,7 @@ def max(input_tensor, dim, keepdims=False):
     """
 
 
-def reduce_max(input_tensor):
+def reduce_max(input_tensor: Tensor) -> Tensor:
     """Returns the maximum of all elements in the input tensor.
 
     Args:
@@ -436,7 +442,7 @@ def reduce_max(input_tensor):
     """
 
 
-def norm(tensor, ord=None, axis=None, keepdims=False):
+def norm(tensor: Tensor, ord: Number | None = None, axis: int | None = None, keepdims: bool = False) -> Tensor:
     """Computes a vector norm.
 
     Due to the incompatibility of different backends, only some vector norms are 
@@ -457,7 +463,7 @@ def norm(tensor, ord=None, axis=None, keepdims=False):
     """
 
 
-def zeros(shape, dtype):
+def zeros(shape: Sequence[int], dtype: dtype) -> Tensor:
     """Creates a tensor with all elements set to zero.
 
     Args:
@@ -469,7 +475,7 @@ def zeros(shape, dtype):
     """
 
 
-def zeros_like(input_tensor):
+def zeros_like(input_tensor: Tensor) -> Tensor:
     """Create a zero tensor with the same shape, dtype and context of the given tensor.
 
     Args:
@@ -480,7 +486,7 @@ def zeros_like(input_tensor):
     """
 
 
-def matmul(x, y):
+def matmul(x: Tensor, y: Tensor) -> Tensor:
     """Compute matrix multiplication for two matrices x and y.
 
     Args:
@@ -492,6 +498,14 @@ def matmul(x, y):
     """
 
 
+@overload
+def sparse_dense_matmul(x: SparseTensor, y: Tensor) -> Tensor: ...
+
+
+@overload
+def sparse_dense_matmul(x: SparseTensor, y: SparseTensor) -> SparseTensor: ...
+
+
 def sparse_dense_matmul(x, y):
     """Compute matrix multiplication of a sparse matrix x and a sparse/dense matrix y.
 
@@ -502,3 +516,4 @@ def sparse_dense_matmul(x, y):
     Returns:
         Tensor: The multiplication result.
     """
+sparse_dense_matmul()
\ No newline at end of file
diff --git a/deepxde/geometry/csg.py b/deepxde/geometry/csg.py
index 7bbcf7b44..e580c6ee4 100644
--- a/deepxde/geometry/csg.py
+++ b/deepxde/geometry/csg.py
@@ -1,13 +1,12 @@
 import numpy as np
 
-from . import geometry
+from .geometry import Geometry
 from .. import config
 
-
-class CSGUnion(geometry.Geometry):
+class CSGUnion(Geometry):
     """Construct an object by CSG Union."""
 
-    def __init__(self, geom1, geom2):
+    def __init__(self, geom1: Geometry, geom2: Geometry):
         if geom1.dim != geom2.dim:
             raise ValueError(
                 "{} | {} failed (dimensions do not match).".format(
@@ -101,10 +100,10 @@ def periodic_point(self, x, component):
         return x
 
 
-class CSGDifference(geometry.Geometry):
+class CSGDifference(Geometry):
     """Construct an object by CSG Difference."""
 
-    def __init__(self, geom1, geom2):
+    def __init__(self, geom1: Geometry, geom2: Geometry):
         if geom1.dim != geom2.dim:
             raise ValueError(
                 "{} - {} failed (dimensions do not match).".format(
@@ -183,10 +182,10 @@ def periodic_point(self, x, component):
         return x
 
 
-class CSGIntersection(geometry.Geometry):
+class CSGIntersection(Geometry):
     """Construct an object by CSG Intersection."""
 
-    def __init__(self, geom1, geom2):
+    def __init__(self, geom1: Geometry, geom2: Geometry):
         if geom1.dim != geom2.dim:
             raise ValueError(
                 "{} & {} failed (dimensions do not match).".format(
diff --git a/deepxde/geometry/geometry.py b/deepxde/geometry/geometry.py
index 7564921bd..804bc16db 100644
--- a/deepxde/geometry/geometry.py
+++ b/deepxde/geometry/geometry.py
@@ -1,37 +1,42 @@
+from __future__ import annotations
+
 import abc
-from typing import Literal
+from typing import Callable, Literal
+from numbers import Number
 
 import numpy as np
+from numpy.typing import NDArray
 
+from ..types import Tensor, TensorOrTensors
 
 class Geometry(abc.ABC):
-    def __init__(self, dim, bbox, diam):
+    def __init__(self, dim: int, bbox: NDArray[np.float_], diam: Number):
         self.dim = dim
         self.bbox = bbox
         self.diam = min(diam, np.linalg.norm(bbox[1] - bbox[0]))
         self.idstr = type(self).__name__
 
     @abc.abstractmethod
-    def inside(self, x):
+    def inside(self, x: NDArray[np.float_]) -> NDArray[np.bool_]:
         """Check if x is inside the geometry (including the boundary)."""
 
     @abc.abstractmethod
-    def on_boundary(self, x):
+    def on_boundary(self, x: NDArray[np.float_]) -> NDArray[np.bool_]:
         """Check if x is on the geometry boundary."""
 
-    def distance2boundary(self, x, dirn):
+    def distance2boundary(self, x: NDArray[np.float_], dirn: Number) -> NDArray[np.float_]:
         raise NotImplementedError(
             "{}.distance2boundary to be implemented".format(self.idstr)
         )
 
-    def mindist2boundary(self, x):
+    def mindist2boundary(self, x: NDArray[np.float_]) -> NDArray[np.float_]:
         raise NotImplementedError(
             "{}.mindist2boundary to be implemented".format(self.idstr)
         )
 
     def boundary_constraint_factor(
-        self, x, smoothness: Literal["C0", "C0+", "Cinf"] = "C0+"
-    ):
+        self, x: NDArray[np.float_], smoothness: Literal["C0", "C0+", "Cinf"] = "C0+"
+    ) -> Tensor:
         """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).
@@ -74,13 +79,13 @@ def boundary_constraint_factor(
             "{}.boundary_constraint_factor to be implemented".format(self.idstr)
         )
 
-    def boundary_normal(self, x):
+    def boundary_normal(self, x: NDArray[np.float_]) -> NDArray[np.float_]:
         """Compute the unit normal at x for Neumann or Robin boundary conditions."""
         raise NotImplementedError(
             "{}.boundary_normal to be implemented".format(self.idstr)
         )
 
-    def uniform_points(self, n, boundary=True):
+    def uniform_points(self, n: int, boundary: bool = True)-> NDArray[np.float_]:
         """Compute the equispaced point locations in the geometry."""
         print(
             "Warning: {}.uniform_points not implemented. Use random_points instead.".format(
@@ -90,10 +95,10 @@ def uniform_points(self, n, boundary=True):
         return self.random_points(n)
 
     @abc.abstractmethod
-    def random_points(self, n, random="pseudo"):
+    def random_points(self, n: int, random: str = "pseudo") -> NDArray[np.float_]:
         """Compute the random point locations in the geometry."""
 
-    def uniform_boundary_points(self, n):
+    def uniform_boundary_points(self, n: int) -> NDArray[np.float_]:
         """Compute the equispaced point locations on the boundary."""
         print(
             "Warning: {}.uniform_boundary_points not implemented. Use random_boundary_points instead.".format(
@@ -103,51 +108,51 @@ def uniform_boundary_points(self, n):
         return self.random_boundary_points(n)
 
     @abc.abstractmethod
-    def random_boundary_points(self, n, random="pseudo"):
+    def random_boundary_points(self, n: int, random: str = "pseudo") -> NDArray[np.float_]:
         """Compute the random point locations on the boundary."""
 
-    def periodic_point(self, x, component):
+    def periodic_point(self, x: NDArray[np.float_], component: int | list[int]) -> NDArray[np.float_]:
         """Compute the periodic image of x for periodic boundary condition."""
         raise NotImplementedError(
             "{}.periodic_point to be implemented".format(self.idstr)
         )
 
-    def background_points(self, x, dirn, dist2npt, shift):
+    def background_points(self, x: NDArray[np.float_], dirn: Number, dist2npt: Callable[[NDArray[np.float_]], int], shift: int) -> NDArray[np.float_]:
         raise NotImplementedError(
             "{}.background_points to be implemented".format(self.idstr)
         )
 
-    def union(self, other):
+    def union(self, other: Geometry) -> Geometry:
         """CSG Union."""
         from . import csg
 
         return csg.CSGUnion(self, other)
 
-    def __or__(self, other):
+    def __or__(self, other: Geometry) -> Geometry:
         """CSG Union."""
         from . import csg
 
         return csg.CSGUnion(self, other)
 
-    def difference(self, other):
+    def difference(self, other: Geometry) -> Geometry:
         """CSG Difference."""
         from . import csg
 
         return csg.CSGDifference(self, other)
 
-    def __sub__(self, other):
+    def __sub__(self, other: Geometry) -> Geometry:
         """CSG Difference."""
         from . import csg
 
         return csg.CSGDifference(self, other)
 
-    def intersection(self, other):
+    def intersection(self, other: Geometry) -> Geometry:
         """CSG Intersection."""
         from . import csg
 
         return csg.CSGIntersection(self, other)
 
-    def __and__(self, other):
+    def __and__(self, other: Geometry) -> Geometry:
         """CSG Intersection."""
         from . import csg
 
diff --git a/deepxde/geometry/geometry_1d.py b/deepxde/geometry/geometry_1d.py
index 98e0eccd1..ac57f39f8 100644
--- a/deepxde/geometry/geometry_1d.py
+++ b/deepxde/geometry/geometry_1d.py
@@ -1,4 +1,5 @@
-from typing import Literal, Union
+from numbers import Number
+from typing import Literal
 
 import numpy as np
 
@@ -10,7 +11,7 @@
 
 
 class Interval(Geometry):
-    def __init__(self, l, r):
+    def __init__(self, l: Number, r: Number):
         super().__init__(1, (np.array([l]), np.array([r])), r - l)
         self.l, self.r = l, r
 
@@ -30,7 +31,7 @@ def boundary_constraint_factor(
         self,
         x,
         smoothness: Literal["C0", "C0+", "Cinf"] = "C0+",
-        where: Union[None, Literal["left", "right"]] = None,
+        where: Literal["left", "right"] | None = None,
     ):
         """Compute the hard constraint factor at x for the boundary.
 
@@ -115,7 +116,7 @@ def uniform_points(self, n, boundary=True):
             self.l, self.r, num=n + 1, endpoint=False, dtype=config.real(np)
         )[1:, None]
 
-    def log_uniform_points(self, n, boundary=True):
+    def log_uniform_points(self, n: int, boundary: bool = True):
         eps = 0 if self.l > 0 else np.finfo(config.real(np)).eps
         l = np.log(self.l + eps)
         r = np.log(self.r + eps)
diff --git a/deepxde/geometry/geometry_2d.py b/deepxde/geometry/geometry_2d.py
index 2a36e7fe8..1db2bd33f 100644
--- a/deepxde/geometry/geometry_2d.py
+++ b/deepxde/geometry/geometry_2d.py
@@ -1,8 +1,11 @@
-__all__ = ["Disk", "Ellipse", "Polygon", "Rectangle", "StarShaped", "Triangle"]
+from __future__ import annotations
 
-from typing import Union, Literal
+__all__ = ["Disk", "Ellipse", "Polygon", "Rectangle", "StarShaped", "Triangle"]
+from numbers import Number
+from typing import Any, Literal
 
 import numpy as np
+from numpy.typing import ArrayLike, NDArray
 from scipy import spatial
 
 from .geometry import Geometry
@@ -10,10 +13,13 @@
 from .sampler import sample
 from .. import backend as bkd
 from .. import config
+from ..types import Tensor
 from ..utils import isclose, vectorize
 
-
 class Disk(Hypersphere):
+    def __init__(self, center, radius):
+        super().__init__(center, radius)
+    
     def inside(self, x):
         return np.linalg.norm(x - self.center, axis=-1) <= self.radius
 
@@ -84,7 +90,7 @@ class Ellipse(Geometry):
             counterclockwise about the center.
     """
 
-    def __init__(self, center, semimajor, semiminor, angle=0):
+    def __init__(self, center: ArrayLike, semimajor: Number, semiminor: Number, angle: Number = 0):
         self.center = np.array(center, dtype=config.real(np))
         self.semimajor = semimajor
         self.semiminor = semiminor
@@ -272,10 +278,10 @@ def random_boundary_points(self, n, random="pseudo"):
 
     def _boundary_constraint_factor_inside(
         self,
-        x,
-        where: Union[None, Literal["left", "right", "bottom", "top"]] = None,
+        x: NDArray[np.float_],
+        where: Literal["left", "right", "bottom", "top"] | None = None,
         smoothness: Literal["C0", "C0+", "Cinf"] = "C0+",
-    ):
+    ) -> Tensor:
         """(Internal use only) Compute the hard constraint factor at `x` for the boundary.
         The points in `x` are assumed to live inside the geometry.
 
@@ -316,9 +322,9 @@ def boundary_constraint_factor(
         self,
         x,
         smoothness: Literal["C0", "C0+", "Cinf"] = "C0+",
-        where: Union[None, Literal["left", "right", "bottom", "top"]] = None,
+        where: Literal["left", "right", "bottom", "top"] | None = None,
         inside: bool = True,
-    ):
+    ) -> Tensor:
         """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).
@@ -380,7 +386,7 @@ def boundary_constraint_factor(
             self.x12_tensor = bkd.as_tensor([self.xmin[0], self.xmax[1]])
             self.x21_tensor = bkd.as_tensor([self.xmax[0], self.xmin[1]])
 
-        dist_left = dist_right = dist_bottom = dist_top = None
+        dist_left = dist_right = dist_bottom = dist_top = 0.0
         if where is None or where == "left":
             dist_left = bkd.abs(
                 bkd.norm(x - self.x11_tensor, axis=-1, keepdims=True)
@@ -421,7 +427,7 @@ def boundary_constraint_factor(
         return dist_left * dist_right * dist_bottom * dist_top
 
     @staticmethod
-    def is_valid(vertices):
+    def is_valid(vertices: NDArray[np.float_]) -> bool:
         """Check if the geometry is a Rectangle."""
         return (
             len(vertices) == 4
@@ -451,7 +457,7 @@ class StarShaped(Geometry):
         coeffs_sin: i-th order coefficients for the i-th sin term (b_i).
     """
 
-    def __init__(self, center, radius, coeffs_cos, coeffs_sin):
+    def __init__(self, center: NDArray[np.float_], radius: Number, coeffs_cos: NDArray[np.float_], coeffs_sin: NDArray[np.float_]):
         self.center = np.array(center, dtype=config.real(np))
         self.radius = radius
         self.coeffs_cos = coeffs_cos
@@ -463,7 +469,7 @@ def __init__(self, center, radius, coeffs_cos, coeffs_sin):
             2 * max_radius,
         )
 
-    def _r_theta(self, theta):
+    def _r_theta(self, theta: Number) -> NDArray[np.float_]:
         """Define the parametrization r(theta) at angles theta."""
         result = self.radius * np.ones(theta.shape)
         for i, (coeff_cos, coeff_sin) in enumerate(
@@ -472,7 +478,7 @@ def _r_theta(self, theta):
             result += coeff_cos * np.cos(i * theta) + coeff_sin * np.sin(i * theta)
         return result
 
-    def _dr_theta(self, theta):
+    def _dr_theta(self, theta: Number) -> NDArray[np.float_]:
         """Evalutate the polar derivative r'(theta) at angles theta"""
         result = np.zeros(theta.shape)
         for i, (coeff_cos, coeff_sin) in enumerate(
@@ -483,7 +489,7 @@ def _dr_theta(self, theta):
             )
         return result
 
-    def inside(self, x):
+    def inside(self, x) -> NDArray[np.bool_]:
         r, theta = polar(x - self.center)
         r_theta = self._r_theta(theta)
         return r_theta >= r
@@ -537,7 +543,7 @@ class Triangle(Geometry):
     vertices will be re-ordered in counterclockwise (right hand rule).
     """
 
-    def __init__(self, x1, x2, x3):
+    def __init__(self, x1: Number, x2: Number, x3: Number):
         self.area = polygon_signed_area([x1, x2, x3])
         # Clockwise
         if self.area < 0:
@@ -684,7 +690,7 @@ def boundary_constraint_factor(
         self,
         x,
         smoothness: Literal["C0", "C0+", "Cinf"] = "C0+",
-        where: Union[None, Literal["x1-x2", "x1-x3", "x2-x3"]] = None,
+        where: Literal["x1-x2", "x1-x3", "x2-x3"] | None = None,
     ):
         """Compute the hard constraint factor at x for the boundary.
 
@@ -739,7 +745,7 @@ def boundary_constraint_factor(
             self.x2_tensor = bkd.as_tensor(self.x2)
             self.x3_tensor = bkd.as_tensor(self.x3)
 
-        diff_x1_x2 = diff_x1_x3 = diff_x2_x3 = None
+        # diff_x1_x2 = diff_x1_x3 = diff_x2_x3 = 0.0
         if where not in ["x1-x3", "x2-x3"]:
             diff_x1_x2 = (
                 bkd.norm(x - self.x1_tensor, axis=-1, keepdims=True)
@@ -779,7 +785,7 @@ class Polygon(Geometry):
             rule).
     """
 
-    def __init__(self, vertices):
+    def __init__(self, vertices: ArrayLike):
         self.vertices = np.array(vertices, dtype=config.real(np))
         if len(vertices) == 3:
             raise ValueError("The polygon is a triangle. Use Triangle instead.")
@@ -813,7 +819,7 @@ def __init__(self, vertices):
         self.normal = clockwise_rotation_90(self.segments.T).T
         self.normal = self.normal / np.linalg.norm(self.normal, axis=1).reshape(-1, 1)
 
-    def inside(self, x):
+    def inside(self, x) -> NDArray[np.bool_]:
         def wn_PnPoly(P, V):
             """Winding number algorithm.
 
@@ -927,7 +933,7 @@ def random_boundary_points(self, n, random="pseudo"):
         return np.vstack(x)
 
 
-def polygon_signed_area(vertices):
+def polygon_signed_area(vertices: ArrayLike) -> Number:
     """The (signed) area of a simple polygon.
 
     If the vertices are in the counterclockwise direction, then the area is positive; if
@@ -941,12 +947,12 @@ def polygon_signed_area(vertices):
     return 0.5 * (np.sum(x[:-1] * y[1:]) - np.sum(x[1:] * y[:-1]))
 
 
-def clockwise_rotation_90(v):
+def clockwise_rotation_90(v: NDArray[np.float_]) -> NDArray[np.float_]:
     """Rotate a vector of 90 degrees clockwise about the origin."""
     return np.array([v[1], -v[0]])
 
 
-def is_left(P0, P1, P2):
+def is_left(P0: ArrayLike, P1: ArrayLike, P2: ArrayLike) -> NDArray[Any]:
     """Test if a point is Left|On|Right of an infinite line.
 
     See: the January 2001 Algorithm "Area of 2D and 3D Triangles and Polygons".
@@ -963,7 +969,7 @@ def is_left(P0, P1, P2):
     return np.cross(P1 - P0, P2 - P0, axis=-1).reshape((-1, 1))
 
 
-def is_rectangle(vertices):
+def is_rectangle(vertices: ArrayLike) -> bool:
     """Check if the geometry is a rectangle.
 
     https://stackoverflow.com/questions/2303278/find-if-4-points-on-a-plane-form-a-rectangle/2304031
@@ -979,7 +985,7 @@ def is_rectangle(vertices):
     return np.allclose(d, np.full(4, d[0]))
 
 
-def is_on_line_segment(P0, P1, P2):
+def is_on_line_segment(P0: ArrayLike, P1: ArrayLike, P2: ArrayLike) -> bool:
     """Test if a point is between two other points on a line segment.
 
     Args:
@@ -1005,7 +1011,7 @@ def is_on_line_segment(P0, P1, P2):
     # or isclose(np.linalg.norm(v12), 0)  # check whether P2 is close to P1
 
 
-def polar(x):
+def polar(x: NDArray[np.float_]) -> tuple[NDArray[np.float_], NDArray[np.float_]]:
     """Get the polar coordinated for a 2d vector in cartesian coordinates."""
     r = np.sqrt(x[:, 0] ** 2 + x[:, 1] ** 2)
     theta = np.arctan2(x[:, 1], x[:, 0])
diff --git a/deepxde/geometry/geometry_3d.py b/deepxde/geometry/geometry_3d.py
index ae08a239d..b52071a52 100644
--- a/deepxde/geometry/geometry_3d.py
+++ b/deepxde/geometry/geometry_3d.py
@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import itertools
-from typing import Union, Literal
+from typing import Literal
 
 import numpy as np
 
@@ -72,9 +74,7 @@ def boundary_constraint_factor(
         self,
         x,
         smoothness: Literal["C0", "C0+", "Cinf"] = "C0+",
-        where: Union[
-            None, Literal["back", "front", "left", "right", "bottom", "top"]
-        ] = None,
+        where: Literal["back", "front", "left", "right", "bottom", "top"] | None = None,
         inside: bool = True,
     ):
         """Compute the hard constraint factor at x for the boundary.
@@ -138,7 +138,6 @@ def boundary_constraint_factor(
             self.xmin_tensor = bkd.as_tensor(self.xmin)
             self.xmax_tensor = bkd.as_tensor(self.xmax)
 
-        dist_l = dist_r = None
         if where not in ["front", "right", "top"]:
             dist_l = bkd.abs(
                 (x - self.xmin_tensor) / (self.xmax_tensor - self.xmin_tensor) * 2
diff --git a/deepxde/geometry/geometry_nd.py b/deepxde/geometry/geometry_nd.py
index a011cd417..db17009d9 100644
--- a/deepxde/geometry/geometry_nd.py
+++ b/deepxde/geometry/geometry_nd.py
@@ -1,7 +1,12 @@
+from __future__ import annotations
+
 import itertools
-from typing import Literal
+from numbers import Number
+from typing import Callable, Literal
 
 import numpy as np
+from numpy.typing import NDArray
+
 from scipy import stats
 from sklearn import preprocessing
 
@@ -9,11 +14,12 @@
 from .sampler import sample
 from .. import backend as bkd
 from .. import config
+from ..types import Tensor
 from ..utils import isclose
 
 
 class Hypercube(Geometry):
-    def __init__(self, xmin, xmax):
+    def __init__(self, xmin: NDArray[np.float_], xmax: NDArray[np.float_]):
         if len(xmin) != len(xmax):
             raise ValueError("Dimensions of xmin and xmax do not match.")
 
@@ -28,19 +34,19 @@ def __init__(self, xmin, xmax):
         )
         self.volume = np.prod(self.side_length)
 
-    def inside(self, x):
+    def inside(self, x: NDArray[np.float_]) -> NDArray[np.bool_]:
         return np.logical_and(
             np.all(x >= self.xmin, axis=-1), np.all(x <= self.xmax, axis=-1)
         )
 
-    def on_boundary(self, x):
+    def on_boundary(self, x: NDArray[np.float_]) -> NDArray[np.bool_]:
         _on_boundary = np.logical_or(
             np.any(isclose(x, self.xmin), axis=-1),
             np.any(isclose(x, self.xmax), axis=-1),
         )
         return np.logical_and(self.inside(x), _on_boundary)
 
-    def boundary_normal(self, x):
+    def boundary_normal(self, x: NDArray[np.float_]) -> NDArray[np.float_]:
         _n = -isclose(x, self.xmin).astype(config.real(np)) + isclose(x, self.xmax)
         # For vertices, the normal is averaged for all directions
         idx = np.count_nonzero(_n, axis=-1) > 1
@@ -53,7 +59,7 @@ def boundary_normal(self, x):
             _n[idx] /= l
         return _n
 
-    def uniform_points(self, n, boundary=True):
+    def uniform_points(self, n: int, boundary: bool = True) -> NDArray[np.float_]:
         dx = (self.volume / n) ** (1 / self.dim)
         xi = []
         for i in range(self.dim):
@@ -81,11 +87,11 @@ def uniform_points(self, n, boundary=True):
             )
         return x
 
-    def random_points(self, n, random="pseudo"):
+    def random_points(self, n: int, random: str = "pseudo") -> NDArray[np.float_]:
         x = sample(n, self.dim, random)
         return (self.xmax - self.xmin) * x + self.xmin
 
-    def random_boundary_points(self, n, random="pseudo"):
+    def random_boundary_points(self, n: int, random: str = "pseudo") -> NDArray[np.float_]:
         x = sample(n, self.dim, random)
         # Randomly pick a dimension
         rand_dim = np.random.randint(self.dim, size=n)
@@ -93,7 +99,7 @@ def random_boundary_points(self, n, random="pseudo"):
         x[np.arange(n), rand_dim] = np.round(x[np.arange(n), rand_dim])
         return (self.xmax - self.xmin) * x + self.xmin
 
-    def periodic_point(self, x, component):
+    def periodic_point(self, x: NDArray[np.float_], component: int | list[int]) -> NDArray[np.float_]:
         y = np.copy(x)
         _on_xmin = isclose(y[:, component], self.xmin[component])
         _on_xmax = isclose(y[:, component], self.xmax[component])
@@ -103,11 +109,11 @@ def periodic_point(self, x, component):
 
     def boundary_constraint_factor(
         self,
-        x,
+        x: NDArray[np.float_],
         smoothness: Literal["C0", "C0+", "Cinf"] = "C0",
-        where: None = None,
+        where: str | None = None,
         inside: bool = True,
-    ):
+    ) -> Tensor:
         """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).
@@ -184,7 +190,7 @@ def boundary_constraint_factor(
 
 
 class Hypersphere(Geometry):
-    def __init__(self, center, radius):
+    def __init__(self, center: Number, radius: Number):
         self.center = np.array(center, dtype=config.real(np))
         self.radius = radius
         super().__init__(
@@ -193,13 +199,13 @@ def __init__(self, center, radius):
 
         self._r2 = radius**2
 
-    def inside(self, x):
+    def inside(self, x: NDArray[np.float_]) -> NDArray[np.bool_]:
         return np.linalg.norm(x - self.center, axis=-1) <= self.radius
 
-    def on_boundary(self, x):
+    def on_boundary(self, x: NDArray[np.float_]) -> NDArray[np.bool_]:
         return isclose(np.linalg.norm(x - self.center, axis=-1), self.radius)
 
-    def distance2boundary_unitdirn(self, x, dirn):
+    def distance2boundary_unitdirn(self, x: NDArray[np.float_], dirn: Number) -> NDArray[np.float_]:
         # https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection
         xc = x - self.center
         ad = np.dot(xc, dirn)
@@ -207,15 +213,15 @@ def distance2boundary_unitdirn(self, x, dirn):
             config.real(np)
         )
 
-    def distance2boundary(self, x, dirn):
+    def distance2boundary(self, x: NDArray[np.float_], dirn: Number) -> NDArray[np.float_]:
         return self.distance2boundary_unitdirn(x, dirn / np.linalg.norm(dirn))
 
-    def mindist2boundary(self, x):
+    def mindist2boundary(self, x: NDArray[np.float_]) -> NDArray[np.float_]:
         return np.amin(self.radius - np.linalg.norm(x - self.center, axis=-1))
 
     def boundary_constraint_factor(
-        self, x, smoothness: Literal["C0", "C0+", "Cinf"] = "C0+"
-    ):
+        self, x: NDArray[np.float_], smoothness: Literal["C0", "C0+", "Cinf"] = "C0+"
+    ) -> Tensor:
         if smoothness not in ["C0", "C0+", "Cinf"]:
             raise ValueError("smoothness must be one of C0, C0+, Cinf")
 
@@ -230,13 +236,13 @@ def boundary_constraint_factor(
             dist = bkd.abs(dist)
         return dist
 
-    def boundary_normal(self, x):
+    def boundary_normal(self, x: NDArray[np.float_]) -> NDArray[np.float_]:
         _n = x - self.center
         l = np.linalg.norm(_n, axis=-1, keepdims=True)
         _n = _n / l * isclose(l, self.radius)
         return _n
 
-    def random_points(self, n, random="pseudo"):
+    def random_points(self, n: int, random: str = "pseudo") -> NDArray[np.float_]:
         # https://math.stackexchange.com/questions/87230/picking-random-points-in-the-volume-of-sphere-with-uniform-probability
         if random == "pseudo":
             U = np.random.rand(n, 1).astype(config.real(np))
@@ -249,7 +255,7 @@ def random_points(self, n, random="pseudo"):
         X = U ** (1 / self.dim) * X
         return self.radius * X + self.center
 
-    def random_boundary_points(self, n, random="pseudo"):
+    def random_boundary_points(self, n: int, random: str = "pseudo") -> NDArray[np.float_]:
         # http://mathworld.wolfram.com/HyperspherePointPicking.html
         if random == "pseudo":
             X = np.random.normal(size=(n, self.dim)).astype(config.real(np))
@@ -259,7 +265,7 @@ def random_boundary_points(self, n, random="pseudo"):
         X = preprocessing.normalize(X)
         return self.radius * X + self.center
 
-    def background_points(self, x, dirn, dist2npt, shift):
+    def background_points(self, x: NDArray[np.float_], dirn: Number, dist2npt: Callable[[NDArray[np.float_]], int], shift: int) -> NDArray[np.float_]:
         dirn = dirn / np.linalg.norm(dirn)
         dx = self.distance2boundary_unitdirn(x, -dirn)
         n = max(dist2npt(dx), 1)
diff --git a/deepxde/geometry/pointcloud.py b/deepxde/geometry/pointcloud.py
index 5644a79ae..49caa1203 100644
--- a/deepxde/geometry/pointcloud.py
+++ b/deepxde/geometry/pointcloud.py
@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import numpy as np
+from numpy.typing import NDArray
 
 from .geometry import Geometry
 from .. import config
@@ -17,7 +20,7 @@ class PointCloud(Geometry):
         boundary_normals: A 2-D NumPy array.
     """
 
-    def __init__(self, points, boundary_points=None, boundary_normals=None):
+    def __init__(self, points: NDArray[np.float_], boundary_points: NDArray[np.float_] | None = None, boundary_normals: NDArray[np.float_] | None = None):
         self.points = np.asarray(points, dtype=config.real(np))
         self.num_points = len(points)
         self.boundary_points = None
diff --git a/deepxde/geometry/sampler.py b/deepxde/geometry/sampler.py
index 7b384eddc..19da42341 100644
--- a/deepxde/geometry/sampler.py
+++ b/deepxde/geometry/sampler.py
@@ -1,12 +1,14 @@
 __all__ = ["sample"]
+from typing import Literal
 
 import numpy as np
+from numpy.typing import NDArray
 import skopt
 
 from .. import config
 
 
-def sample(n_samples, dimension, sampler="pseudo"):
+def sample(n_samples: int, dimension: int, sampler: Literal["pseudo", "LHS", "Halton", "Hammersley", "Sobol"] = "pseudo"):
     """Generate pseudorandom or quasirandom samples in [0, 1]^dimension.
 
     Args:
@@ -23,7 +25,7 @@ def sample(n_samples, dimension, sampler="pseudo"):
     raise ValueError("f{sampler} sampling is not available.")
 
 
-def pseudorandom(n_samples, dimension):
+def pseudorandom(n_samples: int, dimension: int) -> NDArray[np.float_]:
     """Pseudo random."""
     # If random seed is set, then the rng based code always returns the same random
     # number, which may not be what we expect.
@@ -32,7 +34,7 @@ def pseudorandom(n_samples, dimension):
     return np.random.random(size=(n_samples, dimension)).astype(config.real(np))
 
 
-def quasirandom(n_samples, dimension, sampler):
+def quasirandom(n_samples: int, dimension: int, sampler: Literal["LHS", "Halton", "Hammersley", "Sobol"]) -> NDArray[np.float_]:
     # Certain points should be removed:
     # - Boundary points such as [..., 0, ...]
     # - Special points [0, 0, 0, ...] and [0.5, 0.5, 0.5, ...], which cause error in
diff --git a/deepxde/geometry/timedomain.py b/deepxde/geometry/timedomain.py
index 0ee871b6c..6604518c6 100644
--- a/deepxde/geometry/timedomain.py
+++ b/deepxde/geometry/timedomain.py
@@ -1,7 +1,10 @@
 import itertools
+from numbers import Number
 
 import numpy as np
+from numpy.typing import NDArray
 
+from .geometry import Geometry
 from .geometry_1d import Interval
 from .geometry_2d import Rectangle
 from .geometry_3d import Cuboid
@@ -11,17 +14,17 @@
 
 
 class TimeDomain(Interval):
-    def __init__(self, t0, t1):
+    def __init__(self, t0: Number, t1: Number):
         super().__init__(t0, t1)
         self.t0 = t0
         self.t1 = t1
 
-    def on_initial(self, t):
+    def on_initial(self, t: NDArray[np.float_]) -> NDArray[np.bool_]:
         return isclose(t, self.t0).flatten()
 
 
-class GeometryXTime:
-    def __init__(self, geometry, timedomain):
+class GeometryXTime(Geometry):
+    def __init__(self, geometry: Geometry, timedomain: TimeDomain):
         self.geometry = geometry
         self.timedomain = timedomain
         self.dim = geometry.dim + timedomain.dim
diff --git a/deepxde/utils/external.py b/deepxde/utils/external.py
index bca1cc1d1..caa9bc0c0 100644
--- a/deepxde/utils/external.py
+++ b/deepxde/utils/external.py
@@ -1,3 +1,4 @@
+from __future__ import annotations
 """External utilities."""
 
 import csv
@@ -6,6 +7,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+from numpy.typing import NDArray
 import scipy.spatial.distance
 from mpl_toolkits.mplot3d import Axes3D
 from sklearn import preprocessing
@@ -376,7 +378,7 @@ def dat_to_csv(dat_file_path, csv_file_path, columns):
             csv_writer.writerow(row)
 
 
-def isclose(a, b):
+def isclose(a: NDArray[np.float_], b: NDArray[np.float_]) -> NDArray[np.bool_]:
     """A modified version of `np.isclose` for DeepXDE.
 
     This function changes the value of `atol` due to the dtype of `a` and `b`.

From f81f4befa1c14c5d11d4422d1db43ee4ff681eda Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Mon, 4 Sep 2023 14:56:53 -0400
Subject: [PATCH 11/12] testing

---
 pyproject.toml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index a1f7c55bd..f455cf01f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -54,3 +54,8 @@ exclude = ["docker", "docs*", "examples*"]
 
 [tool.setuptools_scm]
 write_to = "deepxde/_version.py"
+
+[mypy]
+python_version = 3.8
+strict = false
+ignore_missing_imports = true
\ No newline at end of file

From 2ef9831741555f4b8f64e006e059415a2c1f19cc Mon Sep 17 00:00:00 2001
From: seetop <siumabon123@gmail.com>
Date: Tue, 5 Sep 2023 12:39:33 -0400
Subject: [PATCH 12/12] bug fix

---
 deepxde/backend/backend.py     |  1 -
 deepxde/geometry/geometry.py   |  2 +-
 deepxde/geometry/pointcloud.py |  2 +-
 deepxde/geometry/timedomain.py |  2 +-
 deepxde/types.py               | 11 +++++++++--
 5 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/deepxde/backend/backend.py b/deepxde/backend/backend.py
index 866522949..883967d11 100644
--- a/deepxde/backend/backend.py
+++ b/deepxde/backend/backend.py
@@ -516,4 +516,3 @@ def sparse_dense_matmul(x, y):
     Returns:
         Tensor: The multiplication result.
     """
-sparse_dense_matmul()
\ No newline at end of file
diff --git a/deepxde/geometry/geometry.py b/deepxde/geometry/geometry.py
index 804bc16db..03f0374cf 100644
--- a/deepxde/geometry/geometry.py
+++ b/deepxde/geometry/geometry.py
@@ -7,7 +7,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from ..types import Tensor, TensorOrTensors
+from ..types import Tensor
 
 class Geometry(abc.ABC):
     def __init__(self, dim: int, bbox: NDArray[np.float_], diam: Number):
diff --git a/deepxde/geometry/pointcloud.py b/deepxde/geometry/pointcloud.py
index 49caa1203..d9883cfd9 100644
--- a/deepxde/geometry/pointcloud.py
+++ b/deepxde/geometry/pointcloud.py
@@ -5,7 +5,7 @@
 
 from .geometry import Geometry
 from .. import config
-from ..data import BatchSampler
+from ..data.sampler import BatchSampler
 from ..utils import isclose
 
 
diff --git a/deepxde/geometry/timedomain.py b/deepxde/geometry/timedomain.py
index 6604518c6..483cee95f 100644
--- a/deepxde/geometry/timedomain.py
+++ b/deepxde/geometry/timedomain.py
@@ -23,7 +23,7 @@ def on_initial(self, t: NDArray[np.float_]) -> NDArray[np.bool_]:
         return isclose(t, self.t0).flatten()
 
 
-class GeometryXTime(Geometry):
+class GeometryXTime():
     def __init__(self, geometry: Geometry, timedomain: TimeDomain):
         self.geometry = geometry
         self.timedomain = timedomain
diff --git a/deepxde/types.py b/deepxde/types.py
index 8cff698e8..159cc4e5a 100644
--- a/deepxde/types.py
+++ b/deepxde/types.py
@@ -1,8 +1,15 @@
 from __future__ import annotations
 from typing import Sequence, TypeVar, Union
 
+# dtype from any backend
+dtype = TypeVar("dtype")
+
+# NN from any backend (Using the `NN` from deepxde is recommended.)
+NN = TypeVar("NN")
+
+# SparseTensor from any backend
+SparseTensor = TypeVar("SparseTensor")
+
 # Tensor from any backend
 Tensor = TypeVar("Tensor")
 TensorOrTensors = Union[Tensor, Sequence[Tensor]]
-SparseTensor = TypeVar("SparseTensor")
-dtype = TypeVar("dtype")
\ No newline at end of file