search

flaport / sax

S + Autograd + XLA :: S-parameter based frequency domain circuit simulations and optimizations using JAX.
https://flaport.github.io/sax
Apache License 2.0
70 stars 17 forks source link

Replace the backend for circuits without backscattering #36

Closed zhetaoj closed 1 month ago

zhetaoj commented 5 months ago

I am looking into the possibility of replacing the backend in the case of no backscattering in the circuit components. Instead of calculating inverse of a matrix, the backend will simply need to do a series of matrix multiplication in the forward direction and gives the results at output ports. This can be much faster for complicated circuit where scattering is weak/absent. Is there an easy way to implement this in SAX? Thanks!

flaport commented 5 months ago

Not sure if something like this is easy to add. But feel free to give it a go.

Have you tried the KLU backend? In my experience it's fast enough for any circuit I've tried so far. Not sure if I see a need for it. Although I would probably accept a PR that adds it.

zhetaoj commented 5 months ago

Thank you @flaport! I was using KLU and it takes ~10s seconds for a large MZI network, but it may be hard to optimize or perform sensitivity analysis. I will try to work on a backend to make it faster.

zhetaoj commented 1 month ago

I have submitted a pull request #39 to implement this backend with a dummy example of a cascade MZI network. It provides decent speed-up compared to klu backend especially when the circuit gets complicated @flaport :)

flaport commented 1 month ago

Thanks @zhetaoj ! Looking at the example notebook the speedup is indeed impressive!