source: trunk/anuga_core/source/anuga_parallel/parallel_shallow_water.py @ 8272

"""Class Parallel_shallow_water_domain -
2D triangular domains for finite-volume computations of
the shallow water equation, with extra structures to allow
communication between other Parallel_domains and itself

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

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

"""

from anuga import Domain

from anuga_parallel.parallel_generic_communications import *
from anuga.abstract_2d_finite_volumes.neighbour_mesh import Mesh

import numpy as num


#Import matplotlib



class Parallel_domain(Domain):

    def __init__(self, coordinates, vertices,
                 boundary=None,
                 full_send_dict=None,
                 ghost_recv_dict=None,
                 number_of_full_nodes=None,
                 number_of_full_triangles=None,
                 geo_reference=None,
                 tri_map=None,
                 inv_tri_map=None): #jj added this

        Domain.__init__(self,
                        coordinates,
                        vertices,
                        boundary,
                        full_send_dict=full_send_dict,
                        ghost_recv_dict=ghost_recv_dict,
                        processor=pypar.rank(),
                        numproc=pypar.size(),
                        number_of_full_nodes=number_of_full_nodes,
                        number_of_full_triangles=number_of_full_triangles,
                        geo_reference=geo_reference) #jj added this
        
        # PETE: Find the number of full nodes and full triangles, this is a temporary fix
        # until the bug with get_number_of_full_[nodes|triangles]() is fixed.

        if number_of_full_nodes is not None:
            self.number_of_full_nodes_tmp = number_of_full_nodes
        else:
            self.number_of_full_nodes_tmp = get_number_of_nodes()

        if number_of_full_triangles is not None:
            self.number_of_full_triangles_tmp = number_of_full_triangles
        else:
            self.number_of_full_triangles_tmp = get_number_of_triangles()

        setup_buffers(self)

        self.global_name = 'domain'
        self.tri_map = tri_map
        self.inv_tri_map = inv_tri_map


    def set_name(self, name):
        """Assign name based on processor number 
        """

        if name.endswith('.sww'):
            name = name[:-4]

        self.global_name = name

        # Call parents method with processor number attached.
        Domain.set_name(self, name + '_P%d_%d' %(self.processor, self.numproc))


    def get_global_name(self):

        return self.global_name


    def update_timestep(self, yieldstep, finaltime):
        """Calculate local timestep
        """

        communicate_flux_timestep(self, yieldstep, finaltime)

        Domain.update_timestep(self, yieldstep, finaltime)



    def update_ghosts(self):

        # We must send the information from the full cells and
        # receive the information for the ghost cells


        communicate_ghosts(self)

    def apply_fractional_steps(self):

        for operator in self.fractional_step_operators:
            operator()

        # PETE: Make sure that there are no deadlocks here

        self.update_ghosts()



    def sww_merge(self, verbose=False):

        if self.processor == 0 and self.numproc > 1:
            import anuga.utilities.sww_merge as merge
            
            merge.sww_merge(self.get_global_name(),self.numproc,verbose)


# =======================================================================
# PETE: NEW METHODS FOR FOR PARALLEL STRUCTURES. Note that we assume the 
# first "number_of_full_[nodes|triangles]" are full [nodes|triangles]
# For full triangles it is possible to enquire self.tri_full_flag == True
# =======================================================================

    def get_number_of_full_triangles(self, *args, **kwargs):
        return self.number_of_full_triangles_tmp

    def get_full_centroid_coordinates(self, *args, **kwargs):
        C = self.mesh.get_centroid_coordinates(*args, **kwargs)
        return C[:self.number_of_full_triangles_tmp, :]

    def get_full_vertex_coordinates(self, *args, **kwargs):
        V = self.mesh.get_vertex_coordinates(*args, **kwargs)
        return V[:3*self.number_of_full_triangles_tmp,:]

    def get_full_triangles(self, *args, **kwargs):
        T = self.mesh.get_triangles(*args, **kwargs)
        return T[:self.number_of_full_triangles_tmp,:]

    def get_full_nodes(self, *args, **kwargs):
        N = self.mesh.get_nodes(*args, **kwargs)
        return N[:self.number_of_full_nodes_tmp,:]

    def get_tri_map(self):
        return self.tri_map

    def get_inv_tri_map(self):
        return self.inv_tri_map

    '''
    Outputs domain triangulation, full triangles are shown in green while ghost triangles are shown in blue.
    The default filename is "domain.png"
    '''
    def dump_triangulation(self, filename="domain.png"):
        # Get vertex coordinates, partition full and ghost triangles based on self.tri_full_flag

        try:
            import matplotlib
            matplotlib.use('Agg')
            import matplotlib.pyplot as plt
            import matplotlib.tri as tri
        except:
            print "Couldn't import module from matplotlib, probably you need to update matplotlib"
            raise

        vertices = self.get_vertex_coordinates()
        full_mask = num.repeat(self.tri_full_flag == 1, 3)
        ghost_mask = num.repeat(self.tri_full_flag == 0, 3)
        
        myid = pypar.rank()
        numprocs = pypar.size()

        if myid == 0:
            fx = {}
            fy = {}
            gx = {}
            gy = {}

            # Proc 0 gathers full and ghost nodes from self and other processors
            fx[0] = vertices[full_mask,0]
            fy[0] = vertices[full_mask,1]
            gx[0] = vertices[ghost_mask,0]
            gy[0] = vertices[ghost_mask,1]
            
            for i in range(1,numprocs):
                fx[i] = pypar.receive(i)
                fy[i] = pypar.receive(i)
                gx[i] = pypar.receive(i)
                gy[i] = pypar.receive(i)

            # Plot full triangles
            for i in range(0, numprocs):
                n = int(len(fx[i])/3)
                            
                triang = num.array(range(0,3*n))
                triang.shape = (n, 3)
                plt.triplot(fx[i], fy[i], triang, 'g-')

            # Plot ghost triangles
            for i in range(0, numprocs):
                n = int(len(gx[i])/3)
                            
                triang = num.array(range(0,3*n))
                triang.shape = (n, 3)
                plt.triplot(gx[i], gy[i], triang, 'b--')

            # Save triangulation to location pointed by filename
            plt.savefig(filename)

        else:
            # Proc 1..numprocs send full and ghost triangles to Proc 0
            pypar.send(vertices[full_mask,0], 0)
            pypar.send(vertices[full_mask,1], 0)
            pypar.send(vertices[ghost_mask,0], 0)
            pypar.send(vertices[ghost_mask,1], 0)