source: inundation/parallel/pmesh_divide.py @ 2094

#########################################################
#
#
#  Read in a data file and subdivide the triangle list
#
#
#  This is only intended as a temporary file, once an
# automatic grid partitioner has been incorporated this
# file will become redundant.
#
#  Authors: Linda Stals and Matthew Hardy, June 2005
#  Modified: Linda Stals, Nov 2005
#            Jack Kelly, Nov 2005
#
#
#########################################################


from math import floor
from Numeric import zeros, Float, Int, reshape

#########################################################
#
# If the triangles list is reordered, the quantities
# assigned to the triangles must also be reorded.
#
# *) quantities contain the quantites in the old ordering
# *) proc_sum[i] contains the number of triangles in
# processor i
# *) tri_index is a map from the old triangle ordering to
# the new ordering, where the new number for triangle
# i is proc_sum[tri_index[i][0]]+tri_index[i][1]
#
# -------------------------------------------------------
#
# *) The quantaties are returned in the new ordering
#
#########################################################

def reorder(quantities, tri_index, proc_sum):

    # find the number triangles

    N = len(tri_index)

    # temporary storage area

    index = zeros(N, Int)
    q_reord = {}

    # find the new ordering of the triangles

    for i in range(N):
        bin = tri_index[i][0]
        bin_off_set = tri_index[i][1]
        index[i] = proc_sum[bin]+bin_off_set

    # reorder each quantity according to the new ordering

    for k in quantities:
        q_reord[k] = zeros((N, 3), Float)
        for i in range(N):
            q_reord[k][index[i]]=quantities[k].vertex_values[i]
    del index

    return q_reord


#########################################################
#
# Do a simple coordinate based division of the grid
#
# *) n_x is the number of subdivisions in the x direction
# *) n_y is the number of subdivisions in the y direction
#
# -------------------------------------------------------
#
# *)  The nodes, triangles, boundary, and quantities are
# returned. triangles_per_proc defines the subdivision.
# The first triangles_per_proc[0] triangles are assigned
# to processor 0, the next triangles_per_proc[1] are
# assigned to processor 1 etc. The boundary and quantites
# are ordered the same way as the triangles
#
#########################################################

def pmesh_divide(domain, n_x = 1, n_y = 1):

    # find the bounding box
    
    x_coord_min = domain.xy_extent[0]
    x_coord_max = domain.xy_extent[2]
    y_coord_min = domain.xy_extent[1]
    y_coord_max = domain.xy_extent[3]

    # find the size of each sub-box

    x_div = (x_coord_max-x_coord_min)/n_x
    y_div = (y_coord_max-y_coord_min)/n_y

    # initialise the lists
    
    tri_list = []
    triangles_per_proc = []
    proc_sum = []
    for i in range(n_x*n_y):
        tri_list.append([])
        triangles_per_proc.append([])
        proc_sum.append([])
        tri_list[i] = []

    # subdivide the triangles depending on which sub-box they sit
    # in (a triangle sits in the sub-box if its first vectex sits
    # in that sub-box)

    tri_index = {}
    N = domain.number_of_elements
    for i in range(N):
        t = domain.triangles[i]

        x_coord = domain.centroid_coordinates[i][0]
        bin_x = int(floor((x_coord-x_coord_min)/x_div))
        if (bin_x == n_x):
            bin_x = n_x-1

        y_coord = domain.centroid_coordinates[i][1]
        bin_y = int(floor((y_coord-y_coord_min)/y_div))
        if (bin_y == n_y):
            bin_y = n_y-1

        bin = bin_y*n_x + bin_x
        tri_list[bin].append(t)
        tri_index[i] = ([bin, len(tri_list[bin])-1])

    # find the number of triangles per processor and order the
    # triangle list so that all of the triangles belonging to
    # processor i are listed before those belonging to processor
    # i+1

    triangles = []
    for i in range(n_x*n_y):
        triangles_per_proc[i] = len(tri_list[i])
        for t in tri_list[i]:
            triangles.append(t)

    # the boundary labels have to changed in accoradance with the
    # new triangle ordering, proc_sum and tri_index help with this

    proc_sum[0] = 0
    for i in range(n_x*n_y-1):
        proc_sum[i+1]=proc_sum[i]+triangles_per_proc[i]

    # relabel the boundary elements to fit in with the new triangle
    # ordering

    boundary = {}
    for b in domain.boundary:
        t =  tri_index[b[0]]
        boundary[proc_sum[t[0]]+t[1], b[1]] = domain.boundary[b]

    quantities = reorder(domain.quantities, tri_index, proc_sum)

    # extract the node list
    
    nodes = domain.coordinates.copy()
    ttriangles = zeros((len(triangles), 3), Int)
    for i in range(len(triangles)):
        ttriangles[i] = triangles[i]
        
    return nodes, ttriangles, boundary, triangles_per_proc, quantities


#########################################################
#
# Do a coordinate based division of the grid using the
# centroid coordinate
#
# *) n_x is the number of subdivisions in the x direction
# *) n_y is the number of subdivisions in the y direction
#
# -------------------------------------------------------
#
# *)  The nodes, triangles, boundary, and quantities are
# returned. triangles_per_proc defines the subdivision.
# The first triangles_per_proc[0] triangles are assigned
# to processor 0, the next triangles_per_proc[1] are
# assigned to processor 1 etc. The boundary and quantites
# are ordered the same way as the triangles
#
#########################################################
def pmesh_divide_steve(domain, n_x = 1, n_y = 1):
    
    # find the bounding box
    x_coord_min = domain.xy_extent[0]
    x_coord_max = domain.xy_extent[2]
    y_coord_min = domain.xy_extent[1]
    y_coord_max = domain.xy_extent[3]


    # find the size of each sub-box

    x_div = (x_coord_max-x_coord_min)/n_x
    y_div = (y_coord_max-y_coord_min)/n_y


    # initialise the lists
    tri_list = []
    triangles_per_proc = []
    proc_sum = []
    for i in range(n_x*n_y):
        tri_list.append([])
        triangles_per_proc.append([])
        proc_sum.append([])
        tri_list[i] = []

    # subdivide the triangles depending on which sub-box they sit
    # in (a triangle sits in the sub-box if its first vectex sits
    # in that sub-box)

    tri_index = {}
    N = domain.number_of_elements
 
    #sort by x coordinate of centroid
    from Numeric import argsort
    sort_order = argsort(argsort(domain.centroid_coordinates[:,0]))

    x_div = float(N)/n_x

    for i in range(N):
        t = domain.triangles[i]
        
        bin = int(floor(sort_order[i]/x_div))
        if (bin == n_x):
            bin = n_x-1

        tri_list[bin].append(t)
        tri_index[i] = ([bin, len(tri_list[bin])-1])

    #print tri_list
    #print tri_index

    # find the number of triangles per processor and order the
    # triangle list so that all of the triangles belonging to
    # processor i are listed before those belonging to processor
    # i+1

    triangles = []
    for i in range(n_x*n_y):
        triangles_per_proc[i] = len(tri_list[i])
        for t in tri_list[i]:
           triangles.append(t)
        

    # the boundary labels have to changed in accoradance with the
    # new triangle ordering, proc_sum and tri_index help with this

    proc_sum[0] = 0
    for i in range(n_x*n_y-1):
        proc_sum[i+1]=proc_sum[i]+triangles_per_proc[i]

    # relabel the boundary elements to fit in with the new triangle
    # ordering

    boundary = {}
    for b in domain.boundary:
        t =  tri_index[b[0]]
        boundary[proc_sum[t[0]]+t[1], b[1]] = domain.boundary[b]

    quantities = reorder(domain.quantities, tri_index, proc_sum)

    # extract the node list
    nodes = domain.coordinates.copy()


    # convert the triangle datastructure to be an array typ
    # this helps with the communication

    ttriangles = zeros((len(triangles), 3), Int)
    for i in range(len(triangles)):
        ttriangles[i] = triangles[i]
        
    return nodes, ttriangles, boundary, triangles_per_proc, quantities




###
#
# Divide the mesh using a call to metis, through pymetis.

from os import sep
from sys import path

path.append('..' + sep + 'pymetis')

from metis import partMeshNodal

def pmesh_divide_metis(domain, n_procs):
    # initialise the lists
    # List, indexed by processor of # triangles.
    triangles_per_proc = []
    # List of lists, indexed by processor of vertex numbers
    tri_list = []
    # List indexed by processor of cumulative total of triangles allocated.
    proc_sum = []
    for i in range(n_procs):
        tri_list.append([])
        triangles_per_proc.append(0)
        proc_sum.append([])

    # Prepare variables for the metis call
    n_tri = len(domain.triangles)
    if n_procs != 1: #Because metis chokes on it...
        n_vert = len(domain.coordinates)
        t_list = domain.triangles.copy()
        t_list = reshape(t_list, (-1,))
    
        # The 1 here is for triangular mesh elements.
        edgecut, epart, npart = partMeshNodal(n_tri, n_vert, t_list, 1, n_procs)
        del edgecut
        del npart
        # print edgecut
        # print npart
    else:
        epart = [0] * n_tri

    # print epart


    # Assign triangles to processes, according to what metis told us.

    # tri_index maps triangle number -> processor, new triangle number
    # (local to the processor)
    tri_index = {}
    triangles = []        
    for i in range(n_tri):
        triangles_per_proc[epart[i]] = triangles_per_proc[epart[i]] + 1
        tri_list[epart[i]].append(domain.triangles[i])
        tri_index[i] = ([epart[i], len(tri_list[epart[i]]) - 1])

    # print tri_list
    # print triangles_per_proc

    # order the triangle list so that all of the triangles belonging
    # to processor i are listed before those belonging to processor
    # i+1

    for i in range(n_procs):
        for t in tri_list[i]:
            triangles.append(t)
            
    # the boundary labels have to changed in accoradance with the
    # new triangle ordering, proc_sum and tri_index help with this

    proc_sum[0] = 0
    for i in range(n_procs - 1):
        proc_sum[i+1]=proc_sum[i]+triangles_per_proc[i]

    # relabel the boundary elements to fit in with the new triangle
    # ordering

    boundary = {}
    for b in domain.boundary:
        t =  tri_index[b[0]]
        boundary[proc_sum[t[0]]+t[1], b[1]] = domain.boundary[b]

    quantities = reorder(domain.quantities, tri_index, proc_sum)

    # extract the node list
    nodes = domain.coordinates.copy()

    # convert the triangle datastructure to be an array type
    # this helps with the communication

    ttriangles = zeros((len(triangles), 3), Int)
    for i in range(len(triangles)):
        ttriangles[i] = triangles[i]
        
    return nodes, ttriangles, boundary, triangles_per_proc, quantities