How to make type-stable

[sofiabelen]

@stepanzh @vvpisarev

1. Passing an element of a vector of functions

In the functions darcy!(::System, ::Parameters) and boundary_velocity!(::System, ::Parameters) at cfd.jl, this is type unstable:

fgas(s) = s^2
fliq(s) = (1 .- s)^2
f = (fgas, fliq)

for k = 1 : 2
  @views darcy!(f=f[k], u=sys.u[:, :, k],
                v=sys.v[:, :, k], P=sys.P,
                s=sys.s, p=p, k=k)
end

I asked for help: https://discourse.julialang.org/t/how-to-make-type-stable-passing-an-element-of-a-vector-of-functions/60068 and the proposed solution works if we didn't have @views. How can I deal with this?

2. Receiving a tuple from a function

The issue is in the function find_pressure!(::System, ::Parameters) at thermo.jl.

This gives seg fault, but is type stable

(sys.P[i, j], sys.s[i, j])::Tuple{T, T} =
            find_pressure(p=p,
            root_finder=newton_raphson,
            ρ₁=sys.ρ[i, j, 1],
            ρ₂=sys.ρ[i, j, 2],
            tait=tait_C₅H₁₂,
            igas=ideal_gas) where T<:AbstractFloat

This is type unstable

sys.P[i, j], sys.s[i, j] =
           find_pressure(p=p,
           root_finder=newton_raphson,
           ρ₁=sys.ρ[i, j, 1],
           ρ₂=sys.ρ[i, j, 2],
           tait=tait_C₅H₁₂,
           igas=ideal_gas)

Problematic part of the output of @code_warntype find_pressure!(sys, p) for the second piece of code:

│   %27 = (%25)(%26)::NamedTuple{(:p, :root_finder, :ρ₁, :ρ₂, :tait, :igas), _A} where _A<:Tuple{Parameters{Float64}, typeof(newton_raphson), Float64, Float64, Any, Any}

I tried explicitely stating all the types in find_pressure (the function being called inside `find_pressure!'):

function find_pressure(; root_finder::Function, ρ₁::T, ρ₂::T,
        tait::TaitEoS, igas::IdealGasEoS,
        p::Parameters, niter::Integer=50,
        eps_x::T=1e-6)::Tuple{T, T} where T<:AbstractFloat

However, this is type stable

This is a simplified version of the original code, and it's type stable. The question is, what is different between this and the original code?

function f(x::Integer, y::Integer)
    return x, y
end

a = zeros(3, 3)
b = zeros(3, 3)

function g()
    for i = 1 : 3
        for j = 1 : 3
            a[i, j], b[i, j] = f(i, j)
        end
    end
end

[stepanzh]

1. Passing an element of a vector of functions

First possible option, try to @view explicitly on objects you want to @view (idk whether it's type stable or not)

darcy!(f=f[k], u=(@view sys.u[:, :, k]),
                v=(@view sys.v[:, :, k]), P=sys.P,
                s=sys.s, p=p, k=k)

Sometimes curly brackets are needed to clarity macro.

Second option. Define type of s argument (I guess, this is the case)

--- fgas(s) = s^2
--- fliq(s) = (1 .- s)^2
+++ fgas(s::Real) = s^2
+++ fliq(s::Real) = (1 - s)^2

Note, that there is not .- in fliq.

To applay scalar function on vector, use broadcast

fgas(5)  # for scalar
fgas.(rand(10))  # for AbstractArray

Third possible option, practicable in future. Instead of passing functions, create functors, like this

struct PermeabilityAir{T}
end

# default type, if not passed
PermeabilityAir() = PermeabilityAir{Float64}()

# actual function (functor)
function (p::PermeabilityAir{T})(s::Real) where {T}
    return s^2
end

fgas = PermeabilityAir()
fgas(0.5)
fgas.(rand(10))

For this example it seems overweighted, but, in future you would like to keep some parameters inside PermeabilityAir. And functors is a proper way to deal with closures.

Also, I very recommend you to replace definitions of f_phase(saturation_gas) with f_phase(saturation_phase). So, instead of using single vector of gas saturation you will keep both: for gas and liquid phase. That's OK. For me, your current definitions of fgas and fliq seem tricky, which would cause semantic errors.

[vvpisarev]

1. Passing an element of a vector of functions

Not sure, looks like the loop doesn't get unrolled automatically. The problem is that f[1] and f[2] have different types, so that the result of a generic loop is technically type-unstable. If none of the @stepanzh's solutions applies, maybe changing the loop to for k in (1, 2) instead of for k in 1:2 has a better chance of unrolling.

2. Receiving a tuple from a function

I guess, the type instability is due to use of non-const globals inside newton_raphson. Those must be defined as

const cache_x₀ = Ref(1e5)
const is_cache = Ref(1) # maybe `true` is more appropriate?

and accessed via is_cache[] and cache_x₀[]. References are required to make the values mutable.

Another option is to define solvers as callable structs and hold the initial values, tolerance etc. within those. E.g.

mutable struct NewtonRaphsonSolver
    x0::Float64
    eps_x::Float64
    maxiter::Int
end

function (nr::NewtonRaphsonSolver)(f, df)
    x₀, eps_x, niter = nr.x0, nr.eps_x, nr.maxiter

    for i = 1 : niter
        @debug x₀ f(x₀)
        x₁ = x₀ - f(x₀) / df(x₀) # `AbstractFloat` type annotation makes any expression with x₁ type-unstable

        ## To avoid domain error: negative values in log
        if x₁ < eps_x
            x₁ = eps_x
        end

        if abs(x₁ - x₀) < eps_x || f(x₁) == 0
            nr.x0 = x₁
            return x₁
        end

        x₀ = x₁
    end

    error("newton raphson method fails to reach accuracy of ",
                 eps_x, " in ", niter, " iterations")
end

# a variant that starts from a user-provided value
function (nr::NewtonRaphsonSolver)(f, df, x₀)
    eps_x, niter = nr.eps_x, nr.maxiter

    for i = 1 : niter
        @debug x₀ f(x₀)
        x₁ = x₀ - f(x₀) / df(x₀)

        ## To avoid domain error: negative values in log
        if x₁ < eps_x
            x₁ = eps_x
        end

        if abs(x₁ - x₀) < eps_x || f(x₁) == 0
            nr.x0 = x₁
            return x₁
        end

        x₀ = x₁
    end

    error("newton raphson method fails to reach accuracy of ",
                 eps_x, " in ", niter, " iterations")
end