In many modern applications, a key bottleneck is the solution of a matrix problem of the form Ax=b where A is a large matrix. In numerical weather prediction, such systems arise as a sub-problem within data assimilation algorithms. In this setting, finding the most likely initial condition with which to initialise a forecast is equivalent to finding the (approximate) solution x.
Due to the enormous size of the problem, we find x via an iterative method where we can only afford a very small number of iterations compared to the size of the problem. It is therefore imperative to design novel preconditioners for the system. Preconditioners allow us to solve an equivalent, but more mathematically tractable problem. In the setting of this case study, we exploit the underlying block structure of the saddle point formulation of the data assimilation problem in a new way to massively accelerate convergence, both in terms of iterations and wallclock times.
Figure 1: Left shows the structure for our matrix A. Currently we assume blocks with the same colour have the same value - apart from the black blocks which are different. Right shows the proposed structure for our preconditioner P - notice it looks almost identical to A but each block has the same value M.
In the saddle point formulation we are interested in, the matrix A has a very special structure (see Figure 1). Here each block is itself a matrix. If all the blocks of the same colour have exactly the same value, we can write A in a Kronecker form. This allows us to rewrite products A with a vector as products of related matrices with tall-skinny matrices - which turns out to be much computationally cheaper (see Figure 2 for an illustration). Unfortunately, the M blocks in A are different. We can’t change A itself (this would mean solving a different linear system), but we can use this idea to design a preconditioner that can be evaluated in a “matrix-oriented” way rather than a “vector-oriented way”.
Figure 2: Shows how a block matrix (where the red blocks are all identical) multiplied with a vector (vector-oriented) can be rewritten as a single block multiplied by the vector rearranged into a matrix (matrix-oriented).
Our new approach uses a preconditioner P where we replace M1, M2, M3 with an average value M (right in Figure 1). We can then rewrite any matrix-vector product Px or P-1x in the matrix-matrix way shown in the bottom of Figure 2. We then replace standard iterative methods (which multiply matrices with vectors) with matrix-oriented iterative methods to solve Ax = b. If we use the same P and A in the vector-oriented or matrix-oriented algorithms we will obtain the same solution, x, in the same number of iterations. However, the computational cost will be much higher for the vector-oriented approach, and using the P shown in Figure 1 in the matrix-vector implementation is not feasible in practice.
Compared to standard matrix-vector methods with other choices of P, we find the matrix-oriented approach with P as shown in Figure 1 yields:
- Improved estimates to the model term (if the Ms in A don’t differ too much). This means P is closer to the original matrix A.
- Dramatic reduction in Iteration counts, as well as computational cost for a serial implementation. This means we can find a solution x in fewer steps, and the overall cost is not a problem for our test case.
The results of this proof-of-concept approach seem very promising! However, there is some work left to do before this will be used in operational weather forecasting centres. Our main challenges are:
- What happens if the other blocks of the system vary (i.e. the red blocks in Figure 1 are different from one another)? These can have different sizes which is a problem for this Kronecker approach.
- We currently ensure competitive computational performance by requiring the full eigenvalue decomposition of M. This is not feasible for operational-scale systems.
We have a few ideas for how to tackle these outstanding issues - stay tuned for the next paper to find out if we manage!
For further details and full numerical experiments:
Palitta, Davide, and Jemima M. Tabeart. "Stein-based preconditioners for weak-constraint 4D-var." Journal of Computational Physics 482 (2023): 112068.
Jemima Tabeart is a Hooke Research Fellow in Oxford Mathematics.