Closure in struct type inference

This comes up a lot in my code style. I’ve been doing this thing which is to store methods inside a struct as closures. (sort of like oo programming I guess?). This is a lot more convenient than storing the captured values.

mutable struct ViralPopulation
    time::Float64 # current time
    fitness_total::Float64 # current fitness. Divided by cap gives φ
    pop::Array{Virus,1} # State of population length is the total population size
    capacity::Int 	# carrying capacity
    r::Float64 	# rate of decay of neutralized viruses
    λ::Float64 	# noise of system
    fitness_function::Function
    mutation_operator::Function
end

The mutation_operator is generated when the struct is initialized in a somewhat expensive operation:

function mutation_operator_closure(ab_profile_list, ν)
    let mut0 = mapreduce(y -> map(x->x[1],y), vcat, ab_profile_list), #relatively expensive mallocs
        mut1 = mapreduce(y -> map(x->x[2],y), vcat, ab_profile_list), 
		tot = sum(mut0) + sum(mut1)
		# this is a let-block to capture a fixed
		function mutate!(vp::ViralPopulation) # attempt to mutate a virus at index ind
            if rand() < tot*ν # if there is a shot at mutation
                ind = rand(1:length(vp.pop)) # choose a random virus
				@inbounds v = vp.pop[ind]
				#mutvec = mut[ (1:2:2*l) .+ bitarray(v.genotype,l)] # get the mutational parameters
                ii = bitmask_toggle_choose(v.genotype, mut0, mut1, tot)
				if ii > 0
					# flippy the dippy and make a new virus
					new_genotype = flip(v.genotype, 2^(ii-1))
					new_fitness = vp.fitness_function(new_genotype)
                    new_escape = vp.escape_function(new_genotype)
                    vp.fitness_total += new_fitness - v.fitness
                    vp.pop[ind] = Virus(new_genotype,new_fitness,new_escape)
				end
			end
        end
        return mutate!
	end
end

In profiling I get what looks like type inference issues in the part of my code when I call v.muation_operator:
closure in profile

Can I make my closure type secure, or do I have to send around the parameters as tuples? What do I need to do to score that last little bit of sweet sweet performance?

You need to parameterise your type-field, as each function is its own type:

struct ViralPopulation{F1,F2}
...
fitness_function::F1
mutation_operator::F2
end
7 Likes

:exploding_head: I’ll try it out.

Q: Does this mean that every instance of ViralPopulation will create a new type? I mean I don’t have that many floating around, but a little worried about what happens when one of them goes out of scope… do types get gc’ed? “Each function is it’s own type.” Holy moly I have so many questions.

Omg, mad genius! Just hitting the rand() now.

Screenshot 2021-06-01 at 13.38.59

Namaste.

4 Likes

Yes, but note that the array there, which is possibly the only memory-hungry data involved, will be shared if you happen to change the functions of your population:

julia> struct A{T}
         x::Vector{Float64}
         f::T
       end

julia> a = A{typeof(sin)}([1,2,3],sin)
A{typeof(sin)}([1.0, 2.0, 3.0], sin)

julia> b = A{typeof(cos)}(a.x,cos)
A{typeof(cos)}([1.0, 2.0, 3.0], cos)

julia> b.x[1] = 5
5

julia> a.x
3-element Vector{Float64}:
 5.0
 2.0
 3.0



I’m not sure that they get GCed. If not, the type defs are small, so unless you got heaps and heaps, that is not an issue. At least it’s not an issue that I encountered.

@mauro3, thanks you saved me days of work, and taught me something cool.

Julia is a deceptively easy language. You can get a cool project working and fast in a couple days and still not fully understand functions and types after a couple years of work.

@lmiq I have to admit, I didn’t realize that pure structs also take arrays as references!

Honestly not 100% sure if in my case I should have the mutable here or not. .time gets updated in place, but I could also use a Ref{Float64}. I sort of feel a much more important distinction isn’t mutable but what isbits returns you know?

1 Like

Probably not. Using a ref, or an array of a single element, are adequate solutions. Or using Setfield. Here is a discussion on the alternatives: Mutable scalar in immutable object: the best alternative?

ps: Array{Virus,1} is the same as Vector{Virus}. Be careful that Virus should be a concrete type here, or that vector would be abstract as well. If Virus is abstract (has parameters, you should better parameterize the population struct as well for the type of Virus).

1 Like

For these purposes, I’d recommend to not use a generic closure at all and define a callable type instead:

struct MutationOperator
    mut0::Vector{Float64}
    mut1::Vector{Float64}
    tot::Float64
    ν::Float64
end

function (op::MutationOperator)(vp::ViralPopulation)
    mut0, mut1, tot, ν = op.mut0, op.mut1, op.tot, op.ν
    if rand() < tot*ν # if there is a shot at mutation
        ind = rand(1:length(vp.pop)) # choose a random virus
        @inbounds v = vp.pop[ind]
        ii = bitmask_toggle_choose(v.genotype, mut0, mut1, tot)
        if ii > 0
            # flippy the dippy and make a new virus
            new_genotype = flip(v.genotype, 2^(ii-1))
            new_fitness = vp.fitness_function(new_genotype)
            new_escape = vp.escape_function(new_genotype)
            vp.fitness_total += new_fitness - v.fitness
            vp.pop[ind] = Virus(new_genotype,new_fitness,new_escape)
        end
    end
end

function MutationOperator(ab_profile_list, ν)
    let mut0 = mapreduce(y -> map(x->x[1],y), vcat, ab_profile_list), #relatively expensive mallocs
        mut1 = mapreduce(y -> map(x->x[2],y), vcat, ab_profile_list), 
        tot = sum(mut0) + sum(mut1)
        return MutationOperator(mut0, mut1, tot, ν)
    end
end
6 Likes

@Vasily_Pisarev that was definitely my reflex (well sort of, I didn’t know you could define functions that way! …but the struct for the dereferenced variables at least) when I thought about how to solve this. But it looks like, at least in this case, this function call is as fast as it can be after the annotation from @mauro3 !

I was worried I’d be paying compile costs every instance, but julia-1.6 @time shows no compilation after the first run and gc is ~1%, and I’m hitting the necessary rand() in profiling 98% of the time the closure is called. It looks like circa 1.6, the compiler understands closures well enough to handle this at top performance with just the right type annotation (and a let block of course). To me it is a more elegant solution.

Yes, that’s because the compiler generates almost the same callable type under the hood for the closure.

OK, but it’s a more informed decision when you know you have more than one option.

Yes, that’s because the compiler generates almost the same callable type under the hood for the closure.

I think I see, yeah, and defining a struct certainly brings that warm fuzzy feeling of knowing exactly what is being allocated. And were I to rewrite my code from scratch, maybe the callable struct is nicer because the action at the function call is more transparent.

Either way it’s a really nice design pattern. You generally want to separate the dynamics from the integration control and confine your parameters to the dynamics. It’s great to see how Julia lets you do that with the flexibility of dynamic dispatch but with no loss of performance with sufficiently typed input. In the past, I passed gigantic structs around because of worries about closure performance. Now I have at least two ways to not do that, so thanks :metal:

4 Likes

Just to chime in with an option three, if you know what the signature of your function will be, you could always use FunctionWrappers.jl.

3 Likes