# Fixing type stability involving functions

**URL:** <https://discourse.julialang.org/t/fixing-type-stability-involving-functions/89258>\
**Category:** Performance\
**Tags:** function, struct, code\_warntype, type-stability\
**Created:** [October 25, 2022, 5:59pm UTC](https://discourse.julialang.org/t/fixing-type-stability-involving-functions/89258 "2022-10-25T17:59:08Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![yhchang96](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yhchang96/32/36499_2.png) [@yhchang96](https://discourse.julialang.org/u/yhchang96)\
**Post date:** [October 25, 2022, 5:59pm UTC](https://discourse.julialang.org/t/fixing-type-stability-involving-functions/89258/1 "2022-10-25T17:59:08Z")

</div>

I have the following codes

```julia
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))
    end

    return c
end

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

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

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)
end

```

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:

```julia
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

```julia
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
Arguments
  #self#::Core.Const(dhdt_1!)
  dh::Vector{Float64}
  h::Vector{Float64}
  par::Test3{var"#61#62", StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}
  t::Int64
Locals
  V::Any
Body::Nothing
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?

---

<div class="post-metadata">

**Author:** ![Vasily\_Pisarev](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/vasily_pisarev/32/7929_2.png) [@Vasily\_Pisarev](https://discourse.julialang.org/u/Vasily_Pisarev)\
**Post date:** [October 25, 2022, 7:48pm UTC](https://discourse.julialang.org/t/fixing-type-stability-involving-functions/89258/2 "2022-10-25T19:48:46Z")

</div>

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.