JuliaMM.jl

##geometric functions
using LinearAlgebra:dot, norm, atan, acos, ×, ⋅
distance(x, y) = norm(x - y)
∠(a, b, c) = acos(-dot((b-a)/norm(b-a), (c-b)/norm(c-b)))

function dihedral_angle(a, b, c, d)
    b1 = b - a
    # b2 = c - b # mdtraj convention
    b2 = b - c  # openmm convention
    b3 = d - c

    c1 = b2 × b3
    c2 = b1 × b2

    p1 = (b1 ⋅ c1) * norm(b2)^2
    p2 = (c1 ⋅ c2)

    return atan(p1, p2)
end


## nonbonded potentials
const ϵ₀ = 5.72765e-4
lennard_jones_potential(r, σ, ϵ) = 4ϵ * ((σ/r)^12 - (σ/r)^6)
coulomb_potential(r, q_prod) = (q_prod / r) / (4π * ϵ₀)


## combining rules
σ_combining_rule(σ₁, σ₂) = (σ₁ + σ₂)/2
ϵ_combining_rule(ϵ₁, ϵ₂) = √abs(ϵ₁ *ϵ₂)
charge_combining_rule(q₁, q₂) = q₁ * q₂


## Standard valence terms
harmonic_bond_potential(r, k, r₀) = (r-r₀)^2 * (k/2)
harmonic_angle_potential(θ, k, θ₀) = (θ-θ₀)^2 * (k/2)
periodic_torsion_potential(θ, k, phase, periodicity) = k * (1 + cos(periodicity * θ - phase))


struct Particle
    ϵ::Float64
    q::Float64
    σ::Float64
end

struct HarmonicBond
    atom₁::Int64
    atom₂::Int64
    k::Float64
    r₀::Float64
end

struct HarmonicAngle
    atom₁::Int64
    atom₂::Int64
    atom₃::Int64
    k::Float64
    θ₀::Float64
end

struct PeriodicTorsion
    atom₁::Int64
    atom₂::Int64
    atom₃::Int64
    atom₄::Int64
    k::Float64
    phase::Float64
    periodicity::Float64
end

struct ValenceModel
    harmonic_bonds::Array{HarmonicBond,1}
    harmonic_angles::Array{HarmonicAngle,1}
    periodic_torsions::Array{PeriodicTorsion,1}
end

using LightXML

xdoc = parse_file("ala2.xml");
xroot = root(xdoc);

ces = collect(child_elements(xroot));

forces = collect(child_elements(ces[4]))

# parse harmonic bond force
harmonic_bond_force = collect(child_elements(forces[1]))[1]
bonds = collect(child_elements(harmonic_bond_force))

function bond_from_attr_dict(attr_dict)
    # note +1, because zero-indexing vs. one-indexing
    p1, p2 = map(tag -> parse(Int64, attr_dict[tag]) + 1, ["p1", "p2"])
    k, d = map(tag -> parse(Float64, attr_dict[tag]), ["k", "d"])

    HarmonicBond(p1, p2, k, d)
end

harmonic_bonds = map(bond_from_attr_dict ∘ attributes_dict, bonds)

# parse harmonic angle force
harmonic_angle_force = collect(child_elements(forces[2]))[1]
angles = collect(child_elements(harmonic_angle_force))

function angle_from_attr_dict(attr_dict)
    p1, p2, p3 = map(tag -> parse(Int64, attr_dict[tag]) + 1, ["p1", "p2", "p3"])
    k, a = map(tag -> parse(Float64, attr_dict[tag]), ["k", "a"])

    HarmonicAngle(p1, p2, p3, k, a)
end

harmonic_angles = map(angle_from_attr_dict ∘ attributes_dict, angles)

# parse periodic torsion force
periodic_torsion_force = collect(child_elements(forces[3]))[1]
torsions = collect(child_elements(periodic_torsion_force))

function torsion_from_attr_dict(attr_dict)
    p1, p2, p3, p4 = map(tag -> parse(Int64, attr_dict[tag])+1, ["p1", "p2", "p3", "p4"])
    k, phase, periodicity = map(tag -> parse(Float64, attr_dict[tag]), ["k", "phase", "periodicity"])

    PeriodicTorsion(p1, p2, p3, p4, k, phase, periodicity)
end

periodic_torsions = map(torsion_from_attr_dict ∘ attributes_dict, torsions)

# (partially) parse nonbonded force
nb_force = collect(child_elements(forces[4]))
particle_entries = collect(child_elements(nb_force[4]))


function particle_from_attr_dict(attr_dict)
    ϵ, q, σ = map(tag -> parse(Float64, attr_dict[tag]), ["eps", "q", "sig"])

    Particle(ϵ, q, σ)
end

particles = map(particle_from_attr_dict ∘ attributes_dict, particle_entries)

function compute_harmonic_bond_potential(X, harmonic_bonds)
    U = 0
    for bond in harmonic_bonds
        r = distance(X[bond.atom₁], X[bond.atom₂])
        U += harmonic_bond_potential(r, bond.k, bond.r₀)
    end
    return U
end

function compute_harmonic_angle_potential(X, harmonic_angles)
    U = 0
    for angle in harmonic_angles
        θ = ∠(X[angle.atom₁], X[angle.atom₂], X[angle.atom₃])
        U += harmonic_angle_potential(θ, angle.k, angle.θ₀)
    end
    return U
end

function compute_periodic_torsion_potential(X, periodic_torsions)
    U = 0
    for torsion in periodic_torsions
        θ = dihedral_angle(X[torsion.atom₁], X[torsion.atom₂], X[torsion.atom₃], X[torsion.atom₄])
        U += periodic_torsion_potential(θ, torsion.k, torsion.phase, torsion.periodicity)
    end
    return U
end

function compute_valence_potential(X, valence_model)
    U_bonds = compute_harmonic_bond_potential(X, valence_model.harmonic_bonds)
    U_angles = compute_harmonic_angle_potential(X, valence_model.harmonic_angles)
    U_torsions = compute_periodic_torsion_potential(X, valence_model.periodic_torsions)
    return U_bonds + U_angles + U_torsions
end

valence_model = ValenceModel(harmonic_bonds, harmonic_angles, periodic_torsions);

# TODO: compute nonbonded sum