InexactError when using complex matrices

[seadra]

The following code fails with an error

using DifferentialEquations, DiffEqFlux, Zygote, SciMLSensitivity, Optimization, OptimizationFlux, OptimizationOptimJL, ComponentArrays, Lux, Random, LinearAlgebra

const T = 10.0;
const ω = π/T;

const id = Matrix{Complex{Float64}}(I,2, 2);
const u0 = id;


const utarget = Matrix{Complex{Float64}}([im 0; 0 -im]);

ann = Lux.Chain(Lux.Dense(1,32), Lux.Dense(32,32,tanh), Lux.Dense(32,1));
rng = Random.default_rng();
ip, st = Lux.setup(rng, ann);

function f_nn(u, p, t)
    local a, _ = ann([t/T],p,st);
    local A = [sin(a[1]) 0.0; 0.0 -a[1]];
    return -(im*A)*u;
end



tspan = (0.0, T)
prob_ode = ODEProblem(f_nn, u0, tspan, ComponentArray(ip));


function loss_adjoint(p)
    local prediction = solve(prob_ode, BS5(), p=p, abstol=1e-7, reltol=1e-7)
    local usol = last(prediction)
    local loss = abs(1.0 - abs(tr(usol*utarget')/2))
    return loss
end

opt_f = Optimization.OptimizationFunction((x, p) -> loss_adjoint(x), Optimization.AutoZygote());
opt_prob = Optimization.OptimizationProblem(opt_f, ComponentArray(ip));
optimized_sol_nn = Optimization.solve(opt_prob, AMSGrad(0.001), maxiters = 100, progress=true);

The error is

┌ Warning: Potential performance improvement omitted. ZygoteVJP tried and failed in the automated AD choice algorithm. To show the stack trace, set SciMLSensitivity.STACKTRACE_WITH_VJPWARN[] = true. To turn off this printing, add `verbose = false` to the `solve` call.
└ @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/bspwn/src/concrete_solve.jl:100
┌ Warning: Potential performance improvement omitted. ReverseDiffVJP tried and failed in the automated AD choice algorithm. To show the stack trace, set SciMLSensitivity.STACKTRACE_WITH_VJPWARN[] = true. To turn off this printing, add `verbose = false` to the `solve` call.
└ @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/bspwn/src/concrete_solve.jl:117

┌ Warning: Potential performance improvement omitted. TrackerVJP tried and failed in the automated AD choice algorithm. To show the stack trace, set SciMLSensitivity.STACKTRACE_WITH_VJPWARN[] = true. To turn off this printing, add `verbose = false` to the `solve` call.
└ @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/bspwn/src/concrete_solve.jl:135
┌ Warning: Reverse-Mode AD VJP choices all failed. Falling back to numerical VJPs
└ @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/bspwn/src/concrete_solve.jl:145


InexactError: Float32(0.23771682913973122 - 0.0036578124321906547im)

I narrowed down the problem to the line [sin(a[1]) 0.0; 0.0 -a[1]];: if we change it to [a[1] 0.0; 0.0 -a[1]]; (as it was in this issue), it works.

[seadra]

A bit tangential regarding those warnings about potential performance improvements, is there a way I can get those performance improvements such a loss function?

[ChrisRackauckas]

Your code doesn't use DiffEqFlux, so just remove it from the using and submit this to SciMLSensitivity in the future.

With SciML/SciMLSensitivity.jl#1064 your code is much faster. You do require complex coefficients in this case. So with that PR and the new patch, the following is really good:

using OrdinaryDiffEq, Zygote, SciMLSensitivity, Optimization, OptimizationOptimisers,
      ComponentArrays, Lux, Random, LinearAlgebra

const T = 10.0;
const ω = π/T;
const id = Matrix{Complex{Float64}}(I,2, 2);
const u0 = id;
const utarget = Matrix{Complex{Float64}}([im 0; 0 -im]);

ann = Lux.Chain(Lux.Dense(1,32), Lux.Dense(32,32,tanh), Lux.Dense(32,1));
rng = Random.default_rng();
ip, st = Lux.setup(rng, ann);

function f_nn(u, p, t)
    local a, _ = ann([t/T],p,st);
    local A = [sin(a[1]) 0.0; 0.0 -a[1]];
    return -(im*A)*u;
end

tspan = (0.0, T)
prob_ode = ODEProblem(f_nn, u0, tspan, ComponentArray{ComplexF32}(ip));

function loss_adjoint(p)
    local prediction = solve(prob_ode, BS5(), p=p, abstol=1e-7, reltol=1e-7)
    local usol = last(prediction)
    local loss = abs(1.0 - abs(tr(usol*utarget')/2))
    return loss
end

opt_f = Optimization.OptimizationFunction((x, p) -> loss_adjoint(x), Optimization.AutoZygote());
opt_prob = Optimization.OptimizationProblem(opt_f, ComponentArray{ComplexF32}(ip));
optimized_sol_nn = Optimization.solve(opt_prob, Adam(0.001), maxiters = 100, progress=true);

[seadra]

Thank you very much for the great progress on supporting complex matrices.

I'm a bit lost on this comment that you made:

You do require complex coefficients in this case.

But I need them to be real. If a[1] becomes complex, the equation corresponds to something unphysical.

[seadra]

Using real coefficients used to work with DiffEqSensitivity.

I just tried using

using OrdinaryDiffEq, Zygote, SciMLSensitivity, Optimization, OptimizationOptimisers,
      ComponentArrays, Lux, Random, LinearAlgebra

const T = 10.0;
const ω = π/T;
const id = Matrix{Complex{Float64}}(I,2, 2);
const u0 = id;
const utarget = Matrix{Complex{Float64}}([im 0; 0 -im]);

ann = Lux.Chain(Lux.Dense(1,32), Lux.Dense(32,32,tanh), Lux.Dense(32,1));
rng = Random.default_rng();
ip, st = Lux.setup(rng, ann);

function f_nn(u, p, t)
    local a, _ = ann([t/T],p,st);
    local A = [sin(a[1]) 0.0; 0.0 -a[1]];
    return -(im*A)*u;
end

tspan = (0.0, T)
prob_ode = ODEProblem(f_nn, u0, tspan, ComponentArray{Float32}(ip));

function loss_adjoint(p)
    local prediction = solve(prob_ode, BS5(), p=p, abstol=1e-7, reltol=1e-7)
    local usol = last(prediction)
    local loss = abs(1.0 - abs(tr(usol*utarget')/2))
    return loss
end

opt_f = Optimization.OptimizationFunction((x, p) -> loss_adjoint(x), Optimization.AutoZygote());
opt_prob = Optimization.OptimizationProblem(opt_f, ComponentArray{Float32}(ip));
optimized_sol_nn = Optimization.solve(opt_prob, Adam(0.001), maxiters = 100, progress=true);

and it seems to have worked.

Is there something that I'm missing?