Skip to content

Commit

Permalink
Merge pull request #1224 from BenHBLiu/master
Browse files Browse the repository at this point in the history
update
  • Loading branch information
cbtxs authored Oct 17, 2024
2 parents 01b840b + 48bcac5 commit 9abf2f4
Show file tree
Hide file tree
Showing 5 changed files with 34 additions and 33 deletions.
3 changes: 2 additions & 1 deletion fealpy/fem/vector_mass_integrator.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,14 +68,15 @@ def assembly_cell_matrix_for_vector_basis_vspace(self, space, index=_S, cellmeas
coef = self.coef
mesh = space.mesh
GD = mesh.geo_dimension()
self.device = mesh.device

if cellmeasure is None:
cellmeasure = mesh.entity_measure('cell', index=index)

NC = len(cellmeasure)
ldof = space.number_of_local_dofs()
if out is None:
D = bm.zeros((NC, ldof, ldof), dtype=space.ftype)
D = bm.zeros((NC, ldof, ldof), device=self.device,dtype=space.ftype)
else:
D = out

Expand Down
24 changes: 12 additions & 12 deletions fealpy/functionspace/bernstein_fe_space.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ class BernsteinFESpace(FunctionSpace, Generic[_MT]):
def __init__(self, mesh: _MT, p: int=1, ctype='C'):
self.mesh = mesh
self.p = p
self.device = mesh.device

assert ctype in {'C', 'D'}
self.ctype = ctype # 空间连续性类型
Expand Down Expand Up @@ -85,19 +86,18 @@ def basis(self, bcs: TensorLike, index: Index=_S, p = None):
ldof = multiIndex.shape[0]

B = bcs
B = bm.ones((NQ, p+1, TD+1), dtype=self.ftype)
B = bm.ones((NQ, p+1, TD+1), device=self.device, dtype=self.ftype)
# B[:, 1:] = bcs[:, None, :]
B = bm.set_at(B,(slice(None),slice(1,None)),bcs[:,None,:])
B = bm.cumprod(B, axis=1)
P = bm.arange(p+1)
P = bm.arange(p+1, device=self.device)
P = bm.set_at(P,(0),1)
P = bm.cumprod(P,axis=0).reshape(1, -1, 1)
B = B/P

# B : (NQ, p+1, TD+1)
# B[:, multiIndex, bm.arange(TD+1).reshape(1, -1)]: (NQ, ldof, TD+1)
phi = P[0, -1, 0]*bm.prod(B[:, multiIndex,
bm.arange(TD+1).reshape(1, -1)], axis=-1)
phi = P[0, -1, 0]*bm.prod(B[:, multiIndex, bm.arange(TD+1).reshape(1, -1)], axis=-1)
return phi[None, :]

@barycentric
Expand Down Expand Up @@ -132,26 +132,26 @@ def grad_basis(self, bcs: TensorLike, index: Index=_S, variable='u',p=None):
ldof = multiIndex.shape[0]

B = bcs
B = bm.ones((NQ, p+1, TD+1), dtype=self.ftype)
B = bm.ones((NQ, p+1, TD+1), device=self.device, dtype=self.ftype)
# B[:, 1:] = bcs[:, None, :]
B = bm.set_at(B,(slice(None),slice(1,None)),bcs[:,None,:])
B = bm.cumprod(B, axis=1)

P = bm.arange(p+1)
P = bm.arange(p+1,device=self.device)
# P[0] = 1
P = bm.set_at(P,(0),1)
P = bm.cumprod(P,axis=0).reshape(1, -1, 1)
B = B/P

F = bm.zeros(B.shape, dtype=self.ftype)
F = bm.zeros(B.shape, device=self.device, dtype=self.ftype)
# F[:, 1:] = B[:, :-1]
F = bm.set_at(F,(slice(None),slice(1,None)),B[:,:-1])
shape = bcs.shape[:-1]+(ldof, TD+1)
R = bm.zeros(shape, dtype=self.ftype)
R = bm.zeros(shape,device=self.device,dtype=self.ftype)
for i in range(TD+1):
idx = list(range(TD+1))
idx.remove(i)
idx = bm.array(idx, dtype=self.itype)
idx = bm.array(idx,device=self.device, dtype=self.itype)
# R[..., i] = bm.prod(B[..., multiIndex[:, idx], idx.reshape(1, -1)],axis=-1)*F[..., multiIndex[:, i], [i]]
R = bm.set_at(R,(...,i),bm.prod(B[..., multiIndex[:, idx], idx.reshape(1, -1)],axis=-1)*F[..., multiIndex[:, i], [i]])
Dlambda = self.mesh.grad_lambda()
Expand All @@ -167,7 +167,7 @@ def hess_basis(self, bcs: TensorLike, index: Index=_S, variable='u'):
g2phi = self.grad_m_basis(bcs, 2, index=index)
TD = self.mesh.top_dimension()
shape = g2phi.shape[:-1] + (TD, TD)
hval = bm.zeros(shape, dtype=self.ftype)
hval = bm.zeros(shape, device=self.device, dtype=self.ftype)
hval = bm.set_at(hval,(...,0,0),g2phi[..., 0])
hval = bm.set_at(hval,(...,0,1),g2phi[..., 1])
hval = bm.set_at(hval,(...,1,0),g2phi[..., 1])
Expand Down Expand Up @@ -222,7 +222,7 @@ def grad_m_basis(self, bcs: TensorLike, m: int, index = _S):
midxp_1 = bm.multi_index_matrix(m, GD) # m 次多重指标

N, N1 = len(symidx), midxp_1.shape[0]
B = bm.zeros((N1, NQ, ldof), dtype=self.ftype)
B = bm.zeros((N1, NQ, ldof), device=self.device, dtype=self.ftype)
symLambdaBeta = bm.zeros((N1, NC, N), dtype=self.ftype)
for beta, Bi, symi in zip(midxp_1, B, symLambdaBeta):
midxp_0 -= beta[None, :]
Expand Down Expand Up @@ -336,7 +336,7 @@ def boundary_interpolate(self, gD: Union[Callable, int, float],
isbdface = self.mesh.boundary_face_flag()
f2d = self.face_to_dof()[isbdface]

isDDof = bm.zeros(gdof, dtype=bm.bool)
isDDof = bm.zeros(gdof, device=self.device, dtype=bm.bool)
# isDDof[f2d] = True
isDDof = bm.set_at(isDDof,(f2d),True)
if callable(gD):
Expand Down
30 changes: 15 additions & 15 deletions fealpy/functionspace/first_nedelec_fe_space_3d.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,14 +51,14 @@ def number_of_global_dofs(self):
def edge_to_dof(self, index=_S):
NE = self.mesh.number_of_edges()
edof = self.number_of_local_dofs(doftype='edge')
return bm.arange(NE*edof).reshape(NE, edof)[index]
return bm.arange(NE*edof, device=self.device).reshape(NE, edof)[index]

def face_to_internal_dof(self, index=_S):
NE = self.mesh.number_of_edges()
NF = self.mesh.number_of_faces()
edof = self.number_of_local_dofs(doftype='edge')
fdof = self.number_of_local_dofs(doftype='face')
return bm.arange(NE*edof, NE*edof+NF*fdof).reshape(NF, fdof)[index]
return bm.arange(NE*edof, NE*edof+NF*fdof,device=self.device).reshape(NF, fdof)[index]

def face_to_dof(self):
p = self.p
Expand Down Expand Up @@ -109,7 +109,7 @@ def cell_to_dof(self):
face = self.mesh.entity('face')
cell = self.mesh.entity('cell')

c2d = bm.zeros((NC, ldof), dtype=bm.int64)
c2d = bm.zeros((NC, ldof), dtype=bm.int64, device=self.device)
# c2d[:, :eldof*6] = e2dof[c2e].reshape(NC, eldof*6)
c2d = bm.set_at(c2d,(slice(None),slice(eldof*6)),e2dof[c2e].reshape(NC, eldof*6))
s = [0, 0, 0, 1, 1, 2]
Expand All @@ -118,12 +118,12 @@ def cell_to_dof(self):
# c2d[flag, eldof*i:eldof*(i+1)] = bm.flip(c2d[flag, eldof*i:eldof*(i+1)],axis=-1)
c2d = bm.set_at(c2d,(flag,slice(eldof*i,eldof*(i+1))),bm.flip(c2d[flag, eldof*i:eldof*(i+1)],axis=-1))
if fldof > 0:
locFace = bm.array([[1, 2, 3], [0, 2, 3], [0, 1, 3], [0, 1, 2]], dtype=self.itype)
locFace = bm.array([[1, 2, 3], [0, 2, 3], [0, 1, 3], [0, 1, 2]],device=self.device, dtype=self.itype)
midx2num = lambda a : (a[:, 1]+a[:, 2])*(1+a[:, 1]+a[:, 2])//2 + a[:, 2]

midx = self.mesh.multi_index_matrix(p-1, 2)
perms = bm.array(list(itertools.permutations([0, 1, 2])))
indices = bm.zeros((6, len(midx)), dtype=self.itype)
perms = bm.array(list(itertools.permutations([0, 1, 2])), device=self.device)
indices = bm.zeros((6, len(midx)),device=self.device, dtype=self.itype)
for i in range(6):
indices[i] = midx2num(midx[:, perms[i]])

Expand Down Expand Up @@ -163,7 +163,7 @@ def is_boundary_dof(self):
ldof = p+1
gdof = self.number_of_global_dofs()

isDDof = bm.zeros(gdof, dtype=bm.bool)
isDDof = bm.zeros(gdof, device=self.device,dtype=bm.bool)
index = self.mesh.boundary_face_index()
face2dof = self.face_to_internal_dof()[index]
isDDof[face2dof] = True
Expand All @@ -173,7 +173,7 @@ def is_boundary_dof(self):
# bdeflag = bm.all(bdnflag[edge], axis=1)
NE = mesh.number_of_edges()
f2e = mesh.face_to_edge()[index]
bdeflag = bm.zeros(NE, dtype=bm.bool)
bdeflag = bm.zeros(NE,device=self.device,dtype=bm.bool)
bdeflag[f2e] = True
index = bm.nonzero(bdeflag)[0]

Expand Down Expand Up @@ -236,7 +236,7 @@ def basis(self, bcs, index=_S):
# edge basis
phi = self.bspace.basis(bcs, p=p)
multiIndex = self.mesh.multi_index_matrix(p, 3)
val = bm.zeros((NC,) + bcs.shape[:-1]+(ldof, 3), dtype=self.ftype)
val = bm.zeros((NC,) + bcs.shape[:-1]+(ldof, 3),device=self.device, dtype=self.ftype)
locEdgeDual = bm.tensor([[2, 3], [1, 3], [1, 2], [0, 3], [0, 2], [0, 1]])
for i in range(6):
flag = bm.all(multiIndex[:, locEdgeDual[i]]==0, axis=1)
Expand Down Expand Up @@ -303,7 +303,7 @@ def face_internal_basis(self, bcs, index=_S):
glambda = mesh.grad_face_lambda()
val = bm.zeros((NF,) + bcs.shape[:-1]+(fdof, 3),device=self.device, dtype=self.ftype)

l = bm.zeros((3, )+bcs[None, :,0, None, None].shape, dtype=self.ftype)
l = bm.zeros((3, )+bcs[None, :,0, None, None].shape, device=self.device,dtype=self.ftype)
l = bm.set_at(l,(0),bcs[None, :,0, None, None])
l = bm.set_at(l,(1),bcs[None, :,1, None, None])
l = bm.set_at(l,(2),bcs[None, :,2, None, None])
Expand All @@ -330,9 +330,9 @@ def face_basis(self, bcs, index=_S):
glambda = mesh.grad_face_lambda()
ledge = bm.array([[1, 2],
[2, 0],
[0, 1]])
[0, 1]],device=self.device)
c2esign = mesh.face_to_edge_sign()
l = bm.zeros((3, )+bcs[None, :,0, None, None].shape, dtype=self.ftype)
l = bm.zeros((3, )+bcs[None, :,0, None, None].shape,device=self.device, dtype=self.ftype)
l = bm.set_at(l,(0),bcs[None, :,0, None, None])
l = bm.set_at(l,(1),bcs[None, :,1, None, None])
l = bm.set_at(l,(2),bcs[None, :,2, None, None])
Expand Down Expand Up @@ -373,7 +373,7 @@ def curl_basis(self, bcs):

c2esign = mesh.cell_to_edge_sign() #(NC, 6, 2)

l = bm.zeros((4, )+bcs[None,:, 0,None, None].shape, dtype=self.ftype)
l = bm.zeros((4, )+bcs[None,:, 0,None, None].shape,device=self.device, dtype=self.ftype)

l = bm.set_at(l,(0),bcs[None, :,0,None, None])
l = bm.set_at(l,(1),bcs[None, :,1,None, None])
Expand Down Expand Up @@ -595,7 +595,7 @@ def set_dirichlet_bc(self, gD, uh, threshold=None, q=None):
p = self.p
mesh = self.mesh
gdof = self.number_of_global_dofs()
isDDof = bm.zeros(gdof, dtype=bm.bool)
isDDof = bm.zeros(gdof, device=self.device, dtype=bm.bool)


index1 = self.mesh.boundary_face_index()
Expand Down Expand Up @@ -626,7 +626,7 @@ def set_dirichlet_bc(self, gD, uh, threshold=None, q=None):
# 边界边界的点
NE = mesh.number_of_edges()
f2e = mesh.face_to_edge()[index1]
bdeflag = bm.zeros(NE, dtype=bm.bool)
bdeflag = bm.zeros(NE, device=self.device, dtype=bm.bool)
bdeflag[f2e] = True
index2 = bm.nonzero(bdeflag)[0]
edge2dof = self.dof.edge_to_dof()[index2]
Expand Down
4 changes: 2 additions & 2 deletions fealpy/mesh/mesh_data_structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,8 +153,8 @@ def cell_to_face(self, index: Index=_S) -> TensorLike:
face2cell = self.face2cell
dtype = self.itype

cell2face = bm.zeros((NC, NFC), dtype=dtype)
arange_tensor = bm.arange(0, NF, dtype=dtype)
cell2face = bm.zeros((NC, NFC),device=bm.get_device(face2cell), dtype=dtype)
arange_tensor = bm.arange(0, NF,device=bm.get_device(face2cell),dtype=dtype)

assert cell2face.dtype == arange_tensor.dtype, f"Data type mismatch: cell2face is {cell2face.dtype}, arange_tensor is {arange_tensor.dtype}"

Expand Down
6 changes: 3 additions & 3 deletions fealpy/mesh/tetrahedron_mesh.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ def face_to_edge_sign(self):
face = self.face
NF = len(face2edge)
NEF = 3
face2edgeSign = bm.zeros((NF, NEF), dtype=bm.bool)
face2edgeSign = bm.zeros((NF, NEF),device=bm.get_device(face), dtype=bm.bool)
n = [1, 2, 0]
for i in range(3):
face2edgeSign[:, i] = (face[:, n[i]] == edge[face2edge[:, i], 0])
Expand Down Expand Up @@ -194,7 +194,7 @@ def grad_lambda(self, index=_S):
node = self.node
cell = self.cell
NC = self.number_of_cells() if index == _S else len(index)
Dlambda = bm.zeros((NC, 4, 3), dtype=self.ftype)
Dlambda = bm.zeros((NC, 4, 3), device=self.device, dtype=self.ftype)
volume = self.entity_measure('cell', index=index)
for i in range(4):
j,k,m = localFace[i]
Expand All @@ -213,7 +213,7 @@ def grad_face_lambda(self, index=_S):
v2 = node[face[..., 1]] - node[face[..., 0]]
GD = self.geo_dimension()
nv = bm.cross(v1, v2)
Dlambda = bm.zeros((NF, 3, GD), dtype=self.ftype)
Dlambda = bm.zeros((NF, 3, GD),device=bm.get_device(face), dtype=self.ftype)

length = bm.linalg.norm(nv, axis=-1, keepdims=True)
n = nv / length
Expand Down

0 comments on commit 9abf2f4

Please sign in to comment.