Matrix-free linear algebra formulations

[valentinsulzer]

Investigate using matrix-free forms for linear algebra. For example, y[2:]-2*y[1:-1]+y[:-2] instead of A@y where A is the [1,-2,1] matrix. Ideally, we could construct the matrix-free form as a post-processing step at the end of discretisation, only for the cases Sparse matrix @ state vector

Early testing suggests this could speed to the solvers.

[martinjrobins]

this would work well, be particularly good for the jax evaluator. Would it go in the simplification code? i.e. https://github.com/pybamm-team/PyBaMM/blob/52e2d964ac7fcff26da459f8b07a1d7b8eaf41c2/pybamm/expression_tree/binary_operators.py#L1245

[valentinsulzer]

I think it should be done after the model has been fully discretized. Having the intermediate matrix-vector form is useful for simplifications. For example, the grad(u) becomes B@y then later div(grad(u)) becomes A@(B@y) which we simplify to (A@B)@y where A@B is now a constant matrix. If we had already rewritten B@y in matrix-free form then it would be harder to figure out what A@B@y should be in matrix-free form.

[martinjrobins]

some issues with the 1st and last elements of the calculation, for the example above the first row is -2*y[0]+y[1] and the last is -2y[-2] + y[-1], so have three expressions you need to concatenate. Going to try this and see if we still get a speedup even with concatenations

[martinjrobins]

sadly I'm not seeing any speed-ups here. @tinosulzer did you want to have a look at this and see if what I implemented is what you were thinking of?, the relevent test is: python tests/unit/test_discretisations/test_discretisation.py TestDiscretise.test_matrix_free_simplification

[valentinsulzer]

Closing for now since it doesn't seem like this works