Skip to content

Commit

Permalink
Added StaggeredNonlinearFEOpearators
Browse files Browse the repository at this point in the history
  • Loading branch information
@JordiManyer
JordiManyer committed Dec 11, 2024
1 parent d3a2e06 commit 6062c1d
Show file tree
Hide file tree
Showing 4 changed files with 227 additions and 105 deletions.
6 changes: 3 additions & 3 deletions src/BlockSolvers/BlockFEOperators.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ end
function BlockArrays.blocks(f::MultiFieldFESpace{<:BlockMultiFieldStyle{NB,SB,P}}) where {NB,SB,P}
block_ranges = MultiField.get_block_ranges(NB,SB,P)
block_spaces = map(block_ranges) do range
(length(range) == 1) ? f[range[1]] : MultiFieldFESpace(f[range])
(length(range) == 1) ? f[range[1]] : MultiFieldFESpace(f.spaces[range])
end
return block_spaces
end
Expand All @@ -104,11 +104,11 @@ function BlockArrays.blocks(f::GridapDistributed.DistributedMultiFieldFESpace{<:
if (length(range) == 1)
space = f[range[1]]
else
global_sf_spaces = f[range]
global_sf_spaces = f.field_fe_space[range]
local_sf_spaces = to_parray_of_arrays(map(local_views,global_sf_spaces))
local_mf_spaces = map(MultiFieldFESpace,local_sf_spaces)
vector_type = GridapDistributed._find_vector_type(local_mf_spaces,gids)
space = MultiFieldFESpace(global_sf_spaces,local_mf_spaces,gids,vector_type)
space = DistributedMultiFieldFESpace(global_sf_spaces,local_mf_spaces,gids,vector_type)
end
space
end
Expand Down
4 changes: 3 additions & 1 deletion src/BlockSolvers/BlockSolvers.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@ using GridapDistributed
using GridapSolvers.MultilevelTools
using GridapSolvers.SolverInterfaces

using Gridap.MultiField: BlockSparseMatrixAssembler

using GridapDistributed: to_parray_of_arrays
using GridapDistributed: DistributedMultiFieldFESpace, DistributedMultiFieldFEFunction

Expand All @@ -33,6 +35,6 @@ export BlockDiagonalSolver
export BlockTriangularSolver

export BlockFEOperator
export StaggeredFEOperator, StaggeredAffineFEOperator, StaggeredFESolver
export StaggeredFEOperator, StaggeredAffineFEOperator, StaggeredNonlinearFEOperator, StaggeredFESolver

end
196 changes: 119 additions & 77 deletions src/BlockSolvers/StaggeredFEOperators.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,20 @@ function MultiField.get_block_ranges(::StaggeredFEOperator{NB,SB}) where {NB,SB}
MultiField.get_block_ranges(NB,SB,Tuple(1:sum(SB)))
end

# TODO: This is type piracy -> move to Gridap
MultiField.num_fields(space::FESpace) = 1

# TODO: We could reuse gids in distributed
function combine_fespaces(spaces::Vector{<:FESpace})
NB = length(spaces)
SB = Tuple(map(num_fields,spaces))
sf_spaces = vcat(map(split_fespace,spaces)...)
MultiFieldFESpace(sf_spaces; style = BlockMultiFieldStyle(NB,SB))
end

split_fespace(space::FESpace) = [space]
split_fespace(space::MultiFieldFESpaceTypes) = [space...]

function get_solution(op::StaggeredFEOperator{NB,SB}, xh::MultiFieldFEFunction, k) where {NB,SB}
r = MultiField.get_block_ranges(op)[k]
if isone(length(r)) # SingleField
Expand Down Expand Up @@ -78,7 +92,6 @@ function Algebra.solve!(xh, solver::StaggeredFESolver{NB}, op::StaggeredFEOperat
xh_k = get_solution(op,xh,k)
op_k = get_operator(op,xhs,k)
xh_k, cache_k = solve!(xh_k,solvers[k],op_k,nothing)
copy!(get_free_dof_values(get_solution(op,xh,k)),get_free_dof_values(xh_k))
xhs, caches, operators = (xhs...,xh_k), (caches...,cache_k), (operators...,op_k)
end
return xh, (caches,operators)
Expand Down Expand Up @@ -127,81 +140,79 @@ struct StaggeredAffineFEOperator{NB,SB} <: StaggeredFEOperator{NB,SB}
trial :: BlockFESpaceTypes{NB,SB}
test :: BlockFESpaceTypes{NB,SB}

@doc """
function StaggeredAffineFEOperator(
biforms :: Vector{<:Function},
liforms :: Vector{<:Function},
trials :: Vector{<:FESpace},
tests :: Vector{<:FESpace},
[assems :: Vector{<:Assembler}]
)
Constructor for a `StaggeredAffineFEOperator` operator, taking in each
equation as a pair of bilinear/linear forms and the corresponding trial/test spaces.
The trial/test spaces can be single or multi-field spaces.
"""
function StaggeredAffineFEOperator(
biforms::Vector{<:Function},
liforms::Vector{<:Function},
trials::Vector{<:FESpace},
tests::Vector{<:FESpace},
assems::Vector{<:Assembler}
biforms :: Vector{<:Function},
liforms :: Vector{<:Function},
trials :: Vector{<:FESpace},
tests :: Vector{<:FESpace},
assems :: Vector{<:Assembler} = map(SparseMatrixAssembler,tests,trials)
)
@assert length(biforms) == length(liforms) == length(trials) == length(tests) == length(assems)
trial = combine_fespaces(trials)
test = combine_fespaces(tests)
NB, SB = length(trials), Tuple(map(num_fields,trials))
new{NB,SB}(biforms,liforms,trials,tests,assems,trial,test)
end
end

"""
@doc """
function StaggeredAffineFEOperator(
biforms::Vector{<:Function},
liforms::Vector{<:Function},
trials::Vector{<:FESpace},
tests::Vector{<:FESpace},
[assems::Vector{<:Assembler}]
)
Constructor for a `StaggeredAffineFEOperator` operator, taking in each
equation as a pair of bilinear/linear forms and the corresponding trial/test spaces.
The trial/test spaces can be single or multi-field spaces.
"""
function StaggeredAffineFEOperator(
biforms::Vector{<:Function},
liforms::Vector{<:Function},
trials::Vector{<:FESpace},
tests::Vector{<:FESpace}
)
assems = map(SparseMatrixAssembler,tests,trials)
return StaggeredAffineFEOperator(biforms,liforms,trials,tests,assems)
biforms :: Vector{<:Function},
liforms :: Vector{<:Function},
trial :: BlockFESpaceTypes{NB,SB,P},
test :: BlockFESpaceTypes{NB,SB,P},
[assem :: BlockSparseMatrixAssembler{NB,NV,SB,P}]
) where {NB,NV,SB,P}
Constructor for a `StaggeredAffineFEOperator` operator, taking in each
equation as a pair of bilinear/linear forms and the global trial/test spaces.
"""
function StaggeredAffineFEOperator(
biforms :: Vector{<:Function},
liforms :: Vector{<:Function},
trial :: BlockFESpaceTypes{NB,SB,P},
test :: BlockFESpaceTypes{NB,SB,P},
assem :: BlockSparseMatrixAssembler{NB,NV,SB,P} = SparseMatrixAssembler(trial,test)
) where {NB,NV,SB,P}
@assert length(biforms) == length(liforms) == NB
@assert P == Tuple(1:sum(SB)) "Permutations not supported"
trials = blocks(trial)
tests = blocks(test)
assems = diag(blocks(assem))
new{NB,SB}(biforms,liforms,trials,tests,assems,trial,test)
end
end

FESpaces.get_trial(op::StaggeredAffineFEOperator) = op.trial
FESpaces.get_test(op::StaggeredAffineFEOperator) = op.test

# Utils

MultiField.num_fields(space::FESpace) = 1

# TODO: We could reuse gids in distributed
function combine_fespaces(spaces::Vector{<:FESpace})
NB = length(spaces)
SB = Tuple(map(num_fields,spaces))

sf_spaces = vcat(map(split_fespace,spaces)...)
MultiFieldFESpace(sf_spaces; style = BlockMultiFieldStyle(NB,SB))
end

split_fespace(space::FESpace) = [space]
split_fespace(space::MultiFieldFESpaceTypes) = [space...]

function get_operator(op::StaggeredFEOperator{NB}, xhs, k) where NB
function get_operator(op::StaggeredAffineFEOperator{NB}, xhs, k) where NB
@assert NB >= k
a(uk,vk) = op.biforms[k](xhs,uk,vk)
l(vk) = op.liforms[k](xhs,vk)
# uhd = zero(op.trials[k])
# A, b = assemble_matrix_and_vector(a,l,op.assems[k],op.trials[k],op.tests[k],uhd)
# return AffineOperator(A,b)
return AffineFEOperator(a,l,op.trials[k],op.tests[k],op.assems[k])
end

function get_operator!(op_k::AffineFEOperator, op::StaggeredFEOperator{NB}, xhs, k) where NB
function get_operator!(op_k::AffineFEOperator, op::StaggeredAffineFEOperator{NB}, xhs, k) where NB
@assert NB >= k
A, b = get_matrix(op_k), get_vector(op_k)
a(uk,vk) = op.biforms[k](xhs,uk,vk)
l(vk) = op.liforms[k](xhs,vk)
uhd = zero(op.trials[k])
assemble_matrix_and_vector!(a,l,A,b,op.assems[k],op.trials[k],op.tests[k],uhd)
return AffineOperator(A,b)
return op_k
end

############################################################################################
Expand All @@ -222,52 +233,83 @@ Such a problem is formulated by a set of residual/jacobian pairs:
"""
struct StaggeredNonlinearFEOperator{NB,SB} <: StaggeredFEOperator{NB,SB}
jacobians :: Vector{<:Function}
residuals :: Vector{<:Function}
jacobians :: Vector{<:Function}
trials :: Vector{<:FESpace}
tests :: Vector{<:FESpace}
assems :: Vector{<:Assembler}
trial :: BlockFESpaceTypes{NB,SB}
test :: BlockFESpaceTypes{NB,SB}

@doc """
function StaggeredNonlinearFEOperator(
res :: Vector{<:Function},
jac :: Vector{<:Function},
trials :: Vector{<:FESpace},
tests :: Vector{<:FESpace}
)
Constructor for a `StaggeredNonlinearFEOperator` operator, taking in each
equation as a pair of residual/jacobian forms and the corresponding trial/test spaces.
The trial/test spaces can be single or multi-field spaces.
"""
function StaggeredNonlinearFEOperator(
res::Vector{<:Function},
jac::Vector{<:Function},
trials::Vector{<:FESpace},
tests::Vector{<:FESpace},
assems::Vector{<:Assembler}
res :: Vector{<:Function},
jac :: Vector{<:Function},
trials :: Vector{<:FESpace},
tests :: Vector{<:FESpace},
assems :: Vector{<:Assembler} = map(SparseMatrixAssembler,tests,trials)
)
@assert length(res) == length(jac) == length(trials) == length(tests) == length(assems)
trial = combine_fespaces(trials)
test = combine_fespaces(tests)
NB, SB = length(trials), Tuple(map(num_fields,trials))
new{NB,SB}(res,jac,trials,tests,assems,trial,test)
end
end

# TODO: Can be compute jacobians from residuals?

"""
@doc """
function StaggeredNonlinearFEOperator(
res::Vector{<:Function},
jac::Vector{<:Function},
trials::Vector{<:FESpace},
tests::Vector{<:FESpace}
)
Constructor for a `StaggeredNonlinearFEOperator` operator, taking in each
equation as a pair of residual/jacobian forms and the corresponding trial/test spaces.
The trial/test spaces can be single or multi-field spaces.
"""
function StaggeredNonlinearFEOperator(
res::Vector{<:Function},
jac::Vector{<:Function},
trials::Vector{<:FESpace},
tests::Vector{<:FESpace}
)
assems = map(SparseMatrixAssembler,tests,trials)
return StaggeredNonlinearFEOperator(res,jac,trials,tests,assems)
res :: Vector{<:Function},
jac :: Vector{<:Function},
trial :: BlockFESpaceTypes{NB,SB,P},
test :: BlockFESpaceTypes{NB,SB,P},
[assem :: BlockSparseMatrixAssembler{NB,NV,SB,P}]
) where {NB,NV,SB,P}
Constructor for a `StaggeredNonlinearFEOperator` operator, taking in each
equation as a pair of bilinear/linear forms and the global trial/test spaces.
"""
function StaggeredNonlinearFEOperator(
res :: Vector{<:Function},
jac :: Vector{<:Function},
trial :: BlockFESpaceTypes{NB,SB,P},
test :: BlockFESpaceTypes{NB,SB,P},
assem :: BlockSparseMatrixAssembler{NB,NV,SB,P} = SparseMatrixAssembler(trial,test)
) where {NB,NV,SB,P}
@assert length(res) == length(jac) == NB
@assert P == Tuple(1:sum(SB)) "Permutations not supported"
trials = blocks(trial)
tests = blocks(test)
assems = diag(blocks(assem))
new{NB,SB}(res,jac,trials,tests,assems,trial,test)
end
end

FESpaces.get_trial(op::StaggeredNonlinearFEOperator) = op.trial
FESpaces.get_test(op::StaggeredNonlinearFEOperator) = op.test

function get_operator(op::StaggeredNonlinearFEOperator{NB}, xhs, k) where NB
@assert NB >= k
jac(uk,duk,dvk) = op.jacobians[k](xhs,uk,duk,dvk)
res(uk,vk) = op.residuals[k](xhs,uk,vk)
return FESpaces.FEOperatorFromWeakForm(res,jac,op.trials[k],op.tests[k],op.assems[k])
end

function get_operator!(
op_k::FESpaces.FEOperatorFromWeakForm, op::StaggeredNonlinearFEOperator{NB}, xhs, k
) where NB
@assert NB >= k
jac(uk,duk,dvk) = op.jacobians[k](xhs,uk,duk,dvk)
res(uk,vk) = op.residuals[k](xhs,uk,vk)
return FESpaces.FEOperatorFromWeakForm(res,jac,op.trials[k],op.tests[k],op.assems[k])
end
Loading

0 comments on commit 6062c1d

Please sign in to comment.