Merge pull request #1197 from brighthe/master

add phi device based on bcs

app/soptx/linear_elasticity/hexahedrom_mesh_pytorch.py

@@ -14,6 +14,7 @@
 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.fem.dirichlet_bc import DirichletBC
 
 from fealpy.experimental.decorator import cartesian
 
@@ -105,22 +106,23 @@ def dirichlet(self, points: TensorLike) -> TensorLike:
         return bm.zeros(points.shape, 
                         dtype=points.dtype, device=points.device)
 
-parser = argparse.ArgumentParser(description="HexahedronMesh 上的任意次 Lagrange 有限元空间的线性弹性问题求解.")
+parser = argparse.ArgumentParser(description="Solve linear elasticity problems \
+                            in arbitrary order Lagrange finite element space on HexahedronMesh.")
 parser.add_argument('--backend', 
                     default='pytorch', type=str,
-                    help='指定计算的后端类型, 默认为 pytorch.')
+                    help='Specify the backend type for computation, default is "pytorch".')
 parser.add_argument('--degree', 
-                    default=1, type=int, 
-                    help='Lagrange 有限元空间的次数, 默认为 1 次.')
+                    default=2, type=int, 
+                    help='Degree of the Lagrange finite element space, default is 1.')
 parser.add_argument('--nx', 
                     default=2, type=int, 
-                    help='x 方向的初始网格单元数, 默认为 2.')
+                    help='Initial number of grid cells in the x direction, default is 2.')
 parser.add_argument('--ny',
                     default=2, type=int,
-                    help='y 方向的初始网格单元数, 默认为 2.')
+                    help='Initial number of grid cells in the y direction, default is 2.')
 parser.add_argument('--nz',
                     default=2, type=int,
-                    help='z 方向的初始网格单元数, 默认为 2.')
+                    help='Initial number of grid cells in the z direction, default is 2.')
 args = parser.parse_args()
 
 pde = BoxDomainPolyUnloaded3d()
@@ -153,7 +155,7 @@ def dirichlet(self, points: TensorLike) -> TensorLike:
     integrator_K = LinearElasticIntegrator(material=linear_elastic_material, q=tensor_space.p+3)
     bform = BilinearForm(tensor_space)
     bform.add_integrator(integrator_K)
-    K = bform.assembly()
+    K = bform.assembly(format='csr')
     tmr.send('stiffness assembly')
 
     integrator_F = VectorSourceIntegrator(source=pde.source, q=tensor_space.p+3)
@@ -168,17 +170,8 @@ def dirichlet(self, points: TensorLike) -> TensorLike:
     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()
+    dbc = DirichletBC(space=tensor_space)
+    K = dbc.apply_matrix(matrix=K, check=True)
     tmr.send('boundary')
 
     uh = tensor_space.function()

fealpy/experimental/mesh/mesh_base.py

@@ -573,7 +573,7 @@ def grad_shape_function(self, bcs: Tuple[TensorLike], p: int=1, *, index: Index=
         n = phi.shape[0]**TD
         ldof = phi.shape[-1]**TD
         shape = (n, ldof, TD)
-        gphi = bm.zeros(shape, dtype=self.ftype)
+        gphi = bm.zeros(shape, dtype=self.ftype, device=bm.get_device(bcs[0]))
 
         if TD == 3:
             gphi0 = bm.einsum('im, jn, ko->ijkmno', dphi, 
@@ -582,13 +582,13 @@ def grad_shape_function(self, bcs: Tuple[TensorLike], p: int=1, *, index: Index=
                               phi).reshape(-1, ldof, 1)
             gphi2 = bm.einsum('im, jn, ko->ijkmno', phi, phi,
                                      dphi).reshape(-1, ldof, 1)
-            gphi = bm.concatenate((gphi0, gphi1, gphi2), axis=-1) 
+            gphi[:] = bm.concatenate((gphi0, gphi1, gphi2), axis=-1)
             if variables == 'x':
                 J = self.jacobi_matrix(bcs, index=index)
                 J = bm.linalg.inv(J)
                 # J^{-T}\nabla_u phi
                 # gphi = bm.einsum('qcmn, qlm -> qcln', J, gphi)
-                gphi = bm.einsum('qcmn, qlm -> cqln', J, gphi)
+                gphi[:] = bm.concatenate((gphi0, gphi1), axis=-1)
                 return gphi
         elif TD == 2:
             gphi0 = bm.einsum('im, jn -> ijmn', dphi, phi).reshape(-1, ldof, 1)
@@ -599,7 +599,7 @@ def grad_shape_function(self, bcs: Tuple[TensorLike], p: int=1, *, index: Index=
                 G = self.first_fundamental_form(J)
                 G = bm.linalg.inv(G)
                 # gphi = bm.einsum('qikm, qimn, qln -> qilk', J, G, gphi)
-                gphi = bm.einsum('qikm, qimn, qln -> iqlk', J, G, gphi)
+                gphi[:] = bm.einsum('qikm, qimn, qln -> iqlk', J, G, gphi)
                 return gphi
         return gphi