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)