def _mpi_wrap_discretization_manage_operators(obj_id):
d = mpi.get_object(obj_id)
operators = {k: mpi.manage_object(v) if v else None for k, v in sorted(d.operators.iteritems())}
functionals = {k: mpi.manage_object(v) if v else None for k, v in sorted(d.functionals.iteritems())}
vectors = {k: mpi.manage_object(v) if v else None for k, v in sorted(d.vector_operators.iteritems())}
products = {k: mpi.manage_object(v) if v else None for k, v in sorted(d.products.iteritems())} if d.products else None
if mpi.rank0:
return operators, functionals, vectors, products
def _mpi_wrap_operator_VectorArrayOperator_manage_array(obj_id):
op = mpi.get_object(obj_id)
array_obj_id = mpi.manage_object(op._array)
subtypes = mpi.comm.gather(op._array.subtype, root=0)
mpi.remove_object(obj_id)
if mpi.rank0:
return array_obj_id, tuple(subtypes)
def process_attribute(v):
if isinstance(v, OperatorInterface):
mpi_range = type(v.range) in mpi_spaces or v.range.id in mpi_spaces
mpi_source = type(v.source) in mpi_spaces or v.source.id in mpi_spaces
if mpi_range or mpi_source:
return _OperatorToWrap(mpi.manage_object(v), mpi_range, mpi_source)
else:
return v
else:
return v
def _mpi_wrap_operator_VectorArrayOperator_manage_array(obj_id, pickle_subtypes):
op = mpi.get_object(obj_id)
array_obj_id = mpi.manage_object(op._array)
subtype = op._array.subtype
if not pickle_subtypes:
subtype = _register_subtype(subtype)
subtypes = mpi.comm.gather(subtype, root=0)
mpi.remove_object(obj_id)
if mpi.rank0:
return array_obj_id, tuple(subtypes)
def _mpi_wrap_operator_VectorArrayOperator_manage_array(obj_id, pickle_local_spaces):
op = mpi.get_object(obj_id)
array_obj_id = mpi.manage_object(op._array)
local_space = op._array.space
if not pickle_local_spaces:
local_space = _register_local_space(local_space)
local_spaces = mpi.comm.gather(local_space, root=0)
mpi.remove_object(obj_id)
if mpi.rank0:
return array_obj_id, tuple(local_spaces)
def _mpi_wrap_operator_VectorArrayOperator_manage_array(
obj_id, pickle_local_spaces):
op = mpi.get_object(obj_id)
array_obj_id = mpi.manage_object(op.array)
local_space = op.array.space
if not pickle_local_spaces:
local_space = _register_local_space(local_space)
local_spaces = mpi.comm.gather(local_space, root=0)
mpi.remove_object(obj_id)
if mpi.rank0:
return array_obj_id, tuple(local_spaces)
def _mpi_wrap_operator_VectorArrayOperator_manage_array(
obj_id, pickle_subtypes):
op = mpi.get_object(obj_id)
array_obj_id = mpi.manage_object(op._array)
subtype = op._array.subtype
if not pickle_subtypes:
subtype = _register_subtype(subtype)
subtypes = mpi.comm.gather(subtype, root=0)
mpi.remove_object(obj_id)
if mpi.rank0:
return array_obj_id, tuple(subtypes)
def process_attribute(v):
if isinstance(v, OperatorInterface):
mpi_range = type(v.range) in mpi_spaces or v.range.id in mpi_spaces
mpi_source = type(
v.source) in mpi_spaces or v.source.id in mpi_spaces
if mpi_range or mpi_source:
return _OperatorToWrap(mpi.manage_object(v), mpi_range,
mpi_source)
else:
return v
else:
return v
def _MPIVectorArray_make_array(cls, subtype=(None, ), count=0, reserve=0):
subtype = subtype[mpi.rank] if len(subtype) > 1 else subtype[0]
if type(subtype) is RegisteredSubtype:
subtype = _subtype_registry[subtype]
obj = cls.make_array(subtype=subtype, count=count, reserve=reserve)
return mpi.manage_object(obj)
def _random_array(dims, length, seed):
np.random.seed(seed + mpi.rank)
dim = dims[mpi.rank] if len(dims) > 1 else dims[0]
array = NumpyVectorArray(np.random.random((length, dim)))
obj_id = mpi.manage_object(array)
return obj_id
def _random_vector(dims, seed):
np.random.seed(seed + mpi.rank)
dim = dims[mpi.rank] if len(dims) > 1 else dims[0]
vector = NumpyVector(np.random.random(dim), copy=False)
obj_id = mpi.manage_object(vector)
return obj_id
def _MPIVector_make_zeros(cls, subtype=(None,)):
subtype = subtype[mpi.rank] if len(subtype) > 1 else subtype[0]
if type(subtype) is RegisteredSubtype:
subtype = _subtype_registry[subtype]
obj = cls.make_zeros(subtype)
return mpi.manage_object(obj)
def _mpi_wrap_operator_LincombOperator_manage_operators(obj_id):
op = mpi.get_object(obj_id)
obj_ids = [mpi.manage_object(o) for o in op.operators]
mpi.remove_object(obj_id)
if mpi.rank0:
return obj_ids
def _MPIVector_make_zeros(cls, subtype=(None,)):
subtype = subtype[mpi.rank] if len(subtype) > 1 else subtype[0]
obj = cls.make_zeros(subtype)
return mpi.manage_object(obj)
def _MPIVectorArray_make_array(cls, subtype=(None,), count=0, reserve=0):
subtype = subtype[mpi.rank] if len(subtype) > 1 else subtype[0]
if type(subtype) is RegisteredSubtype:
subtype = _subtype_registry[subtype]
obj = cls.make_array(subtype=subtype, count=count, reserve=reserve)
return mpi.manage_object(obj)
def _MPIVectorArray_make_array(cls, subtype=(None,), count=0, reserve=0):
subtype = subtype[mpi.rank] if len(subtype) > 1 else subtype[0]
obj = cls.make_array(subtype=subtype, count=count, reserve=reserve)
return mpi.manage_object(obj)
def _MPIVectorSpace_zeros(local_spaces=(None, ), count=0, reserve=0):
local_space = _get_local_space(local_spaces)
obj = local_space.zeros(count=count, reserve=reserve)
return mpi.manage_object(obj)
def _MPIVector_make_zeros(cls, subtype=(None, )):
subtype = subtype[mpi.rank] if len(subtype) > 1 else subtype[0]
if type(subtype) is RegisteredSubtype:
subtype = _subtype_registry[subtype]
obj = cls.make_zeros(subtype)
return mpi.manage_object(obj)
def _MPIVectorSpace_zeros(local_spaces=(None,), count=0, reserve=0):
local_space = _get_local_space(local_spaces)
obj = local_space.zeros(count=count, reserve=reserve)
return mpi.manage_object(obj)
def _MPIOperator_assemble_lincomb(operators, coefficients, name):
operators = [mpi.get_object(op) for op in operators]
return mpi.manage_object(operators[0].assemble_lincomb(operators,
coefficients,
name=name))
def _random_vector(dims, seed):
np.random.seed(seed + mpi.rank)
dim = dims[mpi.rank] if len(dims) > 1 else dims[0]
vector = NumpyVector(np.random.random(dim), copy=False)
obj_id = mpi.manage_object(vector)
return obj_id
def _MPIOperator_assemble_lincomb(operators, coefficients, name):
operators = [mpi.get_object(op) for op in operators]
return mpi.manage_object(operators[0].assemble_lincomb(operators, coefficients, name=name))
def _random_array(dims, length, seed):
np.random.seed(seed + mpi.rank)
dim = dims[mpi.rank] if len(dims) > 1 else dims[0]
array = NumpyVectorSpace.make_array(np.random.random((length, dim)))
obj_id = mpi.manage_object(array)
return obj_id
def _mpi_wrap_operator_LincombOperator_manage_operators(obj_id):
op = mpi.get_object(obj_id)
obj_ids = [mpi.manage_object(o) for o in op.operators]
mpi.remove_object(obj_id)
if mpi.rank0:
return obj_ids