Ah, I see. The problem here is that harmonic_extension and harmonix_extension_trans are local operators. They directly access the locally stored vector values. You can wrap a ParallelMatrix around them to make the code a little more elegant
Code:
hex = ParallelMatrix(harmonic_extension, row_pardods=..., col_pardofs=..., op_type=C2C)
hext = ParallelMatrix(harmonic_extension_trans, row_pardods=..., col_pardofs=..., op_type=D2D)
Now "hex" is an operator that takes a cumulated vector as input and has a cumulated vector as an output. So, when you give it the distributed solution vector from ksp, it automatically cumulates it. You can also write down the entire operation as a single operator:
Code:
hex, hext, aiii = a.harmonic_extension, a.harmonic_extension_trans, a.inner_solve
Id = IdentityMatrix(a.mat.height)
if a.space.mesh.comm.size == 1:
full_precond = ((Id + hex) @ precond @ (Id + hext)) + aiii
else:
Ihex = ParallelMatrix(Id + hex, row_pardofs = a.mat.row_pardofs, col_pardofs = a.mat.row_pardofs, op = ngs.ParallelMatrix.C2C)
Ihext = ParallelMatrix(Id + hext, row_pardofs = self.a.mat.row_pardofs, col_pardofs = self.a.mat.row_pardofs, op = ngs.ParallelMatrix.D2D)
Isolve = ParallelMatrix(aiii, row_pardofs = self.a.mat.row_pardofs, col_pardofs = self.a.mat.row_pardofs, op = ngs.ParallelMatrix.D2C)
full_precond = ( Ihex @ precond @ Ihext ) + aiii
Then you can just write
Code:
gfu.vec.data = full_precond * f.vec
(Nothing wrong with your code, I just find this a little more neat)
I will have a look at the timings.
Best,
Lukas
