Merge pull request #1187 from brighthe/master

add device in hexahedrom_mesh and uniform_mesh_2d

app/soptx/linear_elasticity/hexahedrom_mesh.py → app/soptx/linear_elasticity/hexahedrom_mesh_numpy.py

@@ -122,6 +122,10 @@ def dirichlet(self, points: TensorLike) -> TensorLike:
 args = parser.parse_args()
 
 bm.set_backend(args.backend)
+test = bm.array([1, 2, 3])
+a = bm.get_device(test)
+bm.device('cuda')
+b = bm.get_device(test)
 nx, ny, nz = args.nx, args.ny, args.nz
 mesh = HexahedronMesh.from_box(box=pde.domain(), nx=nx, ny=ny, nz=nz)
 
@@ -138,7 +142,7 @@ def dirichlet(self, points: TensorLike) -> TensorLike:
 next(tmr)
 
 maxit = 4
-errorMatrix = bm.zeros((1, maxit), dtype=bm.float64)
+errorMatrix = bm.zeros((3, maxit), dtype=bm.float64)
 errorType = ['$|| u  - u_h ||_{L2}$', '$|| u -  u_h||_{l2}$', 'boundary']
 NDof = bm.zeros(maxit, dtype=bm.int32)
 for i in range(maxit):
@@ -190,15 +194,13 @@ def dirichlet(self, points: TensorLike) -> TensorLike:
     tmr.send(None)
 
     u_exact = tensor_space.interpolate(pde.solution)
-    # errorMatrix[0, i] = bm.sqrt(bm.sum(bm.abs(uh - u_exact)**2 * (1 / NDof[i])))
+    errorMatrix[0, i] = bm.sqrt(bm.sum(bm.abs(uh - u_exact)**2 * (1 / NDof[i])))
     errorMatrix[0, i] = mesh.error(u=uh, v=pde.solution, q=tensor_space.p+3, power=2)
-    # errorMatrix[2, i] = bm.sqrt(bm.sum(bm.abs(uh[isDDof] - u_exact[isDDof])**2 * (1 / NDof[i])))
-
-    # print("errorMatrix:", errorMatrix)
+    errorMatrix[2, i] = bm.sqrt(bm.sum(bm.abs(uh[isDDof] - u_exact[isDDof])**2 * (1 / NDof[i])))
 
     if i < maxit-1:
         mesh.uniform_refine()
 
 print("errorMatrix:\n", errorMatrix)
-# print("order_l2:\n", bm.log2(errorMatrix[0, :-1] / errorMatrix[0, 1:]))
-print("order_L2:\n ", bm.log2(errorMatrix[0, :-1] / errorMatrix[0, 1:]))
+print("order_l2:\n", bm.log2(errorMatrix[0, :-1] / errorMatrix[0, 1:]))
+print("order_L2:\n ", bm.log2(errorMatrix[1, :-1] / errorMatrix[1, 1:]))

app/soptx/linear_elasticity/hexahedrom_mesh_pytorch.py

@@ -0,0 +1,193 @@
+from fealpy.experimental.backend import backend_manager as bm
+
+from fealpy.experimental.typing import TensorLike
+
+from fealpy.experimental.decorator import cartesian
+
+from fealpy.experimental.mesh import HexahedronMesh
+
+from fealpy.experimental.functionspace import LagrangeFESpace, TensorFunctionSpace
+
+from fealpy.experimental.material.elastic_material import LinearElasticMaterial
+
+from fealpy.experimental.fem.linear_elastic_integrator import LinearElasticIntegrator
+from fealpy.experimental.fem.vector_source_integrator import VectorSourceIntegrator
+from fealpy.experimental.fem.bilinear_form import BilinearForm
+from fealpy.experimental.fem.linear_form import LinearForm
+
+from fealpy.experimental.decorator import cartesian
+
+from fealpy.experimental.sparse import COOTensor
+
+from fealpy.experimental.solver import cg
+
+from app.soptx.soptx.utilfs.timer import timer
+
+import argparse
+
+class BoxDomainPolyUnloaded3d():
+    def __init__(self):
+        pass
+
+    def domain(self):
+        return [0, 1, 0, 1, 0, 1]
+
+    @cartesian
+    def solution(self, points: TensorLike):
+        x = points[..., 0]
+        y = points[..., 1]
+        z = points[..., 2]
+        val = bm.zeros(points.shape, dtype=points.dtype)
+        val[..., 0] = 2*x**3 - 3*x*y**2 - 3*x*z**2
+        val[..., 1] = 2*y**3 - 3*y*x**2 - 3*y*z**2
+        val[..., 2] = 2*z**3 - 3*z*y**2 - 3*z*x**2
+
+        return val
+
+    @cartesian
+    def source(self, points: TensorLike):
+        val = bm.zeros(points.shape, dtype=points.dtype)
+
+        return val
+
+    def dirichlet(self, points: TensorLike) -> TensorLike:
+
+        return self.solution(points)
+
+class BoxDomainPolyLoaded3d():
+    def domain(self):
+        return [0, 1, 0, 1, 0, 1]
+
+    @cartesian
+    def source(self, points: TensorLike):
+        x = points[..., 0]
+        y = points[..., 1]
+        z = points[..., 2]
+        val = bm.zeros(points.shape, dtype=bm.float64)
+        mu = 1
+        factor1 = -400 * mu * (2 * y - 1) * (2 * z - 1)
+        term1 = 3 * (x ** 2 - x) ** 2 * (y ** 2 - y + z ** 2 - z)
+        term2 = (1 - 6 * x + 6 * x ** 2) * (y ** 2 - y) * (z ** 2 - z)
+        val[..., 0] = factor1 * (term1 + term2)
+
+        factor2 = 200 * mu * (2 * x - 1) * (2 * z - 1)
+        term1 = 3 * (y ** 2 - y) ** 2 * (x ** 2 - x + z ** 2 - z)
+        term2 = (1 - 6 * y + 6 * y ** 2) * (x ** 2 - x) * (z ** 2 - z)
+        val[..., 1] = factor2 * (term1 + term2)
+
+        factor3 = 200 * mu * (2 * x - 1) * (2 * y - 1)
+        term1 = 3 * (z ** 2 - z) ** 2 * (x ** 2 - x + y ** 2 - y)
+        term2 = (1 - 6 * z + 6 * z ** 2) * (x ** 2 - x) * (y ** 2 - y)
+        val[..., 2] = factor3 * (term1 + term2)
+
+        return val
+
+    @cartesian
+    def solution(self, points: TensorLike):
+        x = points[..., 0]
+        y = points[..., 1]
+        z = points[..., 2]
+        val = bm.zeros(points.shape, dtype=bm.float64)
+
+        mu = 1
+        val[..., 0] = 200*mu*(x-x**2)**2 * (2*y**3-3*y**2+y) * (2*z**3-3*z**2+z)
+        val[..., 1] = -100*mu*(y-y**2)**2 * (2*x**3-3*x**2+x) * (2*z**3-3*z**2+z)
+        val[..., 2] = -100*mu*(z-z**2)**2 * (2*y**3-3*y**2+y) * (2*x**3-3*x**2+x)
+
+        return val
+
+    def dirichlet(self, points: TensorLike) -> TensorLike:
+
+        return bm.zeros(points.shape, dtype=points.dtype)
+
+parser = argparse.ArgumentParser(description="HexahedronMesh 上的任意次 Lagrange 有限元空间的线性弹性问题求解.")
+parser.add_argument('--backend', 
+                    default='pytorch', type=str,
+                    help='指定计算的后端类型, 默认为 pytorch.')
+parser.add_argument('--degree', 
+                    default=1, type=int, 
+                    help='Lagrange 有限元空间的次数, 默认为 1 次.')
+parser.add_argument('--nx', 
+                    default=2, type=int, 
+                    help='x 方向的初始网格单元数, 默认为 2.')
+parser.add_argument('--ny',
+                    default=2, type=int,
+                    help='y 方向的初始网格单元数, 默认为 2.')
+parser.add_argument('--nz',
+                    default=2, type=int,
+                    help='z 方向的初始网格单元数, 默认为 2.')
+args = parser.parse_args()
+
+pde = BoxDomainPolyUnloaded3d()
+args = parser.parse_args()
+
+bm.set_backend(args.backend)
+
+nx, ny, nz = args.nx, args.ny, args.nz
+mesh = HexahedronMesh.from_box(box=pde.domain(), nx=nx, ny=ny, nz=nz, device='cpu')
+
+p = args.degree
+
+tmr = timer("FEM Solver")
+next(tmr)
+
+maxit = 4
+errorType = ['$|| u  - u_h ||_{L2}$']
+errorMatrix = bm.zeros((len(errorType), maxit), dtype=bm.float64)
+NDof = bm.zeros(maxit, dtype=bm.int32)
+for i in range(maxit):
+    space = LagrangeFESpace(mesh, p=p, ctype='C')
+    tensor_space = TensorFunctionSpace(space, shape=(-1, 3))
+    NDof[i] = tensor_space.number_of_global_dofs()
+
+    linear_elastic_material = LinearElasticMaterial(name='lam1_mu1', 
+                                                lame_lambda=1, shear_modulus=1, 
+                                                hypo='3D')
+    tmr.send('material')
+
+    integrator_K = LinearElasticIntegrator(material=linear_elastic_material, q=tensor_space.p+3)
+    bform = BilinearForm(tensor_space)
+    bform.add_integrator(integrator_K)
+    K = bform.assembly()
+    tmr.send('stiffness assembly')
+
+    integrator_F = VectorSourceIntegrator(source=pde.source, q=tensor_space.p+3)
+    lform = LinearForm(tensor_space)    
+    lform.add_integrator(integrator_F)
+    F = lform.assembly()
+    tmr.send('source assembly')
+
+    uh_bd = bm.zeros(tensor_space.number_of_global_dofs(), dtype=bm.float64)
+    uh_bd, isDDof = tensor_space.boundary_interpolate(gD=pde.dirichlet, uh=uh_bd, threshold=None)
+
+    F = F - K.matmul(uh_bd)
+    F[isDDof] = uh_bd[isDDof]
+
+    indices = K.indices()
+    new_values = bm.copy(K.values())
+    IDX = isDDof[indices[0, :]] | isDDof[indices[1, :]]
+    new_values[IDX] = 0
+
+    K = COOTensor(indices, new_values, K.sparse_shape)
+    index, = bm.nonzero(isDDof)
+    one_values = bm.ones(len(index), **K.values_context())
+    one_indices = bm.stack([index, index], axis=0)
+    K1 = COOTensor(one_indices, one_values, K.sparse_shape)
+    K = K.add(K1).coalesce()
+    tmr.send('boundary')
+
+    uh = tensor_space.function()
+    K = K.tocsr()
+    uh[:] = cg(K, F, maxiter=1000, atol=1e-14, rtol=1e-14)
+    tmr.send('solve(cg)')
+
+    tmr.send(None)
+
+    u_exact = tensor_space.interpolate(pde.solution)
+    errorMatrix[0, i] = mesh.error(u=uh, v=pde.solution, q=tensor_space.p+3, power=2)
+
+    if i < maxit-1:
+        mesh.uniform_refine()
+
+print("errorMatrix:\n", errorType, "\n", errorMatrix)
+print("order_L2:\n ", bm.log2(errorMatrix[0, :-1] / errorMatrix[0, 1:]))

fealpy/experimental/mesh/hexahedron_mesh.py

@@ -353,7 +353,8 @@ def from_one_hexahedron(cls, twist=False):
         return cls(node, cell)
 
     @classmethod
-    def from_box(cls, box=[0, 1, 0, 1, 0, 1], nx=10, ny=10, nz=10, threshold=None):
+    def from_box(cls, box=[0, 1, 0, 1, 0, 1], nx=10, ny=10, nz=10, 
+                threshold=None, *, itype=None, ftype=None, device=None,):
         """
         Generate a hexahedral mesh for a box domain.
 
@@ -363,20 +364,23 @@ def from_box(cls, box=[0, 1, 0, 1, 0, 1], nx=10, ny=10, nz=10, threshold=None):
         @param threshold Optional function to filter cells based on their barycenter coordinates (default: None)
         @return HexahedronMesh instance
         """
+        if itype is None:
+            itype = bm.int32
+        if ftype is None:
+            ftype = bm.float64
+
         shape = (nx+1, ny+1, nz+1)
-        X = bm.linspace(box[0], box[1], nx+1, endpoint=True, dtype=bm.float64)[:, None, None]
-        # X = bm.linspace(box[0], box[1], nx+1, endpoint=True, dtype=bm.float64)
-        # X = X[:, None, None]
-        Y = bm.linspace(box[2], box[3], ny+1, endpoint=True, dtype=bm.float64)[None, :, None]
-        Z = bm.linspace(box[4], box[5], nz+1, endpoint=True, dtype=bm.float64)[None, None, :]
+        X = bm.linspace(box[0], box[1], nx+1, endpoint=True, dtype=ftype, device=device)[:, None, None]
+        Y = bm.linspace(box[2], box[3], ny+1, endpoint=True, dtype=ftype, device=device)[None, :, None]
+        Z = bm.linspace(box[4], box[5], nz+1, endpoint=True, dtype=ftype, device=device)[None, None, :]
         X = bm.broadcast_to(X, shape).reshape(-1, 1)
         Y = bm.broadcast_to(Y, shape).reshape(-1, 1)
         Z = bm.broadcast_to(Z, shape).reshape(-1, 1)
 
         node = bm.concatenate([X, Y, Z], axis=-1)
 
         NN = (nx+1)*(ny+1)*(nz+1)
-        idx = bm.arange(0, NN).reshape(nx+1, ny+1, nz+1)
+        idx = bm.arange(0, NN, device=device).reshape(nx+1, ny+1, nz+1)
         c = idx[:-1, :-1, :-1]
 
         nyz = (ny + 1)*(nz + 1)
@@ -395,11 +399,11 @@ def from_box(cls, box=[0, 1, 0, 1, 0, 1], nx=10, ny=10, nz=10, threshold=None):
             bc = bm.sum(node[cell, :], axis=1)/cell.shape[1]
             isDelCell = threshold(bc)
             cell = cell[~isDelCell]
-            isValidNode = bm.zeros(NN, dtype=bm.bool)
+            isValidNode = bm.zeros(NN, dtype=bm.bool, device=device)
             isValidNode[cell] = True
             node = node[isValidNode]
-            idxMap = bm.zeros(NN, dtype=cell.dtype)
-            idxMap[isValidNode] = bm.arange(isValidNode.sum())
+            idxMap = bm.zeros(NN, dtype=cell.dtype, device=device)
+            idxMap[isValidNode] = bm.arange(isValidNode.sum(), device=device)
             cell = idxMap[cell]
 
         return cls(node, cell)

fealpy/experimental/mesh/uniform_mesh_2d.py

@@ -74,7 +74,7 @@ def __init__(self, extent: tuple[int, int, int, int] = (0, 1, 0, 1),
                 origin: tuple[float, float] = (0.0, 0.0), 
                 ipoints_ordering='yx', 
                 flip_direction=None, 
-                itype=None, ftype=None):
+                *, itype=None, ftype=None, device=None):
         """
         Initializes a 2D uniform structured mesh.
 
@@ -137,14 +137,14 @@ def __init__(self, extent: tuple[int, int, int, int] = (0, 1, 0, 1),
         self.adjusted_edge_mask = self.get_adjusted_edge_mask()
 
         # Specify the counterclockwise drawing
-        self.ccw = bm.array([0, 2, 3, 1], dtype=self.itype)
+        self.ccw = bm.array([0, 2, 3, 1], dtype=self.itype, device=device)
 
         self.edge2cell = self.edge_to_cell()
         self.face2cell = self.edge2cell
         self.cell2edge = self.cell_to_edge()
 
         self.localEdge = bm.array([(0, 2), (1, 3), 
-                                   (0, 1), (2, 3)], dtype=self.itype)   
+                                   (0, 1), (2, 3)], dtype=self.itype, device=device)   
 
 
     # 实体生成方法
@@ -153,14 +153,15 @@ def _get_node(self) -> TensorLike:
         """
         @berif Generate the nodes in a structured mesh.
         """
+        device = self.device
         GD = 2
         nx = self.nx
         ny = self.ny
         box = [self.origin[0], self.origin[0] + nx * self.h[0],
                self.origin[1], self.origin[1] + ny * self.h[1]]
 
-        x = bm.linspace(box[0], box[1], nx + 1, dtype=self.ftype)
-        y = bm.linspace(box[2], box[3], ny + 1, dtype=self.ftype)
+        x = bm.linspace(box[0], box[1], nx + 1, dtype=self.ftype, device=device)
+        y = bm.linspace(box[2], box[3], ny + 1, dtype=self.ftype, device=device)
         xx, yy = bm.meshgrid(x, y, indexing='ij')
 
         node = bm.concatenate((xx[..., None], yy[..., None]), axis=-1)
@@ -177,14 +178,15 @@ def _get_edge(self) -> TensorLike:
         """
         @berif Generate the edges in a structured mesh.
         """
+        device = self.device
         nx = self.nx
         ny = self.ny
         NN = self.NN
         NE = self.NE
 
-        idx = bm.arange(NN, dtype=self.itype).reshape(nx + 1, ny + 1)
+        idx = bm.arange(NN, dtype=self.itype, device=device).reshape(nx + 1, ny + 1)
 
-        edge = bm.zeros((NE, 2), dtype=self.itype)
+        edge = bm.zeros((NE, 2), dtype=self.itype, device=device)
         NE0 = 0
         NE1 = nx * (ny + 1)
         edge = bm.set_at(edge, (slice(NE0, NE1), 0), idx[:-1, :].reshape(-1))
@@ -204,22 +206,23 @@ def _get_cell(self) -> TensorLike:
         """
         @berif Generate the cells in a structured mesh.
         """
+        device = self.device
         nx = self.nx
         ny = self.ny
 
         NN = self.NN
         NC = self.NC
 
-        idx = bm.arange(NN).reshape(nx + 1, ny + 1)
+        idx = bm.arange(NN, device=device).reshape(nx + 1, ny + 1)
 
-        cell = bm.zeros((NC, 4), dtype=self.itype)
+        cell = bm.zeros((NC, 4), dtype=self.itype, device=device)
         c = idx[:-1, :-1]
         cell_0 = c.reshape(-1)
         cell_1 = cell_0 + 1
         cell_2 = cell_0 + ny + 1
         cell_3 = cell_2 + 1
-        cell = bm.concatenate([cell_0[:, None], cell_1[:, None], cell_2[:, None], 
-                            cell_3[:, None]], axis=-1)
+        cell = bm.concatenate([cell_0[:, None], cell_1[:, None], 
+                               cell_2[:, None], cell_3[:, None]], axis=-1)
 
         return cell