# Path: modred/parallel.py
# def get_hostname():
# def get_node_ID():
# def get_num_nodes():
# def get_num_MPI_workers():
# def get_rank():
# def get_num_procs():
# def is_distributed():
# def is_rank_zero():
# def barrier():
# def print_from_rank_zero(msg, output_channel='stdout'):
# def call_from_rank_zero(func, *args, **kwargs):
# def bcast(vals):
# def call_and_bcast(func, *args, **kwargs):
# def find_assignments(tasks, task_weights=None):
# def check_for_empty_tasks(task_assignments):


import os
import numpy as np
import modred as mr
from modred import parallel




# Create directory for output files
out_dir = 'tutorial_ex4_out'
if not os.path.isdir(out_dir):
    os.makedirs(out_dir)

# Define non-uniform grid and corresponding inner product weights
nx = 80
ny = 100
x_grid = 1. - np.cos(np.linspace(0, np.pi, nx))
y_grid = np.linspace(0, 1., ny) ** 2
Y, X = np.meshgrid(y_grid, x_grid)

# Create random snapshot data
num_vecs = 100
snapshots = [
    mr.VecHandlePickle('%s/vec%d.pkl' % (out_dir, i))
    for i in mr.range(num_vecs)]
if parallel.is_rank_zero():
    for i, snap in enumerate(snapshots):
        snap.put(np.sin(X * i) + np.cos(Y * i))

# Path: modred/parallel.py
# def get_hostname():
# def get_node_ID():
# def get_num_nodes():
# def get_num_MPI_workers():
# def get_rank():
# def get_num_procs():
# def is_distributed():
# def is_rank_zero():
# def barrier():
# def print_from_rank_zero(msg, output_channel='stdout'):
# def call_from_rank_zero(func, *args, **kwargs):
# def bcast(vals):
# def call_and_bcast(func, *args, **kwargs):
# def find_assignments(tasks, task_weights=None):
# def check_for_empty_tasks(task_assignments):


# Path: modred/examples/tutorial_ex4.py
import os
import numpy as np
import modred as mr
from modred import parallel





# Create directory for output files
out_dir = 'tutorial_ex4_out'
if not os.path.isdir(out_dir):
    os.makedirs(out_dir)

# Define non-uniform grid and corresponding inner product weights
nx = 80
ny = 100
x_grid = 1. - np.cos(np.linspace(0, np.pi, nx))
y_grid = np.linspace(0, 1., ny) ** 2
Y, X = np.meshgrid(y_grid, x_grid)

# Create random snapshot data
num_vecs = 100
snapshots = [
    mr.VecHandlePickle('%s/vec%d.pkl' % (out_dir, i))
    for i in mr.range(num_vecs)]
if parallel.is_rank_zero():
    for i, snap in enumerate(snapshots):
        snap.put(np.sin(X * i) + np.cos(Y * i))

parallel.barrier()