Brownian Dynamics with cell lists

dbrownian.jl

using Random
using StaticArrays
using LinearAlgebra
using DelimitedFiles
using ThreadPools
using Printf
using CellListMap.PeriodicSystems
import CellListMap.PeriodicSystems: copy_output, reset_output!, reducer

struct Parameters
    ρ::Float64
    ktemp::Float64
    n_particles::Int
end

mutable struct EnergyAndForces
    energy::Float64
    virial::Float64
    forces::Vector{SVector{3,Float64}}
end

function initialize_simulation!(x, y, z, npart, rc, halfdist)
    dist = halfdist * 2.0

    # Define the first positions
    x[1] = -rc + halfdist
    y[1] = -rc + halfdist
    z[1] = -rc + halfdist

    # Create a complete lattice
    for i in 1:(npart - 1)
        x[i + 1] = x[i] + dist
        y[i + 1] = y[i]
        z[i + 1] = z[i]

        if x[i + 1] > rc
            x[i + 1] = -rc + halfdist
            y[i + 1] += dist
            z[i + 1] = z[i]

            if y[i + 1] > rc
                x[i + 1] = -rc + halfdist
                y[i + 1] = -rc + halfdist
                z[i + 1] += dist
            end
        end
    end

    return nothing
end

function pseudohs(rij; lambda=50.0)
    b_param = lambda / (lambda - 1.0)
    a_param = lambda * b_param^(lambda - 1.0)
    ktemp = 1.0
    uij = 0.0
    fij = 0.0

    if rij < b_param
        uij = (a_param / ktemp) * ((1.0 / rij)^lambda - (1.0 / rij)^(lambda - 1.0))
        uij += (1.0 / ktemp)
        fij = lambda * (1.0 / rij)^(lambda + 1.0)
        fij -= (lambda - 1.0) * (1.0 / rij)^lambda
        fij *= -a_param / ktemp
    end

    return uij, fij
end

"Custom copy, reset and reducer functions"
copy_output(x::EnergyAndForces) =
    EnergyAndForces(copy(x.energy), copy(x.virial), copy(x.forces))

function reset_output!(output::EnergyAndForces)
    output.energy = 0.0
    output.virial = 0.0

    for i in eachindex(output.forces)
        output.forces[i] = SVector(0.0, 0.0, 0.0)
    end

    return output
end

function reducer(x::EnergyAndForces, y::EnergyAndForces)
    e_tot = x.energy + y.energy
    vir_tot = x.virial + y.virial
    x.forces .+= y.forces

    return EnergyAndForces(e_tot, vir_tot, x.forces)
end

"Function that updates energy and forces for each pair"
function energy_and_forces!(x, y, i, j, d2, cutoff, output::EnergyAndForces)
    d = sqrt(d2)
    r = y - x
    if d < cutoff
        (uij, fij) = pseudohs(d)
        sumies = @. fij * r / d
        output.virial += dot(sumies, r)
    end
    output.energy += uij
    output.forces[i] = @. output.forces[i] + (fij * r / d)
    output.forces[j] = @. output.forces[j] - (fij * r / d)

    return output
end

function init_system(boxl, cutoff, inter_distance; n_particles=10^3)
    # We can create normal arrays for holding the particles' positions
    x = zeros(n_particles)
    y = zeros(n_particles)
    z = zeros(n_particles)
    initialize_simulation!(x, y, z, n_particles, boxl / 2.0, inter_distance / 2.0)

    # Now we change the arrays to static versions of it
    positions = [@SVector [i, j, k] for (i, j, k) in zip(x, y, z)]
    # Initialize system
    system = PeriodicSystem(;
        xpositions=positions,
        unitcell=[boxl, boxl, boxl],
        cutoff=cutoff,
        output=EnergyAndForces(0.0, 0.0, similar(positions)),
        output_name=:energy_and_forces,
    )

    return system
end

function integrate(positions, forces, ktemp, dt, rng)
    sigma = sqrt(2.0 * dt)
    noise = @SVector randn(rng, 3)
    new_positions = @. positions + (forces * dt / ktemp) + noise * sigma

    return new_positions
end

function simulation(
    params::Parameters, pathname; isave=1_000, eq_steps=100_000, prod_steps=500_000
)
    rng = Random.Xoshiro()
    boxl = cbrt(params.n_particles / params.ρ)
    inter_distance = cbrt(1.0 / params.ρ)
    cutoff = 2.0
    τ = 0.00001

    system = init_system(boxl, cutoff, inter_distance; n_particles=params.n_particles)
    reset_output!(system.energy_and_forces)
    trajectory_file = open(joinpath(pathname, "production.xyz"), "w")
    zfactor = 0.0
    nprom = 0

    for step in 1:eq_steps
        # Compute energy and forces
        reset_output!(system.energy_and_forces)
        map_pairwise!(
            (x, y, i, j, d2, output) ->
                energy_and_forces!(x, y, i, j, d2, system.cutoff, output),
            system,
        )

        for i in eachindex(system.positions, system.energy_and_forces.forces)
            f = system.energy_and_forces.forces[i]
            x = system.positions[i]
            new_x = integrate(x, f, params.ktemp, τ, rng)
            system.positions[i] = new_x
        end

        if mod(step, 10) == 0
            zfactor += system.energy_and_forces.virial
            nprom += 1
        end

        # Every few steps we show the energy and save to disk
        if mod(step, 100) == 0
            ener_part = system.energy_and_forces.energy / params.n_particles
            pressure = 1.0 - (zfactor / (3.0 * nprom * params.n_particles * params.ktemp))
            println("Compressibility: $pressure")
            println("Energy per particle: $ener_part")
        end

        # Save to disk the positions
        if mod(step, 100) == 0
            # Write to file
            println(trajectory_file, params.n_particles)
            println(trajectory_file, "Frame $step")
            for p in system.positions
                writedlm(trajectory_file, [p], " ")
            end
        end
    end

    close(trajectory_file)

    return nothing
end

phi2density(x) = 6.0 * x / pi

function main()
    packings = [0.45]
    densities = phi2density.(packings)

    for (d, e) in zip(densities, packings)
        params = Parameters(d, 1.4737, 20^3)
        # Create a new directory with these parameters
        pathname = joinpath(@__DIR__, "cells_packing=$(@sprintf("%.3f", e))")
        mkpath(pathname)
        simulation(params, pathname; eq_steps=10000, prod_steps=10_000_000)
    end

    return nothing
end

main()