# EnsembleProblem with SteadyStateProblem in Zygote

**URL:** <https://discourse.julialang.org/t/ensembleproblem-with-steadystateproblem-in-zygote/85018>\
**Category:** New to Julia\
**Tags:** question, sciml, differentialequation\
**Created:** [July 30, 2022, 2:28am UTC](https://discourse.julialang.org/t/ensembleproblem-with-steadystateproblem-in-zygote/85018 "2022-07-30T02:28:37Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![kcqpo](https://avatars.discourse-cdn.com/v4/letter/k/a4c791/32.png) [@kcqpo](https://discourse.julialang.org/u/kcqpo)\
**Post date:** [July 30, 2022, 2:28am UTC](https://discourse.julialang.org/t/ensembleproblem-with-steadystateproblem-in-zygote/85018/1 "2022-07-30T02:28:37Z")

</div>

Is there a premade parallel simulation version of SteadyStateProblem (currently using a for loop)? I am assuming there isn’t because EnsembleProblem expects tspan which I don’t provide. My goal is to obtain the steady state solutions with adjoints for training a neural ODE. Thanks!

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 1, 2022, 12:22pm UTC](https://discourse.julialang.org/t/ensembleproblem-with-steadystateproblem-in-zygote/85018/2 "2022-08-01T12:22:42Z")

</div>

> [@kcqpo](#):
>
> I am assuming there isn’t because EnsembleProblem expects tspan which I don’t provide.

It should work. Share an MWE

---

<div class="post-metadata">

**Author:** ![kcqpo](https://avatars.discourse-cdn.com/v4/letter/k/a4c791/32.png) [@kcqpo](https://discourse.julialang.org/u/kcqpo)\
**Post date:** [August 1, 2022, 3:10pm UTC](https://discourse.julialang.org/t/ensembleproblem-with-steadystateproblem-in-zygote/85018/3 "2022-08-01T15:10:34Z")

</div>

```julia
using DifferentialEquations
using SciMLSensitivity
using Plots
using Zygote

# Some model with multiple steady states
function some_model(u, p, t)
    dxdt = p[1]*u[1]*u[2]
    dydt = -p[2]*u[2]
    return [dxdt, dydt]
end

# Specify the initial conditions for each simulation
function prob_func(prob, i, repeat)
    global x0
    remake(prob, u0=x0[i, :])
end

x0 = hcat(rand(10), ones(10)) # Initialize some random initial conditions
foop = [1., 1.] # Parameters
tspan = LinRange(0, 50, 25)

# Do some random gradients just to make sure code works for parametrizing stuff
foo_gs = gradient(foop) do foop
    foo_prob = ODEProblem(some_model, u0, (tspan[1], tspan[end]), foop) # Works
    foo_ensemble = EnsembleProblem(foo_prob, prob_func=prob_func)
    foo_sol = solve(
        foo_ensemble, 
        Tsit5(), 
        EnsembleThreads(), 
        trajectories=size(x0)[1];
        saveat = tspan,
        sensealg = InterpolatingAdjoint(autojacvec=ZygoteVJP())
    )
    Zygote.ignore_derivatives() do
        plot(foo_sol) # Plot to see what is going on
    end

    # Some random function
    foo_val = 0.
    for some_sol in foo_sol.u
        _some_sol_mat = mapreduce(permutedims, vcat, some_sol.u)
        foo_val += sum(abs2, _some_sol_mat)
    end
    @show foo_val
    return foo_val
end

# The following below doesn't work
foo_gs_SS = gradient(foop) do foop
    foo_SS_prob = SteadyStateProblem(some_model, u0, foop)
    foo_SS_ensemble = EnsembleProblem(foo_SS_prob, prob_func=prob_func)
    foo_SS_sol = solve(
        foo_SS_ensemble,
        DynamicSS(Tsit5(), abstol=1e-4, reltol=1e-3, tspan=tspan[end]),
        EnsembleThreads(), 
        trajectories=size(x0)[1];
        sensealg = InterpolatingAdjoint(autojacvec=ZygoteVJP())
    )
    # @show Array(foo_SS_sol.u[1])
    
    # Some random function
    foo_val = 0.
    for some_sol in foo_SS_sol.u
        foo_val += sum(abs2, some_sol)
    end
    @show foo_val
    return foo_val
end

```

I’m not too familiar with the backend of AD packages, but using Zygote seems to break the ensemble problem. Using it without gradients work just fine. The error message is “ERROR: type SteadyStateProblem has no field tspan”.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 1, 2022, 6:56pm UTC](https://discourse.julialang.org/t/ensembleproblem-with-steadystateproblem-in-zygote/85018/4 "2022-08-01T18:56:36Z")

</div>

Thanks! I see. In general Zygote works with ensemble problems (it’s tested with ODEs, SDEs, DDEs, DAEs), but this combination seems to require something extra. Can you open an issue on SciMLSensitivity.jl? I can solve this by the end of the week.

---

<div class="post-metadata">

**Author:** ![kcqpo](https://avatars.discourse-cdn.com/v4/letter/k/a4c791/32.png) [@kcqpo](https://discourse.julialang.org/u/kcqpo)\
**Post date:** [August 2, 2022, 1:56am UTC](https://discourse.julialang.org/t/ensembleproblem-with-steadystateproblem-in-zygote/85018/5 "2022-08-02T01:56:38Z")

</div>

(post deleted by author)

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 5, 2022, 6:49pm UTC](https://discourse.julialang.org/t/ensembleproblem-with-steadystateproblem-in-zygote/85018/6 "2022-08-05T18:49:10Z")

</div>

Solved, in a way. The issue is that for steady state problems you should use SteadyStateAdjoint. [Throw a better error for time-based adjoint on no-time problem by ChrisRackauckas · Pull Request #705 · SciML/SciMLSensitivity.jl · GitHub](https://github.com/SciML/SciMLSensitivity.jl/pull/705) solves this by throwing a very explicit error message saying that.