def run_mpi():
    print "Rank :", rank, "started"
    npix = hp.nside2npix(config.nside_out)
    dim, ind_elements = cov_ut.get_dim(config.pol_type)
    bolo_list = [None, None, None]

    inv_cov_matrix_local = np.zeros((npix, ind_elements), dtype=np.float)
    b_matrix_local = np.zeros((npix, dim), dtype=np.float)
    hitmap_local = np.zeros(npix, dtype=np.float)

    bolo_segment_dict = get_local_bolo_segment_list(rank, size, config.bolo_list, config.segment_list)

    time.sleep(0.1*rank)
    print "Rank :", rank, ", Bolos and Segments :", bolo_segment_dict
    comm.Barrier()

    recon_dir = get_recon_dir()
    if rank == 0:
        make_data_dirs() 

    for bolo_name in bolo_segment_dict.keys():
        print "Rank :", rank, "Bolos class being generated"
        bolo_list = initialise_bolo(bolo_name, bolo_list)
        for segment in bolo_segment_dict[bolo_name]:
            prompter.prompt("Rank : %d doing Bolo : %s and segment : %d" % (rank, bolo_name, segment))
            if config.take_diff_signal:
                signal, v, pol_ang = acquire_difference_signal(bolo_a, bolo_b, segment)
            else:
                signal, v, pol_ang = acquire_signal(bolo, segment)
            if config.subtract_template:
                signal_TEMPLATE = bolo_TEMPLATE.read_timestream(segment, read_list=["signal"])["signal"]
                signal -= estimated_y*signal_TEMPLATE
            print "Rank :", rank, "Bolos signal read"
            hitpix = hp.vec2pix(config.nside_out, v[...,0], v[...,1], v[...,2])
            del v
            cov_ut.get_inv_cov_matrix(hitpix, pol_ang, signal, inv_cov_matrix_local, b_matrix_local, hitmap_local, npix, config.pol_type)
            print "Rank :", rank, "Inverse covariance matrix generated"

    if config.subtract_template:
        del signal_TEMPLATE
    del signal
    del pol_ang
    del hitpix

    inv_cov_matrix_local_segment = distribute_matrix(inv_cov_matrix_local, "cov_matrix")
    del inv_cov_matrix_local
    b_matrix_local_segment = distribute_matrix(b_matrix_local, "b_matrix")
    del b_matrix_local
    hitmap_local_segment = distribute_matrix(hitmap_local, "hitmap")
    del hitmap_local

    cov_matrix_local_segment = cov_ut.get_covariance_matrix(inv_cov_matrix_local_segment, hitmap_local_segment, config.pol_type)

    sky_map_local_segment = cov_ut.get_sky_map(cov_matrix_local_segment, b_matrix_local_segment, hitmap_local_segment, config.pol_type)

    write_segments(hitmap_local_segment, "hitmap", recon_dir) 
    write_segments(inv_cov_matrix_local_segment, "inverse_covariance_matrix", recon_dir) 
    write_segments(cov_matrix_local_segment, "covariance_matrix", recon_dir) 
    write_segments(sky_map_local_segment, "sky_map", recon_dir) 
def distribute_matrix(local_full_matrix, matrix_type):
    npix = hp.nside2npix(config.nside_out)
    dim, ind_elements = cov_ut.get_dim(config.pol_type)
    segment_length = npix/size 

    if matrix_type == "hitmap":
        local_segmented_matrix = np.zeros(segment_length) 
    else:
        inner_dim = {"cov_matrix" : ind_elements, "b_matrix" : dim}
        local_segmented_matrix = np.zeros((segment_length, inner_dim[matrix_type])) 

    for i in range(size):
        start = i*segment_length
        stop = (i+1)*segment_length
        comm.Reduce(local_full_matrix[start:stop], local_segmented_matrix, MPI.SUM, root=i)
        
    return local_segmented_matrix
def accumulate_segments(size):
    recon_dir = get_recon_dir()
    map_segment_dir, hitmap_segment_dir, cov_matrix_segment_dir, inv_cov_matrix_segment_dir = get_segment_dirs()

    npix = hp.nside2npix(config.nside_out)
    npix_segment = npix/size
    dim, ind_elements = cov_ut.get_dim(config.pol_type)

    sky_map = np.empty((dim, npix))

    for i in range(size):
        start = i*npix_segment
        stop = (i+1)*npix_segment
        print os.path.join(recon_dir, map_segment_dir, str(i).zfill(4) + '.npy'), os.path.exists(os.path.join(recon_dir, map_segment_dir, str(i).zfill(4) + '.npy'))
        sky_map_segment = np.load(os.path.join(recon_dir, map_segment_dir, str(i).zfill(4) + '.npy')) 
        sky_map[..., start:stop] = sky_map_segment

    hp.write_map(os.path.join(recon_dir, "sky_map.fits"), sky_map)
    del sky_map
    del sky_map_segment
    shutil.rmtree(os.path.join(recon_dir, map_segment_dir))

    hitmap = np.empty(npix)

    for i in range(size):
        start = i*npix_segment
        stop = (i+1)*npix_segment
        hitmap_segment = np.load(os.path.join(recon_dir, hitmap_segment_dir, str(i).zfill(4) + '.npy')) 
        hitmap[..., start:stop] = hitmap_segment

    hp.write_map(os.path.join(recon_dir, "hitmap.fits"), hitmap)
    del hitmap
    del hitmap_segment
    shutil.rmtree(os.path.join(recon_dir, hitmap_segment_dir))

    inverse_cov_matrix = np.empty((npix, ind_elements))

    for i in range(size):
        start = i*npix_segment
        stop = (i+1)*npix_segment
        inverse_cov_matrix_segment = np.load(os.path.join(recon_dir, inv_cov_matrix_segment_dir, str(i).zfill(4) + '.npy')) 
        inverse_cov_matrix[start:stop] = inverse_cov_matrix_segment

    hp.write_map(os.path.join(recon_dir, "inverse_covariance_maps.fits"), inverse_cov_matrix.T)
    del inverse_cov_matrix
    del inverse_cov_matrix_segment
    shutil.rmtree(os.path.join(recon_dir, inv_cov_matrix_segment_dir))

    cov_matrix = np.empty((npix, ind_elements))

    for i in range(size):
        start = i*npix_segment
        stop = (i+1)*npix_segment
        cov_matrix_segment = np.load(os.path.join(recon_dir, cov_matrix_segment_dir, str(i).zfill(4) + '.npy')) 
        if config.pol_type == "TQU":
            cov_matrix_segment = cov_matrix_segment.reshape((npix_segment, dim**2))[..., np.array([0,1,2,4,5,8])]
        elif config.pol_type =="QU":
            cov_matrix_segment = cov_matrix_segment.reshape((npix_segment, dim**2))[..., np.array([0,1,3])]
        else: 
            cov_matrix_segment = cov_matrix_segment.reshape((npix_segment, dim**2))
        cov_matrix[start:stop] = cov_matrix_segment

    hp.write_map(os.path.join(recon_dir, "covariance_maps.fits"), cov_matrix.T)
    del cov_matrix
    del cov_matrix_segment
    shutil.rmtree(os.path.join(recon_dir, cov_matrix_segment_dir))