def global_flat_indices(dim_data):
"""
Return a list of tuples of indices into the flattened global array.
Parameters
----------
dim_data: dimension dictionary.
Returns
-------
list of 2-tuples of ints.
Each tuple is a (start, stop) interval into the flattened global array.
All selected ranges comprise the indices for this dim_data's sub-array.
"""
# TODO: FIXME: can be optimized when the last dimension is 'n'.
for dd in dim_data:
if dd['dist_type'] == 'n':
dd['start'] = 0
dd['stop'] = dd['size']
glb_shape = tuple(dd['size'] for dd in dim_data)
glb_strides = strides_from_shape(glb_shape)
ranges = [range(dd['start'], dd['stop']) for dd in dim_data[:-1]]
start_ranges = ranges + [[dim_data[-1]['start']]]
stop_ranges = ranges + [[dim_data[-1]['stop']]]
def flatten(idx):
return sum(a * b for (a, b) in zip(idx, glb_strides))
starts = map(flatten, product(*start_ranges))
stops = map(flatten, product(*stop_ranges))
intervals = zip(starts, stops)
return condense(intervals)
def global_flat_indices(dim_data):
"""
Return a list of tuples of indices into the flattened global array.
Parameters
----------
dim_data: dimension dictionary.
Returns
-------
list of 2-tuples of ints.
Each tuple is a (start, stop) interval into the flattened global array.
All selected ranges comprise the indices for this dim_data's sub-array.
"""
# TODO: FIXME: can be optimized when the last dimension is 'n'.
for dd in dim_data:
if dd['dist_type'] == 'n':
dd['start'] = 0
dd['stop'] = dd['size']
glb_shape = tuple(dd['size'] for dd in dim_data)
glb_strides = strides_from_shape(glb_shape)
ranges = [range(dd['start'], dd['stop']) for dd in dim_data[:-1]]
start_ranges = ranges + [[dim_data[-1]['start']]]
stop_ranges = ranges + [[dim_data[-1]['stop']]]
def flatten(idx):
return sum(a * b for (a, b) in zip(idx, glb_strides))
starts = map(flatten, product(*start_ranges))
stops = map(flatten, product(*stop_ranges))
intervals = zip(starts, stops)
return condense(intervals)
def _massage_indices(local_distribution, glb_intervals):
# XXX: TODO: document why we do `-1)+1` below.
local_flat_slices = [(_transform(local_distribution, i[0]),
_transform(local_distribution, i[1] - 1) + 1)
for i in glb_intervals]
return condense(local_flat_slices)
def _massage_indices(local_distribution, glb_intervals):
# XXX: TODO: document why we do `-1)+1` below.
local_flat_slices = [(_transform(local_distribution, i[0]),
_transform(local_distribution, i[1]-1)+1)
for i in glb_intervals]
return condense(local_flat_slices)
