def test_solve_mixed(self, mat, dat):
x = op2.MixedDat(dat.dataset)
op2.solve(mat, x, dat)
b = mat * x
eps = 1.e-12
assert_allclose(dat[0].data_ro, b[0].data_ro, eps)
assert_allclose(dat[1].data_ro, b[1].data_ro, eps)
def test_copy_mixed(self, s, mdat):
"""Copy method on a MixedDat should copy values into given target"""
mdat2 = op2.MixedDat([s, s])
mdat.copy(mdat2)
assert all(all(d.data_ro == d_.data_ro) for d, d_ in zip(mdat, mdat2))
for dat in mdat.data:
dat[:] = -1
assert all(all(d.data_ro != d_.data_ro) for d, d_ in zip(mdat, mdat2))
def test_copy_constructor_mixed(self, mdat):
"""MixedDat copy constructor should copy values"""
mdat2 = op2.MixedDat(mdat)
assert mdat.dataset.set == mdat2.dataset.set
assert all(all(d.data_ro == d_.data_ro) for d, d_ in zip(mdat, mdat2))
for dat in mdat.data:
dat[:] = -1
assert all(all(d.data_ro != d_.data_ro) for d, d_ in zip(mdat, mdat2))
def test_norm_mixed(self):
s = op2.Set(1)
n = op2.Dat(s, [3], np.float64)
o = op2.Dat(s, [4], np.float64)
md = op2.MixedDat([n, o])
assert abs(md.norm - 5) < 1e-12
def make_dat(self, val=None, valuetype=None, name=None, uid=None):
"""Return a newly allocated :class:`pyop2.MixedDat` defined on the
:attr:`dof_dset` of this :class:`MixedFunctionSpace`."""
if val is not None:
assert len(val) == len(self)
else:
val = [None for _ in self]
return op2.MixedDat(s.make_dat(v, valuetype, "%s[cmpt-%d]" % (name, i), utils._new_uid())
for i, (s, v) in enumerate(zip(self._spaces, val)))
def test_norm_mixed(self):
s = op2.Set(1)
n = op2.Dat(s, [3], np.complex128)
o = op2.Dat(s, [4j], np.complex128)
md = op2.MixedDat([n, o])
assert type(md.norm) is float
assert abs(md.norm - 5) < 1e-12
def test_inner_mixed(self):
s = op2.Set(1)
n = op2.Dat(s, [3], np.float64)
o = op2.Dat(s, [4], np.float64)
md = op2.MixedDat([n, o])
n1 = op2.Dat(s, [4], np.float64)
o1 = op2.Dat(s, [5], np.float64)
md1 = op2.MixedDat([n1, o1])
ret = md.inner(md1)
assert abs(ret - 32) < 1e-12
ret = md1.inner(md)
assert abs(ret - 32) < 1e-12
def dat(self, mset, mmap, mdat):
dat = op2.MixedDat(mset)
kernel_code = FunDecl("void", "addone_rhs",
[Decl("double", Symbol("v", (3,))),
Decl("double**", c_sym("d"))],
c_for("i", 3, Incr(Symbol("v", ("i")), FlatBlock("d[i][0]"))))
addone = op2.Kernel(kernel_code, "addone_rhs")
op2.par_loop(addone, mmap.iterset,
dat(op2.INC, mmap[op2.i[0]]),
mdat(op2.READ, mmap))
return dat
def dat(self, mset, mmap, mdat):
dat = op2.MixedDat(mset)
kernel_code = FunDecl("void", "addone_rhs",
[Decl("double", Symbol("v", (3,))),
Decl("double", Symbol("d", (3,)))],
c_for("i", 3, Incr(Symbol("v", ("i")), FlatBlock("d[i]"))),
pred=["static"])
addone = op2.Kernel(kernel_code.gencode(), "addone_rhs")
op2.par_loop(addone, mmap.iterset,
dat(op2.INC, mmap),
mdat(op2.READ, mmap))
return dat
def test_CoW_MixedDat_duplicate_original_changes(self, backend, x, y):
md = op2.MixedDat([x, y])
md_dup = md.duplicate()
x += 1
y += 2
for a, b in zip(md, md_dup):
assert not self.same_data(a, b)
assert numpy.allclose(md_dup.data_ro[0], numpy.arange(nelems))
assert numpy.allclose(md_dup.data_ro[1], 0)
assert numpy.allclose(md.data_ro[0], numpy.arange(nelems) + 1)
assert numpy.allclose(md.data_ro[1], 2)
def test_dat_save_and_load(self, tmpdir, d1, s, mdat):
"""The save method should dump Dat and MixedDat values to
the file 'output', and the load method should read back
those same values from the 'output' file. """
output = tmpdir.join('output').strpath
d1.save(output)
d2 = op2.Dat(s)
d2.load(output)
assert (d1.data_ro == d2.data_ro).all()
mdat.save(output)
mdat2 = op2.MixedDat([d1, d1])
mdat2.load(output)
assert all(all(d.data_ro == d_.data_ro) for d, d_ in zip(mdat, mdat2))
def test_mixed_dat_versioning(self, backend, x, y):
md = op2.MixedDat([x, y])
mdv = md._version
x += 1
assert md._version != mdv
mdv1 = md._version
y += 1
assert md._version != mdv1
assert md._version != mdv
mdv2 = md._version
md.zero()
assert md._version == (0, 0)
y += 2
assert md._version != mdv2
assert md._version != mdv1
assert md._version != mdv
assert md._version != (0, 0)
def test_assemble_mixed_rhs_vector(self, mset, mmap, mvdat):
"""Assemble a simple right-hand side over a mixed space and check result."""
dat = op2.MixedDat(mset ** 2)
assembly = Block(
[Incr(Symbol("v", ("i"), ((2, 0),)), FlatBlock("d[i][0]")),
Incr(Symbol("v", ("i"), ((2, 1),)), FlatBlock("d[i][1]"))], open_scope=True)
kernel_code = FunDecl("void", "addone_rhs_vec",
[Decl("double", Symbol("v", (6,))),
Decl("double**", c_sym("d"))],
c_for("i", 3, assembly))
addone = op2.Kernel(kernel_code, "addone_rhs_vec")
op2.par_loop(addone, mmap.iterset,
dat(op2.INC, mmap[op2.i[0]]),
mvdat(op2.READ, mmap))
eps = 1.e-12
exp = np.kron(list(zip([1.0, 4.0, 6.0, 4.0])), np.ones(2))
assert_allclose(dat[0].data_ro, np.kron(list(zip(rdata(3))), np.ones(2)), eps)
assert_allclose(dat[1].data_ro, exp, eps)
def test_mixed_vec_access(self):
s = op2.Set(1)
ms = op2.MixedSet([s, s])
d = op2.MixedDat(ms)
d.data[0][:] = 1.0
d.data[1][:] = 2.0
with d.vec_ro as v:
assert np.allclose(v.array_r, [1.0, 2.0])
d.data[0][:] = 0.0
d.data[0][:] = 0.0
with d.vec_wo as v:
assert np.allclose(v.array_r, [1.0, 2.0])
v.array[:] = 1
assert d.data[0][0] == 1
assert d.data[1][0] == 1
def mdat(d1):
return op2.MixedDat([d1, d1])
def mvdat(mset):
return op2.MixedDat(op2.Dat(s ** 2, list(zip(rdata(s.size), rdata(s.size)))) for s in mset)
def mdat(mset):
return op2.MixedDat(op2.Dat(s, rdata(s.size)) for s in mset)
def split(self, fields):
from firedrake import replace, as_vector, split
from firedrake import NonlinearVariationalProblem as NLVP
fields = tuple(tuple(f) for f in fields)
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
V = F.arguments()[0].function_space()
# Exposition:
# We are going to make a new solution Function on the sub
# mixed space defined by the relevant fields.
# But the form may refer to the rest of the solution
# anyway.
# So we pull it apart and will make a new function on the
# subspace that shares data.
pieces = [us[i].dat for i in field]
if len(pieces) == 1:
val, = pieces
subu = function.Function(V, val=val)
subsplit = (subu, )
else:
val = op2.MixedDat(pieces)
subu = function.Function(V, val=val)
# Split it apart to shove in the form.
subsplit = split(subu)
# Permutation from field indexing to indexing of pieces
field_renumbering = dict([f, i] for i, f in enumerate(field))
vec = []
for i, u in enumerate(us):
if i in field:
# If this is a field we're keeping, get it from
# the new function. Otherwise just point to the
# old data.
u = subsplit[field_renumbering[i]]
if u.ufl_shape == ():
vec.append(u)
else:
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
# So now we have a new representation for the solution
# vector in the old problem. For the fields we're going
# to solve for, it points to a new Function (which wraps
# the original pieces). For the rest, it points to the
# pieces from the original Function.
# IOW, we've reinterpreted our original mixed solution
# function as being made up of some spaces we're still
# solving for, and some spaces that have just become
# coefficients in the new form.
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp, argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
Vbc = bc.function_space()
if Vbc.parent is not None and isinstance(Vbc.parent.ufl_element(), VectorElement):
index = Vbc.parent.index
else:
index = Vbc.index
cmpt = Vbc.component
# TODO: need to test this logic
if index in field:
if len(field) == 1:
W = V
else:
W = V.sub(field_renumbering[index])
if cmpt is not None:
W = W.sub(cmpt)
bcs.append(type(bc)(W,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=Jp,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)
def split(self, fields):
from firedrake import replace, as_vector, split
from firedrake_ts.ts_solver import DAEProblem
from firedrake.bcs import DirichletBC, EquationBC
fields = tuple(tuple(f) for f in fields)
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
V = F.arguments()[0].function_space()
# Exposition:
# We are going to make a new solution Function on the sub
# mixed space defined by the relevant fields.
# But the form may refer to the rest of the solution
# anyway.
# So we pull it apart and will make a new function on the
# subspace that shares data.
pieces = [us[i].dat for i in field]
if len(pieces) == 1:
(val, ) = pieces
subu = function.Function(V, val=val)
subsplit = (subu, )
else:
val = op2.MixedDat(pieces)
subu = function.Function(V, val=val)
# Split it apart to shove in the form.
subsplit = split(subu)
# Permutation from field indexing to indexing of pieces
field_renumbering = dict([f, i] for i, f in enumerate(field))
vec = []
for i, u in enumerate(us):
if i in field:
# If this is a field we're keeping, get it from
# the new function. Otherwise just point to the
# old data.
u = subsplit[field_renumbering[i]]
if u.ufl_shape == ():
vec.append(u)
else:
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
# So now we have a new representation for the solution
# vector in the old problem. For the fields we're going
# to solve for, it points to a new Function (which wraps
# the original pieces). For the rest, it points to the
# pieces from the original Function.
# IOW, we've reinterpreted our original mixed solution
# function as being made up of some spaces we're still
# solving for, and some spaces that have just become
# coefficients in the new form.
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp,
argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
if isinstance(bc, DirichletBC):
bc_temp = bc.reconstruct(
field=field,
V=V,
g=bc.function_arg,
sub_domain=bc.sub_domain,
method=bc.method,
)
elif isinstance(bc, EquationBC):
bc_temp = bc.reconstruct(field, V, subu, u)
if bc_temp is not None:
bcs.append(bc_temp)
new_problem = DAEProblem(
F,
subu,
problem.udot,
problem.tspan,
bcs=bcs,
J=J,
Jp=Jp,
form_compiler_parameters=problem.form_compiler_parameters,
)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(
type(self)(
new_problem,
mat_type=self.mat_type,
pmat_type=self.pmat_type,
appctx=self.appctx,
transfer_manager=self.transfer_manager,
))
return self._splits.setdefault(tuple(fields), splits)
def test_set_diagonal_invalid_dat(self, backend, mat, mset):
dat = op2.MixedDat(mset**4)
with pytest.raises(TypeError):
mat.set_diagonal(dat)
def mdat(mset):
return op2.MixedDat(mset)
def test_copy_mixed_subset_fails(self, s, mdat):
"""Copy method on a MixedDat does not support subsets"""
with pytest.raises(NotImplementedError):
mdat.copy(op2.MixedDat([s, s]), subset=op2.Subset(s, []))
