GridapDistributed.jl/src/TransientFESpaces.jl at a39a49a7dd013eabce781224c58338ef2fb1c8e3 · oriolcg/GridapDistributed.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
# Functions for transient FE spaces

Fields.evaluate(U::DistributedSingleFieldFESpace,t::Nothing) = U

(U::DistributedSingleFieldFESpace)(t) = U

∂t(U::DistributedSingleFieldFESpace) = HomogeneousTrialFESpace(U)
∂t(U::DistributedMultiFieldFESpace) = MultiFieldFESpace(∂t.(U.field_fe_space))
∂tt(U::DistributedSingleFieldFESpace) = HomogeneousTrialFESpace(U)
∂tt(U::DistributedMultiFieldFESpace) = MultiFieldFESpace(∂tt.(U.field_fe_spaces))

function TransientMultiFieldFESpace(spaces::Vector{<:DistributedSingleFieldFESpace})
  MultiFieldFESpace(spaces)
end

# Functions for transient FE Functions

function ODETools.allocate_jacobian(
  op::TransientFETools.TransientFEOperatorFromWeakForm,
  duh::DistributedCellField,
  cache)
  _matdata_jacobians = TransientFETools.fill_initial_jacobians(op,duh)
  matdata = _vcat_distributed_matdata(_matdata_jacobians)
  allocate_matrix(op.assem_t,matdata)
end

function ODETools.allocate_jacobian(
  op::TransientFETools.TransientFEOperatorFromWeakForm,
  duh::DistributedMultiFieldFEFunction,
  cache)
  _matdata_jacobians = TransientFETools.fill_initial_jacobians(op,duh)
  matdata = _vcat_distributed_matdata(_matdata_jacobians)
  allocate_matrix(op.assem_t,matdata)
end

function ODETools.jacobians!(
  A::AbstractMatrix,
  op::TransientFETools.TransientFEOperatorFromWeakForm,
  t::Real,
  xh::TransientDistributedCellField,
  γ::Tuple{Vararg{Real}},
  cache)
  _matdata_jacobians = TransientFETools.fill_jacobians(op,t,xh,γ)
  matdata = _vcat_distributed_matdata(_matdata_jacobians)
  assemble_matrix_add!(A,op.assem_t, matdata)
  A
end

function _vcat_distributed_matdata(_matdata)
  term_to_cellmat = map_parts(a->a[1],local_views(_matdata[1]))
  term_to_cellidsrows = map_parts(a->a[2],local_views(_matdata[1]))
  term_to_cellidscols = map_parts(a->a[3],local_views(_matdata[1]))
  for j in 2:length(_matdata)
    term_to_cellmat_j = map_parts(a->a[1],local_views(_matdata[j]))
    term_to_cellidsrows_j = map_parts(a->a[2],local_views(_matdata[j]))
    term_to_cellidscols_j = map_parts(a->a[3],local_views(_matdata[j]))
    term_to_cellmat = map_parts((a,b)->vcat(a,b),local_views(term_to_cellmat),local_views(term_to_cellmat_j))
    term_to_cellidsrows = map_parts((a,b)->vcat(a,b),local_views(term_to_cellidsrows),local_views(term_to_cellidsrows_j))
    term_to_cellidscols = map_parts((a,b)->vcat(a,b),local_views(term_to_cellidscols),local_views(term_to_cellidscols_j))
  end
  map_parts( (a,b,c) -> (a,b,c),
    local_views(term_to_cellmat),
    local_views(term_to_cellidsrows),
    local_views(term_to_cellidscols)
  )
end