def apply_dfunc(dfunc, array1, val2):
"""Apply a two argument function to each member of a VTKCompositeDataArray
and another argument The second argument can be a VTKCompositeDataArray, in
which case a one-to-one match between arrays is assumed. Otherwise, the
function is applied to the composite array with the second argument repeated.
VTKArray and numpy arrays are also supported."""
if type(array1) == dsa.VTKCompositeDataArray and type(val2) == dsa.VTKCompositeDataArray:
res = []
for a1, a2 in itertools.izip(array1.Arrays, val2.Arrays):
if a1 is dsa.NoneArray or a2 is dsa.NoneArray:
res.append(dsa.NoneArray)
else:
l = dsa.reshape_append_ones(a1, a2)
res.append(dfunc(l[0], l[1]))
return dsa.VTKCompositeDataArray(res, dataset = array1.DataSet)
elif type(array1) == dsa.VTKCompositeDataArray:
res = []
for a in array1.Arrays :
if a is dsa.NoneArray:
res.append(dsa.NoneArray)
else:
l = dsa.reshape_append_ones(a, val2)
res.append(dfunc(l[0], l[1]))
return dsa.VTKCompositeDataArray(res, dataset = array1.DataSet)
elif array1 is dsa.NoneArray:
return dsa.NoneArray
else:
l = dsa.reshape_append_ones(array1, val2)
return dfunc(l[0], l[1])
def new_dsfunc2(ds):
if type(ds) == dsa.CompositeDataSet:
res = []
for dataset in ds:
res.append(dsfunc(dataset))
return dsa.VTKCompositeDataArray(res, dataset = ds)
else:
return dsfunc(ds)
def apply_ufunc(func, array, args=()):
"""Apply a function to each member of a VTKCompositeDataArray.
VTKArray and numpy arrays are also supported."""
if array is dsa.NoneArray:
return dsa.NoneArray
elif type(array) == dsa.VTKCompositeDataArray:
return dsa.VTKCompositeDataArray(_apply_func2(func, array, args), dataset = array.DataSet)
else:
return func(array)
def new_dsfunc(array, ds=None):
if type(array) == dsa.VTKCompositeDataArray:
res = []
for a in array.Arrays:
if a is dsa.NoneArray:
res.append(dsa.NoneArray)
else:
res.append(dsfunc(a, ds))
return dsa.VTKCompositeDataArray(res, dataset = array.DataSet)
elif array is dsa.NoneArray:
return dsa.NoneArray
else:
return dsfunc(array, ds)
def make_vector(arrayx, arrayy, arrayz=None):
"""Given 2 or 3 scalar arrays, returns a vector array. If only
2 scalars are provided, the third component will be set to 0."""
if type(arrayx) == dsa.VTKCompositeDataArray and type(arrayy) == dsa.VTKCompositeDataArray and (type(arrayz) == dsa.VTKCompositeDataArray or arrayz is None):
res = []
if arrayz is None:
for ax, ay in itertools.izip(arrayx.Arrays, arrayy.Arrays):
if ax is not dsa.NoneArray and ay is not dsa.NoneArray:
res.append(algs.make_vector(ax, ay))
else:
res.append(dsa.NoneArray)
else:
for ax, ay, az in itertools.izip(arrayx.Arrays, arrayy.Arrays, arrayz.Arrays):
if ax is not dsa.NoneArray and ay is not dsa.NoneArray and az is not dsa.NoneArray:
res.append(algs.make_vector(ax, ay, az))
else:
res.append(dsa.NoneArray)
return dsa.VTKCompositeDataArray(res, dataset = arrayx.DataSet)
else:
return algs.make_vector(arrayx, arrayy, arrayz)
def _global_per_block(impl, array, axis=None, controller=None):
if axis > 0:
return impl.op()(array, axis=axis, controller=controller)
try:
dataset = array.DataSet
except AttributeError:
dataset = None
t = type(array)
if t == dsa.VTKArray or t == numpy.ndarray:
from vtk.vtkCommonDataModel import vtkMultiBlockDataSet
array = dsa.VTKCompositeDataArray([array])
ds = vtkMultiBlockDataSet()
ds.SetBlock(0, dataset.VTKObject)
dataset = ds
results = _apply_func2(impl.op2(), array, (axis,))
if controller is None and vtkMultiProcessController is not None:
controller = vtkMultiProcessController.GetGlobalController()
if controller and controller.IsA("vtkMPIController"):
from mpi4py import MPI
comm = vtkMPI4PyCommunicator.ConvertToPython(controller.GetCommunicator())
# First determine the number of components to use
# for reduction
res = dsa.NoneArray
for res in results:
if res is not dsa.NoneArray:
break
max_dims, size = _reduce_dims(res, comm)
# All NoneArrays
if size == 0:
return dsa.NoneArray;
# Next determine the max id to use for reduction
# operations
# Get all ids from dataset, including empty ones.
ids = []
lmax_id = numpy.int32(0)
if dataset is not None:
it = dataset.NewIterator()
it.UnRegister(None)
it.SetSkipEmptyNodes(False)
while not it.IsDoneWithTraversal():
_id = it.GetCurrentFlatIndex()
lmax_id = numpy.max((lmax_id, _id)).astype(numpy.int32)
if it.GetCurrentDataObject() is not None:
ids.append(_id)
it.GoToNextItem()
max_id = numpy.array(0, dtype=numpy.int32)
mpitype = _lookup_mpi_type(numpy.int32)
comm.Allreduce([lmax_id, mpitype], [max_id, mpitype], MPI.MAX)
has_ids = numpy.zeros(max_id+1, dtype=numpy.int32)
for _id in ids:
has_ids[_id] = 1
id_count = numpy.array(has_ids)
comm.Allreduce([has_ids, mpitype], [id_count, mpitype], MPI.SUM)
if numpy.all(id_count <= 1):
return dsa.VTKCompositeDataArray(results, dataset=dataset)
# Now that we know which blocks are shared by more than
# 1 rank. The ones that have a count of 2 or more.
reduce_ids = []
for _id in xrange(len(id_count)):
if id_count[_id] > 1:
reduce_ids.append(_id)
to_reduce = len(reduce_ids)
# If not block is shared, short circuit. No need to
# communicate any more.
if to_reduce == 0:
return dsa.VTKCompositeDataArray(results, dataset=dataset)
# Create the local array that will be used for
# reduction. Set it to a value that won't effect
# the reduction.
lresults = numpy.empty(size*to_reduce)
lresults.fill(impl.default())
# Just get non-empty ids. Doing this again in case
# the traversal above results in a different order.
# We need the same order since we'll use izip below.
if dataset is not None:
it = dataset.NewIterator()
it.UnRegister(None)
ids = []
while not it.IsDoneWithTraversal():
ids.append(it.GetCurrentFlatIndex())
it.GoToNextItem()
# Fill the local array with available values.
for _id, _res in itertools.izip(ids, results):
success = True
try:
loc = reduce_ids.index(_id)
except ValueError:
success = False
if success:
if _res is not dsa.NoneArray:
lresults[loc*size:(loc+1)*size] = _res.flatten()
# Now do the MPI reduction.
rresults = numpy.array(lresults)
mpitype = _lookup_mpi_type(numpy.double)
comm.Allreduce([lresults, mpitype], [rresults, mpitype], impl.mpi_op())
if array is dsa.NoneArray:
return dsa.NoneArray
# Fill in the reduced values.
for i in xrange(to_reduce):
_id = reduce_ids[i]
success = True
try:
loc = ids.index(_id)
except ValueError:
success = False
if success:
if size == 1:
results[loc] = dsa.VTKArray(rresults[i])
else:
results[loc] = rresults[i*size:(i+1)*size].reshape(max_dims)
return dsa.VTKCompositeDataArray(results, dataset=dataset)
