agdestein / IncompressibleNavierStokes.jl

Incompressible Navier-Stokes solver
https://agdestein.github.io/IncompressibleNavierStokes.jl/dev/
MIT License
41 stars 13 forks source link

Error differentiating through the solver #81

Open gbruer15 opened 2 weeks ago

gbruer15 commented 2 weeks ago

Hello! First of all, thank you to the software writer for making this package. There are many examples with lovely pictures.

The documentation site says the package supports differentiating through the solver, but the examples directory doesn't show how to do that.

I tried a very simple test with the Actuator2D example, but it errors just when solving the forward problem, before even starting the backwards gradient calculation.

slightly modified Actuator2D.jl ```julia # # Unsteady actuator case - 2D # # In this example, an unsteady inlet velocity profile at encounters a wind # turbine blade in a wall-less domain. The blade is modeled as a uniform body # force on a thin rectangle. using CairoMakie using IncompressibleNavierStokes # Output directory outdir = joinpath(@__DIR__, "output", "Actuator2D") ispath(outdir) || mkpath(outdir) # A 2D grid is a Cartesian product of two vectors n = 40 x = LinRange(0.0, 10.0, 5n + 1) y = LinRange(-2.0, 2.0, 2n + 1) fig = plotgrid(x, y; figure = (; size = (600, 300))) display(fig) # Boundary conditions boundary_conditions = ( ## x left, x right ( ## Unsteady BC requires time derivatives DirichletBC( (dim, x, y, t) -> sin(π / 6 * sin(π / 6 * t) + π / 2 * (dim() == 1)), (dim, x, y, t) -> (π / 6)^2 * cos(π / 6 * t) * cos(π / 6 * sin(π / 6 * t) + π / 2 * (dim() == 1)), ), PressureBC(), ), ## y rear, y front (PressureBC(), PressureBC()), ) # Actuator body force: A thrust coefficient `Cₜ` distributed over a thin rectangle xc, yc = 2.0, 0.0 # Disk center D = 1.0 # Disk diameter δ = 0.11 # Disk thickness Cₜ = 0.2 # Thrust coefficient cₜ = Cₜ / (D * δ) inside(x, y) = abs(x - xc) ≤ δ / 2 && abs(y - yc) ≤ D / 2 bodyforce(dim, x, y, t) = dim() == 1 ? -cₜ * inside(x, y) : 0.0 # Build setup and assemble operators setup = Setup(x, y; Re = 100.0, boundary_conditions, bodyforce); # Initial conditions (extend inflow) ustart = create_initial_conditions(setup, (dim, x, y) -> dim() == 1 ? 1.0 : 0.0); # Solve unsteady problem function get_state(ustart) state, outputs = solve_unsteady(; setup, ustart, tlims = (0.0, 12.0), method = RKMethods.RK44P2(), Δt = 0.05, processors = ( # rtp = realtimeplotter(; setup, size = (600, 300), nupdate = 5), # log = timelogger(; nupdate = 24), ), ) return state end state = get_state(ustart) # ## Post-process # # We may visualize or export the computed fields # Export to VTK save_vtk(state; setup, filename = joinpath(outdir, "solution")) # We create a box to visualize the actuator. box = ( [xc - δ / 2, xc - δ / 2, xc + δ / 2, xc + δ / 2, xc - δ / 2], [yc + D / 2, yc - D / 2, yc - D / 2, yc + D / 2, yc + D / 2], ) # Plot pressure fig = fieldplot(state; setup, size = (600, 300), fieldname = :pressure) lines!(box...; color = :red) display(fig) # Plot velocity fig = fieldplot(state; setup, size = (600, 300), fieldname = :velocitynorm) lines!(box...; color = :red) display(fig) # Plot vorticity fig = fieldplot(state; setup, size = (600, 300), fieldname = :vorticity) lines!(box...; color = :red) display(fig) # Get gradient using Zygote state, adjoint_action = Zygote.pullback(get_state, ustart); ```
Output, with error message ``` julia> include("examples/Actuator2D.jl") ┌ Warning: Creating new pressure solver for observefield └ @ IncompressibleNavierStokes ~/a/navierstokes/IncompressibleNavierStokes.jl/src/processors.jl:91 ERROR: LoadError: Compiling Tuple{SparseArrays.UMFPACK.var"##umfpack_numeric!#21", Bool, Nothing, typeof(SparseArrays.UMFPACK.umfpack_numeric!), SparseArrays.UMFPACK.UmfpackLU{Float64, Int64}}: ArgumentError: array must be non-empty Stacktrace: [1] macro expansion @ ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [inlined] [2] _pullback(::Zygote.Context{…}, ::SparseArrays.UMFPACK.var"##umfpack_numeric!#21", ::Bool, ::Nothing, ::typeof(SparseArrays.UMFPACK.umfpack_numeric!), ::SparseArrays.UMFPACK.UmfpackLU{…}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:87 [3] umfpack_numeric! @ /my_data/.julia/juliaup/julia-1.11.0-rc3+0.x64.linux.gnu/share/julia/stdlib/v1.11/SparseArrays/src/solvers/umfpack.jl:629 [inlined] [4] _pullback(::Zygote.Context{…}, ::typeof(Core.kwcall), ::@NamedTuple{…}, ::typeof(SparseArrays.UMFPACK.umfpack_numeric!), ::SparseArrays.UMFPACK.UmfpackLU{…}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [5] #lu#7 @ /my_data/.julia/juliaup/julia-1.11.0-rc3+0.x64.linux.gnu/share/julia/stdlib/v1.11/SparseArrays/src/solvers/umfpack.jl:388 [inlined] [6] _pullback(::Zygote.Context{…}, ::SparseArrays.UMFPACK.var"##lu#7", ::Bool, ::Nothing, ::Vector{…}, ::typeof(LinearAlgebra.lu), ::SparseArrays.SparseMatrixCSC{…}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [7] lu @ /my_data/.julia/juliaup/julia-1.11.0-rc3+0.x64.linux.gnu/share/julia/stdlib/v1.11/SparseArrays/src/solvers/umfpack.jl:384 [inlined] [8] _pullback(ctx::Zygote.Context{false}, f::typeof(LinearAlgebra.lu), args::SparseArrays.SparseMatrixCSC{Float64, Int64}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [9] factorize @ /my_data/.julia/juliaup/julia-1.11.0-rc3+0.x64.linux.gnu/share/julia/stdlib/v1.11/SparseArrays/src/linalg.jl:1976 [inlined] [10] _pullback(ctx::Zygote.Context{…}, f::typeof(LinearAlgebra.factorize), args::SparseArrays.SparseMatrixCSC{…}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [11] psolver_direct @ ~/a/navierstokes/IncompressibleNavierStokes.jl/src/pressure.jl:137 [inlined] [12] _pullback(::Zygote.Context{…}, ::typeof(psolver_direct), ::Vector{…}, ::@NamedTuple{…}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [13] psolver_direct @ ~/a/navierstokes/IncompressibleNavierStokes.jl/src/pressure.jl:121 [inlined] [14] default_psolver @ ~/a/navierstokes/IncompressibleNavierStokes.jl/src/pressure.jl:116 [inlined] [15] _pullback(ctx::Zygote.Context{…}, f::typeof(default_psolver), args::@NamedTuple{…}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [16] solve_unsteady @ ~/a/navierstokes/IncompressibleNavierStokes.jl/src/solvers/solve_unsteady.jl:18 [inlined] [17] _pullback(::Zygote.Context{…}, ::typeof(Core.kwcall), ::@NamedTuple{…}, ::typeof(solve_unsteady)) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [18] get_state @ ~/a/navierstokes/IncompressibleNavierStokes.jl/examples/Actuator2D.jl:57 [inlined] [19] _pullback(ctx::Zygote.Context{false}, f::typeof(get_state), args::Tuple{Matrix{Float64}, Matrix{Float64}}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [20] pullback(f::Function, cx::Zygote.Context{false}, args::Tuple{Matrix{Float64}, Matrix{Float64}}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:90 [21] pullback(f::Function, args::Tuple{Matrix{Float64}, Matrix{Float64}}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:88 [22] top-level scope @ ~/a/navierstokes/IncompressibleNavierStokes.jl/examples/Actuator2D.jl:103 [23] include(fname::String) @ Main ./sysimg.jl:38 [24] top-level scope @ REPL[2]:1 in expression starting at /home/gbruer/a/navierstokes/IncompressibleNavierStokes.jl/examples/Actuator2D.jl:103 Some type information was truncated. Use `show(err)` to see complete types. ```

Version info:

agdestein commented 2 weeks ago

Hi, thanks for the feedback.

Differentiating with Zygote requires all operations to be allocating and non-modifying, this can be achieved by using the out-of-place versions of the operators. The function solve_unsteady uses the in-place versions, but you can make a differentiable time-stepping loop as follows (replace your get_state with this version):

# Do timestepping manually with out-of-place stepper
# (use `timestep` instead of `timestep!` to avoid modifying the input)
INS = IncompressibleNavierStokes
psolver = default_psolver(setup)
function get_state(ustart)
    t_end = 12.0
    Δt = 0.05
    method = RKMethods.RK44P2()
    nstep = round(Int, t_end / Δt)
    stepper =
        INS.create_stepper(method; setup, psolver, u = ustart, temp = nothing, t = 0.0)
    for it = 1:nstep
        stepper = INS.timestep(method, stepper, Δt)
    end
    return INS.get_state(stepper)
end

Let me know if this works!

gbruer15 commented 2 weeks ago

Thank you. With that change, the forward solve works. The downsides are it is much slower and perhaps has a memory leak somewhere, although that could be a Zygote issue. (Calling it twice doubled the memory usage and was unaffected by GC.gc())

Now the pullback for DirichletBC seems to not be implemented yet.

Since I don't have a particular problem setup in mind yet, I'll quit working on this for now (and I'll assume the PeriodicBC one works). I appreciate your help!

Actuator2D.jl ```julia # # Unsteady actuator case - 2D # # In this example, an unsteady inlet velocity profile at encounters a wind # turbine blade in a wall-less domain. The blade is modeled as a uniform body # force on a thin rectangle. using CairoMakie using IncompressibleNavierStokes # Output directory outdir = joinpath(@__DIR__, "output", "Actuator2D") ispath(outdir) || mkpath(outdir) # A 2D grid is a Cartesian product of two vectors n = 40 x = LinRange(0.0, 10.0, 5n + 1) y = LinRange(-2.0, 2.0, 2n + 1) # Boundary conditions boundary_conditions = ( ## x left, x right ( ## Unsteady BC requires time derivatives DirichletBC( (dim, x, y, t) -> sin(π / 6 * sin(π / 6 * t) + π / 2 * (dim() == 1)), (dim, x, y, t) -> (π / 6)^2 * cos(π / 6 * t) * cos(π / 6 * sin(π / 6 * t) + π / 2 * (dim() == 1)), ), PressureBC(), ), ## y rear, y front (PressureBC(), PressureBC()), ) # Actuator body force: A thrust coefficient `Cₜ` distributed over a thin rectangle xc, yc = 2.0, 0.0 # Disk center D = 1.0 # Disk diameter δ = 0.11 # Disk thickness Cₜ = 0.2 # Thrust coefficient cₜ = Cₜ / (D * δ) inside(x, y) = abs(x - xc) ≤ δ / 2 && abs(y - yc) ≤ D / 2 bodyforce(dim, x, y, t) = dim() == 1 ? -cₜ * inside(x, y) : 0.0 # Build setup and assemble operators setup = Setup(x, y; Re = 100.0, boundary_conditions, bodyforce); # Initial conditions (extend inflow) ustart = create_initial_conditions(setup, (dim, x, y) -> dim() == 1 ? 1.0 : 0.0); # Solve unsteady problem INS = IncompressibleNavierStokes psolver = default_psolver(setup) function get_state(ustart) t_end = 12.0 Δt = 0.05 method = RKMethods.RK44P2() nstep = round(Int, t_end / Δt) stepper = INS.create_stepper(method; setup, psolver, u = ustart, temp = nothing, t = 0.0) for it = 1:nstep stepper = INS.timestep(method, stepper, Δt) end return INS.get_state(stepper) end # @time state = get_state(ustart) # @time state = get_state(ustart) # 7.038594 seconds (4.56 M allocations: 4.221 GiB, 10.04% gc time, 37.87% compilation time) # 3.871654 seconds (644.46 k allocations: 4.025 GiB, 5.98% gc time) # Get gradient using Zygote # @time state, adjoint_action = Zygote.pullback(get_state, ustart); # @time state, adjoint_action = Zygote.pullback(get_state, ustart); # 32.891385 seconds (41.20 M allocations: 5.919 GiB, 6.28% gc time, 78.29% compilation time) # 25.251067 seconds (1.15 M allocations: 4.080 GiB, 11.74% gc time) @time state, adjoint_action = Zygote.pullback(get_state, ustart); d_ustart = adjoint_action(state) ```
Output with error ``` ulia> include("examples/Actuator2D.jl") 35.876284 seconds (44.16 M allocations: 6.068 GiB, 6.69% gc time, 79.78% compilation time) ERROR: LoadError: MethodError: no method matching apply_bc_u_pullback!(::DirichletBC{…}, ::Tuple{…}, ::Int64, ::Float64, ::@NamedTuple{…}; isright::Bool) The function `apply_bc_u_pullback!` exists, but no method is defined for this combination of argument types. Closest candidates are: apply_bc_u_pullback!(::PeriodicBC, ::Any, ::Any, ::Any, ::Any; isright, kwargs...) @ IncompressibleNavierStokes ~/a/navierstokes/IncompressibleNavierStokes.jl/src/boundary_conditions.jl:292 apply_bc_u_pullback!(::Any, ::Any, ::Any; kwargs...) @ IncompressibleNavierStokes ~/a/navierstokes/IncompressibleNavierStokes.jl/src/boundary_conditions.jl:163 Stacktrace: [1] apply_bc_u_pullback!(φbar::Tuple{…}, t::Float64, setup::@NamedTuple{…}; kwargs::@Kwargs{}) @ IncompressibleNavierStokes ~/a/navierstokes/IncompressibleNavierStokes.jl/src/boundary_conditions.jl:168 [2] apply_bc_u_pullback!(φbar::Tuple{…}, t::Float64, setup::@NamedTuple{…}) @ IncompressibleNavierStokes ~/a/navierstokes/IncompressibleNavierStokes.jl/src/boundary_conditions.jl:163 [3] (::IncompressibleNavierStokes.var"#5#6"{…})(φbar::ChainRulesCore.Tangent{…}) @ IncompressibleNavierStokes ~/a/navierstokes/IncompressibleNavierStokes.jl/src/boundary_conditions.jl:109 [4] (::Zygote.ZBack{IncompressibleNavierStokes.var"#5#6"{…}})(dy::Tuple{Matrix{…}, Matrix{…}}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/chainrules.jl:211 [5] #timestep#358 @ ~/a/navierstokes/IncompressibleNavierStokes.jl/src/time_steppers/step_explicit_runge_kutta.jl:123 [inlined] [6] (::Zygote.Pullback{…})(Δ::@NamedTuple{…}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [7] timestep @ ~/a/navierstokes/IncompressibleNavierStokes.jl/src/time_steppers/step_explicit_runge_kutta.jl:66 [inlined] [8] (::Zygote.Pullback{…})(Δ::@NamedTuple{…}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [9] get_state @ ~/a/navierstokes/IncompressibleNavierStokes.jl/examples/Actuator2D.jl:66 [inlined] [10] (::Zygote.Pullback{Tuple{…}, Any})(Δ::@NamedTuple{u::Tuple{…}, temp::Nothing, t::Float64, n::Int64}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0 [11] (::Zygote.var"#75#76"{Zygote.Pullback{…}})(Δ::@NamedTuple{u::Tuple{…}, temp::Nothing, t::Float64, n::Int64}) @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:91 [12] top-level scope @ ~/a/navierstokes/IncompressibleNavierStokes.jl/examples/Actuator2D.jl:85 [13] include(fname::String) @ Main ./sysimg.jl:38 [14] top-level scope @ REPL[2]:1 in expression starting at /home/gbruer/a/navierstokes/IncompressibleNavierStokes.jl/examples/Actuator2D.jl:85 Some type information was truncated. Use `show(err)` to see complete types. ```
agdestein commented 2 weeks ago

I added the missing chain rules, so now d_ustart = adjoint_action(state) in your example should give gradients on the latest main branch.

An exlanation can be found here.