Performant schelling model in julia

run_schelling.jl

# Showcase
env = Schelling([0.45, 0.45, 0.1], Ba = 0.5, dim = 200)
anim = @animate for i in 1:100
    plot(env; title = i)
    step!(env)
end
gif(anim, "anim.gif", fps = 20)

schelling_julia.jl

import DSP: conv
import Distributions: Categorical
using Random
using Plots

mutable struct Schelling
    A::Array{Int}          # state arrat
    B::Array{Float64}      # neighbour proportion array
    Ba::Float64            # agent preference
    kern::Matrix{Int}      # kernel for convolution
    max_neigh::Matrix{Int} # normalization matrix
end

"""
    Schelling(prop[; p_free, Ba, dim])

Initialize a Schelling agent-based model environment.

# Arguments
- `prop::Vector{Float64}`: the population size proportions
- `free::Float64`: the proportion of free cells
- `Ba::Float64`: neighbour similarity preference (0 < Ba < 1)
- `dim::Int`: the width and height of the environment (always square)

# Examples
```julia-repl
julia> Schelling([0.5, 0.5])
julia> plot(Schelling)
```
"""
function Schelling(prop::Vector{Float64}; free::Float64 = 0.25, Ba::Float64 = 0.5, dim::Int = 30)
    # Generate categorical matrix
    θ = vcat(free / (1-free), prop ./ sum(prop)) ./ (1 + free / (1 - free))
    A_raw = rand(Categorical(θ), dim, dim)

    # transform categorical matrix into exclusive array
    # A[:,:,1] indicates which cells are empty, A[:,:,2:end] are the agents
    A = zeros(Int, (dim, dim, length(θ)))
    for r in 1:dim, c in 1:dim
        A[r, c, A_raw[r, c]] = 1
    end

    # Initialize neighbour array (excludes the free cells)
    B = zeros(Float64, (dim, dim, length(prop)))

    # convolution kernel
    kern = fill(1, (3, 3))
    kern[2,2] = 0
    
    # matrix determining maximum neighbours
    max_neigh = fill(8, (dim, dim))
    max_neigh[:,[1,dim]] .= 5
    max_neigh[[1,dim],:] .= 5
    max_neigh[[1,dim],1] .= 3
    max_neigh[[1,dim],dim] .= 3

    return Schelling(A, B, Ba, kern, max_neigh)
end

"Perform a step (in-place) of the Schelling model"
function step!(env::Schelling)
    # step 1: compute proportion of neighbours like you
    # perform convolution
    C = conv(env.kern, env.A)[2:end-1, 2:end-1, :]
    # update max_neigh with empty neighbours
    empty_neigh = C[:,:,1] .* (1 .- env.A[:,:,1])

    # compute B
    env.B = C[:,:,2:end] .* env.A[:,:,2:end] ./ (env.max_neigh .- empty_neigh)


    # step 2: move 
    empty_ids = findall(isone, env.A[:,:,1])
    len = length(empty_ids)
    unhappy_ids = find_unhappy(env)

    for idx in unhappy_ids
        move_from_idx = CartesianIndex(idx[1], idx[2])
        move_to_idx = popat!(empty_ids, rand(1:len))
        
        # update empty
        env.A[move_to_idx, 1] = 0
        env.A[move_from_idx, 1] = 1
        push!(empty_ids, move_from_idx)

        # update location
        env.A[idx] = 0
        env.A[move_to_idx, idx[3]] = 1
    end
end

"Find the indices of the unhappy agents in a Schelling model"
function find_unhappy(env::Schelling)
    out = []
    R, C, K = size(env.A)
    for k in 2:K, r in 1:R, c in 1:C
        env.A[r, c, k] == 0 && continue
        env.A[r, c, 1] == 1 && continue
        env.B[r, c, k - 1] >= env.Ba && continue
        push!(out, CartesianIndex(r, c, k))
    end
    return out
end

"Plot a Schelling model"
function plot(env::Schelling; kwargs...)
    dd = size(env.A)
    M = zeros(Float64, dd[1:2])
    for i in 2:dd[3]
        M += env.A[:,:,i] .* (i-1)
    end
    M[M .== 0] .= NaN
    col = cgrad(:fall, categorical = true)
    heatmap(M, color = col, aspect_ratio = 1, axis = ([], false), 
            yflip = true, legend = false; kwargs...)
end