def dot(a1, a2): "Returns the dot product of two scalars/vectors of two array of scalars/vectors." if a1 is dsa.NoneArray or a2 is dsa.NoneArray: return dsa.NoneArray if a1.shape[1] != a2.shape[1]: raise RuntimeError, 'Dot product only works with vectors of same dimension.'\ 'Input shapes ' + a1.shape + ' and ' + a2.shape m = a1 * a2 va = dsa.VTKArray(numpy.add.reduce(m, 1)) if a1.DataSet == a2.DataSet: va.DataSet = a1.DataSet return va
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)