force.bend source
force.bend on the hub · documented module
import Baseimport ./geom.bend as Gimport ./protein.bend as P# --- Nonbonded kernels v0: uniform Lennard-Jones + Coulomb, squared cutoffs.# Self pairs (d2 == 0) and pairs beyond c2 contribute nothing.# Convention: for self energy/forces, pass the same reusable list twice,# e.g. lj_total(+xs, xs, eps, sig2, c2) style call sites. That call is a# DOUBLE sum: every unordered pair {i, j} is counted twice (i,j and j,i),# so the physical total is half the answer. The diagonal is dropped by# POSITION (d2 == 0.0), not by index: callers must give distinct atoms# distinct positions, since a coincident distinct pair is silently dropped# instead of diverging.def lj_sq(+x: F32) -> F32: (x * x : F32)def lj_cube(+inv: F32) -> F32: ((inv * inv : F32) * inv : F32)def lj_eval(+u6: F32, +eps: F32) -> F32: ((4.0 * eps : F32) * ((lj_sq(u6) - u6 : F32)) : F32)def lj_from_inv(inv: F32, +eps: F32) -> F32: lj_eval(lj_sq(lj_cube(inv)), eps)# LJ energy from squared distance: 4e*((s2/d2)^12 - (s2/d2)^6), no sqrt.def lj_of_d2(d2: F32, +eps: F32, +sig2: F32) -> F32: lj_from_inv((sig2 / d2 : F32), eps)def lj_self_go(d2: F32, +eps: F32, +sig2: F32, self: Bool) -> F32: match self: case True{}: 0.0 case False{}: lj_of_d2(d2, eps, sig2)def lj_pair_go(+d2: F32, +eps: F32, +sig2: F32, below: Bool) -> F32: match below: case True{}: lj_self_go(d2, eps, sig2, F32.is_eq(d2, 0.0)) case False{}: 0.0def lj_pair_at(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32) -> F32: +d2 = G.Vec3.dist2(pi, pj) lj_pair_go(d2, eps, sig2, F32.is_lt(d2, c2))def lj_head(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, h: P.Atom) -> F32: match h: case P.Atom{serial, elem, +pj}: lj_pair_at(pi, pj, eps, sig2, c2)def lj_row( +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>) -> F32: match ys: case Nil{}: 0.0 case h <> t: (lj_row(pi, eps, sig2, c2, t) + lj_head(pi, eps, sig2, c2, h) : F32)def lj_self_atom( h: P.Atom, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32) -> F32: match h: case P.Atom{serial, elem, +pi}: lj_row(pi, eps, sig2, c2, ys)def lj_total( xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32) -> F32: match xs: case Nil{}: 0.0 case h <> t: (lj_total(t, ys, eps, sig2, c2) + lj_self_atom(h, ys, eps, sig2, c2) : F32)# --- Coulomb with per-atom charges in an aligned List<&2, F32>. ---# The charge lists zip with the atom lists: a shorter charge list silently# drops the remaining atoms (the Nil{} arms answer zero), so callers must# keep charges aligned with atoms; see LAWS coul_trunc_charges.def coul_formula(d2: F32, +qq: F32, +ke: F32) -> F32: ((ke * qq : F32) / F32.sqrt(d2) : F32)def coul_self_go(d2: F32, +qq: F32, +ke: F32, self: Bool) -> F32: match self: case True{}: 0.0 case False{}: coul_formula(d2, qq, ke)def coul_pair_go(+d2: F32, +qq: F32, +ke: F32, below: Bool) -> F32: match below: case True{}: coul_self_go(d2, qq, ke, F32.is_eq(d2, 0.0)) case False{}: 0.0def coul_at(+pi: G.Vec3, +pj: G.Vec3, +qq: F32, +ke: F32, +c2: F32) -> F32: +d2 = G.Vec3.dist2(pi, pj) coul_pair_go(d2, qq, ke, F32.is_lt(d2, c2))def coul_head( +pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, h: P.Atom, qj: F32) -> F32: match h: case P.Atom{serial, elem, +pj}: coul_at(pi, pj, (qi * qj : F32), ke, c2)def coul_row( +pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, ys: List<&2, P.Atom>, qs: List<&2, F32>) -> F32: match ys qs: case Nil{} Nil{}: 0.0 case Nil{} qh <> qt: 0.0 case h <> t Nil{}: 0.0 case h <> t qh <> qt: (coul_row(pi, qi, ke, c2, t, qt) + coul_head(pi, qi, ke, c2, h, qh) : F32)def coul_self( h: P.Atom, +qh: F32, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32) -> F32: match h: case P.Atom{serial, elem, +pi}: coul_row(pi, qh, ke, c2, ys, qs_e)def coul_total( xs: List<&2, P.Atom>, +qs_o: List<&2, F32>, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32) -> F32: match xs qs_o: case Nil{} Nil{}: 0.0 case Nil{} qh <> qt: 0.0 case h <> t Nil{}: 0.0 case h <> t qh <> qt: (coul_total(t, qt, ys, qs_e, ke, c2) + coul_self(h, qh, ys, qs_e, ke, c2) : F32)# --- LJ forces: F_i = 24e*(2*u6^2 - u6)/d2 * (pi - pj), u6 = (sig2/d2)^3^2. ---def lj_k(+u6: F32, +eps: F32, +d2: F32) -> F32: (((24.0 * eps : F32) * (((2.0 * lj_sq(u6) : F32) - u6 : F32)) : F32) / d2 : F32)def lj_fvec(+d: G.Vec3, +k: F32) -> G.Vec3: G.Vec3.scale(k, d)def lj_fpair(+d2: F32, +d: G.Vec3, +eps: F32, +sig2: F32) -> G.Vec3: lj_fvec(d, lj_k(lj_sq(lj_cube((sig2 / d2 : F32))), eps, d2))def lj_fself(+d2: F32, +d: G.Vec3, +eps: F32, +sig2: F32, self: Bool) -> G.Vec3: match self: case True{}: G.V3{0.0, 0.0, 0.0} case False{}: lj_fpair(d2, d, eps, sig2)def lj_fgo(+d2: F32, +d: G.Vec3, +eps: F32, +sig2: F32, below: Bool) -> G.Vec3: match below: case True{}: lj_fself(d2, d, eps, sig2, F32.is_eq(d2, 0.0)) case False{}: G.V3{0.0, 0.0, 0.0}def lj_fat(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32) -> G.Vec3: +d2 = G.Vec3.dist2(pi, pj) lj_fgo(d2, G.Vec3.sub(pi, pj), eps, sig2, F32.is_lt(d2, c2))def lj_fhead(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, h: P.Atom) -> G.Vec3: match h: case P.Atom{serial, elem, +pj}: lj_fat(pi, pj, eps, sig2, c2)def lj_frow( +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>) -> G.Vec3: match ys: case Nil{}: G.V3{0.0, 0.0, 0.0} case h <> t: G.Vec3.add(lj_frow(pi, eps, sig2, c2, t), lj_fhead(pi, eps, sig2, c2, h))def lj_fself_atom( h: P.Atom, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32) -> G.Vec3: match h: case P.Atom{serial, elem, +pi}: lj_frow(pi, eps, sig2, c2, ys)def lj_forces( xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32) -> List<&2, G.Vec3>: match xs: case Nil{}: Nil{} case h <> t: lj_fself_atom(h, ys, eps, sig2, c2) <> lj_forces(t, ys, eps, sig2, c2)# --- Steepest descent: x += F * dt, fuel-bounded. Equal-length xs/fs. ---def sd_move(h: P.Atom, f: G.Vec3, +dt: F32) -> P.Atom: match h f: case P.Atom{serial, elem, +pos} G.V3{+fx, +fy, +fz}: P.Atom{serial, elem, G.Vec3.add(pos, G.Vec3.scale(dt, G.V3{fx, fy, fz}))}def sd_sweep( xs: List<&2, P.Atom>, fs: List<&2, G.Vec3>, +dt: F32) -> List<&2, P.Atom>: match xs fs: case Nil{} Nil{}: Nil{} case Nil{} f <> ft: Nil{} case h <> t Nil{}: h <> t case h <> t f <> ft: sd_move(h, f, dt) <> sd_sweep(t, ft, dt)def minimize( fuel: Nat, +xs: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +dt: F32) -> List<&2, P.Atom>: match fuel: case 0n: xs case 1n+f: minimize(f, sd_sweep(xs, lj_forces(xs, xs, eps, sig2, c2), dt), eps, sig2, c2, dt)# --- Per-element Lennard-Jones: illustrative eps (kcal/mol) and sigma (A)# by element number; unknown elements take a generic (0.20, 3.00). These# are placeholders, not a force field. Lorentz-Berthelot combining:# eps = sqrt(ei*ej), sig2 = ((si+sj)/2)^2.def lj_eps(e: U32) -> F32: match e: case 1: 0.03 case 6: 0.066 case 7: 0.17 case 8: 0.21 case 16: 0.25 case _: 0.2def lj_sig(e: U32) -> F32: match e: case 1: 2.5 case 6: 3.5 case 7: 3.25 case 8: 2.96 case 16: 3.55 case _: 3.0def lj_combine_eps(ei: F32, ej: F32) -> F32: F32.sqrt((ei * ej : F32))def lj_combine_sig2(si: F32, sj: F32) -> F32: lj_sq(((si + sj : F32) / 2.0 : F32))def lj_head_elem(+si: U32, +pi: G.Vec3, +c2: F32, h: P.Atom) -> F32: match h: case P.Atom{serial, +elem, +pj}: lj_pair_at(pi, pj, lj_combine_eps(lj_eps(si), lj_eps(elem)), lj_combine_sig2(lj_sig(si), lj_sig(elem)), c2)def lj_row_elem(+si: U32, +pi: G.Vec3, +ys: List<&2, P.Atom>, +c2: F32) -> F32: match ys: case Nil{}: 0.0 case h <> t: (lj_row_elem(si, pi, t, c2) + lj_head_elem(si, pi, c2, h) : F32)def lj_total_elem(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +c2: F32) -> F32: match xs: case Nil{}: 0.0 case +h <> t: (lj_total_elem(t, ys, c2) + lj_row_elem(P.Atom.elem(h), P.Atom.pos(h), ys, c2) : F32)# --- Coulomb forces: F_i = ke*qi*qj/d2^1.5 * (pi - pj), repulsive for# like charges along +d. Self pairs (d2 == 0) and pairs beyond c2 give zero.def coul_k(+qq: F32, +ke: F32, +d2: F32) -> F32: ((ke * qq : F32) / ((d2 * F32.sqrt(d2) : F32)) : F32)def coul_fpair(+d2: F32, +d: G.Vec3, +qq: F32, +ke: F32) -> G.Vec3: G.Vec3.scale(coul_k(qq, ke, d2), d)def coul_fself(+d2: F32, +d: G.Vec3, +qq: F32, +ke: F32, self: Bool) -> G.Vec3: match self: case True{}: G.V3{0.0, 0.0, 0.0} case False{}: coul_fpair(d2, d, qq, ke)def coul_fgo(+d2: F32, +d: G.Vec3, +qq: F32, +ke: F32, below: Bool) -> G.Vec3: match below: case True{}: coul_fself(d2, d, qq, ke, F32.is_eq(d2, 0.0)) case False{}: G.V3{0.0, 0.0, 0.0}def coul_fat(+pi: G.Vec3, +pj: G.Vec3, +qq: F32, +ke: F32, +c2: F32) -> G.Vec3: +d2 = G.Vec3.dist2(pi, pj) coul_fgo(d2, G.Vec3.sub(pi, pj), qq, ke, F32.is_lt(d2, c2))def coul_fhead(+pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, h: P.Atom, qj: F32) -> G.Vec3: match h: case P.Atom{serial, elem, +pj}: coul_fat(pi, pj, (qi * qj : F32), ke, c2)def coul_frow(+pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, ys: List<&2, P.Atom>, qs: List<&2, F32>) -> G.Vec3: match ys qs: case Nil{} Nil{}: G.V3{0.0, 0.0, 0.0} case Nil{} qh <> qt: G.V3{0.0, 0.0, 0.0} case h <> t Nil{}: G.V3{0.0, 0.0, 0.0} case h <> t qh <> qt: G.Vec3.add(coul_frow(pi, qi, ke, c2, t, qt), coul_fhead(pi, qi, ke, c2, h, qh))def coul_fself_atom(h: P.Atom, +qh: F32, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32) -> G.Vec3: match h: case P.Atom{serial, elem, +pi}: coul_frow(pi, qh, ke, c2, ys, qs_e)def coul_forces(xs: List<&2, P.Atom>, +qs_o: List<&2, F32>, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32) -> List<&2, G.Vec3>: match xs qs_o: case Nil{} Nil{}: Nil{} case Nil{} qh <> qt: Nil{} case h <> t Nil{}: Nil{} case h <> t qh <> qt: coul_fself_atom(h, qh, ys, qs_e, ke, c2) <> coul_forces(t, qt, ys, qs_e, ke, c2)# --- Atomic masses (amu) by element; unknown -> 12.0. ---def mass(e: U32) -> F32: match e: case 1: 1.008 case 6: 12.011 case 7: 14.007 case 8: 15.999 case 15: 30.974 case 16: 32.06 case _: 12.0