feat: Velocity rescaling

Author: 0ffz

CHANGELOG.md

@@ -11,13 +11,14 @@ development.
 ### Added
 
 - Started a changelog based on keep a changelog
-- Mean square velocity calculation via a compute shader that sums log_2 times; plan to extract into a general API in the
-  future
+- **Mean square velocity** calculation via a compute shader that sums log_2 times; plan to extract into a general API in
+  the future
+- **Support for recent projects in web** via local storage, plus "Save as" button to save projects back to disk
+- **Target velocity** via rescaling based on mean square velocity
 - Added more options to graphs displayed in GUI like toggling on/off and clearing
 - Added helper function for drawing graphs for values that update over time
 - Added banner and icon for the project
 - Grid toggle button under visual options
-- Support for recent projects in web via local storage, plus "Save as" button to save projects back to disk
 
 ### Changed
 

particles-config/src/commonMain/kotlin/me/dvyy/particles/dsl/Simulation.kt

@@ -10,6 +10,8 @@ data class Simulation(
     val dT: Double = 0.001,
     val maxVelocity: Double = 20.0,
     val maxForce: Double = 100000.0,
+    val targetVelocity: Double = 0.0,
+    val targetVelocityStrength: Double = 0.0,
     val threeDimensions: Boolean = false,
     val passesPerFrame: Int = 100,
     val size: Size = Size(),

particles-kool/src/commonMain/kotlin/me/dvyy/particles/ParticlesScene.kt

@@ -77,7 +77,10 @@ class ParticlesScene(
         }
 
         offsetsShader.addTo(computePass) // Calculate offsets (start index in particles array for each grid cell)
-        val fieldsPasses = fieldsShader.addTo(computePass) // Run force computations based on particle interactions
+        val fieldsPasses = fieldsShader.addTo(
+            computePass,
+            meanSquareDataShader
+        ) // Run force computations based on particle interactions
         convertShader.addTo(computePass) // Convert particles to different types as needed
 
         // == DATA COLLECTION ==

particles-kool/src/commonMain/kotlin/me/dvyy/particles/compute/data/MeanSquareVelocities.kt

@@ -40,16 +40,32 @@ class MeanSquareVelocities(
             }
         }
     }
+
+    private val outputShader = KslComputeShader("MeanSquareVelocities_output") {
+//        val struct = struct { SimulationStatisticsStruct() }
+//        val total = storage("total", struct)
+        val inputs = storage<KslFloat1>("inputs")
+        val output = storage<KslFloat1>("outputStorage")
+        computeStage(1) {
+            main {
+                output[0.const] = inputs[0.const]
+            }
+        }
+    }
+
+    private var outputBind by outputShader.storage("outputStorage")//.uniformStruct("total", ::SimulationStatisticsStruct)
+
     private var inputs by reduce.storage("inputs")
     private var outputs by reduce.storage("outputs")
     private var total by reduce.uniform1i("total")
 
-    val inputBuffer = Buffers.floats(buffers.count)
-    val outputBuffer = Buffers.floats(buffers.count)
+    private val inputBuffer = Buffers.floats(buffers.count)
+    private val outputBuffer = Buffers.floats(buffers.count)
+    val output = Buffers.floats(1)
 
     private val roundedUp = 1 shl (32 - (buffers.count - 1).countLeadingZeroBits())
     private val iterations = roundedUp.countTrailingZeroBits()
-    val readBack = if (iterations % 2 == 0) inputBuffer else outputBuffer
+    private val readBack = if (iterations % 2 == 0) inputBuffer else outputBuffer
 
     fun addTo(
         pass: ComputePass,
@@ -73,6 +89,10 @@ class MeanSquareVelocities(
                     }
                 }
             }
+            addTask(outputShader.apply {
+                storage("inputs").set(readBack)
+                outputBind = output
+            }, Vec3i(1, 1, 1))
         }
     }
 }

particles-kool/src/commonMain/kotlin/me/dvyy/particles/compute/simulation/FieldsMultiPasses.kt

@@ -2,6 +2,7 @@ package me.dvyy.particles.compute.simulation
 
 import de.fabmax.kool.pipeline.ComputePass
 import me.dvyy.particles.compute.ParticleBuffers
+import me.dvyy.particles.compute.data.MeanSquareVelocities
 import me.dvyy.particles.compute.forces.ForcesDefinition
 import me.dvyy.particles.config.ConfigRepository
 
@@ -32,10 +33,14 @@ class FieldsMultiPasses(
         halfStep.boxMax = configRepo.boxSize
     }
 
-    fun addTo(computePass: ComputePass): List<Pair<ComputePass.Task, ComputePass.Task>> {
+    fun addTo(
+        computePass: ComputePass,
+        velocitiesShader: MeanSquareVelocities,
+    ): List<Pair<ComputePass.Task, ComputePass.Task>> {
         val config = configRepo.config.value
 
         initBuffers()
+        fields.velocityData = velocitiesShader.output
         val passes = buildList {
             repeat(config.simulation.passesPerFrame) { passIndex ->
                 val halfStep = computePass.addTask(halfStep.shader, numGroups = configRepo.numGroups).apply {
@@ -53,6 +58,8 @@ class FieldsMultiPasses(
                             fields.params.set {
                                 maxVelocity.set(simulation.maxVelocity.toFloat())
                                 maxForce.set(simulation.maxForce.toFloat())
+                                targetVelocity.set(simulation.targetVelocity.toFloat())
+                                targetVelocityFixStrength.set(simulation.targetVelocityStrength.toFloat())
                             }
                             val count = configRepo.count
                             fields.count = count

particles-kool/src/commonMain/kotlin/me/dvyy/particles/compute/simulation/FieldsShader.kt

@@ -36,6 +36,7 @@ class FieldsShader(
             val velocities = storage<KslFloat4>("velocities")
             val forces = storage<KslFloat4>("forces")
             val localNeighbours = storage<KslFloat1>("localNeighbours")
+            val velocityData = storage<KslFloat1>("velocityData")
 
             val particleTypes = storage<KslInt1>("particleTypes")
 
@@ -188,10 +189,23 @@ class FieldsShader(
                 `if`(length(nextVelocity) gt params.maxVelocity.ksl) {
                     nextVelocity set normalize(nextVelocity) * params.maxVelocity.ksl
                 }
+                val target = params.targetVelocity.ksl
+                val totalSqrtVelocities = float1Var(velocityData[0.const])
+                val average = totalSqrtVelocities / count.toFloat1()
+//                val halved = totalSqrtVelocities/2f.const
+//                val degreesOfFreedom = 2f.const
+                val strength = params.targetVelocityFixStrength.ksl
+                // nudge particles towards target velocity
+                nextVelocity set nextVelocity * sqrt(
+                    1f.const + (dT * strength) * ((target) / max(
+                        average,
+                        0.1f.const
+                    ) - 1f.const)
+                )
 
                 forces[id] = float4Value(nextForce, 0f)
                 velocities[id] = float4Value(nextVelocity, 0f)
-                localNeighbours[id] = localCount
+                localNeighbours[id] = totalSqrtVelocities
             }
         }
 
@@ -214,6 +228,7 @@ class FieldsShader(
     var localNeighbours by shader.storage("localNeighbours")
     var forces by shader.storage("forces")
     var particleTypes by shader.storage("particleTypes")
+    var velocityData by shader.storage("velocityData")
 }
 //
 //context(scope: KslScopeBuilder)

particles-kool/src/commonMain/kotlin/me/dvyy/particles/compute/simulation/SimulationParametersStruct.kt

@@ -6,4 +6,6 @@ import de.fabmax.kool.util.Struct
 class SimulationParametersStruct : Struct("SimulationParametersStruct", MemoryLayout.Std140) {
     val maxVelocity = float1()
     val maxForce = float1()
+    val targetVelocity = float1()
+    val targetVelocityFixStrength = float1()
 }

particles-kool/src/commonMain/kotlin/me/dvyy/particles/helpers/Buffers.kt

@@ -17,9 +17,9 @@ object Buffers {
     }
 
     fun randomPosition(size: Vec3f) = Vec4f(
-        Random.Default.nextDouble(size.x.toDouble()).toFloat(),
-        Random.Default.nextDouble(size.y.toDouble()).toFloat(),
-        if (size.z == 0f) 0f else Random.Default.nextDouble(size.z.toDouble()).toFloat(),
+        Random.nextDouble(size.x.toDouble()).toFloat(),
+        Random.nextDouble(size.y.toDouble()).toFloat(),
+        if (size.z == 0f) 0f else Random.nextDouble(size.z.toDouble()).toFloat(),
         0f
     )
 

particles-kool/src/commonMain/kotlin/me/dvyy/particles/ui/viewmodels/ParticlesViewModel.kt

@@ -52,6 +52,12 @@ class ParticlesViewModel(
                 UiConfigurable.Slider("dT", state.dT, 0f, 0.01f, precision = 4) {
                     updateState { copy(dT = it.toDouble()) }
                 },
+                UiConfigurable.Slider("Target Velocity", state.targetVelocity, 0f, 100f) {
+                    updateState { copy(targetVelocity = it.toDouble()) }
+                },
+                UiConfigurable.Slider("Targetting Strength", state.targetVelocityStrength, 0f, 100f) {
+                    updateState { copy(targetVelocityStrength = it.toDouble()) }
+                },
                 UiConfigurable.Slider("Max Velocity", state.maxVelocity, 0f, 100f) {
                     updateState { copy(maxVelocity = it.toDouble()) }
                 },
@@ -89,8 +95,8 @@ class ParticlesViewModel(
     }
 
     suspend fun readbackMeanSquareVelocity() {
-        val result = Float32Buffer(buffers.count)
-        meanSquareData.readBack.downloadData(result)
+        val result = Float32Buffer(1)
+        meanSquareData.output.downloadData(result)
         val msqV = result[0] / buffers.count
         msqvOverTime.pushNewValueRight(msqV)
         meanSquareVelocity.update { msqV }