Open navidcy opened 2 years ago
Not yet, but that's a good idea. If you have a few example matrices you could share, I could take a look at the best way to cover a variety of cases. One thing to keep in mind is that geometric multi-grid is only appropriate for matrices from structured grids. So there is no point trying to deal with general matrices. For that, you want algebraic multi-grid.
Let me give a bit more background. There is a discussion here: https://github.com/CliMA/Oceananigans.jl/pull/2396 In particular, we want to apply multigrid solver to the equation mentioned at https://github.com/CliMA/Oceananigans.jl/pull/2396#issuecomment-1115850783.
We don't have the matrix A
in hand but rather we have a function that returns the operation of the matrix after applied to x
. But then we can get the matrix by applying all the unit vectors to this function. (See discussion a bit above the comment.)
And thus I was wondering: with the matrix A in hand how do we then proceed calling, e.g., st = mg_state(A, zero(x), b)
? We need to convert A
into a FieldMatrix
. That's where we got bit stuck..
Ah. That's why you said linear operator instead of matrix. So you were hoping to do something "matrix-free" like the IterativeSovlers.jl library. While it would be easy to use a linear map on any giving level, I don't think a linear map is sufficient to construct the restricted linear map one level down. Again, that's the whole idea of a FieldMatrix
- it embeds the topology of the problem to make this restriction as simple as possible.
From the link, it looks like @glwagner has the code to construct the explicit matrix coefficients. So the only issue is filling the FieldMatrix.L
array with his Ax,Ay
and FieldMatrix.D
with his diag
. Two complications I'm seeing are
Thanks for this response! The idea was to construct the sparse matrix itself and then not have to reside to "matrix-free" iterative solvers.
Let me have a look and get back to you!
I don't think a linear map is sufficient to construct the restricted linear map one level down.
Can we envision using a linear operator at the outermost level, but then constructing matrices that represent the restricted linear map at lower levels? For some problems that might produce more efficient code and save memory.
@navidcy Note that FieldMatrix
is an extremely sparse representation. It's basically a wrapper around the "small" arrays Gregory defined in that example. It's even more memory efficient than a general sparse matrix representation since the memory layout is fixed.
@glwagner On a CPU, the inner pressure loop is typically compute restricted, not memory restricted. Maybe on a GPU, this balance would shift and computing the matrix on the fly would be worthwhile. Another potential advantage of the linear operator approach would be the ability to handle a wider variety of matrix layouts with very general code.
However, it would be very slow to define the operator at level-n as A_n=A_0 R^n
where A_0
is the top-level linear operator and R
is a restriction operator. Even on a GPU, I can't imagine that being performant. I'm sure it would be possible, for certain simple cases, to define a multi-level linear operator from the start, ie A(n)
and this could potentially be faster in some cases.
On a CPU, the inner pressure loop is typically compute restricted, not memory restricted. Maybe on a GPU, this balance would shift and computing the matrix on the fly would be worthwhile.
Hmm, yes this makes sense. I agree there's little motivation from a "memory savings" standpoint to have such a capability. But maybe my logic was a little cloudy... here are more detailed thoughts on function-based vs. matrix-based operators:
For some of the problem we work on, it's easier to work with function-based linear operators, because in the context of solving PDEs, typical software comes built-in with powerful methods for building function representations of linear operators. This is the case for Oceananigans; and we coded a function-based linear operator to represent the "left hand side" of a particular implicit equation (in this case, it's an equation representing the evolution of the free surface in a hydrostatic model). We can easily code boundary conditions into these operators too (like periodic boundary conditions).
There's a whole hierarchy of models one might consider for free surface evolution (linear, nonlinear) and also different coordinate systems (ie coordinate systems that move with the flow). Since we already need to build infrastructure to build function-based operators for these cases, it's natural to use those same tools to build a function-based linear operator for the left hand side of the free surface equation. Not only is this often "easy" (eg one line in our case), it creates the possibility that our code is more generic. So, if we implement a new coordinate system (for example), or implement a more accurate immersed boundary representation of a complex domain, we can reuse the same function for the linear operator. We don't have to re-derive (by hand?) a matrix based representation when we change the fundamental numerics or coordinate system. (But maybe there are ways to build functions that return matrix representations in a similarly generic way? Not sure.)
But this means that to use matrix-based solvers we then need a way to bridge the divide between the "function" representation of a linear operator and the matrix representation.
As @navidcy mentioned, one way to achieve this is to compute the matrix representation by evaluating the function operator for a set of unit vectors. The problem is this requires N^2 evaluations of the operator, which is expensive even for a preprocessing step (and may not be feasible for very very large problems?)
So I made that propose (without understanding any of the underlying details) thinking that, perhaps, there is a way to "only" compute matrix representations for the lower levels (which have lower resolution, and is therefore cheaper?) while "keeping" the function-based representation at the outermost level. But perhaps not (maybe this depends on what "matrix properties" of the function-based representation need to be computed...)
And I guess the final note is that for a conjugate gradient iteration we don't need the matrix representation, since the algorithm requires only that we compute the action of a linear operator on some vector (leading eg to a "matrix-free" algorithm).
So, a capability for constructing multigrid preconditioners given function-based linear operators would be "very convenient" for accelerating matrix-free conjugate gradient methods --- I guess.
But I will definitely defer to your judgement @weymouth because I am not the expert!
I agree that having the capability to swap out the underlying function as needed without ever worrying about matrix implementation problems is a big advantage. This is my favorite part of a linear-operator approach.
I'm not seeing the appeal of this mixed approach since I don't see how it helps you avoid the N^2
cost of constructing the matrix. You need to construct it before you restrict it. If you have a way of restricting the operator without forming the matrix first, then we don't need the mixed approach. We can just run a V-cycle using the operator for each level.
If you have a way of restricting the operator without forming the matrix first, then we don't need the mixed approach.
Ah ok, that makes sense! We won't end up in a scenario where we can restrict a function operator, but "only once"...
Where is restriction defined for FieldMatrix
?
Huh, ok. So this is a non-local operation. We can certainly write this as a function (eg just a(I, i)
), but unless we do something "special" it will not be performant.
To make this performant, one has to write this in a way that can be "inlined" by the compiler. There's a few techniques for that. One "easy" way is that if you can write this with matrix algebra, then you can use StaticArrays
(basically, tuples that look like arrays, with matrix algebra, which allows the compiler to "unroll" small matrix operations, with some cost paid at compile time). There are other ways it can be achieved. Most likely, it requires wrapping the linear operation in an object with a static type parameter that indicates to the compiler the size of up(I, i)
(which is what we need to unroll that loop).
I don't understand the index notation / convention being used so forgive me if I muck this up (we are using j, i
rather than i, j
?). But if we want to start down this road I would create a utility with this pseudocode:
function restricted_operation(i, j, parent_operation, args...)
a_Kj = 0 # the a[K, j] component of a
for k in up(K, i)
a_Kj += parent_operation(k, j, args...)
end
return a_Kj
end
This is composable, so one can write
double_restriction(i, j, parent_operation, args...) = restricted_operation(i, j, child_operation, restricted_operation, args...)
for double refinement.
Then performance optimization focuses on child_operation
.
I'll look at this more carefully tomorrow.
For now, I'll try to clear up notation. I
is the CartesianIndex
on level n+1
and i
is the face/direction index. (So Ax
in your previous example is a[:,1]
, Ay
is a[:,2]
etc.) J ∈ up(I,i)
are the set of CartesianIndices
on level n
which contribute to the matrix coefficient. The up
function is defined just a little further up: https://github.com/weymouth/GeometricMultigrid.jl/blob/4e09502c86a0da1941544a784bc709acb8737052/src/MultiGrid.jl#L26
It appears that the Poisson equation is hard coded to be the linear operator https://github.com/weymouth/GeometricMultigrid.jl/blob/4e09502c86a0da1941544a784bc709acb8737052/src/MultiGrid.jl#L44
Is this for simplicity? Could we instead use the restriction and interpolation matrix technique detailed here so we can work with arbitrary matrices?
Yes. As I said in the "readme" the current solver is for Poisson only, but this is only because I don't (personally) have another use case.
The handout you have linked to shows one restriction operator, but the options are infinite. I've implemented one such operator in the code, but it's specific to the Poisson equation. If you supply the discrete matrix A_0
, then it's easy to get A_n = A_0 R^n
for each level. But again, R
is not unique! My restriction operator for the Poisson eq ensures the data structure matches on every level, since this let me write efficient code. However, it won't be true for a general A_0
and R
that all A_n
have the same data structure, so if you wanted to go this way, I guess we would need to switch to a general SPA representation.
Linking this to the discussion above - this data-structure stuff would all be avoided if we used a linear-operator approach, but I'm not sure how to get the operator at level n
from the top-level operator alone. Clearly, this would be an advantage of the operator-approach if we could figure it out.
I can't seem to figure out how I can construct a
FieldMatrix
if I have in hand the matrix (or sparsematrix) versions of operatorA
. Is there a way to do that?cc @Elise-palethorpe