Complex numbers in PyBaMM

[brosaplanella]

Description

At the moment we do not have a class to deal with complex numbers in PyBaMM. Python already offers the functionality so might be easy to bring it in, but not sure if CasADI is able to deal with complex numbers.

Motivation

This would enable including EIS simulations (and data fitting) in the future.

Possible Implementation

No response

Additional context

No response

[brosaplanella]

@pybamm-team/maintainers what do you think?

[martinjrobins]

sadly I don't think Casadi supports complex numbers:

from casadi import *

x = MX.sym("x")
f = x**2 + 10
f = Function('f',[x],\
           [x**2 + 10])
print(f(2))
print(f(2 + 3j))

14
Traceback (most recent call last):
  File "/home/mrobins/git/PyBaMM/test.py", line 8, in <module>
    print(f(2 + 3j))
  File "/home/mrobins/git/PyBaMM/env/lib/python3.10/site-packages/casadi/casadi.py", line 23363, in __call__
    ret = self.call(args)
  File "/home/mrobins/git/PyBaMM/env/lib/python3.10/site-packages/casadi/casadi.py", line 20021, in call
    return _casadi.Function_call(self, *args)
NotImplementedError: Wrong number or type of arguments for overloaded function 'Function_call'.
  Possible prototypes are:
    call(self,dict:DM,bool,bool)
    call(self,[DM],bool,bool)
    call(self,[SX],bool,bool)
    call(self,dict:SX,bool,bool)
    call(self,dict:MX,bool,bool)
    call(self,[MX],bool,bool)
  You have: '(Function,(complex))'

[martinjrobins]

Looks like Jax does tho: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.complex.html

[martinjrobins]

Longer term I think it would be good to move to an MLIR-based backend rather than casadi, this will enable complex numbers and gpu support in conjunction with the more robust sundials solver (rather than the jax solver).

[brosaplanella]

In the EIS case the CasADI issue would be easy to bypass, as the system to solve would be algebraic (it is in the frequency domain) so we could get away with other solvers.