solving for Implicit operators

[simonbyrne]

We need a way to handle part of the operators implicitly. To reduce the scope of this problem, initially we will only consider:

• linear operators
• on a single field (or where the same stencil is applied to multiple fields, e.g. advection)
• we will initially require users to construct these manually

Rough plan:

1. Add a MatrixOperator: this would wrap an AbstractMatrix object (typically a Bidiagonal or Tridiagonal, could expand to BandedMatrices.jl objects for more complicated stencils).
2. Define the "forward" stencil (stencil_interior, stencil_left_boundary, etc.): mainly so that we can check it is correct.
3. Add a FactorizedMatrixOperator: this would wrap the result of factorize of the underlying matrix.
4. Define ldiv!:
   • for scalar-valued fields, just pull out the vector and call that
   • for struct-valued fields we would need to reshape the underlying field into the appropriate matrix
5. Figure out how to integrate with the implicit solvers in OrdinaryDiffEq (I think you should be able to pass it as a jac argument)

Once we have this in place, we can tackle the more ambitious stuff:

• automatically constructing the matrix
• linearising non-linear operators
• handling multiple fields at once
• incorporating into non-linear solvers
  some of this I think we can do using autodiff.

[valeriabarra]

Sounds like a good plan.
For point 1, we might also want to consider a BlockedMatrixOperator for performance purposes

[simonbyrne]

From Tapio:

What we need is the following:

1. Solve for linear vertically propagating sound and gravity waves, which means solving (linearly) for density, u_3, and energy (or potential temperature). Start with a pedestrian linear solver that solves for these 3 variables.
2. A more efficient solver solves for just one variable, which can be u_3 or pressure (Schur complement). This means solving a second-order equation that we need to derive by hand. Other variables then follow by backsubstitution. Toby has worked on the PDE-level version of this equation. There’s more discussion on equations for vertical velocity in the COSMO documentation, section 3.3.2. Important for your planning is that this involves solution of a separate PDE for fast waves, not just solution for specific terms in the original equations.
3. Once 1. and 2. are in place, we can use the linear solution as a pre-conditioner for a fully implicit nonlinear solver. Experience so far is that AD here is not faster or better than finite differences for Jacobians, which are the standard way this is handled.

[simonbyrne]

The decision was made today to reduce the priority of this in favor of CG work

[simonbyrne]

After some discussion with @bischtob, I think the way to think of it is as follows:

1. The timestepper ultimately wants a way to solve a problem of the form: "find Y such that Y - γ fᵢ(Y) = Y₀", where fᵢ is the implicit part.
   • if fᵢ is linear, then this is a linear problem (I - γ L) Y = Y₀.
   • This roughly corresponds to Wfact option for DiffEqFunction objects
2. The "naive" way to do this is to form L directly: we can pass this as a jac argument, or factorize it and pass I - γ L as Wfact.
3. The Schur complement / pressure solve requires introducing an additional variable p, which satisfies a wave equation (I - γ^2 G) p = γ * h(p₀, ...), and then the remainder of the variables can be computed as Y = Y₀ + γ K p. In this case, Wfact should be an object which
   • forms the tridiagonal matrix of the operator (I - γ^2 G) (can be done in advance)
   • defines an ldiv! method which :
     a. solves for p.
     b. computes Y from p.

[valeriabarra]

After some discussion with @bischtob, I think the way to think of it is as follows:

1. The timestepper ultimately wants a way to solve a problem of the form: "find Y such that Y - γ fᵢ(Y) = Y₀", where fᵢ is the implicit part.

Hum... I think this form does not really follow the canonical form. Is Y₀ meant to be an Initial Condition here? Are we considering the explicit terms in this formulation? I thought what we used previously in our tutorials (see eq. (3) here) and design docs was put differently

[simonbyrne]

Is Y₀ meant to be an Initial Condition here? Are we considering the explicit terms in this formulation?

No, it is going to be the current state + explicit part (e.g. explicit RK stages). I was just trying to write it in a general-ish form to clarify what the timestepper actually requires.

[simonbyrne]

See some old notes Thomas and I wrote here: https://clima.github.io/TimeMachine.jl/dev/background/ark/

[jakebolewski]

X-Ref: CliMA/ClimateMachine.jl#2255 (comment)

[simonbyrne]

OrdinaryDiffEq.jl v6 has changed its interface such that our existing approach no longer works: see this comment: CliMA/ClimateMachine.jl#2255 (comment)

From what I understand, the main changes:

1. We no longer use Wfact: instead we should use jac (and jac_prototype), which should return a custom DiffEqLinearOperator
2. We need an update_coefficients function
3. We need to specify a custom linear solver

@dennisYatunin can you expand on how we do all that?

[dennisYatunin]

As far as I currently understand the OrdinaryDiffEq interface, we should do the following:

1. Our custom W representation should be made into a subtype of AbstractDiffEqLinearOperator. This will be passed to ODEFunction through the jac_prototype argument. I think the resulting object will then be wrapped by the WOperator object, though there might be a way to avoid that.
2. In order to satisfy the AbstractDiffEqOperator interface, we must also define a method for update_coefficients! that applies to our new operator. This should be the same as our old (::CustomWRepresentation)(...) function, and it gets called whenever the Jacobian is updated; e.g., on every Newton iteration.
3. In principle, we could stop here, in which case I think that the linear solver for our new operator would default to GMRES. In order to ensure that we use the Schur complement solve for which our new operator is specialized, we must define a new solver. I initially thought I understood this part when I read some older documentation. However, the documentation has since changed, and I need to think more carefully about this last step (as well as the previous steps).

[simonbyrne]

@ChrisRackauckas Can you give a few pointers on how this should work? For example, suppose that we have:

# our internal representation of our Jacobian operator
struct CustomJacobian
 ...
end
# function which updates the CustomJacobian at the current state
function jupdate!(J::CustomJacobian, u, p, t)
  ...
end
# function which solves (I - alpha * J) * x = b
function wsolve!(x, J::CustomJacobian, alpha::Real, b) 
  ...
end

how should we incorporate this into the new framework?

[ChrisRackauckas]

As a starting point, I would highly recommend looking at the new large-scale PDE tutorial. That will give the vocabulary of what the new structure is like and we can use that as a starting point.

https://diffeq.sciml.ai/stable/tutorials/advanced_ode_example/

Our custom W representation should be made into a subtype of AbstractDiffEqLinearOperator. This will be passed to ODEFunction through the jac_prototype argument. I think the resulting object will then be wrapped by the WOperator object, though there might be a way to avoid that.

I'm not sure what a custom W representation would do. Essentially it should just hold M, gamma, J, and W_transform for whether a given algorithm understands the setup as gamma*M - J or M - J/gamma (this is dependent on the ODE solver. That's all that WOperator does. It holds those lazy, defines a conversion to concrete via convert(AbstractMatrix,W) which algorithms can use if they need it, and defines mul!. Factorization algorithms call convert, Jacobian-free algorithms don't. In the preconditioners, you can see that something like incomplete-LU preconditioners so the user takes the WOperator and calls convert on it, while matrix-free preconditioners like a geometric multigrid would not need to (and that would then make it never concretize). Note that WOperator also has a jacvec operation defined on it so that matrix multiplication is faster than matrix generation, using either the finite difference or forward single chunk J*v trick. With this trick, I don't think there's ever a reason for a user to supply a Jacobian action, but you can with a jac_prototype being a DiffEqOperator.

Details about W will get added to https://diffeq.sciml.ai/stable/features/linear_nonlinear/ shortly, as we now expect users to interact with it in the preconditioner interface so we need to define the fields and actions.

Add a MatrixOperator: this would wrap an AbstractMatrix object (typically a Bidiagonal or Tridiagonal, could expand to BandedMatrices.jl objects for more complicated stencils). Add a FactorizedMatrixOperator: this would wrap the result of factorize of the underlying matrix.

What you're describing here is equivalent to DiffEqArrayOperator, though with the lazy W you probably don't need to define a DiffEqArrayOperator, as if you are using a factorization method it will just use the array, and if you are using a Krylov method it will just not make the array and instead use the jacvec, and only in the preconditioner it will have the WOperator which will then only concretize if and when you choose.

Our custom W representation should be made into a subtype of AbstractDiffEqLinearOperator. This will be passed to ODEFunction through the jac_prototype argument. I think the resulting object will then be wrapped by the WOperator object, though there might be a way to avoid that.

Mostly correct. Instead of a custom W representation, you just define a custom J representation, so it knows what to do with the Jacobian. If that's a DiffEqOperator, then it keeps the pieces separate in the WOperator anyways, and you just need update_coefficients! for it to know how to update in the Newton-Krylov iterations.

In principle, we could stop here, in which case I think that the linear solver for our new operator would default to GMRES. In order to ensure that we use the Schur complement solve for which our new operator is specialized, we must define a new solver. I initially thought I understood this part when I read some older documentation. However, the documentation has since changed, and I need to think more carefully about this last step (as well as the previous steps).

Note that Pardiso's solvers have a Schur complement solve setup with MKL goodies, so you might want to look into that? It's not enabled with the LinearSolve.jl's defaults, but I saw it in the Pardiso docs so it shouldn't be hard to expose. It should just be the flip of a few options.

how should we incorporate this into the new framework?

So yeah, onto the action items for @simonbyrne. The steps would be to:

1. Create a DiffEqOperator following the interface https://diffeq.sciml.ai/stable/features/diffeq_operator/ and pass that as the jac_prototype. Then DiffEq will internally use your mul!, update_coefficients!, and convert operations internally.
2. OrdinaryDiffEq.jl uses the SciML LinearProblem interface for specifying linear solves. So in theory you just need to define:

linsolvecache = SciMLBase.init(prob::LinearProblem,alg::YourAlgorithm)
linearsolution = SciMLBase.solve(linearcache::YourLinearCacheType;kwargs...)

and the caching API interface functions (http://linearsolve.sciml.ai/dev/basics/CachingAPI/). But, given that most of the shuttling like "don't refactor" and "don't redo symbolic factorization" is all very common between algorithms, I think doing a "from scratch" implementation of a linear solve is probably not the right way to do it. Instead, we should write some documentation in LinearSolve.jl for how to easily add new linear solver algorithms using the internal machinery. I'll go write that right now. Essentially you provide two functions:

1. init_cacheval for how to init algorithm-specific caches.
2. a solve dispatch.

For example, the KLU.jl wrapper is quite clear: https://github.com/SciML/LinearSolve.jl/blob/main/src/factorization.jl#L97-L99 and https://github.com/SciML/LinearSolve.jl/blob/main/src/factorization.jl#L101-L130. That's all, and that will make it cache symbolic factorizations, factorizations, etc. So let me get that written up.

We can jump on a call afterwards to discuss it a bit more. You're the first to define outside linear solvers on the new interface so in theory it all works but in practice I may need to make a few things less bumpy.

[vpuri3]

looks like a lot of the functionality OP seeks is available in SciMLOperators.jl

[vpuri3]

take a look at how I'm doing BVPs in code: https://github.com/vpuri3/PDEInterfaces.jl/tree/master/src/BoundaryConditions

In principle with implicit solvers, you are solving a BVP at every time step. I'm still cleaning up the BC interface but i think you can get away with passing in a BC respecting rank sufficient SciMLOperator to the implicit part of your SplitODEProblem. everything should just work :))