def _flatten_nested_1(form_or_block_form, flattened_block_form, block_function_space):
is_zero = _is_zero(form_or_block_form)
assert is_zero or isinstance(form_or_block_form, (array, Form, list))
if is_zero:
pass
elif isinstance(form_or_block_form, Form):
args = _extract_arguments(form_or_block_form)
test_block_index = None
test_block_function_space = None
for arg in args:
assert arg.number() == 0
if test_block_index is not None:
assert test_block_index == arg.block_index(), "Test functions corresponding to different blocks appear in the same form."
assert test_block_function_space == arg.block_function_space(), "Test functions defined in different block function spaces appear in the same form."
else:
test_block_index = arg.block_index()
test_block_function_space = arg.block_function_space()
assert test_block_index is not None
assert test_block_function_space is not None
if hasattr(block_function_space[0], "is_block_subspace"):
assert test_block_index in block_function_space[0].components_to_sub_components, "Block function space and test block index are not consistent on the sub space."
test_block_index = block_function_space[0].components_to_sub_components[test_block_index]
assert test_block_function_space == block_function_space[0].parent_block_function_space
else:
assert test_block_function_space == block_function_space[0]
flattened_block_form[test_block_index] += form_or_block_form
elif isinstance(form_or_block_form, (array, list)):
assert _get_block_form_rank(form_or_block_form) == 1
for block_form_I_nested in form_or_block_form:
_flatten_nested_1(block_form_I_nested, flattened_block_form, block_function_space)
else:
raise AssertionError("Invalid case in _flatten_nested_1")
def _block_form_preprocessing(block_form,
block_function_space=None,
block_form_rank=None):
assert isinstance(block_form, (array, list))
if block_form_rank is None:
block_form_rank = _get_block_form_rank(block_form)
assert block_form_rank is not None, \
"A block form rank should be provided when assemblying a zero block vector/matrix."
assert block_form_rank in (1, 2)
if block_form_rank == 2:
# Extract BlockFunctionSpace from the current form, if required
if not block_function_space:
assert not all([_is_zero(block_form_I_J) for block_form_I in block_form for block_form_I_J in block_form_I]), \
"A BlockFunctionSpace should be provided when assemblying a zero block matrix."
block_function_space = _extract_block_function_space_2(block_form)
assert len(block_function_space) == 2
assert block_function_space[0] is not None
assert block_function_space[1] is not None
block_function_space = [
block_function_space[0], block_function_space[1]
] # convert from dict to list
else:
assert isinstance(block_function_space, list)
assert len(block_function_space) == 2
assert isinstance(block_function_space[0], BlockFunctionSpace)
assert isinstance(block_function_space[1], BlockFunctionSpace)
# Flatten nested blocks, if any
block_form = block_flatten_nested(block_form, block_function_space)
# Return preprocessed data
return (block_form, block_function_space, block_form_rank)
elif block_form_rank == 1:
# Extract BlockFunctionSpace from the current form, if required
if not block_function_space:
assert not all([_is_zero(block_form_I) for block_form_I in block_form]), \
"A BlockFunctionSpace should be provided when assemblying a zero block vector."
block_function_space = _extract_block_function_space_1(block_form)
assert len(block_function_space) == 1
assert block_function_space[0] is not None
block_function_space = [block_function_space[0]
] # convert from dict to list
else:
assert isinstance(block_function_space, BlockFunctionSpace)
block_function_space = [block_function_space]
# Flatten nested blocks, if any
block_form = block_flatten_nested(block_form, block_function_space)
# Return preprocessed data
return (block_form, block_function_space, block_form_rank)
def _block_form_preprocessing(block_form,
block_function_space=None,
block_form_rank=None):
assert isinstance(block_form, (array, list))
if block_form_rank is None:
block_form_rank = _get_block_form_rank(block_form)
assert block_form_rank is not None, \
"A block form rank should be provided when assemblying a zero block vector/matrix."
assert block_form_rank in (1, 2)
if block_form_rank is 2:
# Extract BlockFunctionSpace from the current form, if required
if not block_function_space:
assert not all([_is_zero(block_form_I_J) for block_form_I in block_form for block_form_I_J in block_form_I]), \
"A BlockFunctionSpace should be provided when assemblying a zero block matrix."
block_function_space = _extract_block_function_space_2(block_form)
assert len(block_function_space) == 2
assert block_function_space[0] is not None
assert block_function_space[1] is not None
block_function_space = [
block_function_space[0], block_function_space[1]
] # convert from dict to list
else:
assert isinstance(block_function_space, list)
assert len(block_function_space) is 2
assert isinstance(block_function_space[0], BlockFunctionSpace)
assert isinstance(block_function_space[1], BlockFunctionSpace)
# Flatten nested blocks, if any
block_form = block_flatten_nested(block_form, block_function_space)
# ... and compute size accordingly
if block_function_space[0] == block_function_space[1]:
_assert_flattened_form_2_is_square(block_form)
N = len(block_form)
M = len(block_form[0])
# Replace zero blocks, if any
replaced_block_form = empty((N, M), dtype=object)
for I in range(N):
for J in range(M):
replaced_block_form[I, J] = block_replace_zero(
block_form, (I, J), block_function_space)
# Return preprocessed data
return (replaced_block_form, block_function_space, block_form_rank)
elif block_form_rank is 1:
# Extract BlockFunctionSpace from the current form, if required
if not block_function_space:
assert not all([_is_zero(block_form_I) for block_form_I in block_form]), \
"A BlockFunctionSpace should be provided when assemblying a zero block vector."
block_function_space = _extract_block_function_space_1(block_form)
assert len(block_function_space) == 1
assert block_function_space[0] is not None
block_function_space = [block_function_space[0]
] # convert from dict to list
else:
assert isinstance(block_function_space, BlockFunctionSpace)
block_function_space = [block_function_space]
# Flatten nested blocks, if any
block_form = block_flatten_nested(block_form, block_function_space)
# ... and compute size accordingly
N = len(block_form)
# Replace zero blocks, if any
replaced_block_form = empty((N, ), dtype=object)
for I in range(N):
replaced_block_form[I] = block_replace_zero(
block_form, (I, ), block_function_space)
# Return preprocessed data
return (replaced_block_form, block_function_space, block_form_rank)
