Mentions légales du service

Skip to content
Snippets Groups Projects
genfem.py 39.68 KiB
#!/usr/bin/env python
# Copyright 2016 INRIA
#
# louis.poirel@inria.fr
#
# This file is part of the genfem python package.
#
# This software is governed by the CeCILL-C license under French law and
# abiding by the rules of distribution of free software.  You can  use,
# modify and/ or redistribute the software under the terms of the CeCILL-C
# license as circulated by CEA, CNRS and INRIA at the following URL
# "http://www.cecill.info".
#
# As a counterpart to the access to the source code and  rights to copy,
# modify and redistribute granted by the license, users are provided only
# with a limited warranty  and the software's author,  the holder of the
# economic rights,  and the successive licensors  have only  limited
# liability.
#
# In this respect, the user's attention is drawn to the risks associated
# with loading,  using,  modifying and/or developing or reproducing the
# software by the user in light of its specific status of free software,
# that may mean  that it is complicated to manipulate,  and  that  also
# therefore means  that it is reserved for developers  and  experienced
# professionals having in-depth computer knowledge. Users are therefore
# encouraged to load and test the software's suitability as regards their
# requirements in conditions enabling the security of their systems and/or
# data to be ensured and,  more generally, to use and operate it in the
# same conditions as regards security.
#
# The fact that you are presently reading this means that you have had
# knowledge of the CeCILL-C license and that you accept its terms.

"""This module provides a fast and easy way to generate finite
element matrices. It is possible to generate a mesh and an element
matrix and assemble the matrix. """

from __future__ import print_function
import numpy as np
import scipy.sparse as ssp
import sympy
import itertools
import functools


# First, we create a mesh
def nodes(shape, **kwargs):
    """Returns an array of global node indexes corresponding to a domain
    or a subdomain.

    Parameters
    ----------
    shape : array or list
        [nX, nY, ...] is a list or array containing the number of
        nodes in each direction of the whole domain

    lim : list
        list of [imin, imax] arrays, delimiting the local
        subdomain. The selected nodes will be such that their index in
        dimension k is between imin and imax INCLUDED.

    wrap_lim : boolean
        if True, negative indexes are counted from the other side (-1
        is the last index). If False, when imin>imax, returns a void
        array

    Returns
    -------
    node : array of node indexes

    Examples
    --------
    >>> nodes((3,4))
    array([[ 0,  3,  6,  9],
           [ 1,  4,  7, 10],
           [ 2,  5,  8, 11]])
    >>> nodes((3,4), lim=[[0,1],[1,3]])
    array([[ 3,  6,  9],
           [ 4,  7, 10]])
    >>> nodes((3,4), lim=[[0,-1], [0,-2]])
    array([[0, 3, 6],
           [1, 4, 7],
           [2, 5, 8]])
    >>> nodes((3,4), lim=[[0,-1], [-2,-1]])
    array([[ 6,  9],
           [ 7, 10],
           [ 8, 11]])
    >>> nodes(3)
    array([0, 1, 2])
    >>> nodes(3, lim=[[0,-1]], wrap_lim=False)
    array([], dtype=int64)

    """
    # cleaning up the input
    shape = np.array(shape, ndmin=1).tolist()
    lim = kwargs.setdefault("lim", [[0, n-1] for n in shape])
    lim = np.array(lim)
    assert len(lim) == len(shape)
    for k, (imin, imax) in enumerate(lim):
        if kwargs.setdefault("wrap_lim", True):
            if imax < 0:
                lim[k, 1] += shape[k]
            if imin < 0:
                lim[k, 0] += shape[k]
        else:
            if imin > imax:
                return np.array([], dtype=np.integer, ndmin=len(shape))
    # computing the indices
    node = np.zeros([imax+1-imin for imin, imax in lim], dtype=np.integer)
    grid = np.ix_(*[range(imin, imax+1) for imin, imax in lim])
    node += grid[0]
    for g, l in zip(grid[1:], np.cumprod(shape)):
        node += g*l
    return node


def mesh_rect(shape, **kwargs):

    """Returns the elements delimited by lim in the domain of size
    shape. Each column has the indices of all nodes of one element.

    Parameters
    ----------
    shape : array or list
        [nX, nY, ...] is a list or array containing the number of
        nodes in each direction of the whole domain

    Additional keyword arguments are passed to nodes.

    Returns
    -------
    tElem : numpy array
        tElem[i,k] is the index of the i-th node of the k-th element.

    Examples
    --------
    >>> mesh_rect((3,4))
    array([[ 0,  3,  6,  1,  4,  7],
           [ 3,  6,  9,  4,  7, 10],
           [ 1,  4,  7,  2,  5,  8],
           [ 4,  7, 10,  5,  8, 11]])
    >>> mesh_rect((3,4), lim=[[0,1],[1,3]])
    array([[ 3,  6],
           [ 6,  9],
           [ 4,  7],
           [ 7, 10]])
    >>> mesh_rect((3,4), lim=[[0,-1],[0,-2]])
    array([[0, 3, 1, 4],
           [3, 6, 4, 7],
           [1, 4, 2, 5],
           [4, 7, 5, 8]])
    >>> mesh_rect((3,4), lim=[[0,-1],[-2,-1]])
    array([[ 6,  7],
           [ 9, 10],
           [ 7,  8],
           [10, 11]])
    >>> mesh_rect(3)
    array([[0, 1],
           [1, 2]])
    """
    shape = np.array(shape, ndmin=1).tolist()
    # in each direction, nElem = nNode-1
    node = nodes(shape, **kwargs)[[np.s_[:-1]]*len(shape)].ravel()
    # the nodes of element e are N(e) = e+N(0)
    neighbors = nodes(shape, lim=[[0, 1]]*len(shape)).ravel()
    return np.vstack(node+v for v in neighbors)


# Then, we compute the element matrix
def elem_rect_(dim=2, dtype=np.float64):

    """Returns the element stiffness and mass matrices in sparse format
    for the heat equation in dimension d, computed using sympy

    Parameters
    ----------
    dim : int
        dimension of the problem, optional, default: 2

    dtype : type
        type of the data in the returned matrix, optional, default:
        numpy.float64

    Returns
    -------
    Ae : numpy array
        Element stiffness matrix

    Me : numpy array
        Element mass matrix

    Examples:
    ---------
    >>> elem_rect(1)
    (array([[ 1., -1.],
           [-1.,  1.]]), array([[ 0.33333333,  0.16666667],
           [ 0.16666667,  0.33333333]]))
    >>> er2 = elem_rect(2)
    >>> er2[0]
    array([[ 0.66666667, -0.16666667, -0.16666667, -0.33333333],
           [-0.16666667,  0.66666667, -0.33333333, -0.16666667],
           [-0.16666667, -0.33333333,  0.66666667, -0.16666667],
           [-0.33333333, -0.16666667, -0.16666667,  0.66666667]])
    >>> er2[1]
    array([[ 0.11111111,  0.05555556,  0.05555556,  0.02777778],
           [ 0.05555556,  0.11111111,  0.02777778,  0.05555556],
           [ 0.05555556,  0.02777778,  0.11111111,  0.05555556],
           [ 0.02777778,  0.05555556,  0.05555556,  0.11111111]])
    >>> for i in range(1, 4):
    ...     er = elem_rect(i)
    ...     er_ = elem_rect_(i)
    ...     np.any(np.round(er[0] - er_[0], 15))
    ...     np.any(np.round(er[1] - er_[1], 15))
    ... # doctest: +NORMALIZE_WHITESPACE
    False
    False
    False
    False
    False
    False
    """
    # symbolic parameters x, y
    X = sympy.symbols('x:'+str(dim))
    # Polynom base of the element  1, x, y, x*y, ...
    base = [np.multiply.reduce(e)
            for e in itertools.product(*[[sympy.S.One, x] for x in X])]
    # coordinates of the nodes [0, 0], [0, 1], [1, 0], [1, 1]
    elem = [e for e in itertools.product(*[[0, 1] for x in X])]
    # value of the polynoms on the nodes
    V = np.array([[P.subs(zip(X, node)) for node in elem]
                  for P in base], dtype=dtype)
    # Lagrange interpolation polynomials = basis functions
    U = np.linalg.inv(V).dot(base)
    gradU = [[sympy.diff(u, x) for x in X] for u in U]
    Ae = np.array([[functools.reduce(lambda f, x: sympy.integrate(f,
                                                                  (x, 0, 1)),
                                     X, np.dot(du, dv))
                    for du in gradU] for dv in gradU], dtype=dtype)
    Me = np.array([[functools.reduce(lambda f, x: sympy.integrate(f,
                                                                  (x, 0, 1)),
                                     X, u*v)
                    for u in U] for v in U], dtype=dtype)
    return Ae, Me


def elem_rect(dim=2, dtype=np.float64):

    """Returns the element stiffness and mass matrices in sparse format
    for the heat equation in dimension d, using cache result if dim < 4

    Parameters
    ----------
    dim : int
        dimension of the problem, optional, default: 2

    dtype : type
        type of the data in the returned matrix, optional, default:
        numpy.float64

    Returns
    -------
    Ae : numpy array
        Element stiffness matrix

    Me : numpy array
        Element mass matrix

    Examples:
    ---------
    >>> elem_rect(1)
    (array([[ 1., -1.],
           [-1.,  1.]]), array([[ 0.33333333,  0.16666667],
           [ 0.16666667,  0.33333333]]))
    >>> er2 = elem_rect(2)
    >>> er2[0]
    array([[ 0.66666667, -0.16666667, -0.16666667, -0.33333333],
           [-0.16666667,  0.66666667, -0.33333333, -0.16666667],
           [-0.16666667, -0.33333333,  0.66666667, -0.16666667],
           [-0.33333333, -0.16666667, -0.16666667,  0.66666667]])
    >>> er2[1]
    array([[ 0.11111111,  0.05555556,  0.05555556,  0.02777778],
           [ 0.05555556,  0.11111111,  0.02777778,  0.05555556],
           [ 0.05555556,  0.02777778,  0.11111111,  0.05555556],
           [ 0.02777778,  0.05555556,  0.05555556,  0.11111111]])
    >>> for i in range(1, 4):
    ...     er = elem_rect(i)
    ...     er_ = elem_rect_(i)
    ...     np.any(np.round(er[0] - er_[0], 15))
    ...     np.any(np.round(er[1] - er_[1], 15))
    ... # doctest: +NORMALIZE_WHITESPACE
    False
    False
    False
    False
    False
    False
    """
    if dim == 1:
        return (np.array([[1., -1.],
                          [-1., 1.]]),
                np.array([[2.,  1.],
                          [1.,  2.]])/6.)
    elif dim == 2:
        return (np.array([[4., -1., -1., -2.],
                          [-1.,  4., -2., -1.],
                          [-1., -2.,  4., -1.],
                          [-2., -1., -1.,  4.]])/6.,
                np.array([[4.,  2.,  2.,  1.],
                          [2.,  4.,  1.,  2.],
                          [2.,  1.,  4.,  2.],
                          [1.,  2.,  2.,  4.]])/36.)
    elif dim == 3:
        return (np.array([[4., -0., -0., -1., -0., -1., -1., -1.],
                          [-0.,  4., -1.,  0., -1.,  0., -1., -1.],
                          [-0., -1.,  4.,  0., -1., -1.,  0., -1.],
                          [-1.,  0.,  0.,  4., -1., -1., -1.,  0.],
                          [-0., -1., -1., -1.,  4., -0., -0., -1.],
                          [-1.,  0., -1., -1., -0.,  4., -1.,  0.],
                          [-1., -1.,  0., -1., -0., -1.,  4.,  0.],
                          [-1., -1., -1.,  0., -1.,  0.,  0.,  4.]])/12.,
                np.array([[8.,  4.,  4.,  2.,  4.,  2.,  2.,  1.],
                          [4.,  8.,  2.,  4.,  2.,  4.,  1.,  2.],
                          [4.,  2.,  8.,  4.,  2.,  1.,  4.,  2.],
                          [2.,  4.,  4.,  8.,  1.,  2.,  2.,  4.],
                          [4.,  2.,  2.,  1.,  8.,  4.,  4.,  2.],
                          [2.,  4.,  1.,  2.,  4.,  8.,  2.,  4.],
                          [2.,  1.,  4.,  2.,  4.,  2.,  8.,  4.],
                          [1.,  2.,  2.,  4.,  2.,  4.,  4.,  8.]])/216.)
    else:
        return elem_rect_(dim, dtype)


# Then, we assemble the matrix and RHS
def assemble_matrix(tElem, Ae, Me, dirichlet_boundary_indices=[], dirichlet_boundary_value=0, source=1., K=1, 
                    ndof=0,rhs_sparse=False, symmetry=False, dirac_source_index=-1, dirac_source_value=None):
    """Assemble the problem given by Ae, Me on the mesh tElem, with a
    Dirichlet bc. on boundary, Neumann elsewhere and a dirac source at the center of the domain.

    Parameters
    ----------
    tElem : numpy array
        tElem[i,k] is the index of the i-th node of the k-th element.
        Nodes and elements indexes vary in the first dimension (x)
        first.

    Ae : numpy array
        Element stiffness matrix

    Me : numpy array
        Element mass matrix

    dirichlet_boundary_indices : list
        indices of the Dirichlet bc nodes

    dirichlet_boundary_value : number
        value of the solution on boundary, optional, default 0

    dirac_source_index : number
        index of the dirac source node (located at the center)

    dirac_source_value : number
        value of the dirac source that will be imposed at 
        the center of the domain, optional, default None

    source : numpy array
        source term of the equation, optional, default: 1

    K : numpy array or number
        if a number, the constant conductivity, optional, default: 1
        if an array, K[k] is the conductivity coefficient of element
        telem[:,k], optional, default: 1

    ndof : integer
        number of nodes (order of the matrix), optional, default:
        computed from tElem

    rhs_sparse : boolean
        if True, b is returned as a sparse matrix. Else, it is a numpy
        array.

    symmetry : boolean
        if True, only the upper triangular part of A is returned

    Returns
    -------
    A : scipy sparse matrix
        A is the assembled fem matrix

    b : numpy array or sparse matrix
        b is the assembled rhs, see rhs_sparse

    Examples
    --------
    >>> shape = 3,
    >>> source = -2.
    >>> tElem = mesh_rect(shape)
    >>> boundary = nodes(shape, lim=[[0,0]]+[[0,-1]]*(len(shape)-1)).ravel()
    >>> value = (shape[0]-1)**2*np.ones_like(boundary)
    >>> Ae, Me = elem_rect(len(shape))
    >>> A, b = assemble_matrix(tElem, Ae, Me, boundary, value, source)
    >>> A.toarray().round(2)
    array([[ 1.,  0.,  0.],
           [ 0.,  2., -1.],
           [ 0., -1.,  1.]])
    >>> b
    array([ 4.,  2., -1.])
    >>> A_upper, b = assemble_matrix(tElem, Ae, Me, boundary, value, source,
    ...                              symmetry=True)
    >>> A_upper.toarray().round(2)
    array([[ 1.,  0.,  0.],
           [ 0.,  2., -1.],
           [ 0.,  0.,  1.]])
    >>> b
    array([ 4.,  2., -1.])
    >>> import scipy.sparse.linalg as sla
    >>> sla.spsolve(A,b).round(2)[nodes(shape)]
    array([ 4.,  1.,  0.])
    >>> A, b = assemble_matrix(tElem, Ae, Me, K=[1,10])
    >>> A.toarray().round(2)
    array([[  1.,  -1.,   0.],
           [ -1.,  11., -10.],
           [  0., -10.,  10.]])
    """
    if not ndof:
        ndof = tElem.max() + 1
    all_conditions = []
    # Check if the current subdomain has one or more Dirichlet boundary nodes
    if len(dirichlet_boundary_indices) > 0:
        all_conditions = np.append(all_conditions,dirichlet_boundary_indices)
    # Check if the current subdomain has the node where the dirac source is imposed
    if dirac_source_index >= 0:
        all_conditions = np.append(all_conditions,dirac_source_index)
    elem_boundary = np.in1d(tElem, all_conditions).reshape(tElem.shape)    
    Ae = ssp.coo_matrix(Ae)
    Ae.data[np.abs(Ae.data) < np.finfo(float).eps] = 0
    Ae.eliminate_zeros()
    K = np.asarray(K)
    if not K.shape:
        K = K*np.ones(tElem.shape[1])
    # Check if the current subdomain has one or more Dirichlet boundary nodes
    if len(dirichlet_boundary_indices) > 0:
        dirichlet_boundary_value = np.asarray(dirichlet_boundary_value)
        if not dirichlet_boundary_value.shape:
            dirichlet_boundary_value = dirichlet_boundary_value*np.ones_like(dirichlet_boundary_indices)
    if tElem.shape[1]:
        II = []
        JJ = []
        VV = []
        II_Dir = []
        JJ_Dir = []
        VV_Dir = []
        for i, j, v in zip(Ae.row, Ae.col, Ae.data):
            is_not_boundary = np.all(elem_boundary[[i, j]] == 0, axis=0)
            if i <= j or not symmetry:
                II.extend(tElem[i, is_not_boundary])
                JJ.extend(tElem[j, is_not_boundary])
                VV.extend(K[is_not_boundary]*v)
            is_boundary = (is_not_boundary == 0)
            II_Dir.extend(tElem[i, is_boundary])
            JJ_Dir.extend(tElem[j, is_boundary])
            VV_Dir.extend(K[is_boundary]*v)
        A = ssp.csc_matrix((VV, (II, JJ)), shape=(ndof, ndof))
        A_Dir = ssp.csc_matrix((VV_Dir, (II_Dir, JJ_Dir)), shape=(ndof, ndof))
    A = A + ssp.csc_matrix(([1]*len(all_conditions), (all_conditions, all_conditions)),
                           shape=(ndof, ndof))
    # Computation of b
    if type(source) != np.ndarray:
        source *= np.ones(ndof)
    b = np.zeros_like(source,dtype=Ae)
    Me = ssp.coo_matrix(Me)
    Me.data[np.abs(Me.data) < np.finfo(float).eps] = 0
    Me.eliminate_zeros()
    if tElem.shape[1]:
        for i, j, v in zip(Me.row, Me.col, Me.data):
            b[tElem[i]] += v*source[tElem[j]]
    # We update b to take into account the Dirichlet bc
    # Check if the current subdomain has one or more Dirichlet boundary nodes
    if len(dirichlet_boundary_indices) > 0:
        b -= A_Dir[:,dirichlet_boundary_indices].dot(dirichlet_boundary_value)
        b[dirichlet_boundary_indices] = dirichlet_boundary_value
    # Check if the current subdomain has the node where the dirac source is imposed
    if dirac_source_index >= 0:
        b -= A_Dir[:,dirac_source_index].dot([dirac_source_value])
        b[dirac_source_index] = dirac_source_value
    if rhs_sparse:
        b = ssp.coo_matrix(b.reshape(-1, 1))
    return A, b

# Then, we assemble the surface mass matrix at the interface between subdomains
def assemble_surface_mass_matrix_at_interfaces(shape,limits,glob2loc,interface,A_d):
    """ Assemble the surface mass matrix at the interfaces between subdomains

    Parameters
    ----------
    shape : array or list
        [nX, nY, ...] is a list or array containing the number of
        nodes in each direction of the whole domain
        first.
    limits : list
        list of [imin, imax] arrays, delimiting the local
        subdomain. The selected nodes will be such that their index in
        dimension k is between imin and imax INCLUDED.
    glob2loc : function
        glob2loc(loc2glob[I]) = loc2glob[glob2loc(I)].
    interface : array
        list of global indexes of the interface nodes
    A_d: DistMatrix
        the distributed assembled fem matrix

    Returns
    -------
    M : scipy sparse matrix defined considering only the interface nodes
        M is the assembled surface mass matrix at the inteface between subdomains

    """
    Ae2, Me2 = elem_rect(len(shape)-1)
    
    faceSize = len(shape)
    if faceSize == 3:
        faceSize = 4
        
    # global and local element nodes 
    elements_glo = mesh_rect(shape,lim=limits)
    elements_loc = glob2loc(elements_glo)
    
    nOfDomains = len(A_d.dd.neighbors)
    II = []
    JJ = []
    VV = []                    
    
    # loop over neigbor subdoamins
    for iDomain in range(nOfDomains):
        # nodes in common with a subdomain
        indices = list(A_d.dd.neighbors.items())[iDomain][1]
        mask = np.in1d(elements_loc,indices).reshape(elements_loc.shape)
        maxim = max(sum(mask))
        if maxim == faceSize:
            # list of faces in common
            faces = elements_loc[mask].reshape(faceSize,-1)
            for iFace in range(faces.shape[1]):
                # rows and columns of the face mass matrix
                rows = np.repeat(faces[:,iFace],faceSize).reshape(faceSize,-1)
                cols = np.repeat(faces[:,iFace],faceSize).reshape(faceSize,-1).transpose()
                II.extend(rows.reshape(-1))
                JJ.extend(cols.reshape(-1))
                VV.extend( Me2.reshape(-1))
    # Mass matrix defined considering all the subdomain nodes
    M  = ssp.csc_matrix((VV, (II, JJ)), shape=A_d.local.shape)

    # Mass matrix defined considering only the interface nodes 
    n  = A_d.local.shape[0]
    nG = len(interface)
    RG = ssp.csc_matrix((np.ones_like(glob2loc(interface)),
                        (range(nG), glob2loc(interface))),
                        shape=(nG, n))
    M = RG.dot(M.dot(RG.T))

    return M



# We can build the matrix on a subdomain
def subdomain(shape, nDom, rank):
    """Returns the connectivity information of a subdomain

    Parameters
    ----------
    shape : array or list
        [nX, nY, ...] is a list or array containing the number of nodes in each
        direction of the whole domain

    nDom : array or list
        [nX, nY, ...] is a list or array containing the number of subdomains in
        each direction of the whole domain

    rank : integer
        rank of the domain/MPI process

    Returns
    -------

    lim : list
        list of [imin, imax] arrays, delimiting the local
        subdomain. The selected nodes will be such that their index in
        dimension k is between imin and imax INCLUDED.

    neighbors : list of integers
        Ranks of the neighbors

    interfaces : list of arrays
        interfaces[i] is the array of indexes in interface that are
        shared with neighbors[i]


    Examples
    --------
    >>> shape = [5,5]
    >>> nDom = [3,2]
    >>> for rank in range(np.multiply.reduce(nDom)):
    ...     (rank,) + subdomain(shape, nDom, rank)
    ... # doctest: +NORMALIZE_WHITESPACE
    (0, array([[0, 1],
               [0, 2]]),
        [3, 1, 4],
        [array([10, 11]),
         array([ 1,  6, 11]),
         array([11])])
    (1, array([[1, 2],
              [0, 2]]),
        [0, 3, 4, 2, 5],
         [array([ 1,  6, 11]),
          array([11]),
          array([11, 12]),
          array([ 2,  7, 12]),
          array([12])])
    (2, array([[2, 4],
              [0, 2]]),
        [1, 4, 5],
        [array([ 2,  7, 12]),
         array([12]),
         array([12, 13, 14])])
    (3, array([[0, 1],
              [2, 4]]),
        [0, 1, 4],
        [array([10, 11]),
         array([11]),
         array([11, 16, 21])])
    (4, array([[1, 2],
              [2, 4]]),
        [0, 3, 1, 2, 5],
        [array([11]),
         array([11, 16, 21]),
         array([11, 12]),
         array([12]),
         array([12, 17, 22])])
    (5, array([[2, 4],
              [2, 4]]),
        [1, 4, 2],
        [array([12]),
         array([12, 17, 22]),
         array([12, 13, 14])])
    >>> shape = 5
    >>> nDom = 3
    >>> for rank in range(np.multiply.reduce(nDom)):
    ...     (rank,) + subdomain(shape, nDom, rank)
    ...
    (0, array([[0, 1]]), [1], [array([1])])
    (1, array([[1, 2]]), [0, 2], [array([1]), array([2])])
    (2, array([[2, 4]]), [1], [array([2])])

    """
    nDom = np.array(nDom, ndmin=1)
    shape = np.array(shape, ndmin=1)
    domain_index = np.array(np.unravel_index(rank, nDom, order='F'))
    limits = np.vstack([(domain_index*(shape-1)//nDom),
                        (domain_index+1)*(shape-1)//nDom]).T
    limits[limits[:, 0] <= 0, 0] = 0
    limits[limits[:, 1] >= shape, 1] -= 1
    neighbors = []
    interfaces = []
    for it in itertools.product(*[range(i-1, i+2) for i in domain_index]):
        it = np.array(it)
        if all(it >= 0) and all(it < nDom) and not all(it == domain_index):
            neighbors.append(np.ravel_multi_index(it, nDom, order='F'))
            lim_ = np.vstack([limits[range(len(shape)),
                                     np.where(it <= domain_index, 0, 1)],
                              limits[range(len(shape)),
                                     np.where(it < domain_index, 0, 1)]]).T
            interfaces.append(np.sort(nodes(shape, lim=lim_).ravel()))
    return limits, neighbors, interfaces


def local_ordering(shape, limits, neighbors, interfaces):
    """Computes a local ordering for the subdomain

    Parameters
    ----------
    shape : array or list
        [nX, nY, ...] is a list or array containing the number of nodes in each
        direction of the whole domain

    lim : list
        list of [imin, imax] arrays, delimiting the local
        subdomain. The selected nodes will be such that their index in
        dimension k is between imin and imax INCLUDED.

    neighbors : list of integers
        Ranks of the neighbors

    interfaces : list of arrays
        interfaces[i] is the array of global indexes in interface that
        are shared with neighbors[i]

    Returns
    -------

    loc2glob : list of integers
        Mapping from local to global ordering (if I is a set of
        indices in local ordering, loc2glob[I] is in global ordering)

    glob2loc : function
        glob2loc(loc2glob[I]) = loc2glob[glob2loc(I)]

    interfaces : list of arrays
        interfaces[i] is the new array of local indexes in interface
        that are shared with neighbors[i]

    interface : array
        list of global indexes of the interface nodes

    Examples
    --------
    >>> shape = [5,5]
    >>> nDom = [3,2]
    >>> rank = 2
    >>> limits, neighbors, interfaces = subdomain(shape, nDom, rank)
    >>> loc2glob, glob2loc, interfaces, interface = local_ordering(shape,
    ...                                                            limits,
    ...                                                            neighbors,
    ...                                                            interfaces)
    >>> print(loc2glob)
    [ 3  4  8  9  2  7 12 13 14]
    >>> print(glob2loc(loc2glob))
    [0 1 2 3 4 5 6 7 8]
    >>> print(interfaces)
    [[0, 1, 2], [2], [2, 3, 4]]
    >>> print(interface)
    [ 2  7 12 13 14]
    """
    dom = nodes(shape, lim=limits).ravel()
    interface = np.sort(np.unique(np.concatenate(interfaces)))
    interior = np.setdiff1d(dom, interface)
    loc2glob = np.hstack((interior, interface))
    ind_sort = np.argsort(loc2glob)

    def glob2loc(I):
        return ind_sort[np.searchsorted(loc2glob, I, sorter=ind_sort)]
    interfaces = [[i for i in range(len(interface))
                   if interface[i] in interfaces[j]]
                  for j in range(len(interfaces))]
    return loc2glob, glob2loc, interfaces, interface


# We provide an example
def stratified_helmholtz(shape, K1=1, K2=1, nLineK=3, rhs_sparse=False,
                         symmetry=False, loc2glob=None, glob2loc=None,
                         dirichlet_boundary_value=0, dirac_source_value=None,
                         source=1, w=1.0, **kwargs):
    """Create the fem matrix of the Helmholtz equation in a stratified medium
    Parameters
    ----------
    shape : array or list
        [nX, nY, ...] is a list or array containing the number of nodes in each
        direction of the whole domain

    K1 : number
        conductivity of the first layer, optional, default: 1

    K2 : number
        conductivity of the second layer, optional, default: 2

    nLineK : integer
        width of each layer, optional, default: 3

    rhs_sparse : boolean
        if True, b is returned as a sparse matrix. Else, it is a numpy array.

    symmetry : boolean
        if True, only the upper triangular part of A is returned

    loc2glob : list of integers
        Mapping from local to global ordering (if I is a set of
        indices in local ordering, loc2glob[I] is in global ordering)
        If present, the matrix is returned in local ordering

    glob2loc : function
        glob2loc(loc2glob[I]) = loc2glob[glob2loc(I)]. If loc2glob is
        given and not glob2loc, glob2loc is computed from loc2glob

    dirichlet_boundary_value : number
        value of the solution on boundary, optional, default 0

    dirac_source_value : number
        value of the dirac source that will be imposed at 
        the center of the domain, optional, default None

    source : numpy array
        source term of the equation, optional, default: 1

    w:
        is the wavenumber of the Helmholtz equation

    Additional keyword arguments are passed to mesh_rect

    Returns
    -------
    A : scipy sparse matrix
        A is the assembled fem matrix

    b : numpy array
        b is the assembled rhs, see rhs_sparse

    """
    shape = np.array(shape, ndmin=1)
    tElem_glob = mesh_rect(shape, **kwargs)
    returnGlobal = (loc2glob is None)
    if loc2glob is None:
        loc2glob = np.unique(tElem_glob)

        def glob2loc(I):
            return np.searchsorted(loc2glob, I)
    elif glob2loc is None:
        ind_sort = np.argsort(loc2glob)

        def glob2loc(I):
            return ind_sort[np.searchsorted(loc2glob, I, sorter=ind_sort)]
    tElem_loc = glob2loc(tElem_glob)
    Ae, Me = elem_rect(len(shape))
    K_vec = np.array([K1]*nLineK+[K2]*nLineK)
    K = K_vec[np.unravel_index(tElem_glob[0], shape,
                               order='F')[min(1, len(shape)-1)] % len(K_vec)]			       
    dirichlet_boundary_indices_loc = []
    dirac_source_index_loc = -1
    # Search if the currrent subdomain has one or more Dirichlet boundary nodes
    # when a value for this kind of boundary is given as an input
    if dirichlet_boundary_value is not None:
        dirichlet_boundary_indices_glob = nodes(shape, lim=[[0, 0]]+[[0, -1]]*(len(shape)-1)).ravel()
        dirichlet_boundary_indices_glob = dirichlet_boundary_indices_glob[(np.in1d(dirichlet_boundary_indices_glob, loc2glob))]
        dirichlet_boundary_indices_loc  = glob2loc(dirichlet_boundary_indices_glob)
    # Search if the current subdomain has the node where the dirac source is imposed
    # when a value for this kind of source is given as an input
    if dirac_source_value is not None:
        center_index = int(shape[0]/2)
        if len(shape) > 1: center_index += int(shape[1]/2)*shape[0]
        if len(shape) > 2: center_index += int(shape[2]/2)*shape[0]*shape[1]
        dirac_source_index_glob = center_index
        if np.in1d(dirac_source_index_glob, loc2glob) == True:
            dirac_source_index_loc  = glob2loc(dirac_source_index_glob)
    A, b = assemble_matrix(tElem_loc, Ae - w*w*Me, Me,dirichlet_boundary_indices_loc, dirichlet_boundary_value, source, K=K, 
                           rhs_sparse=rhs_sparse, symmetry=False, dirac_source_index=dirac_source_index_loc,dirac_source_value=dirac_source_value)

    A.data[np.abs(A.data) < np.finfo(float).eps] = 0
    A.eliminate_zeros()
    if symmetry:
        A = ssp.triu(A)
    if returnGlobal:
        A = A.tocoo()
        A._shape = np.multiply.reduce(shape), np.multiply.reduce(shape)
        A.col = loc2glob[A.col]
        A.row = loc2glob[A.row]
        if rhs_sparse:
            b = b.tocoo()
            b._shape = np.multiply.reduce(shape), 1
            b.row = loc2glob[b.row]
        else:
            b_ = np.zeros((A.shape[0], 1), dtype=A)
            b_[loc2glob, 0] = b
            b = b_
    else:
        b = b.reshape((-1, 1))
    return A, b

def stratified_heat(shape, K1=1, K2=1, nLineK=3, rhs_sparse=False,
                    symmetry=False, loc2glob=None, glob2loc=None,
                    dirichlet_boundary_value=0, source=1, **kwargs):
    """Create the fem matrix of the heat equation in a stratified medium
    For that, we consider the heat equation as a particular case of 
    the Helmholtz equation

    Parameters
    ----------
    shape : array or list
        [nX, nY, ...] is a list or array containing the number of nodes in each
        direction of the whole domain

    K1 : number
        conductivity of the first layer, optional, default: 1

    K2 : number
        conductivity of the second layer, optional, default: 2

    nLineK : integer
        width of each layer, optional, default: 3

    rhs_sparse : boolean
        if True, b is returned as a sparse matrix. Else, it is a numpy array.

    symmetry : boolean
        if True, only the upper triangular part of A is returned

    loc2glob : list of integers
        Mapping from local to global ordering (if I is a set of
        indices in local ordering, loc2glob[I] is in global ordering)
        If present, the matrix is returned in local ordering

    glob2loc : function
        glob2loc(loc2glob[I]) = loc2glob[glob2loc(I)]. If loc2glob is
        given and not glob2loc, glob2loc is computed from loc2glob

    dirichlet_boundary_value : number
        value of the solution on boundary, optional, default 0

    source : numpy array
        source term of the equation, optional, default: 1

    Additional keyword arguments are passed to mesh_rect

    Returns
    -------
    A : scipy sparse matrix
        A is the assembled fem matrix

    b : numpy array
        b is the assembled rhs, see rhs_sparse

    Examples
    --------
    >>> shape = 3,7
    >>> A,b = stratified_heat(shape, K1=1, K2=10, nLineK=2)
    >>> import scipy.sparse.linalg as sla
    >>> sla.spsolve(A.tocsc(),b).round(2)[nodes(shape)]
    array([[ 0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ,  0.  ],
           [ 1.08,  0.92,  0.31,  0.27,  0.31,  0.92,  1.08],
           [ 1.4 ,  1.18,  0.42,  0.36,  0.42,  1.18,  1.4 ]])
    >>> stratified_heat((3,), lim=[[1,2]])[0].A
    array([[ 0.,  0.,  0.],
           [ 0.,  1., -1.],
           [ 0., -1.,  1.]])
    >>> shape = 5
    >>> nDom = 2
    >>> loc2glob = ([0, 1, 2], [2, 3, 4])
    >>> for rank in range(np.multiply.reduce(nDom)):
    ...     limits, neighbors, interfaces = subdomain(shape, nDom, rank)
    ...     A,b = stratified_heat(shape, lim=limits, rhs_sparse=True)
    ...     A.A
    ...     b.A
    ...     A,b = stratified_heat(shape, lim=limits, rhs_sparse=False)
    ...     b
    ...     A,b = stratified_heat(shape, lim=limits, loc2glob=loc2glob[rank])
    ...     A.A
    ...     b
    ...
    array([[ 1.,  0.,  0.,  0.,  0.],
           [ 0.,  2., -1.,  0.,  0.],
           [ 0., -1.,  1.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.]])
    array([[ 0. ],
           [ 1. ],
           [ 0.5],
           [ 0. ],
           [ 0. ]])
    array([[ 0. ],
           [ 1. ],
           [ 0.5],
           [ 0. ],
           [ 0. ]])
    array([[ 1.,  0.,  0.],
           [ 0.,  2., -1.],
           [ 0., -1.,  1.]])
    array([[ 0. ],
           [ 1. ],
           [ 0.5]])
    array([[ 0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  1., -1.,  0.],
           [ 0.,  0., -1.,  2., -1.],
           [ 0.,  0.,  0., -1.,  1.]])
    array([[ 0. ],
           [ 0. ],
           [ 0.5],
           [ 1. ],
           [ 0.5]])
    array([[ 0. ],
           [ 0. ],
           [ 0.5],
           [ 1. ],
           [ 0.5]])
    array([[ 1., -1.,  0.],
           [-1.,  2., -1.],
           [ 0., -1.,  1.]])
    array([[ 0.5],
           [ 1. ],
           [ 0.5]])
    """
    A,b = stratified_helmholtz(shape=shape, K1=K1, K2=K2, nLineK=nLineK, rhs_sparse=rhs_sparse,
                    symmetry=symmetry, loc2glob=loc2glob, glob2loc=glob2loc,
                    dirichlet_boundary_value=dirichlet_boundary_value, source=source, w=0.0, **kwargs)
    return A, b

from ddmpy import DomainDecomposition, DistMatrix, DistVector

def subdomains_baton(N=24, dim=3):
    return [N, ] + [1, ]*(dim - 1)

def subdomains_cuboid(N=24, dim=3):
    assert(N%3 == 0)
    nDom = subdomains_baton(min(N, 3), dim)
    N //= 3
    while N > 1:
        assert(N%2 == 0)
        N //= 2
        nDom[-1] *= 2
        nDom = np.sort(nDom)[::-1]
    return nDom

def generate_example(geom="cuboid", N=24, dim=3,
                     local_shape=[30, 30, 30],
                     K1=1, K2=1, nLineK=3,
                     random_rhs=True, rank=0):
    if geom == "cuboid":
        nDom = subdomains_cuboid(N, dim)
    elif geom == "baton":
        nDom = subdomains_baton(N, dim)
    global_shape = [ls*n + 1 for ls, n in zip(local_shape, nDom)]
    limits, neighbors, interfaces = subdomain(global_shape, nDom, rank)
    loc2glob, glob2loc, interfaces, interface = local_ordering(
        global_shape, limits, neighbors, interfaces)
    Ai, bi = stratified_heat(global_shape, lim=limits,
                             K1=K1, K2=K2, nLineK=nLineK,
                             loc2glob=loc2glob, glob2loc=glob2loc)
    nei = {n: np.array(i) + len(bi) - len(interface)
           for n, i in zip(neighbors, interfaces)}
    dd = DomainDecomposition(Ai.shape[0], nei)
    A_d = DistMatrix(Ai, dd)
    ni = Ai.shape[0]
    if random_rhs:
        xi = np.random.RandomState(42+rank).rand(ni).reshape([ni, 1])
        x_s = DistVector(dd.D * xi, dd, assemble=True)
        b_d = A_d.dot(x_s)
    else:
        x_s = None
        b_d = DistVector(bi, dd, assemble=True)
    return A_d, b_d, x_s, loc2glob


# Main function for standalone execution
def entry_point():
    import argparse
    import scipy.io as sio
    import os.path
    parser = argparse.ArgumentParser(description="Generates a fem problem. "
                                     "Run for instance: "
                                     "genfem.py --shape 7 --K1 1 --K2 10")
    parser.add_argument('--shape', '-s', nargs='*', type=int, default=[10],
                        help="global number of nodes in each direction")
    parser.add_argument('--nDom', '-n', nargs='*', type=int,
                        help="number of subdomains in each direction")
    parser.add_argument('--K1', '-k', type=float, default=1,
                        help="conductivity of the even layers, default: 1")
    parser.add_argument('--K2', '-K', type=float, default=1,
                        help="conductivity of the odd layers, default: 1")
    parser.add_argument('--nLineK', '-l', type=int, default=3,
                        help=" width of each layer, default: 3")
    parser.add_argument('--local', action='store_true',
                        help=" Gives the matrix and rhs in local ordering")
    parser.add_argument('--mpi', action='store_true',
                        help=" computes the rank from the MPI_RANK")
    parser.add_argument('--sym', action='store_true',
                        help=" stores the matrix as a symmetric matrix")
    parser.add_argument('--test', '-t', action='store_true',
                        help=" runs the tests")

    p = parser.parse_args()
    if p.test:
        import doctest
        doctest.testmod(verbose=True)
        return
    shape, nDom, K1, K2, nLineK = p.shape, p.nDom, p.K1, p.K2, p.nLineK
    if nDom is None:
        nDom = np.ones_like(shape)
    if len(nDom) != len(shape) or any(np.subtract(shape, nDom) < 1):
        raise argparse.ArgumentTypeError("nDom=" + str(nDom) +
                                         " and shape=" + str(shape) +
                                         " are incompatible")
    if p.mpi:
        from mpi4py import MPI
        comm = MPI.COMM_WORLD
        ranks = range(comm.Get_rank(),
                      np.multiply.reduce(nDom),
                      comm.Get_size())
    else:
        ranks = range(np.multiply.reduce(nDom))
    symmetry = "symmetric" if p.sym else "general"
    for rank in ranks:
        if os.path.isfile("local_rhs" + str(rank+1) + ".mtx"):
            continue
        else:
            print(rank)
        limits, neighbors, interfaces = subdomain(shape, nDom, rank)
        A, b = stratified_heat(shape, K1, K2, nLineK,
                               lim=limits, rhs_sparse=True)
        if p.local:
            Ri, interfaces, interface = local_ordering(shape, limits,
                                                       neighbors, interfaces)
            A = Ri.dot(A.dot(Ri.T))
            b = Ri.dot(b)
            myInterface = " ".join(str(i+1) for i in interface)
            indexVi = " ".join(str(i) for i in neighbors)
            ptrIndexVi = " ".join(str(i+1) for i in
                                  np.cumsum(list(map(len, interfaces))))
            indexIntrf = " ".join(str(i+1) for i in
                                  np.hstack(interfaces))
            dom = "ndof= {!s}\n".format(b.shape[0]) + \
                  "sizeIntrf= {!s}\n".format(len(interface)) + \
                  "myInterface= {!s}\n".format(myInterface) + \
                  "\n" + \
                  "nbVi= {!s}\n".format(len(neighbors)) + \
                  "indexVi= {!s}\n".format(indexVi) + \
                  "ptrIndexVi= 1 {!s}\n".format(ptrIndexVi) + \
                  "indexIntrf= {!s}\n".format(indexIntrf)
            with open("maphys_local_domain" + str(rank+1) + ".dom", 'w') as f:
                print(dom, file=f)
        sio.mmwrite("local_matrix" + str(rank+1) + ".mtx",
                    A, symmetry=symmetry)
        sio.mmwrite("local_rhs" + str(rank+1) + ".mtx", b)


if __name__ == '__main__':
    entry_point()