def test_c_size(self):
dim_dict = {'dist_type': 'c',
'size': 42,
'proc_grid_size': 2,
'proc_grid_rank': 0,
'start': 0}
dist = Distribution.from_global_dim_data(self.context, (dim_dict,))
ddpr = dist.get_dim_data_per_rank()
shapes = metadata_utils.shapes_from_dim_data_per_rank(ddpr)
self.assertEqual(shapes, [(21,), (21,)])
def create_distribution_plot_and_documentation(context, params):
"""Create an array distribution plot and the related .rst documentation."""
def shape_text(shape):
""" Get a text string describing the array shape. """
# Always want to display at least N X M.
if len(shape) == 1:
shape = (1, shape[0])
shape_labels = ['%d' % (s) for s in shape]
shape_text = ' X '.join(shape_labels)
return shape_text
title = params['title']
labels = params['labels']
shape = params['shape']
grid_shape = params.get('grid_shape', None)
text = params.get('text', None)
dist = params.get('dist', None)
dimdata = params.get('dimdata', None)
filename = params.get('filename', None)
skip = params.get('skip', False)
if skip:
return
# Create array, either from dist or dimdata.
if dist is not None:
distribution = Distribution(context,
shape,
dist=dist,
grid_shape=grid_shape)
elif dimdata is not None:
distribution = Distribution.from_global_dim_data(context, dimdata)
else:
raise ValueError('Must provide either dist or dimdata.')
array = context.empty(distribution)
# Fill the array. This is slow but not a real problem here.
value = 0.0
if len(shape) == 1:
for i in range(shape[0]):
array[i] = value
value += 1.0
elif len(shape) == 2:
for row in range(shape[0]):
for col in range(shape[1]):
array[row, col] = value
value += 1.0
elif len(shape) == 3:
for i in range(shape[0]):
for j in range(shape[1]):
for k in range(shape[2]):
array[i, j, k] = value
value += 1.0
else:
# TODO: Even better would be to generalize this to any dimensions.
raise ValueError('Array must be 1, 2, or 3 dimensional.')
# Get all process grid coordinates.
# This is duplicating work in print_array_documentation(),
# but it is needed for the local array plots.
def _get_process_coords(local_arr):
return local_arr.cart_coords
process_coords = context.apply(_get_process_coords, (array.key, ),
targets=array.targets)
# Plot title and axis labels.
plot_title = title + ' ' + shape_text(shape) + '\n'
if len(shape) == 1:
# add more space for cramped plot.
plot_title += '\n'
xlabel = 'Axis 1, %s' % (labels[1])
ylabel = 'Axis 0, %s' % (labels[0])
# Documentation title and text description.
doc_title = title
dist_text = ' X '.join(["'%s'" % (label) for label in labels])
# Choose 'a' vs 'an' appropriately.
if title[0] in 'aeiouAEIOU':
article = 'an'
else:
article = 'a'
doc_text = 'A (%s) array, with %s %s (%s) distribution over a (%s) process grid.' % (
shape_text(shape), article, title, dist_text,
shape_text(array.grid_shape))
if text is not None:
doc_text = doc_text + "\n\n" + text
# Filenames for array plots.
global_plot_filename = filename
local_plot_filename = None
if global_plot_filename is not None:
root, ext = os.path.splitext(global_plot_filename)
local_plot_filename = root + '_local' + ext
# Create plot.
if len(shape) in [1, 2]:
plotting.plot_array_distribution(
array,
process_coords,
title=plot_title,
xlabel=xlabel,
ylabel=ylabel,
legend=True,
global_plot_filename=global_plot_filename,
local_plot_filename=local_plot_filename)
else:
# Not plottable, avoid writing links to missing plots.
global_plot_filename = None
local_plot_filename = None
# Print documentation.
print_array_documentation(context,
array,
title=doc_title,
text=doc_text,
global_plot_filename=global_plot_filename,
local_plot_filename=local_plot_filename)
def create_distribution_plot_and_documentation(context, params):
"""Create an array distribution plot and the related .rst documentation."""
def shape_text(shape):
""" Get a text string describing the array shape. """
# Always want to display at least N X M.
if len(shape) == 1:
shape = (1, shape[0])
shape_labels = ['%d' % (s) for s in shape]
shape_text = ' X '.join(shape_labels)
return shape_text
title = params['title']
labels = params['labels']
shape = params['shape']
grid_shape = params.get('grid_shape', None)
text = params.get('text', None)
dist = params.get('dist', None)
dimdata = params.get('dimdata', None)
filename = params.get('filename', None)
skip = params.get('skip', False)
if skip:
return
# Create array, either from dist or dimdata.
if dist is not None:
distribution = Distribution(context, shape, dist=dist,
grid_shape=grid_shape)
elif dimdata is not None:
distribution = Distribution.from_global_dim_data(context, dimdata)
else:
raise ValueError('Must provide either dist or dimdata.')
array = context.empty(distribution)
# Fill the array. This is slow but not a real problem here.
value = 0.0
if len(shape) == 1:
for i in range(shape[0]):
array[i] = value
value += 1.0
elif len(shape) == 2:
for row in range(shape[0]):
for col in range(shape[1]):
array[row, col] = value
value += 1.0
elif len(shape) == 3:
for i in range(shape[0]):
for j in range(shape[1]):
for k in range(shape[2]):
array[i, j, k] = value
value += 1.0
else:
# TODO: Even better would be to generalize this to any dimensions.
raise ValueError('Array must be 1, 2, or 3 dimensional.')
# Get all process grid coordinates.
# This is duplicating work in print_array_documentation(),
# but it is needed for the local array plots.
def _get_process_coords(local_arr):
return local_arr.cart_coords
process_coords = context.apply(_get_process_coords,
(array.key,),
targets=array.targets)
# Plot title and axis labels.
plot_title = title + ' ' + shape_text(shape) + '\n'
if len(shape) == 1:
# add more space for cramped plot.
plot_title += '\n'
xlabel = 'Axis 1, %s' % (labels[1])
ylabel = 'Axis 0, %s' % (labels[0])
# Documentation title and text description.
doc_title = title
dist_text = ' X '.join(["'%s'" % (label) for label in labels])
# Choose 'a' vs 'an' appropriately.
if title[0] in 'aeiouAEIOU':
article = 'an'
else:
article = 'a'
doc_text = 'A (%s) array, with %s %s (%s) distribution over a (%s) process grid.' % (
shape_text(shape), article, title, dist_text, shape_text(array.grid_shape))
if text is not None:
doc_text = doc_text + "\n\n" + text
# Filenames for array plots.
global_plot_filename = filename
local_plot_filename = None
if global_plot_filename is not None:
root, ext = os.path.splitext(global_plot_filename)
local_plot_filename = root + '_local' + ext
# Create plot.
if len(shape) in [1, 2]:
plotting.plot_array_distribution(
array,
process_coords,
title=plot_title,
xlabel=xlabel,
ylabel=ylabel,
legend=True,
global_plot_filename=global_plot_filename,
local_plot_filename=local_plot_filename)
else:
# Not plottable, avoid writing links to missing plots.
global_plot_filename = None
local_plot_filename = None
# Print documentation.
print_array_documentation(
context,
array,
title=doc_title,
text=doc_text,
global_plot_filename=global_plot_filename,
local_plot_filename=local_plot_filename)