source: inundation/ga/storm_surge/parallel/parallel_advection.py @ 1460

import sys
from os import sep
sys.path.append('..'+sep+'pyvolution')

"""Class Parallel_Domain -
2D triangular domains for finite-volume computations of
the advection equation, with extra structures to allow
communication between other Parallel_Domains and itself

This module contains a specialisation of class Domain from module advection.py

Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
Geoscience Australia, 2004
"""

from advection import *
Advection_Domain = Domain
from Numeric import zeros, Float, Int, ones, allclose, array
import pypar


class Parallel_Domain(Advection_Domain):

    def __init__(self, coordinates, vertices, boundary = None,
                 full_send_dict = None, ghost_recv_dict = None,
                 velocity = None):

        self.processor = pypar.rank()
        self.numproc   = pypar.size()
        #print 'Processor %d'%self.processor
        #velocity = [(self.processor+1),0.0]

        #print 'velocity',velocity

        Advection_Domain.__init__(self, coordinates, vertices, boundary, velocity)

        N = self.number_of_elements

        self.processor = pypar.rank()
        self.numproc   = pypar.size()

        self.full_send_dict  = full_send_dict
        self.ghost_recv_dict = ghost_recv_dict

        #print self.full_send_dict
        #print self.ghost_recv_dict

    def check_integrity(self):
        Advection_Domain.check_integrity(self)

        msg = 'Will need to check global and local numbering'
        assert self.conserved_quantities[0] == 'stage', msg



    def update_timestep(self, yieldstep, finaltime):

        # Calculate local timestep
        Advection_Domain.update_timestep(self, yieldstep, finaltime)

        # For some reason it looks like pypar only reduces numeric arrays
        # hence we need to create some dummy arrays for communication
        ltimestep = ones( 1, Float )
        ltimestep[0] = self.timestep

        gtimestep = zeros( 1, Float) # Buffer for results

        pypar.raw_reduce(ltimestep, gtimestep, pypar.MIN, 0)
        pypar.broadcast(gtimestep,0)
        #pypar.Barrier()

        self.timestep = gtimestep[0]



    def update_ghosts(self):

        # We must send the information from the full cells and
        # receive the information for the ghost cells
        # We have a dictionary of lists with ghosts expecting updates from
        # the separate processors

        from weave import converters

        stage_cv = self.quantities['stage'].centroid_values

        # update of non-local ghost cells
        for iproc in range(self.numproc):

            if iproc == self.processor:
                #Send data from iproc processor to other processors
                for send_proc in self.full_send_dict:
                    if send_proc != iproc:
                        Idf  = self.full_send_dict[send_proc][0]
                        Xout = self.full_send_dict[send_proc][1]
                        N = len(Xout)

                        """
                        # Original python Code
                        for i in range(N):
                            Xout[i] = stage_cv[Idf[i]]
                        """
                        code1 = """
                        for (int i=0; i<N ; i++){
                            Xout(i) = stage_cv(Idf(i));
                        }
                        """
                        weave.inline(code1, ['stage_cv','Idf','Xout','N'],
                                     type_converters = converters.blitz, compiler='gcc');

                        pypar.send(Xout,send_proc)
                        #print 'Processor %d Sending to Processor %d'%(self.processor,send_proc)
            else:
                #Receive data from the iproc processor
                if  self.ghost_recv_dict.has_key(iproc):
                    Idg = self.ghost_recv_dict[iproc][0]
                    X   = self.ghost_recv_dict[iproc][1]

                    X = pypar.receive(iproc,X)
                    N = len(X)
                    #print 'Processor %d receiving from Processor %d'%(self.processor,iproc)
                    for i in range(N):
                        stage_cv[Idg[i]] = X[i]
            pypar.barrier()

        #local update of ghost cells
        iproc = self.processor
        if self.full_send_dict.has_key(iproc):
            Idf  = self.full_send_dict[iproc][0]
            #print Idf
            Idg  = self.ghost_recv_dict[iproc][0]
            N = len(Idg)
            #print Idg
            for i in range(N):
                #print i,Idg[i],Idf[i]
                stage_cv[Idg[i]] = stage_cv[Idf[i]]


#        if self.ghosts is not None:
#            stage_cv = self.quantities['stage'].centroid_values
#            for triangle in self.ghosts:
#                stage_cv[triangle] = stage_cv[self.ghosts[triangle]]


    def write_time(self):
        if self.min_timestep == self.max_timestep:
            print 'Processor %d, Time = %.4f, delta t = %.8f, steps=%d (%d)'\
                  %(self.processor, self.time, self.min_timestep, self.number_of_steps,
                    self.number_of_first_order_steps)
        elif self.min_timestep > self.max_timestep:
            print 'Processor %d, Time = %.4f, steps=%d (%d)'\
                  %(self.processor, self.time, self.number_of_steps,
                    self.number_of_first_order_steps)
        else:
            print 'Processor %d, Time = %.4f, delta t in [%.8f, %.8f], steps=%d (%d)'\
                  %(self.processor, self.time, self.min_timestep,
                    self.max_timestep, self.number_of_steps,
                    self.number_of_first_order_steps)



    def evolve(self, yieldstep = None, finaltime = None):
        """Specialisation of basic evolve method from parent class
        """

        #Initialise real time viz if requested
        if self.time == 0.0:
            pass

        #Call basic machinery from parent class
        for t in Advection_Domain.evolve(self, yieldstep, finaltime):

            #Pass control on to outer loop for more specific actions
            yield(t)




def parallel_rectangular(m, n, len1=1.0, len2=1.0, origin = (0.0, 0.0)):


    """Setup a rectangular grid of triangles
    with m+1 by n+1 grid points
    and side lengths len1, len2. If side lengths are omitted
    the mesh defaults to the unit square, divided between all the
    processors

    len1: x direction (left to right)
    len2: y direction (bottom to top)
    """

    from config import epsilon
    from Numeric import zeros, Float, Int

    processor = pypar.rank()
    numproc   = pypar.size()



    delta1 = float(len1)/m
    delta2 = float(len2)/n

    #Calculate number of points
    Np = (m+1)*(n+1)

    class VIndex:

        def __init__(self, n,m):
            self.n = n
            self.m = m

        def __call__(self, i,j):
            return j+i*(self.n+1)

    class EIndex:

        def __init__(self, n,m):
            self.n = n
            self.m = m

        def __call__(self, i,j):
            return 2*(j+i*self.n)


    I = VIndex(n,m)
    E = EIndex(n,m)

    points = zeros( (Np,2), Float)

    for i in range(m+1):
        for j in range(n+1):

            points[I(i,j),:] = [i*delta1 + origin[0], j*delta2 + origin[1]]

    #Construct 2 triangles per rectangular element and assign tags to boundary
    #Calculate number of triangles
    Nt = 2*m*n


    elements = zeros( (Nt,3), Int)
    boundary = {}
    Idgl = []
    Xgl  = []
    Idfl = []
    Xfl  = []
    Idgr = []
    Xgr  = []
    Idfr = []
    Xfr  = []

    full_send_dict = {}
    ghost_recv_dict = {}
    nt = -1
    for i in range(m):
        for j in range(n):

            i1 = I(i,j+1)
            i2 = I(i,j)
            i3 = I(i+1,j+1)
            i4 = I(i+1,j)

            #Lower Element
            nt = E(i,j)
            if i == m-1:
                #print 'nt =',nt
                Idgr.append(nt)
                Idfr.append(E(1,j))
            if i == 0:
                Idgl.append(nt)
                Idfl.append(E(m-2,j))

            if i == m-1:
                boundary[nt, 2] = 'right'
            if j == 0:
                boundary[nt, 1] = 'bottom'
            elements[nt,:] = [i4,i3,i2]

            #Upper Element
            nt = E(i,j)+1
            if i == m-1:
                Idgr.append(nt)
                Idfr.append(E(1,j)+1)
            if i == 0:
                Idgl.append(nt)
                Idfl.append(E(m-2,j)+1)

            if i == 0:
                boundary[nt, 2] = 'left'
            if j == n-1:
                boundary[nt, 1] = 'top'
            elements[nt,:] = [i1,i2,i3]

    Idfl = array(Idfl,Int)
    Idgl = array(Idgl,Int)
    Xfl  = zeros(Idfl.shape,Float)
    Xgl  = zeros(Idgl.shape,Float)

    Idfr = array(Idfr,Int)
    Idgr = array(Idgr,Int)
    Xfr  = zeros(Idfr.shape,Float)
    Xgr  = zeros(Idgr.shape,Float)

    #print Idf
    #print Idg
    full_send_dict[(processor-1)%numproc]  = [Idfl, Xfl]
    ghost_recv_dict[(processor-1)%numproc] = [Idgl, Xgl]
    full_send_dict[(processor+1)%numproc]  = [Idfr, Xfr]
    ghost_recv_dict[(processor+1)%numproc] = [Idgr, Xgr]

    return  points, elements, boundary, full_send_dict, ghost_recv_dict



def rectangular_periodic(m, n, len1=1.0, len2=1.0, origin = (0.0, 0.0)):


    """Setup a rectangular grid of triangles
    with m+1 by n+1 grid points
    and side lengths len1, len2. If side lengths are omitted
    the mesh defaults to the unit square.

    len1: x direction (left to right)
    len2: y direction (bottom to top)

    Return to lists: points and elements suitable for creating a Mesh or
    FVMesh object, e.g. Mesh(points, elements)
    """

    from config import epsilon
    from Numeric import zeros, Float, Int

    delta1 = float(len1)/m
    delta2 = float(len2)/n

    #Calculate number of points
    Np = (m+1)*(n+1)

    class VIndex:

        def __init__(self, n,m):
            self.n = n
            self.m = m

        def __call__(self, i,j):
            return j+i*(self.n+1)

    class EIndex:

        def __init__(self, n,m):
            self.n = n
            self.m = m

        def __call__(self, i,j):
            return 2*(j+i*self.n)


    I = VIndex(n,m)
    E = EIndex(n,m)

    points = zeros( (Np,2), Float)

    for i in range(m+1):
        for j in range(n+1):

            points[I(i,j),:] = [i*delta1 + origin[0], j*delta2 + origin[1]]

    #Construct 2 triangles per rectangular element and assign tags to boundary
    #Calculate number of triangles
    Nt = 2*m*n


    elements = zeros( (Nt,3), Int)
    boundary = {}
    ghosts = {}
    nt = -1
    for i in range(m):
        for j in range(n):

            i1 = I(i,j+1)
            i2 = I(i,j)
            i3 = I(i+1,j+1)
            i4 = I(i+1,j)

            #Lower Element
            nt = E(i,j)
            if i == m-1:
                ghosts[nt] = E(1,j)
            if i == 0:
                ghosts[nt] = E(m-2,j)

            if j == n-1:
                ghosts[nt] = E(i,1)

            if j == 0:
                ghosts[nt] = E(i,n-2)

            if i == m-1:
                boundary[nt, 2] = 'right'
            if j == 0:
                boundary[nt, 1] = 'bottom'
            elements[nt,:] = [i4,i3,i2]

            #Upper Element
            nt = E(i,j)+1
            if i == m-1:
                ghosts[nt] = E(1,j)+1
            if i == 0:
                ghosts[nt] = E(m-2,j)+1

            if j == n-1:
                ghosts[nt] = E(i,1)+1

            if j == 0:
                ghosts[nt] = E(i,n-2)+1

            if i == 0:
                boundary[nt, 2] = 'left'
            if j == n-1:
                boundary[nt, 1] = 'top'
            elements[nt,:] = [i1,i2,i3]

    #bottom left
    nt = E(0,0)
    nf = E(m-2,n-2)
    ghosts[nt]   = nf
    ghosts[nt+1] = nf+1

    #bottom right
    nt = E(m-1,0)
    nf = E(1,n-2)
    ghosts[nt]   = nf
    ghosts[nt+1] = nf+1

    #top left
    nt = E(0,n-1)
    nf = E(m-2,1)
    ghosts[nt]   = nf
    ghosts[nt+1] = nf+1

    #top right
    nt = E(m-1,n-1)
    nf = E(1,1)
    ghosts[nt]   = nf
    ghosts[nt+1] = nf+1

    return points, elements, boundary, ghosts

def rectangular_periodic_lr(m, n, len1=1.0, len2=1.0, origin = (0.0, 0.0)):


    """Setup a rectangular grid of triangles
    with m+1 by n+1 grid points
    and side lengths len1, len2. If side lengths are omitted
    the mesh defaults to the unit square.

    len1: x direction (left to right)
    len2: y direction (bottom to top)

    Return to lists: points and elements suitable for creating a Mesh or
    Domain object, e.g. Mesh(points, elements)
    """

    from config import epsilon
    from Numeric import zeros, Float, Int

    delta1 = float(len1)/m
    delta2 = float(len2)/n

    #Calculate number of points
    Np = (m+1)*(n+1)

    class VIndex:

        def __init__(self, n,m):
            self.n = n
            self.m = m

        def __call__(self, i,j):
            return j+i*(self.n+1)

    class EIndex:

        def __init__(self, n,m):
            self.n = n
            self.m = m

        def __call__(self, i,j):
            return 2*(j+i*self.n)


    I = VIndex(n,m)
    E = EIndex(n,m)

    points = zeros( (Np,2), Float)

    for i in range(m+1):
        for j in range(n+1):

            points[I(i,j),:] = [i*delta1 + origin[0], j*delta2 + origin[1]]

    #Construct 2 triangles per rectangular element and assign tags to boundary
    #Calculate number of triangles
    Nt = 2*m*n


    elements = zeros( (Nt,3), Int)
    boundary = {}
    ghosts = {}
    nt = -1
    for i in range(m):
        for j in range(n):

            i1 = I(i,j+1)
            i2 = I(i,j)
            i3 = I(i+1,j+1)
            i4 = I(i+1,j)

            #Lower Element
            nt = E(i,j)
            if i == m-1:
                ghosts[nt] = E(1,j)
            if i == 0:
                ghosts[nt] = E(m-2,j)

            if i == m-1:
                boundary[nt, 2] = 'right'
            if j == 0:
                boundary[nt, 1] = 'bottom'
            elements[nt,:] = [i4,i3,i2]

            #Upper Element
            nt = E(i,j)+1
            if i == m-1:
                ghosts[nt] = E(1,j)+1
            if i == 0:
                ghosts[nt] = E(m-2,j)+1

            if i == 0:
                boundary[nt, 2] = 'left'
            if j == n-1:
                boundary[nt, 1] = 'top'
            elements[nt,:] = [i1,i2,i3]


    return points, elements, boundary, ghosts