QuantumOptics.jl icon indicating copy to clipboard operation
QuantumOptics.jl copied to clipboard

Get the `DifferentialEquations.jl` problem object

Open Krastanov opened this issue 3 years ago • 16 comments

There are a lot of extra tools that DifferentialEquations.jl provides that I want to be able to use (e.g. sensitivity analysis and optimal control, which I would prefer to do directly, instead of having them wrapped into something specific to QuantumOptics.jl).

Looking at some of the timeevolution objects, it seems those objects are not really abstracted away and easily available to the user. How receptive are the developers of QuantumOptics to introducing:

  • A standard way to get the ODE function (e.g., make dschroedinger from https://github.com/qojulia/QuantumOptics.jl/blob/master/src/schroedinger.jl public, stable, and documented)?
  • A standard way to get the ODE problem object (e.g., explicitly create a problem object in https://github.com/qojulia/QuantumOptics.jl/blob/master/src/schroedinger.jl#L49 and optionally return it, instead of integrating it)

If this is acceptable, I might be able to volunteer making these changes myself. It would drastically expand the problems in which I would be using QuantumOptics.

Krastanov avatar Apr 01 '21 17:04 Krastanov

I really like this idea. I've actually been thinking about this myself for a while now, but never got around to actually doing it. I'm not sure yet what the best way to do this is though.

  • ODEProblems: We could implement a constructor for each of the timeevolution functions, e.g. SchroedingerProblem, which is internally constructed when schroedinger is called, but which is also part of the public API. This gives users easy access to the ODEProblem. One question I don't have an answer to yet are the custom container types (such as Ket) that are used to store the results. Currently, we use a SavingCallback which wraps the vector in a Ket. This callback needs to be passed to solve though, so if you construct a SchroedingerProblem you input a Ket but get a Vector as output if you just call solve on it. Not sure how to deal with this. One option might be to unwrap everything internally so you just end up with pure Julia vectors and matrices. So if you solve with schroedinger you get Kets, but if you construct the SchroedingerProblem you get Vectors. This would also give you as much "freedom" as possible, but you can't use other QuantumOptics functions on the output. I'm not sure how to deal with this yet. Any ideas?

  • ODEFunctions: Depending on how we do ODEProblems this might be already given because you can access the ODEFunction from prob.f. Otherwise I think it's a good idea to make dschroedinger, dmaster, etc. part of the API.

It would be awesome if you could take a stab at this! Let me know if you need help.

david-pl avatar Apr 02 '21 07:04 david-pl

Concerning the wrapping in Ket: that might not be necessary, as the DiffEq library is supposed to work on any object that defines a few simple interfaces. Are there any immediate argument against that?

Krastanov avatar Apr 02 '21 14:04 Krastanov

"Simple interfaces" might have overstated it a bit, but here are a few more details https://discourse.julialang.org/t/juliadiffeq-with-custom-types/16294

Krastanov avatar Apr 02 '21 17:04 Krastanov

Are there any immediate argument against that?

Getting DiffEq to work directly on QO types would be ideal. Not sure how easy it is to do that. It's easy enough thinking about the simple matrix*vector Schrödinger equation, but there are other cases such as semiclassical and also lazy operators such as the FFTOperator which aren't represented by matrices.

david-pl avatar Apr 06 '21 09:04 david-pl

Just for reference, after #306 gets merged you can just do something like this:

using QuantumOptics
using OrdinaryDiffEq

b = SpinBasis(1//2)
ψ0 = spindown(b)
H = DenseOperator(sigmax(b))

const Hdata = H.data
function f(du,u,p,t)
    QuantumOptics.timeevolution.dschroedinger(u, Hdata, du)
    return nothing
end

prob = ODEProblem(f,ψ0.data,(0.0,10.0))
sol = solve(prob,Tsit5())

That will give you an ODEProblem with standard Julia arrays. The upside is that you don't have to worry about performance here. The downsides are that you can't directly work with QO types afterwards (though you can just do something simple like ψt = [Ket(b,u) for u ∈ sol.u]). It also won't work with non-array data types (lazy operators, FFTs), but you can just rewrap the state u as Ket before calling dschroedinger to get around that (that's how it's done internally).

david-pl avatar Jun 18 '21 09:06 david-pl

Would it be possible to change the signature of dschroedinger(u, Data, du) to something more standard, like dscroedinger!(du, Hdata, u)?

PhilipVinc avatar Jun 18 '21 10:06 PhilipVinc

@PhilipVinc certainly. It's not part of the public API yet, so changing it should be fine. Maybe we should add it to the API afterwards. Also, if we go for this option it might make sense to add some convenience function that defines the derivative function needed by DiffEq for you, rewrapping things as Ket if necessary (for non-array types).

david-pl avatar Jun 18 '21 11:06 david-pl

After #312 gets merged, the syntax will be

function f(du,u,p,t)
    timeevolution.dschroedinger!(du, Hdata, u)
    return nothing
end

It's actually rather simple to work with non-data operators as well. You can just use the recast! function (which is also used internally). For example:

b1 = SpinBasis(1//2)
b = b1^2
ψ0 = tensor([spindown(b1) for i=1:2]...)

const Hlazy = LazyTensor(b,[1,2],(sigmax(b1),sigmax(b1)))
const dpsi = Ket(b)
const psi = Ket(b)
function f(du,u,p,t)
    timeevolution.recast!(dpsi,du)
    timeevolution.recast!(psi,u)
    timeevolution.dschroedinger!(dpsi,Hlazy,psi)
    timeevolution.recast!(du,dpsi)
    return nothing
end

prob = ODEProblem(f,ψ0.data,(0.0,10.0))
sol = solve(prob,Tsit5())

That's a pretty solid way to do this as it works with any type QO.jl implements and it will be just as fast as the functions from timeevolution. The only difference to calling timeevolution now is that there is an additional SavingCallback rewrapping the result in a Ket. I'm not sure if you even want that if you want to work on the ODEProblem directly.

The only thing that might still be missing now is a neat way to define such an ODEProblem.

@Krastanov What do you think? I'd actually prefer this approach here as we don't have to worry about performance - as opposed to what we started in https://github.com/qojulia/QuantumOpticsBase.jl/pull/16

david-pl avatar Jul 15 '21 14:07 david-pl

I think it is important for both features to be implemented. Julia is supposed to provide zero cost abstractions and these recast! and dscroedinger! functions, while extremely useful, look more like some C code. In particular, I am writing some code that on the fly decides the best way to simulate something (e.g. jumping from stabilizer formalism, to pseudo-probability stabilizers, to kets and density matrices). I want it to work seamlessly with QuantumOptics and to be natural to write ODEProblems without knowing the internals of QuantumOptics (I claim recast! is not natural and it is a sign of code that would be difficult to interoperate with and that implementation details are leaking out instead of being abstracted).

So, while the example from the last comment seems like a great improvement, it falls short from a truly natural interface like:

b1 = SpinBasis(1//2)
b = b1^2
ψ0 = tensor([spindown(b1) for i=1:2]...)

const Hlazy = LazyTensor(b,[1,2],(sigmax(b1),sigmax(b1)))
const dpsi = Ket(b)
const psi = Ket(b)
function f(du,u,p,t)
    timeevolution.dschroedinger!(du,Hlazy,u)
    return nothing
end

prob = ODEProblem(f,ψ0,(0.0,10.0))
sol = solve(prob,Tsit5())

On the other hand, my complaint can be sidestepped by simply implementing multiple methods for dschroedinger!. Especially given that I have not had time to work on qojulia/QuantumOpticsBase.jl#16

What is your opinion on this suggestion:

  1. implementing the dschroedinger! you mentioned hence needing recast when creating an ODEProblem
  2. implementing a second dschroedinger! method that works on arrays, just as a convenience (probably not used much, but a good intermediary method) hence not needing recast, but still needing ψ0.data in the ODEProblem
  3. finishing the broadcasts so that ψ0.data stops being necessary in ODEProblem and the output is an array of quantum objects instead of naked arrays. That would also enable the use of the built-in saveat callback, instead of the custom one implemented in this library.

The reason I believe point 3 is important, is because it drastically simplifies the ODE code in this library, which makes it much easier to use the advanced features of the ODE solvers. One extremely important example is autodifferentiation: currently it is extremely painful to get autodifferentiation to work with QuantumOptics, which makes optimal control tasks unnecessarily difficult to write (exactly the thing that julia is promising should not be a problem). Another reason is so that we can run QuantumOptics on a GPU. These and many other interoperability features will become possible once the interfaces for point 3 are implemented.

Krastanov avatar Jul 15 '21 17:07 Krastanov

I am excited to see #312 merged. If my point 3 above works, it would just make recast! unnecessary (which would make it easier for autodiff and GPU use and saveat callbacks).

Krastanov avatar Jul 15 '21 17:07 Krastanov

  1. This will work after #312 is merged, so yes I'm all for that.
  2. Type constraints have been removed somewhat, so dschroedinger! will work on any type that implements a mul! method. You'd only need to use recast! if e.g. your Hamiltonian cannot be represented by a matrix (which was the point of my previous example). If your H is a matrix you can just do f(du,u,p,t) = timeevolution.dschroedinger!(du,H,u). Or do you mean implementing a method that works on mixed types such as dschroedinger!(du::Vector,H::Operator,u::Vector)?
  3. Just to be very clear here: I am all for this approach. It would really be the ideal solution. The reason behind not already having it is that all my attempts to get broadcasting as efficient as with pure Julia vectors failed so far.

As a side note: things should already run on GPU. The .data type can be anything, so CuArrays can also be used. Dispatch should take care of the rest so that the appropriate mul! methods are called. Autodiff is another story, though. As you say, this won't just work on QO types.

david-pl avatar Jul 16 '21 11:07 david-pl

Hi!

What is missing for this to be merged? Is there any way I can help with it?

golanor avatar May 12 '22 09:05 golanor

Is there a way to do the same for a time-dependent Hamiltonian? If I run something similar:

using QuantumOptics
using OrdinaryDiffEq

fock = FockBasis(n_photon)
a = destroy(fock); at = create(fock);
n = number(fock); I = identityoperator(fock)
ψ0 = fockstate(fock, 1)
H = TimeDependentSum([1.0, fcomp(-1), fcomp(1)],[omega0*n, -0.1*(a + at), -0.1*(a + at)])

const D = timeevolution.schroedinger_dynamic_function(H)

function f(du,u,p,t)
    timeevolution.dschroedinger_dynamic!(du, D, u, t)
    return nothing
end

prob = ODEProblem(f,ψ0.data,(0.0,10.0))
sol = solve(prob,Tsit5())

I get the MethodError:

MethodError: no method matching mul!(::Vector{ComplexF64}, ::TimeDependentSum{...})

oameye avatar Aug 01 '23 18:08 oameye

@oameye You're missing this part where we wrap the update function in an outer function that "recasts" the raw vector to a Ket.

amilsted avatar Aug 02 '23 00:08 amilsted

I think we may want to consider having our own "Problem" type rather like Bloqade is doing.

That way we can have a solve method that takes care of messy things like recasting.

amilsted avatar Aug 02 '23 00:08 amilsted

@amilsted Thank you! Looking into the source code worked out for me. I needed this to utilise the ability of threading with OrdinaryDiffEq.EnsembleProblem. Ones this feature is implemented, it would be cool to explain the functionanlity in the docs.

oameye avatar Aug 02 '23 10:08 oameye