Fixing type stability involving functions

I have the following codes

using Parameters

function spatiotemporal(x, t, a_coefs::AbstractMatrix, b_coefs::AbstractMatrix)

    ns = size(a_coefs,1) - 1
    nt = Int((size(a_coefs,2) - 1)/2)

    c = zero(eltype(a_coefs))

    for m = 0:ns, r = -nt:nt
        a, b = a_coefs[m+1,r+nt+1], b_coefs[m+1,r+nt+1]
        c += a*cos(2*pi*(m*x + r*t)) + b*sin(2*pi*(m*x + r*t))

    return c

@with_kw struct Test3{T1,T2}
    x::T2 = range(0, 1, length=129)[1:end-1]
    ns::Int64 = 10
    nt::Int64 = 10

function dhdt_1!(dh, h, par, t)
    if isa(par.V, Function)
        V = par.V(t)
        V = par.V
    dh .= V .+ 1.0
    return nothing

function costFun!(C, par)
    ns, nt = par.ns, par.nt
    n_total = (ns+1)*(2*nt+1)
    a_coefs = reshape(C[1:n_total], ns+1, 2*nt+1)
    b_coefs = reshape(C[n_total+1:end], ns+1, 2*nt+1)
    par2 = @set par.V = t -> spatiotemporal.(par1.x, t, (a_coefs,), (b_coefs,))
    # solves ODE using dhdt_1!(dh, h, par2, t)

Basically, I’m doing an optimization over the cost function costFun! where in each iteration, the unknown variable C which contains the Fourier coefficients is used to construct a function stored in par2.V via the general function spatiotemporal. The struct par2 which now contains the newly created function with the updated Fourier coefficients is used to solve a stiff ODE, where the results will be used to compute the cost function.

However, it looks like there are type stability issues with par.V. With the following:

a_coefs = [1. 2. 3.; 3. 4. 5.]
b_coefs = [1. 2. 3.; 4. 5. 6.]

h0 = ones(128)
dh = similar(dh)
par1 = Test3(V = t->1.0)
par2 = @set par1.V = t -> spatiotemporal.(par1.x, t, (a_coefs,), (b_coefs,))

Using par1 is good after checking with @code_warntype dhdt_1!(dh, h0, par1, 0.0). However, for par2 where par2.V is a function constructed using a_coefs and b_coefs, @code_warntype dhdt_1!(dh, h0, par2, 0.0) gives

MethodInstance for dhdt_1!(::Vector{Float64}, ::Vector{Float64}, ::Test3{var"#61#62", StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, ::Int64)
  from dhdt_1!(dh, h, par, t) in Main at In[102]:10
  par::Test3{var"#61#62", StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}
1 ─      Core.NewvarNode(:(V))
│   %2 = Base.getproperty(par, :V)::Core.Const(var"#61#62"())
│   %3 = (%2 isa Main.Function)::Core.Const(true)
│        Core.typeassert(%3, Core.Bool)
│   %5 = Base.getproperty(par, :V)::Core.Const(var"#61#62"())
│        (V = (%5)(t))
└──      goto #3
2 ─      Core.Const(:(V = Base.getproperty(par, :V)))
3 ┄ %9 = Base.broadcasted(Main.:+, V, 1.0)::Any
│        Base.materialize!(dh, %9)
└──      return Main.nothing

V is type Any, so it looks like it can’t infer the type of a_coefs and b_coefs. In my actual problem, a_coefs and b_coefs will come from C passed to the cost function.

So are there any ways to fix this type stability issues?

If the whole a_coefs... par2 = ... block is wrapped into a local scope, I have no type instability warning on my machine.
The issue is, I think, that the closure in par2 uses global variables.