Fix higher dimensional FunctionOperator size

[ChrisRackauckas]

If a function operator takes in an MxNxO tensor and spits out an MxNxO tensor, then it's an MNO x MNO operator, not an MxM operator. It can be interpreted in the vec(input) -> vec(output) sense and thus length is the right measure to use. This fixes downstream use cases like:

using OrdinaryDiffEq, LinearAlgebra, Symbolics, LinearSolve

const N = 32
const xyd_brusselator = range(0,stop=1,length=N)
brusselator_f(x, y, t) = (((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.
limit(a, N) = a == N+1 ? 1 : a == 0 ? N : a
function brusselator_2d_loop(du, u, p, t)
  A, B, alpha, dx = p
  alpha = alpha/dx^2
  @inbounds for I in CartesianIndices((N, N))
    i, j = Tuple(I)
    x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
    ip1, im1, jp1, jm1 = limit(i+1, N), limit(i-1, N), limit(j+1, N), limit(j-1, N)
    du[i,j,1] = alpha*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1]) +
                B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
    du[i,j,2] = alpha*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2]) +
                A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
    end
end
p = (3.4, 1., 10., step(xyd_brusselator))

function init_brusselator_2d(xyd)
  N = length(xyd)
  u = zeros(N, N, 2)
  for I in CartesianIndices((N, N))
    x = xyd[I[1]]
    y = xyd[I[2]]
    u[I,1] = 22*(y*(1-y))^(3/2)
    u[I,2] = 27*(x*(1-x))^(3/2)
  end
  u
end
u0 = init_brusselator_2d(xyd_brusselator)
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,u0,(0.,11.5),p)

using Symbolics
du0 = copy(u0)
jac_sparsity = Symbolics.jacobian_sparsity((du,u)->brusselator_2d_loop(du,u,p,0.0),du0,u0)

f = ODEFunction(brusselator_2d_loop;jac_prototype=float.(jac_sparsity))
prob_ode_brusselator_2d_sparse = ODEProblem(f,u0,(0.,11.5),p)

sol1 = solve(prob_ode_brusselator_2d,KenCarp47(linsolve=LUFactorization()),save_everystep=false, reltol=1e-12, abstol=1e-12)
sol2 = solve(prob_ode_brusselator_2d,KenCarp47(linsolve=KrylovJL_GMRES(), concrete_jac = true),save_everystep=false, reltol=1e-12, abstol=1e-12)
sol3 = solve(prob_ode_brusselator_2d,KenCarp47(linsolve=KrylovJL_GMRES()),save_everystep=false, reltol=1e-12, abstol=1e-12)
@test norm(sol1[end] - sol2[end]) < 0.0002
@test norm(sol1[end] - sol3[end]) < 0.0002
@test norm(sol2[end] - sol3[end]) == 0.0

[codecov[bot]]

Codecov Report

Merging #198 (8b8f195) into master (24609db) will decrease coverage by 61.11%. The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #198       +/-   ##
==========================================
- Coverage   68.51%   7.40%   -61.11%     
==========================================
  Files          11      11               
  Lines        1515    1512        -3     
==========================================
- Hits         1038     112      -926     
- Misses        477    1400      +923     

Impacted Files	Coverage Δ
src/func.jl	44.75% <100.00%> (-35.92%)	:arrow_down:

... and 9 files with indirect coverage changes

:mega: We’re building smart automated test selection to slash your CI/CD build times. Learn more

[vpuri3]

input/output can be (N, K) matrices, where the operator is being applied coulmn-wise

This is the correct implementation of brusselator https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/test/interface/preconditioners.jl#L8-L52

In OrdinaryDiffEq, if u0 is a matrix, each column is a separate ODE that is being evolved. For that reason, the brusselator u0 should be vec'd before being passed to OrdinaryDiffEq

https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/test/interface/preconditioners.jl#L51

Inside the ODEFunction, the you can reshape u, du to be matrices

https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/test/interface/preconditioners.jl#L17-L18

[vpuri3]

with that in mind, a SciMLOperator with size (M, N) has been designed to act on AbstractVecOrMats with size(u, 1) = N. That way, we can take advantage of BLAS in fitting K batches in an (N, K) matrix. That's why size is set to (size(output, 1), size(input, 1)), because the second dimension refers to the batch.

[vpuri3]

Tests in the above Brusselator example are passing for me locally with master branch of ODE, SciMLOperators with changes mentioned in the above comment

[ChrisRackauckas]

This is the correct implementation of brusselator https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/test/interface/preconditioners.jl#L8-L52

No it's not. It should allow for a higher dimensional tensor. Did you change that test? I don't remember it being vecd and it shouldn't be. That test is supposed to be finding this bug.

In OrdinaryDiffEq, if u0 is a matrix, each column is a separate ODE that is being evolved.

No it's not. That can be true by the way that the ODE is defined, but it can also be a matrix ODE. You only have separate ODEs if the Jacobian has a block diagonal sparsity pattern, which is not guaranteed.

Inside the ODEFunction, the you can reshape u, du to be matrices

You shouldn't need to, otherwise lots of code would be broken.

with that in mind, a SciMLOperator with size (M, N) has been designed to act on AbstractVecOrMats with size(u, 1) = N. That way, we can take advantage of BLAS in fitting K batches in an (N, K) matrix. That's why size is set to (size(output, 1), size(input, 1)), because the second dimension refers to the batch.

If that's the case then we'd need to be diligent about vec and reshaping inside of the OrdinaryDiffEq code, otherwise many PDE tutorials will break. A lot of PDE codes will have issues right now because of this, so we need to fast track the bug fix.

[vpuri3]

Ok, thanks explaining. If ODE.jl design allows for matrix ODEs and AbstractArray{N} ODEs, we need a way to differentiate the batch dimension. Batch dimension needs to be determined only for FunctionOperator, as it is equal to 2 for all others, eg. MatrixOperator applies the same matrix to u[:, i] for all i, but the operation is fast thanks to BLAS.

I'll add a keyword argument to FunctionOperator for batch_dim and internally do size computation involving dims 1:batch_dim-1.

Being able to apply SciMLOperators on (N, K) matrices, where K is batch size, is pretty fundamental to scimlops. Not just because BLAS mul!, * are much faster than broadcasting, but also because the lazy tensor product operator internally reshapes AbstractVectors to AbstractMatrix for fast evaluation.

[vpuri3]

@ChrisRackauckas LMK if that makes sense before I follow up with a PR

[ChrisRackauckas]

I'll add a keyword argument to FunctionOperator for batch_dim and internally do size computation involving dims 1:batch_dim-1.

That makes sense. Default to not batching but allow for it.