diff --git a/Project.toml b/Project.toml index dd13e39..5b2e688 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "SimplexCellLists" uuid = "fbe69cd6-f244-4ce4-a90b-8979d86c9ea4" -version = "0.2.3-dev" +version = "0.3.0-dev" authors = ["Nathan Zimmerberg"] [workspace] diff --git a/README.md b/README.md index 30e5069..111e39a 100644 --- a/README.md +++ b/README.md @@ -41,7 +41,7 @@ Neighbor lists and collide forces are configured with a `CollidePolicy`. The policy decides which objects and pairs participate in the neighbor lists (`filter_object`, `filter_pair`), how per-object parameters combine into per-pair parameters (`mix_params`), and which force law applies to each edge -(`nl_edge_forces!`). +(`nl_edge_force`). `DefaultCollidePolicy` stores a stiffness and collision layer masks per object, and applies a soft repulsive potential `E = k/2 * (L - d)²` when two @@ -79,5 +79,5 @@ reused across steps. `setup_neighbors_naive!` is a reference implementation of A custom policy subtypes `CollidePolicy{ObjectParams, PairParams}` with its own parameter types and methods for `filter_object`, `filter_pair`, and -`mix_params`, and can define its own `nl_edge_forces!` methods to change the +`mix_params`, and can define its own `nl_edge_force` methods to change the force law. diff --git a/benchmark/force-bench.jl b/benchmark/force-bench.jl index f53f722..aa14206 100644 --- a/benchmark/force-bench.jl +++ b/benchmark/force-bench.jl @@ -19,29 +19,44 @@ function count_interacting(pos, nl_list, policy) end end +# Each system has copies of its positions, which the benchmark runs in turn. +# When the same positions are run over and over, the branch predictor learns +# which groups of edges have one in range, if there are fewer than about 6000 +# such groups, and the time is too low. So the ring is jittered into copies, +# as its positions would change in a simulation. The isosphere systems have +# enough groups already. systems = [ - ("ring", Ring, Ring.pos), - ("isosphere", Isosphere, Isosphere.pos), + ("ring", Ring, [Ring.pos .+ 0.3f0 .* randn(Xoshiro(s), eltype(Ring.pos), length(Ring.pos)) for s in 1:32]), + ("isosphere", Isosphere, [Isosphere.pos]), # The isosphere is so regular that no edges are in range, so also move its vertices randomly. - ("isosphere jittered", Isosphere, Isosphere.pos .+ 3f0 .* randn(Xoshiro(42), eltype(Isosphere.pos), length(Isosphere.pos))), + ("isosphere jittered", Isosphere, [Isosphere.pos .+ 3f0 .* randn(Xoshiro(42), eltype(Isosphere.pos), length(Isosphere.pos))]), ] -for (name, sys, pos) in systems - local nl = NeighborLists(sys.policy) - setup_neighbors_grid!(nl, pos, sys.inputs) - println("$name: $(length(pos)) beads") +for (name, sys, copies) in systems + local nls = map(copies) do pos + nl = NeighborLists(sys.policy) + setup_neighbors_grid!(nl, pos, sys.inputs) + nl + end + local K = length(copies) + println("$name: $(length(first(copies))) beads", K > 1 ? ", $K copies, edge counts are means" : "") for f in (:PTNL, :CCNL, :LLNL) - local list = getfield(nl, f) - isempty(list) && continue - println(" $f: $(length(list)) edges, $(count_interacting(pos, list, sys.policy)) interacting") + isempty(getfield(first(nls), f)) && continue + local nedges = sum(nl -> length(getfield(nl, f)), nls) + local ninteracting = sum(((pos, nl),) -> count_interacting(pos, getfield(nl, f), sys.policy), zip(copies, nls)) + println(" $f: $(round(Int, nedges/K)) edges, $(round(Int, ninteracting/K)) interacting") end for (FE, calc_type) in ( (ForceEnergyFixedPoint{30, 30}, Float64), (ForceEnergyFixedPoint{30, 30}, Float32), (ForceEnergyFloat64, Float64), ) - local fe = FE(length(pos)) + local fe = FE(length(first(copies))) + # Positions in the calculation type. + local pos_calc = map(pos -> map(p -> SimplexCellLists.Vec3{calc_type}(p), pos), copies) + # Each sample is one call on the next copy. + local k = Ref(0) print(" collide_forces! $FE $calc_type:") - display(@b (zero_force_energy!($fe); collide_forces!($fe, $pos, $nl, $calc_type)) seconds=2) + display(@b (k[] = mod1(k[] + 1, $K)) (zero_force_energy!($fe); collide_forces!($fe, $pos_calc[_], $nls[_], $calc_type)) evals=1 seconds=2) end end diff --git a/src/SimplexCellLists.jl b/src/SimplexCellLists.jl index 3d34249..c9bf70a 100644 --- a/src/SimplexCellLists.jl +++ b/src/SimplexCellLists.jl @@ -157,6 +157,7 @@ end include("collide-forces.jl") export collide_forces! public nl_edge_forces! +public nl_edge_force # Precompile the collide forces for the default policy let NL = NeighborLists{DefaultCollidePolicy, DefaultPairParams} diff --git a/src/collide-forces.jl b/src/collide-forces.jl index 543088c..9898864 100644 --- a/src/collide-forces.jl +++ b/src/collide-forces.jl @@ -1,364 +1,417 @@ """ - nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge, policy::CollidePolicy, calc_type)::Nothing + nl_edge_force(policy::CollidePolicy, x::Simplex{N,T}, y::Simplex{M,T}, (d2, r, s, t), L::T, params)::Tuple{T, Simplex{N,T}, Simplex{M,T}} -Add the collide force and energy of a single neighbor-list edge to `force_energy`. +Return `(energy, fx, fy)`, the collide energy of one neighbor-list edge and +the forces on the vertices of `x` and `y`. -Methods dispatch on the index part types of `edge` and on `policy`, so a custom -`CollidePolicy` with its own `PairParams` can define its own force laws. +- `x` and `y` are the vertex positions of `edge.a` and `edge.b`, relative to + the first vertex of `x`, converted to the calculation type `T`. `x[1]` is + always `zero(Vec3{T})`, positive zero. +- `(d2, r, s, t)` is `dist_sqr_r_s_t(x, y)`, see [`dist_sqr_r_s_t`](@ref). +- `L` is `edge.L` converted to `T`, and `params` is `edge.params`. + +The hook is only called for edges in range, `d2 < L^2`. + +The kind of edge is given by `N` and `M`: 1 for a point, 2 for a line +segment, 3 for a triangle. + +Methods dispatch on `policy`, so a custom `CollidePolicy` with its own +`PairParams` can define its own force laws. Methods should be `@inline`, +branch free, and do no memory access, so the forces of many edges can be +vectorized. """ -function nl_edge_forces! end - -Base.@propagate_inbounds function nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge{PointIdxPart, PointIdxPart, DefaultPairParams}, policy::DefaultCollidePolicy, calc_type::T) where T - x1 = pos[edge.a.i] - x2 = pos[edge.b.i] - d = map(calc_type, x1 - x2) - d2cl = sum(abs2, d) - L0 = calc_type(edge.L) - if d2cl < L0^2 - L = _sqrt_fast(d2cl) - k = calc_type(edge.params.k) - ΔL = L0 - L - E = k*calc_type(1//2)*ΔL^2 - F = (k*ΔL*inv(L))*d - add_bead_force!(force_energy, edge.a.i, F) - add_bead_force!(force_energy, edge.b.i, -F) - add_energy!(force_energy, E) - end - nothing -end +function nl_edge_force end -Base.@propagate_inbounds function nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge{PointIdxPart, CLineIdxPart, DefaultPairParams}, policy::DefaultCollidePolicy, calc_type::T) where T - x1 = pos[edge.a.i] - y1 = pos[edge.b.i] - y2 = pos[edge.b.i + UInt32(1)] - r = map(calc_type, y2 - y1) - p = map(calc_type, x1 - y1) - t = clamp01nan(dot3(p, r)*inv(dot3(r, r))) - d = t*r - p - d2cl = sum(abs2, d) - L0 = calc_type(edge.L) - if d2cl < L0^2 - L = _sqrt_fast(d2cl) - k = calc_type(edge.params.k) - ΔL = L0 - L - E = k*calc_type(1//2)*ΔL^2 - F = (k*ΔL*inv(L))*d - add_bead_force!(force_energy, edge.a.i, -F) - add_bead_force!(force_energy, edge.b.i, (oneunit(t)-t)*F) - add_bead_force!(force_energy, edge.b.i + UInt32(1), t*F) - add_energy!(force_energy, E) - end - nothing +@inline function nl_edge_force(policy::DefaultCollidePolicy, x::Simplex{N, T}, y::Simplex{M, T}, (d2, r, s, t), L::T, params::DefaultPairParams) where {N, M, T} + k = T(params.k) + d = _sqrt_fast(d2) + ΔL = L - d + # The force on the closest point of `y`, and its negative on `x`, spread + # over the vertices by their barycentric weights. + F = scale3(k*ΔL*inv(d), r) + wx = SVector(one(T) - sum(s), s...) + wy = SVector(one(T) - sum(t), t...) + k*T(1//2)*ΔL^2, map(w -> scale3(-w, F), wx), map(w -> scale3(w, F), wy) end -Base.@propagate_inbounds function nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge{PointIdxPart, LineIdxPart, DefaultPairParams}, policy::DefaultCollidePolicy, calc_type::T) where T - x1 = pos[edge.a.i] - y1 = pos[edge.b.i] - y2 = pos[edge.b.j] - r = map(calc_type, y2 - y1) - p = map(calc_type, x1 - y1) - t = clamp01nan(dot3(p, r)*inv(dot3(r, r))) - d = t*r - p - d2cl = sum(abs2, d) - L0 = calc_type(edge.L) - if d2cl < L0^2 - L = _sqrt_fast(d2cl) - k = calc_type(edge.params.k) - ΔL = L0 - L - E = k*calc_type(1//2)*ΔL^2 - F = (k*ΔL*inv(L))*d - add_bead_force!(force_energy, edge.a.i, -F) - add_bead_force!(force_energy, edge.b.i, (oneunit(t)-t)*F) - add_bead_force!(force_energy, edge.b.j, t*F) - add_energy!(force_energy, E) - end - nothing +# Energy, and forces on x[2], y[1] and y[2], without the `k` factor, of the +# spring of rest length `L` between the points at `s` on `x` and `t` on `y`, +# separated by `r`. +@inline function _line_line_u_f(L::T, s::T, t::T, r) where T + d = norm_fast(r) + # A boundary critical point can be out of range, where the spring must + # not pull, and the closest point can be by rounding. + ΔL = relu(L - d) + F = scale3(ΔL*inv(d), r) + T(1//2)*ΔL^2, scale3(-s, F), scale3(one(T) - t, F), scale3(t, F) end -""" -Core helper function for line-line collide force calculation. -Returns (do_force::Bool, u, fp1, fq0, fq1) where: -- do_force: whether forces should be applied -- u: energy (without k factor) -- fp1: force on P1 (without k factor) -- fq0: force on Q0 (without k factor) -- fq1: force on Q1 (without k factor) -Force on P0 can be computed as -(fp1+fq0+fq1) -""" -@inline function line_line_force_core(P0, P1, Q0, Q1, calc_type::T, switchover_scale_unitless, L0) where T - P = map(calc_type, P1-P0) - Q = map(calc_type, Q1-Q0) - P0mQ0 = map(calc_type, P0-Q0) +# The closest point alone gives a discontinuous force between nearly parallel +# line segments, so it is blended with the four boundary critical points. +@inline function nl_edge_force(policy::DefaultCollidePolicy, x::LineSeg{T}, y::LineSeg{T}, (d2, r, s, t), L::T, params::DefaultPairParams) where T + P = x[2] - x[1] + Q = y[2] - y[1] + Q0mP0 = y[1] - x[1] a = dot3(P, P) b = dot3(P, Q) c = dot3(Q, Q) - d = dot3(P, P0mQ0) - e = dot3(Q, P0mQ0) - f = dot3(P0mQ0, P0mQ0) - Δ = a*c - b^2 - #assuming both segments are not zero length - #critical points - s0 = clamp01nan(-d/a) - s1 = clamp01nan((b-d)/a) - t0 = clamp01nan(e/c) - t1 = clamp01nan((b+e)/c) - sbar = clamp01nan((b*e - c*d)/Δ) - tbar = clamp01nan((a*e - b*d)/Δ) - r(s,t) = a*s^2 - 2b*s*t + c*t^2 + 2d*s - 2e*t - mins = sbar - mint = tbar - minr = r(sbar,tbar) - if r(s0,0) < minr - minr = r(s0,0) - mins = s0 - mint = zero(a) - end - if r(s1,1) < minr - minr = r(s1,1) - mins = s1 - mint = one(a) - end - if r(0,t0) < minr - minr = r(0,t0) - mins = zero(a) - mint = t0 - end - if r(1,t1) < minr - minr = r(1,t1) - mins = one(a) - mint = t1 + d = dot3(P, Q0mP0) + e = dot3(Q, Q0mP0) + # `‖P×Q‖^2`, which doesn't cancel for nearly parallel line segments like + # `a*c - b*b` does. `cross3_accurate`, as the rounding of `cross3` is + # about `eps*p*q`, which would switch over parallel line segments. + n = cross3_accurate(P, Q) + Δ = dot3(n, n) + # The closest points on `y` of the end points of `x`, the end point + # candidates of `dist_sqr_r_s_t`, and the closest points on `x` of the end + # points of `y`, its `s` at `t = 0` and `t = 1`. + s0 = clamp01nan(d/a) + s1 = clamp01nan((b + d)/a) + t0 = clamp01nan(-e/c) + t1 = clamp01nan((b - e)/c) + ucl, fp1cl, fq0cl, fq1cl = _line_line_u_f(L, s[1], t[1], r) + # The separation between the points at `s` on `x` and `t` on `y` is + # `Q0mP0 + t*Q - s*P`. + us0, fp1s0, fq0s0, fq1s0 = _line_line_u_f(L, s0, zero(T), fma3(-s0, P, Q0mP0)) + us1, fp1s1, fq0s1, fq1s1 = _line_line_u_f(L, s1, one(T), fma3(-s1, P, y[2] - x[1])) + ut0, fp1t0, fq0t0, fq1t0 = _line_line_u_f(L, zero(T), t0, fma3(t0, Q, Q0mP0)) + ut1, fp1t1, fq0t1, fq1t1 = _line_line_u_f(L, one(T), t1, fma3(t1, Q, Q0mP0 - P)) + uother = us0 + us1 + ut0 + ut1 + fp1other = fp1s0 + fp1s1 + fp1t0 + fp1t1 + fq0other = fq0s0 + fq0s1 + fq0t0 + fq0t1 + fq1other = fq1s0 + fq1s1 + fq1t0 + fq1t1 + + switchover_scale = T(policy.switchover_scale_unitless)*inv(L)^2 # units of 1/length^2 + p = _sqrt_fast(a) + q = _sqrt_fast(c) + pq = p*q + # `1/(p*q)`, or zero when a line segment has zero length, or is so short + # that the inverse would overflow. Such a line segment has no direction, + # so it is taken as parallel to the other, which is the limit of the + # energy as its length goes to zero. The force from the switchover has no + # such limit: the gradient of the switchover along the short line segment + # stays about `‖Q‖*sin²θ*switchover_scale`, so the force depends on its + # direction. Here that force is zero. + ipq = _vifelse(pq > floatmin(T), inv(pq), zero(T)) + # Smoothly switch from 0 for parallel lines to exactly 1 for x ≥ 3, where + # x = Δ/(p*q)*switchover_scale, with continuous first and second + # derivatives: switchover = 1 - (1 - x/3)^3 for x < 3. + Δipq = Δ*ipq + w = relu(1 - Δipq*switchover_scale*T(1//3)) + switchover = 1 - w^3 + d_switchover = w^2 # derivative of switchover with respect to x + u = switchover*ucl + T(1//2)*(1-switchover)*uother + # The gradients of `Δ/(p*q)` are `(2*Q×n - Δ/a*P)/(p*q)` for `P` and + # `(2*n×P - Δ/c*Q)/(p*q)` for `Q`. Both terms are `O(‖n‖)`, so they are + # exactly zero for parallel line segments. The expanded + # `2*(c*P - b*Q) = 2*Q×n` cancels to a rounding of `‖Q‖*‖P‖`. `Δipq*ipq` + # is `sin²θ`, which doesn't overflow like `ipq^2` can. + sin2 = Δipq*ipq + uswitched = (ucl + -T(1//2)*uother) + fp1 = switchover*fp1cl + T(1//2)*(1-switchover)*fp1other + d_switchover_dp1 = scale3(-switchover_scale*d_switchover, fma3(-sin2*(c*ipq), P, scale3(2*ipq, cross3(Q, n)))) + fp1 += d_switchover_dp1*uswitched + + fq0 = switchover*fq0cl + T(1//2)*(1-switchover)*fq0other + fq1 = switchover*fq1cl + T(1//2)*(1-switchover)*fq1other + d_switchover_dq0 = scale3(switchover_scale*d_switchover, fma3(-sin2*(a*ipq), Q, scale3(2*ipq, cross3(n, P)))) + d_switchover_dq1 = - d_switchover_dq0 + fq0 += d_switchover_dq0*uswitched + fq1 += d_switchover_dq1*uswitched + + k = T(params.k) + # The force on x[1] balances the others, by translation invariance. + k*u, SVector(-k*(fp1+fq0+fq1), k*fp1), SVector(k*fq0, k*fq1) +end + +_vertex_idxs(x::PointIdxPart) = (x.i,) +_vertex_idxs(x::CLineIdxPart) = (x.i, x.i + UInt32(1)) +_vertex_idxs(x::LineIdxPart) = (x.i, x.j) +_vertex_idxs(x::TriangleIdxPart) = (x.i, x.j, x.k) + +# The vertex positions `(x, y)` of `edge.a` and `edge.b`, relative to the first +# vertex of `edge.a`: subtracted in the position type, then converted to `T`. +# `x[1]` is exactly `zero(Vec3{T})`. +Base.@propagate_inbounds function _edge_simplexes(pos, edge::NeighborListEdge, ::Type{T}) where T + xa = load_positions(pos, edge.a) + N = length(xa) + x = SVector{N, Vec3{T}}(ntuple(i -> i == 1 ? zero(Vec3{T}) : map(T, xa[i] - xa[1]), Val(N))) + y = map(p -> map(T, p - xa[1]), load_positions(pos, edge.b)) + x, y +end + +# Add the forces `fx` and `fy` on the vertices of the index parts `a` and +# `b`, and the energy `E`, to `force_energy`. +Base.@propagate_inbounds function _add_edge_force_energy!(force_energy::ForceEnergy, a, b, E, fx, fy) + for (n, idx) in enumerate(_vertex_idxs(a)) + add_bead_force!(force_energy, idx, fx[n]) end - d2cl = relu(f+minr) - if d2cl < L0^2 - #return the energy and force on each point give s and t - # doesn't have the e.k factor - function u_f(s,t) - local Pcl = P0mQ0 + s*P - local Qcl = t*Q - local Rcl = Qcl - Pcl - local L = norm_fast(Rcl) - local ΔL = relu(L0 - L) - local u = calc_type(1//2)*ΔL^2 - local FP = (-ΔL*inv(L))*Rcl - (u, s*FP, -(1-t)*FP, -t*FP) - end - switchover_scale = calc_type(switchover_scale_unitless)*inv(L0)^2 # units of 1/length^2 - # Smoothly switch from 0 for parallel lines to exactly 1 for x ≥ 3, where - # x = Δ/sqrt(a*c)*switchover_scale, with continuous first and second - # derivatives: switchover = 1 - (1 - x/3)^3 for x < 3. - w = relu(1 - Δ/sqrt(a*c)*switchover_scale*calc_type(1//3)) - switchover = 1 - w^3 - d_switchover = w^2 # derivative of switchover with respect to x - #P0 is set as origin for now, subtract forces later to get final fp0 - ucl, fp1cl, fq0cl, fq1cl = u_f(mins,mint) - us0, fp1s0, fq0s0, fq1s0 = u_f(s0,0) - us1, fp1s1, fq0s1, fq1s1 = u_f(s1,1) - ut0, fp1t0, fq0t0, fq1t0 = u_f(0,t0) - ut1, fp1t1, fq0t1, fq1t1 = u_f(1,t1) - uother = +( - us0, - us1, - ut0, - ut1, - ) - fp1other = +( - fp1s0, - fp1s1, - fp1t0, - fp1t1, - ) - fq0other = +( - fq0s0, - fq0s1, - fq0t0, - fq0t1, - ) - fq1other = +( - fq1s0, - fq1s1, - fq1t0, - fq1t1, - ) - u = switchover*ucl + calc_type(1//2)*(1-switchover)*uother - p = sqrt(a) - q = sqrt(c) - g = b/(p*q) - uswitched = (ucl + -calc_type(1//2)*uother) - fp1 = switchover*fp1cl + calc_type(1//2)*(1-switchover)*fp1other - d_switchover_dp1 = -switchover_scale*d_switchover*((q/p+b*g/a)*P - 2*g*Q) - fp1 += d_switchover_dp1*uswitched - - fq0 = switchover*fq0cl + calc_type(1//2)*(1-switchover)*fq0other - fq1 = switchover*fq1cl + calc_type(1//2)*(1-switchover)*fq1other - d_switchover_dq0 = switchover_scale*d_switchover*((p/q+b*g/c)*Q - 2*g*P) - d_switchover_dq1 = - d_switchover_dq0 - fq0 += d_switchover_dq0*uswitched - fq1 += d_switchover_dq1*uswitched - - return true, u, fp1, fq0, fq1 - else - return false, zero(L0), zero(P), zero(P), zero(P) + for (n, idx) in enumerate(_vertex_idxs(b)) + add_bead_force!(force_energy, idx, fy[n]) end + add_energy!(force_energy, E) end -Base.@propagate_inbounds function nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge{CLineIdxPart, CLineIdxPart, DefaultPairParams}, policy::DefaultCollidePolicy, calc_type::T) where T - P0 = pos[edge.a.i] - P1 = pos[edge.a.i + UInt32(1)] - Q0 = pos[edge.b.i] - Q1 = pos[edge.b.i + UInt32(1)] - L0 = calc_type(edge.L) - do_force, u, fp1, fq0, fq1 = line_line_force_core(P0, P1, Q0, Q1, calc_type, policy.switchover_scale_unitless, L0) - if do_force - ke = calc_type(edge.params.k) - add_bead_force!(force_energy, edge.a.i, -ke*(fp1+fq0+fq1)) - add_bead_force!(force_energy, edge.a.i + UInt32(1), ke*fp1) - - add_bead_force!(force_energy, edge.b.i, ke*fq0) - add_bead_force!(force_energy, edge.b.i + UInt32(1), ke*fq1) - - add_energy!(force_energy, u*ke) +""" + nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge, policy::CollidePolicy, calc_type::Type{T})::Nothing where T + +Add the collide force and energy of a single neighbor-list edge to +`force_energy`, doing the calculations in the floating-point type `calc_type`. + +This is the scalar path of [`collide_forces!`](@ref) for one edge: the range +check `dist_sqr < L^2`, then the force law [`nl_edge_force`](@ref). To change +the force law, define `nl_edge_force` methods, not methods of this function. +""" +Base.@propagate_inbounds function nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge, policy::CollidePolicy, calc_type::Type{T}) where T + x, y = _edge_simplexes(pos, edge, T) + drst = dist_sqr_r_s_t(x, y) + L = T(edge.L) + if drst[1] < L^2 + E, fx, fy = nl_edge_force(policy, x, y, drst, L, edge.params) + _add_edge_force_energy!(force_energy, edge.a, edge.b, E, fx, fy) end nothing end -Base.@propagate_inbounds function nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge{CLineIdxPart, LineIdxPart, DefaultPairParams}, policy::DefaultCollidePolicy, calc_type::T) where T - P0 = pos[edge.a.i] - P1 = pos[edge.a.i + UInt32(1)] - Q0 = pos[edge.b.i] - Q1 = pos[edge.b.j] - L0 = calc_type(edge.L) - do_force, u, fp1, fq0, fq1 = line_line_force_core(P0, P1, Q0, Q1, calc_type, policy.switchover_scale_unitless, L0) - if do_force - ke = calc_type(edge.params.k) - add_bead_force!(force_energy, edge.a.i, -ke*(fp1+fq0+fq1)) - add_bead_force!(force_energy, edge.a.i + UInt32(1), ke*fp1) - - add_bead_force!(force_energy, edge.b.i, ke*fq0) - add_bead_force!(force_energy, edge.b.j, ke*fq1) - - add_energy!(force_energy, u*ke) - end - nothing +# Split 1:n into nthreads contiguous chunks with sizes differing by at most +# one, and return chunk `chunk`. Chunks past the end are empty. +function _chunk_range(n, chunk, nthreads) + q, rem = divrem(n, nthreads) + start = (chunk - 1)*q + min(chunk - 1, rem) + 1 + stop = chunk*q + min(chunk, rem) + start:stop end -Base.@propagate_inbounds function nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge{LineIdxPart, LineIdxPart, DefaultPairParams}, policy::DefaultCollidePolicy, calc_type::T) where T - P0 = pos[edge.a.i] - P1 = pos[edge.a.j] - Q0 = pos[edge.b.i] - Q1 = pos[edge.b.j] - L0 = calc_type(edge.L) - do_force, u, fp1, fq0, fq1 = line_line_force_core(P0, P1, Q0, Q1, calc_type, policy.switchover_scale_unitless, L0) - if do_force - ke = calc_type(edge.params.k) - add_bead_force!(force_energy, edge.a.i, -ke*(fp1+fq0+fq1)) - add_bead_force!(force_energy, edge.a.j, ke*fp1) - - add_bead_force!(force_energy, edge.b.i, ke*fq0) - add_bead_force!(force_energy, edge.b.j, ke*fq1) - - add_energy!(force_energy, u*ke) +# The scalar reference path, one edge at a time: works for any `calc_type`. +function nl_forces_scalar!(force_energy::ForceEnergy, pos, nl, policy::CollidePolicy, calc_type::Type{T}, chunk, nthreads) where T + for i in _chunk_range(length(nl), chunk, nthreads) + nl_edge_forces!(force_energy, pos, nl[i], policy, calc_type) end nothing end +function nl_forces!(force_energy::ForceEnergy, pos, nl, policy::CollidePolicy, calc_type::Type{T}, chunk, nthreads) where T + nl_forces_scalar!(force_energy, pos, nl, policy, calc_type, chunk, nthreads) +end + +const COLLIDE_SIMD_WIDTH = 8 + +# Return the simplices as one simplex whose coordinates have a lane per +# simplex. +@inline function _simplex_lanes(simplices::Tuple{SVector{N}, Vararg{SVector{N}, M}}) where {N, M} + SVector(ntuple(Val(N)) do v + SVector(ntuple(Val(3)) do axis + SIMD.Vec(ntuple(l -> simplices[l][v][axis], Val(M + 1))) + end) + end) +end + +_nvertices(::Type{PointIdxPart}) = 1 +_nvertices(::Type{CLineIdxPart}) = 2 +_nvertices(::Type{LineIdxPart}) = 2 +_nvertices(::Type{TriangleIdxPart}) = 3 + #= -Based on https://www.geometrictools.com/Documentation/DistancePoint3Triangle3.pdf -// David Eberly, Geometric Tools, Redmond WA 98052 -// Copyright (c) 1998-2022 -// Distributed under the Boost Software License, Version 1.0. -// https://www.boost.org/LICENSE_1_0.txt -// https://www.geometrictools.com/License/Boost/LICENSE_1_0.txt -// Version: 6.0.2022.01.06 +A structure of arrays buffer of `W` neighbor-list edges with index parts of +types `A` and `B`, and params of type `P`, in the calculation type `T`. Lane +`l` holds an edge's positions `x[2:N]` and `y`, its `L`, its index parts, and +its params. Its forces `fx[2:N]` and `fy` are stored over `x[2:N]` and `y`, +so only `fx[1]` and the energy need columns of their own. + +The `T` values share one `Memory`, a column of `W` lanes per coordinate, so +LLVM can tell from the indexes that they don't overlap. With separate +buffers it needs runtime checks, and then finds vectorizing the force loop +not worth it. The getters and setters are `@inline` and unchecked. =# -Base.@propagate_inbounds function nl_edge_forces!(force_energy::ForceEnergy, pos, edge::NeighborListEdge{PointIdxPart, TriangleIdxPart, DefaultPairParams}, policy::DefaultCollidePolicy, calc_type::T) where T - P = pos[edge.a.i] - B = pos[edge.b.i] - E0 = map(calc_type, pos[edge.b.j] - B) - E1 = map(calc_type, pos[edge.b.k] - B) - BP = map(calc_type, B - P) - a = dot3(E0, E0) - b = dot3(E0, E1) - c = dot3(E1, E1) - d = dot3(E0, BP) - e = dot3(E1, BP) - f = dot3(BP, BP) - Δ = a*c - b^2 - invΔ = inv(Δ) - sbar = clamp01nan((b*e - c*d)*invΔ) - tbar = clamp01nan((-a*e + b*d)*invΔ) - - # this ensures sbar and tbar are in the domain - # if they are out of the domain, the min distance is on a boundary - outside = sbar+tbar > one(sbar) - sbar = ifelse(outside,zero(sbar),sbar) - tbar = ifelse(outside,zero(tbar),tbar) - s0 = clamp01nan(-d/a) - t0 = clamp01nan(-e/c) - sd = clamp01nan(-(b-c+d-e)/(a-2b+c)) - td = one(sd) - sd - r(s,t) = a*s^2 + 2b*s*t + c*t^2 + 2d*s + 2e*t - - mins = sbar - mint = tbar - minr = r(sbar,tbar) - if r(s0,0) < minr - minr = r(s0,0) - mins = s0 - mint = zero(a) +struct EdgeBuffer{T, W, A, B, P} + data::Memory{T} + a::Memory{A} + b::Memory{B} + params::Memory{P} +end + +# The first columns of the positions `x[2:N]` and `y`, `L`, the force +# `fx[1]`, and the energy, and the number of columns. +function _columns(::Type{<:EdgeBuffer{T, W, A, B}}) where {T, W, A, B} + N = _nvertices(A) + M = _nvertices(B) + x = 0 + y = x + 3*(N - 1) + L = y + 3*M + fx1 = L + 1 + E = fx1 + 3 + (; x, y, L, fx1, E, total = E + 1) +end + +function EdgeBuffer{T, W, A, B, P}() where {T, W, A, B, P} + EdgeBuffer{T, W, A, B, P}( + Memory{T}(undef, W*_columns(EdgeBuffer{T, W, A, B, P}).total), + Memory{A}(undef, W), + Memory{B}(undef, W), + Memory{P}(undef, W), + ) +end + +# Store `x[2:N]` in lane `l`. `x[1]` is not stored. +@inline function _set_x!(b::EdgeBuffer{T, W, A}, l, x) where {T, W, A} + c = _columns(typeof(b)) + for v in 2:_nvertices(A), axis in 1:3 + @inbounds b.data[l + W*(c.x + 3*(v - 2) + axis - 1)] = x[v][axis] + end +end + +@inline function _set_y!(b::EdgeBuffer{T, W, A, B}, l, y) where {T, W, A, B} + c = _columns(typeof(b)) + for v in 1:_nvertices(B), axis in 1:3 + @inbounds b.data[l + W*(c.y + 3*(v - 1) + axis - 1)] = y[v][axis] + end +end + +# Load `x` from lane `l`, with `x[1]` exactly `zero(Vec3{T})`. +@inline function _get_x(b::EdgeBuffer{T, W, A}, l) where {T, W, A} + c = _columns(typeof(b)) + SVector{_nvertices(A), Vec3{T}}(ntuple(Val(_nvertices(A))) do v + v == 1 && return zero(Vec3{T}) + SVector(ntuple(axis -> @inbounds(b.data[l + W*(c.x + 3*(v - 2) + axis - 1)]), Val(3))) + end) +end + +@inline function _get_y(b::EdgeBuffer{T, W, A, B}, l) where {T, W, A, B} + c = _columns(typeof(b)) + SVector{_nvertices(B), Vec3{T}}(ntuple(Val(_nvertices(B))) do v + SVector(ntuple(axis -> @inbounds(b.data[l + W*(c.y + 3*(v - 1) + axis - 1)]), Val(3))) + end) +end + +# Store `edge` in lane `l`, with its relative positions `x` and `y`, where +# `x[1]` is zero and not stored. +@inline function _set_edge!(b::EdgeBuffer{T, W}, l, x, y, edge::NeighborListEdge) where {T, W} + _set_x!(b, l, x) + _set_y!(b, l, y) + @inbounds b.data[l + W*_columns(typeof(b)).L] = T(edge.L) + @inbounds b.a[l] = edge.a + @inbounds b.b[l] = edge.b + @inbounds b.params[l] = edge.params +end + +@inline _get_L(b::EdgeBuffer{T, W}, l) where {T, W} = @inbounds b.data[l + W*_columns(typeof(b)).L] +@inline _get_a(b::EdgeBuffer, l) = @inbounds b.a[l] +@inline _get_b(b::EdgeBuffer, l) = @inbounds b.b[l] +@inline _get_params(b::EdgeBuffer, l) = @inbounds b.params[l] + +# Copy the edge in lane `src` to lane `dst`. +@inline function _copy_edge!(b::EdgeBuffer{T, W}, dst, src) where {T, W} + for j in 0:_columns(typeof(b)).L + @inbounds b.data[dst + W*j] = b.data[src + W*j] end - if r(0,t0) < minr - minr = r(0,t0) - mins = zero(a) - mint = t0 + @inbounds b.a[dst] = b.a[src] + @inbounds b.b[dst] = b.b[src] + @inbounds b.params[dst] = b.params[src] +end + +# Store the energy and forces of the edge in lane `l`, over its positions. +@inline function _set_energy_forces!(b::EdgeBuffer{T, W}, l, E, fx, fy) where {T, W} + c = _columns(typeof(b)) + _set_x!(b, l, fx) + _set_y!(b, l, fy) + for axis in 1:3 + @inbounds b.data[l + W*(c.fx1 + axis - 1)] = fx[1][axis] end - if r(sd,td) < minr - minr = r(sd,td) - mins = sd - mint = td + @inbounds b.data[l + W*c.E] = E +end + +@inline function _get_energy_forces(b::EdgeBuffer{T, W}, l) where {T, W} + c = _columns(typeof(b)) + fx1 = SVector(ntuple(axis -> @inbounds(b.data[l + W*(c.fx1 + axis - 1)]), Val(3))) + (@inbounds b.data[l + W*c.E]), Base.setindex(_get_x(b, l), fx1, 1), _get_y(b, l) +end + +# Compute the forces of the `W` buffered edges, and add the first `n` to +# `force_energy`. +@inline function _flush_forces!(force_energy, policy, b::EdgeBuffer{T, W}, n) where {T, W} + # Loaded to local variables first, so the force loop only loads from `b.data`. + params = ntuple(l -> _get_params(b, l), Val(W)) + # A plain loop over the lanes that LLVM vectorizes. + for l in 1:W + x = _get_x(b, l) + y = _get_y(b, l) + E, fx, fy = nl_edge_force(policy, x, y, dist_sqr_r_s_t(x, y), _get_L(b, l), params[l]) + _set_energy_forces!(b, l, E, fx, fy) end - d2cl = relu(f+minr) - L0 = calc_type(edge.L) - if d2cl < L0^2 - # dvec is the vector from P to closest point on triangle - dvec = BP + mins*E0 + mint*E1 - L = _sqrt_fast(d2cl) - k = calc_type(edge.params.k) - ΔL = L0 - L - E = k*calc_type(1//2)*ΔL^2 - F = (k*ΔL*inv(L))*dvec - add_bead_force!(force_energy, edge.a.i, -F) - # Forces on triangle vertices - # Closest point on triangle is B + mins*E0 + mint*E1 = (1-mins-mint)*B + mins*(B+E0) + mint*(B+E1) - # = (1-mins-mint)*y[1] + mins*y[2] + mint*y[3] - w0 = oneunit(mins) - mins - mint - add_bead_force!(force_energy, edge.b.i, w0*F) - add_bead_force!(force_energy, edge.b.j, mins*F) - add_bead_force!(force_energy, edge.b.k, mint*F) - add_energy!(force_energy, E) + for l in 1:n + E, fx, fy = _get_energy_forces(b, l) + @inbounds _add_edge_force_energy!(force_energy, _get_a(b, l), _get_b(b, l), E, fx, fy) end - nothing end -function nl_forces!(force_energy::ForceEnergy, pos, nl, policy::CollidePolicy, calc_type::T, chunk, nthreads) where T - # Split 1:length(nl) into nthreads contiguous chunks with sizes differing - # by at most one; chunks past the end are empty. - q, r = divrem(length(nl), nthreads) - start = (chunk - 1)*q + min(chunk - 1, r) + 1 - stop = chunk*q + min(chunk, r) - for i in start:stop - nl_edge_forces!(force_energy, pos, nl[i], policy, calc_type) +#= +The vectorized path for `Float32` and `Float64`, with the same results as the +scalar path, as it calls the same functions on the same values. + +The range check is done `W` edges at a time, with `dist_sqr` on `SIMD.Vec` +lanes, which is bit identical to the scalar `dist_sqr_r_s_t` by its contract. +Each edge in range is copied to the next lane of a structure of arrays +buffer, without `x[1]`, which is zero. When `W` are buffered, their forces +are computed in a plain loop over the lanes, then added to `force_energy` in +edge order. The last edges, fewer than `W`, take the scalar path. + +The `Memory` buffers never leave the function, so LLVM keeps them on the +stack, as long as every function using them is inlined and they are under +about 2000 bytes. +=# +function nl_forces!(force_energy::ForceEnergy, pos, nl::AbstractVector{NeighborListEdge{EA, EB, P}}, policy::CollidePolicy, calc_type::Type{T}, chunk, nthreads) where {EA, EB, P, T <: Union{Float32, Float64}} + W = COLLIDE_SIMD_WIDTH + buf = EdgeBuffer{T, W, EA, EB, P}() + n = 0 + range = _chunk_range(length(nl), chunk, nthreads) + ngroups = length(range) ÷ W + for k in first(range):W:(first(range) + W*ngroups - 1) + lanes = ntuple(l -> @inbounds(nl[k + l - 1]), Val(W)) + simplexes = map(edge -> @inbounds(_edge_simplexes(pos, edge, T)), lanes) + x = _simplex_lanes(map(first, simplexes)) + y = _simplex_lanes(map(last, simplexes)) + # Converted to `T` as a vector. Converting each lane made LLVM load + # them with a slow gather instruction. + L = convert(SIMD.Vec{W, T}, SIMD.Vec(map(edge -> edge.L, lanes))) + hits = SIMD.bitmask(dist_sqr(x, y) < L^2) + while !iszero(hits) + i = k + trailing_zeros(hits) + hits &= hits - one(hits) + edge = @inbounds nl[i] + xi, yi = @inbounds _edge_simplexes(pos, edge, T) + n += 1 + _set_edge!(buf, n, xi, yi, edge) + if n == W + _flush_forces!(force_energy, policy, buf, W) + n = 0 + end + end + end + if n > 0 + # Fill the unused lanes with copies of the last edge, so they hold + # valid values. + for l in (n + 1):W + _copy_edge!(buf, l, n) + end + _flush_forces!(force_energy, policy, buf, n) + end + # Padding the last group instead, with an index clamp, made LLVM + # vectorize the address computations and load the positions with slow + # gather instructions. + for i in (first(range) + W*ngroups):last(range) + @inbounds nl_edge_forces!(force_energy, pos, nl[i], policy, calc_type) end nothing end """ - collide_forces!(force_energy::ForceEnergy, pos, s::NeighborLists, calc_type; chunk=1, nthreads=1)::Nothing + collide_forces!(force_energy::ForceEnergy, pos, s::NeighborLists, calc_type::Type{T}; chunk=1, nthreads=1)::Nothing where T Add the collide forces and energy of every neighbor-list edge in `s` to `force_energy`, doing the calculations in the floating-point type `calc_type`. -Each edge's contribution is computed by [`nl_edge_forces!`](@ref), so a custom -`CollidePolicy` can change the force laws. +Each edge in range, with `dist_sqr < L^2`, gets its force and energy from +[`nl_edge_force`](@ref), so a custom `CollidePolicy` can change the force laws. For multithreading, split the work into `nthreads` chunks: thread `t` calls `collide_forces!` with `chunk=t` and the same `nthreads`, accumulating into @@ -366,7 +419,7 @@ its own `force_energy`, and the per-thread accumulators are merged afterwards with [`combine_force_energy!`](@ref). Every edge is visited exactly once across the chunks. """ -function collide_forces!(force_energy::ForceEnergy, pos, s::NeighborLists, calc_type::T; chunk=1, nthreads=1) where T +function collide_forces!(force_energy::ForceEnergy, pos, s::NeighborLists, calc_type::Type{T}; chunk=1, nthreads=1) where T nl_forces!(force_energy, pos, s.PPNL, s.policy, calc_type, chunk, nthreads) nl_forces!(force_energy, pos, s.PCNL, s.policy, calc_type, chunk, nthreads) nl_forces!(force_energy, pos, s.PLNL, s.policy, calc_type, chunk, nthreads) diff --git a/src/neighbor-lists.jl b/src/neighbor-lists.jl index fceab96..b7d5523 100644 --- a/src/neighbor-lists.jl +++ b/src/neighbor-lists.jl @@ -310,7 +310,7 @@ A single potentially-colliding pair of simplices in a [`NeighborLists`](@ref). - `a::A`: index part of the first simplex, one of `PointIdxPart`, `CLineIdxPart`, or `LineIdxPart`. Triangles never initiate an interaction, so `A` is never `TriangleIdxPart`. - `b::B`: index part of the second simplex, one of `PointIdxPart`, `CLineIdxPart`, `LineIdxPart`, or `TriangleIdxPart`. - `L::Float32`: the interaction distance for the pair, i.e. the sum of the two simplices' radii. -- `params::PairParams`: the mixed pair parameters produced by the policy's `mix_params`, passed to `nl_edge_forces!`. +- `params::PairParams`: the mixed pair parameters produced by the policy's `mix_params`, passed to `nl_edge_force`. """ struct NeighborListEdge{A, B, PairParams} a::A # IdxPart diff --git a/test/Project.toml b/test/Project.toml index 33140f6..43f385d 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -6,6 +6,7 @@ JuMP = "4076af6c-e467-56ae-b986-b466b2749572" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Rotations = "6038ab10-8711-5258-84ad-4b1120ba62dc" +SIMD = "fdea26ae-647d-5447-a871-4b548cad5224" SimplexCellLists = "fbe69cd6-f244-4ce4-a90b-8979d86c9ea4" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/test-collide-forces.jl b/test/test-collide-forces.jl index 8516d89..663490e 100644 --- a/test/test-collide-forces.jl +++ b/test/test-collide-forces.jl @@ -1,5 +1,6 @@ using Test using StaticArrays +using Random: Xoshiro, randperm using Rotations: RotMatrix using SimplexCellLists using SimplexCellLists: dist_sqr, nl_edge_forces!, NeighborListEdge @@ -177,6 +178,58 @@ end end end +@testset "Float32 deep overlap of large simplexes" begin + # The force laws use `d2` from `dist_sqr_r_s_t` as is, so it must stay + # accurate in `Float32` when the simplexes are much larger than their + # separation, here 0.01. + x = SA[1.0, 0.0, 0.0] + y = SA[0.0, 1.0, 0.0] + z = SA[0.0, 0.0, 1.0] + w = x/2 + sqrt(3)/2*y + for (positions, a, b) in ( + ([-37x, 63x, 0.01z - 41w, 0.01z + 59w], LineIdxPart(1, 2), LineIdxPart(3, 4)), + ([0.01z + 3x + 2y, -23x - 7y, 31x - 9y, 2x + 27y], PointIdxPart(1), TriangleIdxPart(2, 3, 4)), + ) + edge = NeighborListEdge(a, b, 6f0, DefaultPairParams(10f0)) + fe32 = ForceEnergyFloat64(length(positions)) + nl_edge_forces!(fe32, positions, edge, DefaultCollidePolicy(), Float32) + fe64 = ForceEnergyFloat64(length(positions)) + nl_edge_forces!(fe64, positions, edge, DefaultCollidePolicy(), Float64) + @test all(f -> all(isfinite, f), get_force(fe32)) + @test get_energy(fe32) ≈ get_energy(fe64) rtol=1e-4 + @test flat_force(fe32) ≈ flat_force(fe64) rtol=1e-3 + end +end + +@testset "$T parallel diagonal line segments" for T in (Float32, Float64) + # `P×P` with `cross3` is a rounding residual, not zero, which switched + # over in `Float32` and gave about 21% too little force. + positions = [SA[0.0, 0.0, 0.0], SA[100.1, 100.1, 0.0], SA[0.0, 0.0, 1e-6], SA[100.1, 100.1, 1e-6]] + edge = NeighborListEdge(LineIdxPart(1, 2), LineIdxPart(3, 4), 2f-6, DefaultPairParams(1f0)) + fe = ForceEnergyFloat64(length(positions)) + nl_edge_forces!(fe, positions, edge, DefaultCollidePolicy(), T) + @test get_energy(fe) ≈ 1e-12 rtol=1e-4 + @test get_force(fe) ≈ [SA[0, 0, -1e-6], SA[0, 0, -1e-6], SA[0, 0, 1e-6], SA[0, 0, 1e-6]] rtol=1e-4 +end + +@testset "$T no force along parallel line segments" for T in (Float32, Float64) + # The line segments are parallel and separated along z, so the forces + # must point along z. The switchover gradient written as `c*P - b*Q` + # cancels to a rounding of `‖P‖*‖Q‖` even when `P×Q` is zero, which gave + # a force of 3.8e-6 along the line segments, in y, in `Float32`, larger + # than the 1e-6 force along z that pushes them apart. + for (Q1, f) in ( + (SA[3.0, 26.0, 1e-6], [-1, -1, 1, 1]*1e-6), + (SA[6.0, 52.0, 1e-6], [-1, -0.5, 1.25, 0.25]*1e-6), + ) + positions = [SA[0.0, 0.0, 0.0], SA[3.0, 26.0, 0.0], SA[0.0, 0.0, 1e-6], Q1] + edge = NeighborListEdge(LineIdxPart(1, 2), LineIdxPart(3, 4), 2f-6, DefaultPairParams(1f0)) + fe = ForceEnergyFloat64(length(positions)) + nl_edge_forces!(fe, positions, edge, DefaultCollidePolicy(), T) + @test get_force(fe) ≈ [SA[0, 0, fz] for fz in f] rtol=1e-4 + end +end + @testset "collide_forces! integration" begin policy = DefaultCollidePolicy() # 3×3×4 grid, spacing 0.4: with radius 0.3 overlaps are guaranteed, but @@ -219,6 +272,12 @@ end @test get_energy(fe_manual) == get_energy(fe) @test get_force(fe_manual) == get_force(fe) + # the scalar fallback works for other calc types + fe_big = DebugForceEnergy{BigFloat}(length(positions)) + collide_forces!(fe_big, positions, nl, BigFloat) + @test get_energy(fe_big) ≈ get_energy(fe) + @test get_force(fe_big) ≈ get_force(fe) + # the chunked path visits every edge exactly once for nthreads in (1, 2, 3, 7, 100) local fe_chunked = ForceEnergyFloat64(length(positions)) @@ -238,4 +297,74 @@ end @test iszero(get_energy(fe_empty)) @test all(iszero, get_force(fe_empty)) end + +const NL_FIELDS = (:PPNL, :PCNL, :PLNL, :PTNL, :CCNL, :CLNL, :LLNL) + +make_part(::Type{PointIdxPart}, v) = PointIdxPart(v[1]) +make_part(::Type{CLineIdxPart}, v) = CLineIdxPart(v[1]) +make_part(::Type{LineIdxPart}, v) = LineIdxPart(v[1], v[2]) +make_part(::Type{TriangleIdxPart}, v) = TriangleIdxPart(v[1], v[2], v[3]) + +# `nedges` edges of every type between random beads, with no bead shared by +# the two simplexes of an edge. +function random_neighbor_lists(rng, policy, nbeads, nedges, make_params) + nl = NeighborLists(policy) + for f in NL_FIELDS + list = getfield(nl, f) + A = fieldtype(eltype(list), :a) + B = fieldtype(eltype(list), :b) + for _ in 1:nedges + # Odd beads, so the second bead of a `CLineIdxPart` isn't used + # by another part. + v = UInt32.(2 .* randperm(rng, nbeads ÷ 2)[1:6] .- 1) + push!(list, NeighborListEdge(make_part(A, v[1:3]), make_part(B, v[4:6]), 3*rand(rng, Float32), make_params())) + end + end + nl +end + +# Check that `collide_forces!` gives exactly the forces and energy of the +# scalar path, and return the energy. +function check_matches_scalar(pos, nl, T; nthreads) + fe = ForceEnergyFloat64(length(pos)) + fe_scalar = ForceEnergyFloat64(length(pos)) + for chunk in 1:nthreads + collide_forces!(fe, pos, nl, T; chunk, nthreads) + for f in NL_FIELDS + SimplexCellLists.nl_forces_scalar!(fe_scalar, pos, getfield(nl, f), nl.policy, T, chunk, nthreads) + end + end + @test get_energy(fe) == get_energy(fe_scalar) + @test get_force(fe) == get_force(fe_scalar) + get_energy(fe) +end + +# A policy with params that aren't isbits. +struct BoxedPairParams + k::Vector{Float32} +end +struct BoxedPolicy <: SimplexCellLists.CollidePolicy{DefaultObjectParams, BoxedPairParams} end +@inline function SimplexCellLists.nl_edge_force(::BoxedPolicy, x, y, drst, L, params::BoxedPairParams) + SimplexCellLists.nl_edge_force(DefaultCollidePolicy(), x, y, drst, L, DefaultPairParams(params.k[1])) +end + +@testset "vectorized path matches the scalar path" begin + rng = Xoshiro(1234) + nbeads = 300 + # No group, one group, and several flushes with a last group and edges + # left over. + for P in (Float64, Float32), nedges in (0, 5, 8, 203) + pos = [SVector{3, P}(4 .* rand(rng, 3)) for _ in 1:nbeads] + nl = random_neighbor_lists(rng, DefaultCollidePolicy(), nbeads, nedges, () -> DefaultPairParams(10*rand(rng, Float32))) + for T in (Float32, Float64), nthreads in (1, 3) + E = check_matches_scalar(pos, nl, T; nthreads) + nedges == 203 && @test E > 0 + end + end + pos = [SVector{3, Float64}(4 .* rand(rng, 3)) for _ in 1:nbeads] + boxed_nl = random_neighbor_lists(rng, BoxedPolicy(), nbeads, 203, () -> BoxedPairParams([10*rand(rng, Float32)])) + for T in (Float32, Float64) + @test check_matches_scalar(pos, boxed_nl, T; nthreads=1) > 0 + end +end nothing