Ejemplo n.º 1
0
class SurfaceTriangleMesh():
    def __init__(self, mesh, surface, p=1, scale=None):
        """
        Initial a object of Surface Triangle Mesh.

        Parameters
        ----------
        self :
            Surface Triangle Mesh Object
        mesh :
            mesh object, represents a triangulation with flat triangle faces.
        surface :
            The continuous surface which was represented as a level set
            function.
        p : int
            The degree of the Lagrange space

        Returns
        -------

        See Also
        --------

        Notes
        -----
        """
        self.mesh = mesh
        self.p = p
        self.space = LagrangeFiniteElementSpace(mesh, p)
        self.surface = surface
        self.ds = mesh.ds
        self.scale = scale
        if scale is not None:
            self.mesh.node *= scale

        if scale is None:
            self.node, d = self.surface.project(
                self.space.interpolation_points())
        else:
            self.node, d = self.surface.project(
                self.space.interpolation_points() / scale)
            self.node *= scale

        self.meshtype = 'stri'
        self.ftype = mesh.ftype
        self.itype = mesh.itype
        self.nodedata = {}
        self.celldata = {}

    def vtk_cell_type(self):
        return 69

    def project(self, p):
        if self.scale is None:
            return self.surface.project(p)
        else:
            p, d = self.surface.project(p / self.scale)
            return p * self.scale, d * self.scale

    def integrator(self, k, etype='cell'):
        return TriangleQuadrature(k)

    def entity(self, etype=2):
        if etype in {'cell', 2}:
            return self.ds.cell
        elif etype in {'edge', 1}:
            return self.ds.edge
        elif etype in {'node', 0}:
            return self.mesh.node
        else:
            raise ValueError("`entitytype` is wrong!")

    def entity_measure(self, etype=2):
        p = self.p
        if etype in {'cell', 2}:
            return self.area(p + 1)
        elif etype in {'edge', 'face', 1}:
            return self.mesh.entity_measure('edge')
        else:
            raise ValueError("`entitytype` is wrong!")

    def entity_barycenter(self, etype=2):
        p = self.p
        return self.mesh.entity_barycenter(etype=etype)

    def number_of_nodes(self):
        return self.node.shape[0]

    def number_of_edges(self):
        return self.mesh.ds.NE

    def number_of_cells(self):
        return self.mesh.ds.NC

    def geo_dimension(self):
        return self.node.shape[1]

    def top_dimension(self):
        return 2

    def jacobi_matrix(self, bc, index=np.s_[:]):
        mesh = self.mesh
        cell2dof = self.space.dof.cell2dof

        grad = self.space.grad_basis(bc, index=index)
        # the tranpose of the jacobi matrix between S_h and K
        Jh = mesh.jacobi_matrix(index=index)

        # the tranpose of the jacobi matrix between S_p and S_h
        Jph = np.einsum('ijm, ...ijk->...imk', self.node[cell2dof[index], :],
                        grad)

        # the transpose of the jacobi matrix between S_p and K
        Jp = np.einsum('...ijk, imk->...imj', Jph, Jh)
        grad = np.einsum('ijk, ...imk->...imj', Jh, grad)
        return Jp, grad

    def normal(self, bc, index=None):
        Js, _, ps = self.surface_jacobi_matrix(bc, index=index)
        n = np.cross(Js[..., 0, :], Js[..., 1, :], axis=-1)
        return n, ps

    def surface_jacobi_matrix(self, bc, index=None):
        Jp, grad = self.jacobi_matrix(bc, index=index)
        ps = self.bc_to_point(bc, index=index)
        Jsp = self.surface.jacobi_matrix(ps)
        Js = np.einsum('...ijk, ...imk->...imj', Jsp, Jp)
        return Js, grad, ps

    def bc_to_point(self, bc, index=None):
        phi = self.space.basis(bc)
        cell2dof = self.space.dof.cell2dof
        if index is None:
            bcp = np.einsum('...ij, ijk->...ik', phi, self.node[cell2dof, :])
        else:
            bcp = np.einsum('...ij, ijk->...ik', phi,
                            self.node[cell2dof[index], :])

        bcp, _ = self.project(bcp)
        return bcp

    def area(self, q=3):
        integrator = self.integrator(q)
        bcs, ws = integrator.quadpts, integrator.weights
        Jp, _ = self.jacobi_matrix(bcs)
        n = np.cross(Jp[..., 0, :], Jp[..., 1, :], axis=-1)
        l = np.sqrt(np.sum(n**2, axis=-1))
        a = np.einsum('i, ij->j', ws, l) / 2.0
        return a

    def add_plot(self,
                 plot,
                 nodecolor='w',
                 edgecolor='k',
                 cellcolor=[0.5, 0.9, 0.45],
                 aspect='equal',
                 linewidths=1,
                 markersize=50,
                 showaxis=False,
                 showcolorbar=False,
                 cmap='rainbow'):

        if isinstance(plot, ModuleType):
            fig = plot.figure()
            fig.set_facecolor('white')
            axes = fig.gca()
        else:
            axes = plot
        return show_mesh_2d(axes,
                            self.mesh,
                            nodecolor=nodecolor,
                            edgecolor=edgecolor,
                            cellcolor=cellcolor,
                            aspect=aspect,
                            linewidths=linewidths,
                            markersize=markersize,
                            showaxis=showaxis,
                            showcolorbar=showcolorbar,
                            cmap=cmap)

    def find_node(self,
                  axes,
                  node=None,
                  index=None,
                  showindex=False,
                  color='r',
                  markersize=100,
                  fontsize=24,
                  fontcolor='k'):

        if node is None:
            node = self.node

        if (index is None) and (showindex is True):
            index = np.array(range(node.shape[0]))

        find_node(axes,
                  node,
                  index=index,
                  showindex=showindex,
                  color=color,
                  markersize=markersize,
                  fontsize=fontsize,
                  fontcolor=fontcolor)

    def find_edge(self,
                  axes,
                  index=None,
                  showindex=False,
                  color='g',
                  markersize=150,
                  fontsize=24,
                  fontcolor='k'):

        find_entity(axes,
                    self.mesh,
                    entity='edge',
                    index=index,
                    showindex=showindex,
                    color=color,
                    markersize=markersize,
                    fontsize=fontsize,
                    fontcolor=fontcolor)

    def find_cell(self,
                  axes,
                  index=None,
                  showindex=False,
                  color='y',
                  markersize=200,
                  fontsize=24,
                  fontcolor='k'):

        find_entity(axes,
                    self.mesh,
                    entity='cell',
                    index=index,
                    showindex=showindex,
                    color=color,
                    markersize=markersize,
                    fontsize=fontsize,
                    fontcolor=fontcolor)
Ejemplo n.º 2
0
class HuZhangFiniteElementSpace():
    """
    Hu-Zhang Mixed Finite Element Space.
    """
    def __init__(self, mesh, p):
        self.space = LagrangeFiniteElementSpace(mesh, p)  # the scalar space
        self.mesh = mesh
        self.p = p
        self.dof = self.space.dof
        self.dim = self.space.GD
        self.init_orth_matrices()
        self.init_cell_to_dof()

    def init_orth_matrices(self):
        """
        Initialize the othogonal symetric matrix basis.
        """
        mesh = self.mesh
        gdim = self.geo_dimension()

        NE = mesh.number_of_edges()
        if gdim == 2:
            idx = np.array([(0, 0), (1, 1), (0, 1)])
            self.TE = np.zeros((NE, 3, 3), dtype=np.float)
            self.T = np.array([[(1, 0), (0, 0)], [(0, 0), (0, 1)],
                               [(0, 1), (1, 0)]])
        elif gdim == 3:
            idx = np.array([(0, 0), (1, 1), (2, 2), (1, 2), (0, 2), (0, 1)])
            self.TE = np.zeros((NE, 6, 6), dtype=np.float)
            self.T = np.array([[(1, 0, 0), (0, 0, 0), (0, 0, 0)],
                               [(0, 0, 0), (0, 1, 0), (0, 0, 0)],
                               [(0, 0, 0), (0, 0, 0), (0, 0, 1)],
                               [(0, 0, 0), (0, 0, 1), (0, 1, 0)],
                               [(0, 0, 1), (0, 0, 0), (1, 0, 0)],
                               [(0, 1, 0), (1, 0, 0), (0, 0, 0)]])

        t = mesh.edge_unit_tagent()
        _, _, frame = np.linalg.svd(
            t[:, np.newaxis, :])  # get the axis frame on the edge by svd
        frame[:, 0, :] = t
        for i, (j, k) in enumerate(idx):
            self.TE[:,
                    i] = (frame[:, j, idx[:, 0]] * frame[:, k, idx[:, 1]] +
                          frame[:, j, idx[:, 1]] * frame[:, k, idx[:, 0]]) / 2
        self.TE[:, gdim:] *= np.sqrt(2)

        if gdim == 3:
            NF = mesh.number_of_faces()
            self.TF = np.zeros((NF, 6, 6), dtype=np.float)
            n = mesh.face_unit_normal()
            _, _, frame = np.linalg.svd(
                n[:, np.newaxis, :])  # get the axis frame on the edge by svd
            frame[:, 0, :] = n
            for i, (j, k) in enumerate(idx):
                self.TF[:, i] = (
                    frame[:, j, idx[:, 0]] * frame[:, k, idx[:, 1]] +
                    frame[:, j, idx[:, 1]] * frame[:, k, idx[:, 0]]) / 2

            self.TF[:, gdim:] *= np.sqrt(2)

    def __str__(self):
        return "Hu-Zhang mixed finite element space!"

    def number_of_global_dofs(self):
        """
        """
        p = self.p
        gdim = self.geo_dimension()
        tdim = self.tensor_dimension()

        mesh = self.mesh

        NC = mesh.number_of_cells()
        NN = mesh.number_of_nodes()
        gdof = tdim * NN

        if p > 1:
            edof = p - 1
            NE = mesh.number_of_edges()
            gdof += (tdim - 1) * edof * NE  # 边内部连续自由度的个数
            E = mesh.number_of_edges_of_cells()  # 单元边的个数
            gdof += NC * E * edof  # 边内部不连续自由度的个数

        if p > 2:
            fdof = (p + 1) * (p + 2) // 2 - 3 * p  # 面内部自由度的个数
            if gdim == 2:
                gdof += tdim * fdof * NC
            elif gdim == 3:
                NF = mesh.number_of_faces()
                gdof += 3 * fdof * NF  # 面内部连续自由度的个数
                F = mesh.number_of_faces_of_cells()  # 每个单元面的个数
                gdof += 3 * F * fdof * NC  # 面内部不连续自由度的个数

        if (p > 3) and (gdim == 3):
            ldof = self.dof.number_of_local_dofs()
            V = mesh.number_of_nodes_of_cells()  # 单元顶点的个数
            cdof = ldof - E * edof - F * fdof - V
            gdof += tdim * cdof * NC
        return gdof

    def number_of_local_dofs(self):
        tdim = self.tensor_dimension()
        ldof = self.dof.number_of_local_dofs()
        return tdim * ldof

    def cell_to_dof(self):
        return self.cell2dof

    def init_cell_to_dof(self):
        """
        构建局部自由度到全局自由度的映射矩阵

        Returns
        -------
        cell2dof : ndarray with shape (NC, ldof*tdim)
            NC: 单元个数
            ldof: p 次标量空间局部自由度的个数
            tdim: 对称张量的维数
        """
        mesh = self.mesh
        NN = mesh.number_of_nodes()
        NE = mesh.number_of_edges()
        NC = mesh.number_of_cells()

        gdim = self.geo_dimension()
        tdim = self.tensor_dimension()  # 张量维数
        p = self.p
        dof = self.dof  # 标量空间自由度管理对象

        c2d = dof.cell2dof[..., np.newaxis]
        ldof = dof.number_of_local_dofs()  # ldof : 标量空间单元上自由度个数
        cell2dof = np.zeros((NC, ldof, tdim),
                            dtype=np.int)  # 每个标量自由度变成 tdim 个自由度

        dofFlags = self.dof_flags_1()  # 把不同类型的自由度区分开来
        idx, = np.nonzero(dofFlags[0])  # 局部顶点自由度的编号
        cell2dof[:, idx, :] = tdim * c2d[:, idx] + np.arange(tdim)

        base0 = 0
        base1 = 0
        idx, = np.nonzero(dofFlags[1])  # 边内部自由度的编号
        if len(idx) > 0:
            base0 += NN  # 这是标量编号的新起点
            base1 += tdim * NN  # 这是张量自由度编号的新起点
            #  0号局部自由度对应的是切向不连续的自由度, 留到后面重新编号
            cell2dof[:, idx,
                     1:] = base1 + (tdim - 1) * (c2d[:, idx] -
                                                 base0) + np.arange(tdim - 1)

        idx, = np.nonzero(dofFlags[2])
        if len(idx) > 0:
            edof = p - 1
            base0 += edof * NE
            base1 += (tdim - 1) * edof * NE
            if gdim == 2:
                cell2dof[:, idx, :] = base1 + tdim * (c2d[:, idx] -
                                                      base0) + np.arange(tdim)
            elif gdim == 3:
                # 1, 2, 3 号局部自由度对应切向不连续的张量自由度, 留到后面重新编号
                # TODO: check it is right
                cell2dof[:, idx.reshape(-1, 1),
                         np.array([0, 4, 5])] = base1 + (tdim - 3) * (
                             c2d[:, idx] - base0) + np.arange(tdim - 3)

        fdof = (p + 1) * (p + 2) // 2 - 3 * p  # 边内部自由度
        if gdim == 3:
            idx, = np.nonzero(dofFlags[3])
            if len(idx) > 0:
                NF = mesh.number_of_faces()
                base0 += fdof * NF
                base1 += (tdim - 3) * fdof * NF
                cell2dof[:, idx, :] = base1 + tdim * (c2d[:, idx] -
                                                      base0) + np.arange(tdim)
            cdof = ldof - 4 * fdof - 6 * edof - 4  # 单元内部自由度
        else:
            cdof = fdof

        idx, = np.nonzero(dofFlags[1])
        if len(idx) > 0:
            base1 += tdim * cdof * NC
            cell2dof[:, idx, 0] = base1 + np.arange(NC * len(idx)).reshape(
                NC, len(idx))

        if gdim == 3:
            base1 += NC * len(idx)
            idx, = np.nonzero(dofFlags[2])
            print(idx)
            if len(idx) > 0:
                cell2dof[:, idx.reshape(-1, 1),
                         np.array([1, 2, 3])] = base1 + np.arange(
                             NC * len(idx) * 3).reshape(NC, len(idx), 3)

        self.cell2dof = cell2dof.reshape(NC, -1)

    def geo_dimension(self):
        return self.dim

    def tensor_dimension(self):
        dim = self.dim
        return dim * (dim - 1) // 2 + dim

    def interpolation_points(self):
        return self.dof.interpolation_points()

    def dof_flags(self):
        """ 对标量空间中的自由度进行分类, 分为边内部自由度, 面内部自由度(如果是三维空间的话)及其它自由度 

        Returns
        -------

        isOtherDof : ndarray, (ldof,)
            除了边内部和面内部自由度的其它自由度
        isEdgeDof : ndarray, (ldof, 3) or (ldof, 6) 
            每个边内部的自由度
        isFaceDof : ndarray, (ldof, 4)
            每个面内部的自由度
        -------

        """
        dim = self.geo_dimension()
        dof = self.dof

        isPointDof = dof.is_on_node_local_dof()
        isEdgeDof = dof.is_on_edge_local_dof()
        isEdgeDof[isPointDof] = False

        isEdgeDof0 = np.sum(isEdgeDof, axis=-1) > 0  #
        isOtherDof = (~isEdgeDof0)  # 除了边内部自由度之外的其它自由度
        # dim = 2: 包括点和面内部自由度
        # dim = 3: 包括点, 面内部和体内部自由度
        if dim == 2:
            return isOtherDof, isEdgeDof
        elif dim == 3:
            isFaceDof = dof.is_on_face_local_dof()
            isFaceDof[isPointDof, :] = False
            isFaceDof[isEdgeDof0, :] = False

            isFaceDof0 = np.sum(isFaceDof, axis=-1) > 0
            isOtherDof = isOtherDof & (~isFaceDof0)  # 三维情形下, 从其它自由度中除去面内部自由度

            return isOtherDof, isEdgeDof, isFaceDof
        else:
            raise ValueError('`dim` should be 2 or 3!')

    def dof_flags_1(self):
        """ 
        对标量空间中的自由度进行分类, 分为:
            点上的自由由度
            边内部的自由度
            面内部的自由度
            体内部的自由度

        Returns
        -------

        """
        gdim = self.geo_dimension()  # the geometry space dimension
        dof = self.dof
        isPointDof = dof.is_on_node_local_dof()
        isEdgeDof = dof.is_on_edge_local_dof()
        isEdgeDof[isPointDof] = False
        isEdgeDof0 = np.sum(isEdgeDof, axis=-1) > 0
        if gdim == 2:
            return isPointDof, isEdgeDof0, ~(isPointDof | isEdgeDof0)
        elif gdim == 3:
            isFaceDof = dof.is_on_face_local_dof()
            isFaceDof[isPointDof, :] = False
            isFaceDof[isEdgeDof0, :] = False

            isFaceDof0 = np.sum(isFaceDof, axis=-1) > 0
            return isPointDof, isEdgeDof0, isFaceDof0, ~(
                isPointDof | isEdgeDof0 | isFaceDof0)
        else:
            raise ValueError('`dim` should be 2 or 3!')

    def basis(self, bc, cellidx=None):
        """

        Parameters
        ----------
        bc : ndarray with shape (NQ, dim+1)
            bc[i, :] is i-th quad point
        cellidx : ndarray
            有时我我们只需要计算部分单元上的基函数
        Returns
        -------
        phi : ndarray with shape (NQ, NC, ldof*tdim, 3 or 6)
            NQ: 积分点个数
            NC: 单元个数
            ldof: 标量空间的单元自由度个数
            tdim: 对称张量的维数
        """
        mesh = self.mesh

        gdim = self.geo_dimension()
        tdim = self.tensor_dimension()

        if cellidx is None:
            NC = mesh.number_of_cells()
            cell2edge = mesh.ds.cell_to_edge()
        else:
            NC = len(cellidx)
            cell2edge = mesh.ds.cell_to_edge()[cellidx]

        phi0 = self.space.basis(bc)  # the shape of phi0 is (NQ, ldof)
        shape = list(phi0.shape)
        shape.insert(-1, NC)
        shape += [tdim, tdim]
        # The shape of `phi` is (NQ, NC, ldof, tdim, tdim), where
        #   NQ : the number of quadrature points
        #   NC : the number of cells
        #   ldof : the number of dofs in each cell
        #   tdim : the dimension of symmetric tensor matrix
        phi = np.zeros(shape, dtype=np.float)

        dofFlag = self.dof_flags()
        # the dof on the vertex and the interior of the cell
        isOtherDof = dofFlag[0]
        idx, = np.nonzero(isOtherDof)
        if len(idx) > 0:
            phi[..., idx[..., np.newaxis],
                range(tdim),
                range(tdim)] = phi0[..., np.newaxis, idx, np.newaxis]

        isEdgeDof = dofFlag[1]
        for i, isDof in enumerate(isEdgeDof.T):
            phi[..., isDof, :, :] = np.einsum('...j, imn->...ijmn',
                                              phi0[..., isDof],
                                              self.TE[cell2edge[:, i]])

        if gdim == 3:
            if cellidx is None:
                cell2face = mesh.ds.cell_to_face()
            else:
                cell2face = mesh.ds.cell_to_face()[cellidx]
            isFaceDof = dofFlag[2]
            for i, isDof in enumerate(isFaceDof.T):
                phi[..., isDof, :, :] = np.einsum('...j, imn->...ijmn',
                                                  phi0[..., isDof],
                                                  self.TF[cell2face[:, i]])
        # The shape of `phi` should be (NQ, NC, ldof*tdim, tdim)?
        shape = phi.shape[:-3] + (-1, tdim)
        return phi.reshape(shape)

    def div_basis(self, bc, cellidx=None):
        mesh = self.mesh

        gdim = self.geo_dimension()
        tdim = self.tensor_dimension()

        # the shape of `gphi` is (NQ, NC, ldof, gdim)
        gphi = self.space.grad_basis(bc, cellidx=cellidx)
        shape = list(gphi.shape)
        shape.insert(-1, tdim)
        # the shape of `dphi` is (NQ, NC, ldof, tdim, gdim)
        dphi = np.zeros(shape, dtype=np.float)

        dofFlag = self.dof_flags()
        # the dof on the vertex and the interior of the cell
        isOtherDof = dofFlag[0]
        dphi[..., isOtherDof, :, :] = np.einsum('...ijm, kmn->...ijkn',
                                                gphi[...,
                                                     isOtherDof, :], self.T)

        if cellidx is None:
            cell2edge = mesh.ds.cell_to_edge()
        else:
            cell2edge = mesh.ds.cell_to_edge()[cellidx]
        isEdgeDof = dofFlag[1]
        for i, isDof in enumerate(isEdgeDof.T):
            VAL = np.einsum('ijk, kmn->ijmn', self.TE[cell2edge[:, i]], self.T)
            dphi[..., isDof, :, :] = np.einsum('...ikm, ijmn->...ikjn',
                                               gphi[..., isDof, :], VAL)

        if gdim == 3:
            if cellidx is None:
                cell2face = mesh.ds.cell_to_face()
            else:
                cell2face = mesh.ds.cell_to_face()[cellidx]
            isFaceDof = dofFlag[2]
            for i, isDof in enumerate(isFaceDof.T):
                VAL = np.einsum('ijk, kmn->ijmn', self.TF[cell2face[:, i]],
                                self.T)
                dphi[..., isDof, :, :] = np.einsum('...ikm, ijmn->...ikjn',
                                                   gphi[..., isDof, :], VAL)

        # The new shape of `dphi` is `(NQ, NC, ldof*tdim, gdim)`, where
        shape = dphi.shape[:-3] + (-1, gdim)
        return dphi.reshape(shape)

    def value(self, uh, bc, cellidx=None):
        phi = self.basis(bc, cellidx=cellidx)
        cell2dof = self.cell_to_dof()
        tdim = self.tensor_dimension()
        if cellidx is None:
            uh = uh[cell2dof]
        else:
            uh = uh[cell2dof[cellidx]]
        phi = np.einsum('...jk, kmn->...jmn', phi, self.T)
        val = np.einsum('...ijmn, ij->...imn', phi, uh)
        return val

    def div_value(self, uh, bc, cellidx=None):
        dphi = self.div_basis(bc, cellidx=cellidx)
        cell2dof = self.cell_to_dof()
        tdim = self.tensor_dimension()
        if cellidx is None:
            uh = uh[cell2dof]
        else:
            uh = uh[cell2dof[cellidx]]
        val = np.einsum('...ijm, ij->...im', dphi, uh)
        return val

    def interpolation(self, u):

        mesh = self.mesh
        gdim = self.geo_dimension()
        tdim = self.tensor_dimension()

        if gdim == 2:
            idx = np.array([(0, 0), (1, 1), (0, 1)])
        elif gdim == 3:
            idx = np.array([(0, 0), (1, 1), (2, 2), (1, 2), (0, 2), (0, 1)])

        ipoint = self.dof.interpolation_points()
        c2d = self.dof.cell2dof
        val = u(ipoint)[c2d]

        ldof = self.dof.number_of_local_dofs()
        cell2dof = self.cell2dof.reshape(-1, ldof, tdim)

        uI = Function(self)
        dofFlag = self.dof_flags()
        isOtherDof = dofFlag[0]
        idx0, = np.nonzero(isOtherDof)
        uI[cell2dof[:, idx0, :]] = val[:, idx0][..., idx[:, 0], idx[:, 1]]

        isEdgeDof = dofFlag[1]
        cell2edge = self.mesh.ds.cell_to_edge()
        for i, isDof in enumerate(isEdgeDof.T):
            TE = np.einsum('ijk, kmn->ijmn', self.TE[cell2edge[:, i]], self.T)
            uI[cell2dof[:, isDof, :]] = np.einsum('ikmn, ijmn->ikj',
                                                  val[:, isDof, :, :], TE)

        if gdim == 3:
            cell2face = mesh.ds.cell_to_face()
            isFaceDof = dofFlag[2]
            for i, isDof in enumerate(isFaceDof.T):
                TF = np.einsum('ijk, kmn->ijmn', self.TF[cell2face[:, i]],
                               self.T)
                uI[cell2dof[:, isDof, :]] = np.einsum('ikmn, ijmn->ikj',
                                                      val[...,
                                                          isDof, :, :], TF)
        return uI

        def function(self, dim=None):
            f = Function(self)

        return f

    def array(self, dim=None):
        gdof = self.number_of_global_dofs()
        return np.zeros(gdof, dtype=np.float)
Ejemplo n.º 3
0
class ParallelTwoFluidsWithGeostressSimulator():
    """

    Notes
    -----
    这是一个并行模拟两种流体和地质力学耦合的程序, 这里只是线性系统是并行的  

    * S_0: 流体 0 的饱和度,一般是水
    * S_1: 流体 1 的饱和度,可以为气或油
    * v: 总速度
    * p: 压力
    * u: 岩石位移  

    饱和度 P0 间断元,单元边界用迎风格式
    速度 RT0 元
    压强 P0 元
    岩石位移 P1 连续元离散

    目前, 模型
    * 忽略了毛细管压强和重力作用
    * 饱和度用分片线性间断元求解, 非线性的迎风格式

    渐近解决方案:
    1. Picard 迭代
    2. 气的可压性随着压强的变化而变化
    3. 考虑渗透率随着孔隙度的变化而变化 
    4. 考虑裂缝,裂缝用自适应网格加密,设设置更大的渗透率实现

    体积模量:  K = E/3/(1 - 2*nu) = lambda + 2/3* mu

    Develop
    ------

    """
    def __init__(self, mesh, args, ctx):
        self.ctx = ctx
        self.args = args # 模拟相关参数

        NT = int((args.T1 - args.T0)/args.DT)
        self.timeline = UniformTimeLine(args.T0, args.T1, NT)
        self.mesh = mesh 

        if self.ctx.myid == 0:
            self.GD = mesh.geo_dimension()
            if self.GD == 2:
                self.vspace = RaviartThomasFiniteElementSpace2d(self.mesh, p=0) # 速度空间
            elif self.GD == 3:
                self.vspace = RaviartThomasFiniteElementSpace3d(self.mesh, p=0)

            self.pspace = self.vspace.smspace # 压力和饱和度所属的空间, 分片常数
            self.cspace = LagrangeFiniteElementSpace(self.mesh, p=1) # 位移空间

            # 上一时刻物理量的值
            self.v = self.vspace.function() # 速度函数
            self.p = self.pspace.function() # 压力函数
            self.s = self.pspace.function() # 水的饱和度函数 默认为0, 初始时刻区域内水的饱和度为0
            self.u = self.cspace.function(dim=self.GD) # 位移函数
            self.phi = self.pspace.function() # 孔隙度函数, 分片常数

            # 当前时刻物理量的值, 用于保存临时计算出的值, 模型中系数的计算由当前时刻
            # 的物理量的值决定
            self.cv = self.vspace.function() # 速度函数
            self.cp = self.pspace.function() # 压力函数
            self.cs = self.pspace.function() # 水的饱和度函数 默认为0, 初始时刻区域内水的饱和度为0
            self.cu = self.cspace.function(dim=self.GD) # 位移函数
            self.cphi = self.pspace.function() # 孔隙度函数, 分片常数

            # 初值
            self.s[:] = self.mesh.celldata['fluid_0']
            self.p[:] = self.mesh.celldata['pressure_0']  # MPa
            self.phi[:] = self.mesh.celldata['porosity_0'] # 初始孔隙度 

            self.s[:] = self.mesh.celldata['fluid_0']
            self.cp[:] = self.p # 初始地层压力
            self.cphi[:] = self.phi # 当前孔隙度系数

            # 源汇项 
            self.f0 = self.cspace.function()
            self.f1 = self.cspace.function()

            self.f0[:] = self.mesh.nodedata['source_0'] # 流体 0 的源汇项
            self.f1[:] = self.mesh.nodedata['source_1'] # 流体 1 的源汇项

            # 一些常数矩阵和向量

            # 速度散度矩阵, 速度方程对应的散度矩阵, (\nabla\cdot v, w) 
            self.B = self.vspace.div_matrix()

            # 压力方程对应的位移散度矩阵, (\nabla\cdot u, w) 位移散度矩阵
            # * 注意这里利用了压力空间分片常数, 线性函数导数也是分片常数的事实
            c = self.mesh.entity_measure('cell')
            c *= self.mesh.celldata['biot']

            val = self.mesh.grad_lambda() # (NC, TD+1, GD)
            val *= c[:, None, None]
            pc2d = self.pspace.cell_to_dof()
            cc2d = self.cspace.cell_to_dof()
            pgdof = self.pspace.number_of_global_dofs()
            cgdof = self.cspace.number_of_global_dofs()
            I = np.broadcast_to(pc2d, shape=cc2d.shape)
            J = cc2d 
            self.PU0 = csr_matrix(
                    (val[..., 0].flat, (I.flat, J.flat)), 
                    shape=(pgdof, cgdof)
                    )
            self.PU1 = csr_matrix(
                    (val[..., 1].flat, (I.flat, J.flat)),
                    shape=(pgdof, cgdof)
                    )

            if self.GD == 3:
                self.PU2 = csr_matrix(
                        (val[..., 2].flat, (I.flat, J.flat)),
                        shape=(pgdof, cgdof)
                        )

            # 线弹性矩阵的右端向量
            self.FU = np.zeros(self.GD*cgdof, dtype=np.float64)
            self.FU[0*cgdof:1*cgdof] -= [email protected]
            self.FU[1*cgdof:2*cgdof] -= [email protected]

            if self.GD == 3:
                self.FU[2*cgdof:3*cgdof] -= [email protected]

            # 初始应力和等效应力项
            sigma = self.mesh.celldata['stress_0'] + self.mesh.celldata['stress_eff']# 初始应力和等效应力之和
            self.FU[0*cgdof:1*cgdof] -= [email protected]
            self.FU[1*cgdof:2*cgdof] -= [email protected]
            if self.GD == 3:
                self.FU[2*cgdof:3*cgdof] -= [email protected]

    def recover(self, val):
        """

        Notes
        -----
        给定一个分片常数的量, 恢复为分片连续的量
        """

        mesh = self.mesh
        cell = self.mesh.entity('cell')
        NC = mesh.number_of_cells()
        NN = mesh.number_of_nodes()

        w = 1/self.mesh.entity_measure('cell') # 恢复权重
        w = np.broadcast_to(w[:, None], shape=cell.shape)

        r = np.zeros(NN, dtype=np.float64)
        d = np.zeros(NN, dtype=np.float64)

        np.add.at(d, cell, w)
        np.add.at(r, cell, val[:, None]*w)

        return r/d

    def add_time(self, n):
        """

        Notes
        ----

        增加 n 步的计算时间 
        """

        if self.ctx.myid == 0:
            self.timeline.add_time(n)


    def pressure_coefficient(self):

        """

        Notes
        -----
        计算当前物理量下的压强质量矩阵系数
        """

        b = self.mesh.celldata['biot'] 
        K = self.mesh.celldata['K']
        c0 = self.mesh.meshdata['fluid_0']['compressibility']
        c1 = self.mesh.meshdata['fluid_1']['compressibility']

        s = self.cs # 流体 0 当前的饱和度 (NQ, NC)
        phi = self.cphi # 当前的孔隙度

        val = phi*s*c0  
        val += phi*(1 - s)*c1 # 注意这里的 co 是常数, 但在水气混合物条件下应该依赖于压力
        val += (b - phi)/K
        return val

    def saturation_pressure_coefficient(self):
        """

        Notes
        -----
        计算当前物理量下, 饱和度方程中, 压强项对应的系数
        """

        b = self.mesh.celldata['biot'] 
        K = self.mesh.celldata['K']
        c0 = self.mesh.meshdata['fluid_0']['compressibility']

        s = self.cs # 当前流体 0 的饱和度
        phi = self.cphi # 当前孔隙度

        val  = (b - phi)/K
        val += phi*c0
        val *= s 
        return val

    def flux_coefficient(self):
        """
        Notes
        -----

        计算**当前**物理量下, 速度方程对应的系数

        计算通量的系数, 流体的相对渗透率除以粘性系数得到流动性系数 
        k_0/mu_0 是气体的流动性系数 
        k_1/mu_1 是油的流动性系数

        这里假设压强的单位是 MPa 

        1 d = 9.869 233e-13 m^2 = 1000 md
        1 cp = 1 mPa s = 1e-9 MPa*s = 1e-3 Pa*s

        """

        mu0 = self.mesh.meshdata['fluid_0']['viscosity'] # 单位是 1 cp = 1 mPa.s
        mu1 = self.mesh.meshdata['fluid_1']['viscosity'] # 单位是 1 cp = 1 mPa.s 

        # 岩石的绝对渗透率, 这里考虑了量纲的一致性, 压强单位是 Pa
        k = self.mesh.celldata['permeability']*9.869233e-4 

        s = self.cs # 流体 0 当前的饱和度系数

        lam0 = self.mesh.fluid_relative_permeability_0(s) # TODO: 考虑更复杂的饱和度和渗透的关系 
        lam0 /= mu0

        lam1 = self.mesh.fluid_relative_permeability_1(s) # TODO: 考虑更复杂的饱和度和渗透的关系 
        lam1 /= mu1

        val = 1/(lam0 + lam1)
        val /=k 

        return val

    def fluid_fractional_flow_coefficient_0(self):
        """

        Notes
        -----

        计算**当前**物理量下, 饱和度方程中, 主流体的的流动性系数
        """


        s = self.cs # 流体 0 当前的饱和度系数

        mu0 = self.mesh.meshdata['fluid_0']['viscosity'] # 单位是 1 cp = 1 mPa.s
        lam0 = self.mesh.fluid_relative_permeability_0(s) 
        lam0 /= mu0

        lam1 = self.mesh.fluid_relative_permeability_1(s) 
        mu1 = self.mesh.meshdata['fluid_1']['viscosity'] # 单位是 1 cp = 1 mPa.s 
        lam1 /= mu1 

        val = lam0/(lam0 + lam1)
        return val


    def get_total_system(self):
        """
        Notes
        -----
        构造整个系统

        二维情形:
        x = [s, v, p, u0, u1]

        A = [[   S, None,   SP,  SU0,  SU1]
             [None,    V,   VP, None, None]
             [None,   PV,    P,  PU0,  PU1]
             [None, None,  UP0,  U00,  U01]
             [None, None,  UP1,  U10,  U11]]
        F = [FS, FV, FP, FU0, FU1]

        三维情形:

        x = [s, v, p, u0, u1, u2]
        A = [[   S, None,   SP,  SU0,  SU1,  SU2]
             [None,    V,   VP, None, None, None]
             [None,   PV,    P,  PU0,  PU1,  PU2]
             [None, None,  UP0,  U00,  U01,  U02]
             [None, None,  UP1,  U10,  U11,  U12]
             [None, None,  UP2,  U20,  U21,  U22]]
        F = [FS, FV, FP, FU0, FU1, FU2]

        """

        GD = self.GD
        A0, FS, isBdDof0 = self.get_saturation_system(q=2)
        A1, FV, isBdDof1 = self.get_velocity_system(q=2)
        A2, FP, isBdDof2 = self.get_pressure_system(q=2)

        if GD == 2:
            A3, A4, FU, isBdDof3 = self.get_dispacement_system(q=2)
            A = bmat([A0, A1, A2, A3, A4], format='csr')
            F = np.r_['0', FS, FV, FP, FU]
        elif GD == 3:
            A3, A4, A5, FU, isBdDof3 = self.get_dispacement_system(q=2)
            A = bmat([A0, A1, A2, A3, A4, A5], format='csr')
            F = np.r_['0', FS, FV, FP, FU]
        isBdDof = np.r_['0', isBdDof0, isBdDof1, isBdDof2, isBdDof3]
        return A, F, isBdDof

    def get_saturation_system(self, q=2):
        """
        Notes
        ----
        计算饱和度方程对应的离散系统

        [   S, None,   SP,  SU0,  SU1]

        [   S, None,   SP,  SU0,  SU1, SU2]
        """

        GD = self.GD
        qf = self.mesh.integrator(q, etype='cell')
        bcs, ws = qf.get_quadrature_points_and_weights()
        cellmeasure = self.mesh.entity_measure('cell')

        # SP 是对角矩阵 
        c = self.saturation_pressure_coefficient() # (NC, )
        c *= cellmeasure 
        SP = diags(c, 0)

        # S 质量矩阵组装, 对角矩阵
        c = self.cphi[:]*cellmeasure
        S = diags(c, 0)

        # SU0, SU1, 饱和度方程中的位移散度对应的矩阵
        b = self.mesh.celldata['biot']
        val = self.mesh.grad_lambda() # (NC, TD+1, GD)
        c = self.cs[:]*b # (NC, ), 注意用当前的水饱和度
        c *= cellmeasure 
        val *= c[:, None, None]

        pgdof = self.pspace.number_of_global_dofs() # 压力空间自由度个数
        cgdof = self.cspace.number_of_global_dofs() # 连续空间自由度个数

        pc2d = self.pspace.cell_to_dof()
        cc2d = self.cspace.cell_to_dof()

        I = np.broadcast_to(pc2d, shape=cc2d.shape)
        J = cc2d 

        SU0 = csr_matrix(
                (val[..., 0].flat, (I.flat, J.flat)),
                shape=(pgdof, cgdof)
                )
        SU1 = csr_matrix(
                (val[..., 1].flat, (I.flat, J.flat)),
                shape=(pgdof, cgdof)
                )

        if GD == 3:
            SU2 = csr_matrix(
                    (val[..., 2].flat, (I.flat, J.flat)),
                    shape=(pgdof, cgdof)
                    )



        # 右端矩阵
        dt = self.timeline.current_time_step_length()
        FS = self.f0.value(bcs) # (NQ, NC)
        FS *= ws[:, None]
        FS = np.sum(FS, axis=0)
        FS *= cellmeasure
        FS *= dt

        FS += [email protected] # 上一时间步的压强 
        FS += [email protected] # 上一时间步的饱和度
        FS += [email protected][:, 0] # 上一时间步的位移 x 分量
        FS += [email protected][:, 1] # 上一个时间步的位移 y 分量 
        if GD == 3:
            FS += [email protected][:, 2] # 上一个时间步的位移 z 分量 


        # 用当前时刻的速度场, 构造迎风格式
        face2cell = self.mesh.ds.face_to_cell()
        isBdFace = face2cell[:, 0] == face2cell[:, 1]

        qf = self.mesh.integrator(2, 'face') # 边或面上的积分公式
        bcs, ws = qf.get_quadrature_points_and_weights()
        facemeasure = self.mesh.entity_measure('face') 

        # 边的定向法线,它是左边单元的外法线,右边单元内法线。
        fn = self.mesh.face_unit_normal() 
        val0 = np.einsum('qfm, fm->qf', self.cv.face_value(bcs), fn) # 当前速度和法线的内积

        # 流体 0 的流动分数, 与水的饱和度有关, 如果饱和度为0, 则为 0
        F0 = self.fluid_fractional_flow_coefficient_0()
        val1 = F0[face2cell[:, 0]] # 边的左边单元的水流动分数
        val2 = F0[face2cell[:, 1]] # 边的右边单元的水流动分数 
        val2[isBdFace] = 0.0 # 边界的贡献是 0,没有流入和流出

        flag = val0 < 0.0 # 对于左边单元来说,是流入项
                           # 对于右边单元来说,是流出项

        # 左右单元流入流出的绝对量是一样的
        val = val0*val1[None, :] # val0 >= 0.0, 左边单元是流出
                                 # val0 < 0.0, 左边单元是流入
        val[flag] = (val0*val2[None, :])[flag] # val0 >= 0, 右边单元是流入
                                               # val0 < 0, 右边单元是流出

        b = np.einsum('q, qf, f->f', ws, val, facemeasure)
        b *= dt
        np.subtract.at(FS, face2cell[:, 0], b)  

        isInFace = ~isBdFace # 只处理内部边
        np.add.at(FS, face2cell[isInFace, 1], b[isInFace])  

        isBdDof = np.zeros(pgdof, dtype=np.bool_) 

        if GD == 2:
            return [   S, None,   SP,  SU0,  SU1], FS, isBdDof
        elif GD == 3:
            return [   S, None,   SP,  SU0,  SU1, SU2], FS, isBdDof

    def get_velocity_system(self, q=2):
        """
        Notes
        -----
        计算速度方程对应的离散系统.


         [None,    V,   VP, None, None]

         [None,    V,   VP, None, None, None]
        """

        GD = self.GD
        dt = self.timeline.current_time_step_length()
        cellmeasure = self.mesh.entity_measure('cell')
        qf = self.mesh.integrator(q, etype='cell')
        bcs, ws = qf.get_quadrature_points_and_weights()

        # 速度对应的矩阵  V
        c = self.flux_coefficient()
        c *= cellmeasure
        phi = self.vspace.basis(bcs)
        V = np.einsum('q, qcin, qcjn, c->cij', ws, phi, phi, c, optimize=True)

        c2d = self.vspace.cell_to_dof()
        I = np.broadcast_to(c2d[:, :, None], shape=V.shape)
        J = np.broadcast_to(c2d[:, None, :], shape=V.shape)

        gdof = self.vspace.number_of_global_dofs()
        V = csr_matrix(
                (V.flat, (I.flat, J.flat)), 
                shape=(gdof, gdof)
                )

        # 压强矩阵
        VP = -self.B

        # 右端向量, 0 向量
        FV = np.zeros(gdof, dtype=np.float64)
        isBdDof = self.vspace.dof.is_boundary_dof()

        if GD == 2:
            return [None, V, VP, None, None], FV, isBdDof 
        elif GD == 3:
            return [None, V, VP, None, None, None], FV, isBdDof 

    def get_pressure_system(self, q=2):
        """
        Notes
        -----
        计算压强方程对应的离散系统

        这里组装矩阵时, 利用了压强是分片常数的特殊性 

        [  Noe,  PV,   P, PU0,   PU1]

        [  None, PV,   P, PU0,   PU1, PU2]
        """

        GD = self.GD
        dt = self.timeline.current_time_step_length()
        cellmeasure = self.mesh.entity_measure('cell')
        qf = self.mesh.integrator(q, etype='cell')
        bcs, ws = qf.get_quadrature_points_and_weights()

        PV = dt*self.B.T

        # P 是对角矩阵, 利用分片常数的
        c = self.pressure_coefficient() # (NQ, NC)
        c *= cellmeasure 
        P = diags(c, 0)

        # 组装压力方程的右端向量
        # * 这里利用了压力空间基是分片常数
        FP = self.f1.value(bcs) + self.f0.value(bcs) # (NQ, NC)
        FP *= ws[:, None]

        FP = np.sum(FP, axis=0)
        FP *= cellmeasure
        FP *= dt

        FP += [email protected] # 上一步的压力向量
        FP += [email protected][:, 0] 
        FP += [email protected][:, 1]

        if GD == 3:
            FP += [email protected][:, 2]

        gdof = self.pspace.number_of_global_dofs()
        isBdDof = np.zeros(gdof, dtype=np.bool_)

        if GD == 2:
            return [None, PV, P, self.PU0, self.PU1], FP, isBdDof
        elif GD == 3:
            return [None, PV, P, self.PU0, self.PU1, self.PU2], FP, isBdDof


    def get_dispacement_system(self, q=2):
        """
        Notes
        -----
        计算位移方程对应的离散系统, 即线弹性方程系统.
        
        GD == 2:
         [None, None, UP0, U00,   U01]
         [None, None, UP1, U10,   U11]

        GD == 3:
         [None, None, UP0, U00, U01, U02]
         [None, None, UP1, U10, U11, U12]
         [None, None, UP2, U20, U21, U22]

        """


        GD = self.GD
        # 拉梅参数 (lambda, mu)
        lam = self.mesh.celldata['lambda']
        mu = self.mesh.celldata['mu']
        U = self.linear_elasticity_matrix(lam, mu, format='list')

        isBdDof = self.cspace.dof.is_boundary_dof()

        if GD == 2:
            return (
                    [None, None, -self.PU0.T, U[0][0], U[0][1]], 
                    [None, None, -self.PU1.T, U[1][0], U[1][1]], 
                    self.FU, np.r_['0', isBdDof, isBdDof]
                    )
        elif GD == 3:
            return (
                    [None, None, -self.PU0.T, U[0][0], U[0][1], U[0][2]], 
                    [None, None, -self.PU1.T, U[1][0], U[1][1], U[1][2]], 
                    [None, None, -self.PU2.T, U[2][0], U[2][1], U[2][2]], 
                    self.FU, np.r_['0', isBdDof, isBdDof, isBdDof]
                    )

    def linear_elasticity_matrix(self, lam, mu, format='csr', q=None):
        """
        Notes
        -----
        注意这里的拉梅常数是单元分片常数
        lam.shape == (NC, ) # MPa
        mu.shape == (NC, ) # MPa
        """

        GD = self.GD
        if GD == 2:
            idx = [(0, 0), (0, 1),  (1, 1)]
            imap = {(0, 0):0, (0, 1):1, (1, 1):2}
        elif GD == 3:
            idx = [(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)]
            imap = {(0, 0):0, (0, 1):1, (0, 2):2, (1, 1):3, (1, 2):4, (2, 2):5}
        A = []

        qf = self.cspace.integrator if q is None else self.mesh.integrator(q, 'cell')
        bcs, ws = qf.get_quadrature_points_and_weights()
        grad = self.cspace.grad_basis(bcs) # (NQ, NC, ldof, GD)


        # 分块组装矩阵
        gdof = self.cspace.number_of_global_dofs()
        cellmeasure = self.cspace.cellmeasure
        for k, (i, j) in enumerate(idx):
            Aij = np.einsum('i, ijm, ijn, j->jmn', ws, grad[..., i], grad[..., j], cellmeasure)
            A.append(Aij)

        if GD == 2:
            C = [[None, None], [None, None]]
            D = mu[:, None, None]*(A[0] + A[2]) 
        elif GD == 3:
            C = [[None, None, None], [None, None, None], [None, None, None]]
            D = mu[:, None, None]*(A[0] + A[3] + A[5])

        
        cell2dof = self.cspace.cell_to_dof() # (NC, ldof)
        ldof = self.cspace.number_of_local_dofs()
        NC = self.mesh.number_of_cells()
        shape = (NC, ldof, ldof)
        I = np.broadcast_to(cell2dof[:, :, None], shape=shape)
        J = np.broadcast_to(cell2dof[:, None, :], shape=shape)

        for i in range(GD):
            Aii = D + (mu + lam)[:, None, None]*A[imap[(i, i)]] 
            C[i][i] = csr_matrix((Aii.flat, (I.flat, J.flat)), shape=(gdof, gdof))
            for j in range(i+1, GD):
                Aij = lam[:, None, None]*A[imap[(i, j)]] + mu[:, None, None]*A[imap[(i, j)]].swapaxes(-1, -2)
                C[i][j] = csr_matrix((Aij.flat, (I.flat, J.flat)), shape=(gdof, gdof)) 
                C[j][i] = C[i][j].T 

        if format == 'csr':
            return bmat(C, format='csr') # format = bsr ??
        elif format == 'bsr':
            return bmat(C, format='bsr')
        elif format == 'list':
            return C

    def solve_linear_system_0(self):
        # 构建总系统

        if self.ctx.myid == 0:
            A, F, isBdDof = self.get_total_system()

            # 处理边界条件, 这里是 0 边界
            gdof = len(isBdDof)
            bdIdx = np.zeros(gdof, dtype=np.int_)
            bdIdx[isBdDof] = 1 
            Tbd = diags(bdIdx)
            T = diags(1-bdIdx)
            A = T@A@T + Tbd
            F[isBdDof] = 0.0

            # 求解
            self.ctx.set_centralized_sparse(A)
            x = F.copy()
            self.ctx.set_rhs(x) # Modified in place

        self.ctx.run(job=6) # Analysis + Factorization + Solve

        if self.ctx.myid == 0:
            vgdof = self.vspace.number_of_global_dofs()
            pgdof = self.pspace.number_of_global_dofs()
            cgdof = self.cspace.number_of_global_dofs()
            
            start = 0
            end = pgdof
            s = x[start:end]

            start = end
            end += vgdof
            v = x[start:end]

            start = end
            end += pgdof
            p = x[start:end]

            start = end
            u = x[start:]

            return s, v, p, u



    def solve_linear_system_1(self):
        """
        Notes
        -----
        构造整个系统

        二维情形:
        x = [s, v, p, u0, u1]

        A = [[   S, None,   SP,  SU0,  SU1]
             [None,    V,   VP, None, None]
             [None,   PV,    P,  PU0,  PU1]
             [None, None,  UP0,  U00,  U01]
             [None, None,  UP1,  U10,  U11]]
        F = [FS, FV, FP, FU0, FU1]

        三维情形:

        x = [s, v, p, u0, u1, u2]
        A = [[   S, None,   SP,  SU0,  SU1,  SU2]
             [None,    V,   VP, None, None, None]
             [None,   PV,    P,  PU0,  PU1,  PU2]
             [None, None,  UP0,  U00,  U01,  U02]
             [None, None,  UP1,  U10,  U11,  U12]
             [None, None,  UP2,  U20,  U21,  U22]]
        F = [FS, FV, FP, FU0, FU1, FU2]

        """
        if self.ctx.myid == 0:
            vgdof = self.vspace.number_of_global_dofs()
            pgdof = self.pspace.number_of_global_dofs()
            cgdof = self.cspace.number_of_global_dofs()

            GD = self.GD
            A0, FS, isBdDof0 = self.get_saturation_system(q=2)
            A1, FV, isBdDof1 = self.get_velocity_system(q=2)
            A2, FP, isBdDof2 = self.get_pressure_system(q=2)

            if GD == 2:
                A3, A4, FU, isBdDof3 = self.get_dispacement_system(q=2)
                A = bmat([A1[1:], A2[1:], A3[1:], A4[1:]], format='csr')
                F = np.r_['0', FV, FP, FU]
            elif GD == 3:
                A3, A4, A5, FU, isBdDof3 = self.get_dispacement_system(q=2)
                A = bmat([A1[1:], A2[1:], A3[1:], A4[1:], A5[1:]], format='csr')
                F = np.r_['0', FV, FP, FU]

            isBdDof = np.r_['0', isBdDof1, isBdDof2, isBdDof3]

            gdof = len(isBdDof)
            bdIdx = np.zeros(gdof, dtype=np.int_)
            bdIdx[isBdDof] = 1 
            Tbd = diags(bdIdx)
            T = diags(1-bdIdx)
            A = T@A@T + Tbd
            F[isBdDof] = 0.0

        #[   S, None,   SP,  SU0,  SU1, SU2]
            self.ctx.set_centralized_sparse(A)
            x = F.copy()
            self.ctx.set_rhs(x) # Modified in place

        self.ctx.run(job=6) # Analysis + Factorization + Solve

        if self.ctx.myid == 0:

            v = x[0:vgdof]
            p = x[vgdof:vgdof+pgdof]
            u = x[vgdof+pgdof:]

            s = FS
            s -= A0[2]@p 
            s -= A0[3]@u[0*cgdof:1*cgdof] 
            s -= A0[4]@u[1*cgdof:2*cgdof] 
            if GD == 3:
                s -= A0[5]@u[2*cgdof:3*cgdof]
            s /= A0[0].diagonal()

            return s, v, p, u
        else:
            return None, None, None, None


    def picard_iteration(self):
        """
        """

        e0 = 1.0
        k = 0
        while e0 > 1e-10: 

            s, v, p, u = self.solve_linear_system_1()

            if self.ctx.myid == 0:
                e0 = 0.0
                e0 += np.sum((self.cs - s)**2)
                self.cs[:] = s 

                e0 += np.sum((self.cp - p)**2)
                self.cp[:] = p 

                self.cv[:] = v 
                self.cu.T.flat[:] = u

                e0 = np.sqrt(e0) # 误差的 l2 norm
                print(e0)

            e0 = self.ctx.comm.bcast(e0, root=0)
            k += 1
            if k >= self.args.npicard: 
                if self.ctx.myid == 0:
                    print('picard iteration arrive max iteration with error:', e0)
                break

        if self.ctx.myid == 0:
            # 更新解
            self.s[:] = self.cs
            self.v[:] = self.cv
            self.p[:] = self.cp
            self.u[:] = self.cu

    def update_mesh_data(self):
        """

        Notes
        -----
        更新 mesh 中的数据
        """
        GD = self.GD
        bc = np.array((GD+1)*[1/(GD+1)], dtype=np.float64)

        mesh = self.mesh

        v = self.v 
        p = self.p
        s = self.s
        u = self.u


        # 单元中心的流体速度
        val = v.value(bc)
        if GD == 2:
            val = np.concatenate((val, np.zeros((val.shape[0], 1), dtype=val.dtype)), axis=1)
        mesh.celldata['velocity'] = val 

        # 分片常数的压强
        val = self.recover(p[:])
        mesh.nodedata['pressure'] = val

        # 分片常数的饱和度
        val = self.recover(s[:])
        mesh.nodedata['fluid_0'] = val

        val = self.recover(1 - s)
        mesh.nodedata['fluid_1'] = val 

        # 节点处的位移
        if GD == 2:
            u = np.concatenate((u[:], np.zeros((u.shape[0], 1), dtype=u.dtype)), axis=1)
        mesh.nodedata['displacement'] = u[:] 

        # 增加应变和应力的计算

        NC = self.mesh.number_of_cells()
        s0 = np.zeros((NC, 3, 3), dtype=np.float64)
        s1 = np.zeros((NC, 3, 3), dtype=np.float64)

        s = u.grad_value(bc) # (NC, GD, GD)
        s0[:, 0:GD, 0:GD] = s + s.swapaxes(-1, -2)
        s0[:, 0:GD, 0:GD] /= 2

        lam = self.mesh.celldata['lambda']
        mu = self.mesh.celldata['mu']

        s1[:, 0:GD, 0:GD] = s0[:, 0:GD, 0:GD]
        s1[:, 0:GD, 0:GD] *= mu[:, None, None]
        s1[:, 0:GD, 0:GD] *= 2 

        s1[:, range(GD), range(GD)] += (lam*s0.trace(axis1=-1, axis2=-2))[:, None]

        s1[:, range(GD), range(GD)] += mesh.celldata['stress_0'][:, None]

        s1[:, range(GD), range(GD)] -= (mesh.celldata['biot']*(p - mesh.celldata['pressure_0']))[:, None]

        mesh.celldata['strain'] = s0.reshape(NC, -1)
        mesh.celldata['stress'] = s1.reshape(NC, -1)



    def run(self, writer=None):
        """

        Notes
        -----

        计算所有时间层物理量。
        """

        args = self.args

        timeline = self.timeline
        dt = timeline.current_time_step_length()

        if (self.ctx.myid == 0) and (writer is not None):
            n = timeline.current
            fname = args.output + str(n).zfill(10) + '.vtu'
            self.update_mesh_data()
            writer(fname, self.mesh)

        while not timeline.stop():
            if self.ctx.myid == 0:
                ct = timeline.current_time_level()/3600/24 # 天为单位
                print('当前时刻为第', ct, '天')

            self.picard_iteration()
            timeline.current += 1
            if timeline.current%args.step == 0:
                if (self.ctx.myid == 0) and (writer is not None):
                    n = timeline.current
                    fname = args.output + str(n).zfill(10) + '.vtu'
                    self.update_mesh_data()
                    writer(fname, self.mesh)

        if (self.ctx.myid == 0) and (writer is not None):
            n = timeline.current
            fname = args.output + str(n).zfill(10) + '.vtu'
            self.update_mesh_data()
            writer(fname, self.mesh)
Ejemplo n.º 4
0
from fealpy.mesh import MeshFactory as MF
from fealpy.mesh import TriangleMesh
from fealpy.mesh.core import multi_index_matrix2d
from fealpy.functionspace import LagrangeFiniteElementSpace

mesh = MF.one_triangle_mesh(meshtype='equ')

space = LagrangeFiniteElementSpace(mesh, p=5)

n = 50
bcs = multi_index_matrix2d(n) / n  # ldof = (n+1)(n+2)/2
ps = mesh.bc_to_point(bcs).reshape(-1, 2)

val = space.basis(bcs)  # (NQ, 1, ldof)

val = space.grad_basis(bcs)  # (NQ, NC, ldof, GD)

val = val[:, 0, 0]

fig = plt.figure()
axes = fig.gca()
mesh.add_plot(axes)
mesh.find_node(axes, node=ps, markersize=20)

fig = plt.figure()
axes = fig.add_subplot(projection='3d')
axes.plot_trisurf(ps[:, 0],
                  ps[:, 1],
                  val,
                  linewidth=0.2,
                  antialiased=True,
Ejemplo n.º 5
0
from fealpy.tools.show import showmultirate, show_error_table
n = 10
p = 1
mf = MeshFactory()
mesh = mf.boxmesh2d([0,1,0,1], nx=n, ny=n,meshtype='tri')
pde = PDE()

space = LagrangeFiniteElementSpace(mesh, p=p)
gdof = space.number_of_global_dofs()
cell2dof = space.cell_to_dof()


cellmeasure = mesh.entity_measure('cell')#计算每个单元的面积
qf = mesh.integrator(1, 'cell')#积分点
bcs, ws = qf.get_quadrature_points_and_weights()
gphi = space.grad_basis(bcs)#(1,2,3,2)一个积分点,2个单元,3个基函数,每个基函数梯度两个分量

a = np.array([(10.0,-1.0),(-1.0,2.0)],dtype=np.float64)
dgphi = np.einsum('mn,ijkn->ijkm',a,gphi)#(2,3,2)
A = np.einsum('i, ijkl, ijml, j->jkm', ws,dgphi, gphi, cellmeasure)#(2,3,3)
I = np.broadcast_to(cell2dof[:,:,None], shape=A.shape)
J = np.broadcast_to(cell2dof[:,None,:],shape=A.shape)
A = csr_matrix((A.flat,(I.flat,J.flat)),shape = (gdof,gdof))#(gdof,gdof)刚度矩阵
# A = space.stiff_matrix(c=a)
# print(A.toarray())
#  #质量矩阵组装
# phi = space.basis(bcs)#(1,1,3)一个积分点,每个单元上的重心坐标
@cartesian
def r(p):
    x = p[...,0]
    y = p[...,1]
Ejemplo n.º 6
0
NC = mesh.number_of_cells()
print("网格中节点、边和单元的个数分别为:", NN, NE, NC)

print('创建拉格朗日有限元空间...')
space = LagrangeFiniteElementSpace(mesh, p=p)
ldof = space.number_of_local_dofs()
gdof = space.number_of_global_dofs()

print('拉格朗日空间的次数为:', p)
print('每个单元上的局部自由度个数:', ldof)
print('每个单元上的全局自由度个数', gdof)

print('计算空间基函数在每个单元重心坐标点处的值...')
bc = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)  # (3, 3)
ps = mesh.bc_to_point(bc)  # (NQ, NC, 2)

phi = space.basis(bc)
gphi = space.grad_basis(bc)

print('重心坐标数组:', bc)
print('bc.shape:', bc.shape)
print('phi.shape:', phi.shape)
print('gphi.shape:', gphi.shape)

fig = plt.figure()
axes = fig.gca()
mesh.add_plot(axes)
mesh.find_node(axes, showindex=True)
mesh.find_cell(axes, showindex=True)
plt.show()