Skip to content

Commit f599241

Browse files
committed
[orx-force-2d] Add ManyNodeForce
1 parent 0d0570d commit f599241

3 files changed

Lines changed: 243 additions & 0 deletions

File tree

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
package org.openrndr.extra.force2d
2+
3+
import org.openrndr.extra.force2d.quadtree.QuadTreeNode
4+
import org.openrndr.math.Vector2
5+
import org.openrndr.math.Vector2.Axis
6+
import org.openrndr.shape.Rectangle
7+
import kotlin.math.sqrt
8+
import kotlin.random.Random
9+
10+
private fun Rectangle.outerSquare(): Rectangle {
11+
return if (this.majorAxis == Axis.X) {
12+
Rectangle(x, y, width, width)
13+
} else {
14+
Rectangle(x, y, height, height)
15+
}
16+
}
17+
18+
19+
class ManyNodeForce(val body: Body): Force {
20+
21+
var theta = 1.0
22+
val random = Random(0)
23+
lateinit private var quadTree: QuadTreeNode
24+
25+
var strength: (Node) -> Double = { -0.0 }
26+
27+
private var strengths = DoubleArray(body.nodes.size)
28+
29+
30+
override suspend fun initializeFrame(body: Body) {
31+
body.updateBounds()
32+
val bounds = body.bounds.outerSquare()
33+
quadTree = QuadTreeNode(bounds.x, bounds.y, bounds.x + bounds.width, bounds.y + bounds.height)
34+
for (i in body.nodes.indices) {
35+
strengths[i] = strength(body.nodes[i])
36+
}
37+
38+
for ((index, node) in body.nodes.withIndex()) {
39+
quadTree.add(node, strengths[index])
40+
}
41+
42+
}
43+
44+
private fun applyForce(node: Node, delta: Vector2, distance: Double, distance2: Double, strength: Double, alpha: Double) {
45+
node.velocity += delta * (strength * alpha / (distance2 * distance))
46+
}
47+
private fun accumulateForce(node: Node, octNode: QuadTreeNode, theta2: Double, distanceMin2: Double, distanceMax2: Double, alpha: Double) {
48+
val delta = octNode.centerOfMass - node.position
49+
val d2 = delta.dot(delta)
50+
val w = octNode.xmax - octNode.xmin
51+
52+
if (w * w / theta2 < d2) {
53+
if (d2 < distanceMax2) {
54+
if (d2 < distanceMin2) {
55+
applyForce(node, delta, sqrt(d2), distanceMin2, octNode.totalStrength, alpha)
56+
} else {
57+
applyForce(node, delta, sqrt(d2), d2, octNode.totalStrength, alpha)
58+
}
59+
}
60+
} else {
61+
if (octNode.children != null) {
62+
for (child in octNode.children!!) {
63+
if (child != null) {
64+
accumulateForce(node, child, theta2, distanceMin2, distanceMax2, alpha)
65+
}
66+
}
67+
} else if (octNode.node != null && octNode.node != node) {
68+
if (d2 < distanceMax2) {
69+
var dist2 = d2
70+
var dist = sqrt(d2)
71+
if (d2 < 0.00001) { // Jitter for coincident nodes
72+
val rx = (random.nextDouble() - 0.5) * 0.01
73+
val ry = (random.nextDouble() - 0.5) * 0.01
74+
val jitteredDelta = Vector2(rx, ry)
75+
dist2 = jitteredDelta.dot(jitteredDelta)
76+
dist = sqrt(dist2)
77+
applyForce(node, jitteredDelta, dist, maxOf(dist2, distanceMin2), octNode.totalStrength, alpha)
78+
} else {
79+
applyForce(node, delta, dist, maxOf(dist2, distanceMin2), octNode.totalStrength, alpha)
80+
}
81+
}
82+
}
83+
}
84+
}
85+
86+
override suspend fun apply(body: Body, dt: Double) {
87+
88+
val theta2 = theta * theta
89+
for (i in body.nodes.indices) {
90+
val node = body.nodes[i]
91+
accumulateForce(node, quadTree, theta2, 0.0, 360.0*360.0, 1.0)
92+
}
93+
94+
}
95+
}
96+
97+
fun Body.manyNodeForce(configure: ManyNodeForce.() -> Unit) {
98+
forces.add(ManyNodeForce(this).apply(configure))
99+
}
Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package org.openrndr.extra.force2d.quadtree
2+
3+
import org.openrndr.extra.force2d.Node
4+
import org.openrndr.math.Vector2
5+
import org.openrndr.math.Vector3
6+
7+
internal class QuadTreeNode(
8+
var xmin: Double, var ymin: Double,
9+
var xmax: Double, var ymax: Double
10+
) {
11+
var centerOfMass: Vector2 = Vector2.ZERO
12+
var totalStrength: Double = 0.0
13+
var children: Array<QuadTreeNode?>? = null
14+
var node: Node? = null
15+
16+
fun add(forceNode: Node, strengthValue: Double = 1.0) {
17+
if (node == null && children == null) {
18+
node = forceNode
19+
totalStrength = strengthValue
20+
centerOfMass = forceNode.position
21+
return
22+
}
23+
24+
if (children == null) {
25+
children = arrayOfNulls(4)
26+
val existingNode = node!!
27+
val existingStrength = totalStrength // This should be retrieved/passed if we want per-node strength
28+
// Re-add existing node to children
29+
node = null
30+
totalStrength = 0.0
31+
centerOfMass = Vector2.ZERO
32+
addToChildren(existingNode, existingStrength)
33+
}
34+
35+
addToChildren(forceNode, strengthValue)
36+
updateCenterOfMass(forceNode, strengthValue)
37+
}
38+
39+
private fun addToChildren(forceNode: Node, strengthValue: Double) {
40+
val xm = (xmin + xmax) / 2.0
41+
val ym = (ymin + ymax) / 2.0
42+
43+
val i = if (forceNode.position.x >= xm) 1 else 0
44+
val j = if (forceNode.position.y >= ym) 1 else 0
45+
val idx = i + (j shl 1)
46+
47+
if (children!![idx] == null) {
48+
val nxMin = if (i == 1) xm else xmin
49+
val nxMax = if (i == 1) xmax else xm
50+
val nyMin = if (j == 1) ym else ymin
51+
val nyMax = if (j == 1) ymax else ym
52+
children!![idx] = QuadTreeNode(nxMin, nyMin, nxMax, nyMax)
53+
}
54+
children!![idx]!!.add(forceNode, strengthValue)
55+
}
56+
57+
private fun updateCenterOfMass(forceNode: Node, strengthValue: Double) {
58+
val newTotalStrength = totalStrength + strengthValue
59+
if (newTotalStrength != 0.0) {
60+
centerOfMass = (centerOfMass * totalStrength + forceNode.position * strengthValue) / newTotalStrength
61+
} else {
62+
// If total strength is 0, we just use the last position to avoid NaN,
63+
// though Barnes-Hut with 0 strength nodes is unusual.
64+
centerOfMass = forceNode.position
65+
}
66+
totalStrength = newTotalStrength
67+
}
68+
}
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import kotlinx.coroutines.runBlocking
2+
import org.openrndr.application
3+
import org.openrndr.color.ColorRGBa
4+
import org.openrndr.extra.force2d.Body
5+
import org.openrndr.extra.force2d.ForceSimulation
6+
import org.openrndr.extra.force2d.Link
7+
import org.openrndr.extra.force2d.Node
8+
import org.openrndr.extra.force2d.linkLengthConstraint
9+
import org.openrndr.extra.force2d.manyNodeForce
10+
import org.openrndr.extra.force2d.rectangularBoundsConstraint
11+
import org.openrndr.extra.kdtree.kdTree
12+
import org.openrndr.extra.noise.scatter
13+
import org.openrndr.math.Vector2
14+
import org.openrndr.shape.LineSegment
15+
16+
/**
17+
* Demonstrates a simple force graph layout
18+
*/
19+
fun main() {
20+
application {
21+
configure {
22+
width = 720
23+
height = 720
24+
}
25+
program {
26+
val pts = drawer.bounds.offsetEdges(-30.0).scatter(7.0)
27+
val index = pts.mapIndexed { index, vector2 -> vector2 to index }.toMap()
28+
val kd = pts.kdTree()
29+
30+
val nodes = pts.map {
31+
Node(it, it, Vector2.ZERO, radius = 4.0)
32+
}
33+
34+
val links = pts.flatMapIndexed { source, it ->
35+
kd.findKNearest(it, 2).map {
36+
val idx = index.getValue(it)
37+
Link(source, idx)
38+
}
39+
}
40+
41+
val body = Body(nodes, links = links).apply {
42+
linkLengthConstraint {
43+
compliance = 1.0
44+
iterations = 1
45+
for (i in nodes.indices) {
46+
// restLengths[i] = 30.0
47+
}
48+
49+
}
50+
manyNodeForce {
51+
52+
theta = 1.0
53+
strength = { -3.0 }
54+
}
55+
rectangularBoundsConstraint {
56+
bounds = drawer.bounds.offsetEdges(-10.0)
57+
}
58+
}
59+
60+
val sim = ForceSimulation(mutableListOf(body))
61+
62+
extend {
63+
runBlocking {
64+
sim.simulate(1.0 / 60.0, 10)
65+
}
66+
67+
drawer.stroke = ColorRGBa.PINK
68+
drawer.lineSegments(body.links.map {
69+
LineSegment(body.nodes[it.source].position, body.nodes[it.target].position)
70+
})
71+
72+
drawer.circles(body.nodes.map { it.position }, 4.0)
73+
}
74+
}
75+
}
76+
}

0 commit comments

Comments
 (0)