example: add a pendulum simulation (#9992)

pull/10016/head
Jordan Bonecutter 2021-05-05 06:09:38 -07:00 committed by GitHub
parent 4ac751d773
commit 2c4a59f367
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 366 additions and 0 deletions

View File

@ -0,0 +1,366 @@
// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
// sim.v * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
// created by: jordan bonecutter * * * * * * * * * * * * * * * * * * *
// jpbonecutter@gmail.com * * * * * * * * * * * * * * * * * * * * * *
// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
//
// I wrote the pendulum simulator to learn V, I think it could be a
// good addition to the examples directory.
// Essentially, the pendulum sim runs a simulation of a pendulum with
// a metallic tip swinging over three magnets.
// I run this simulation with the initial position at each pixel in an
// image and color the pixel according to the magnet over which it
// finally rests.
// I used some fun features in V like coroutines, channels,
// struct embedding, mutability, methods, and the like.
import math
import os
import term
import runtime
// customisable through setting VJOBS
const parallel_workers = runtime.nr_jobs()
const width = 800
const height = 600
struct Vec3D {
x f64
y f64
z f64
}
fn (v Vec3D) add(v2 Vec3D) Vec3D {
return Vec3D{
x: v.x + v2.x
y: v.y + v2.y
z: v.z + v2.z
}
}
fn (v Vec3D) dot(v2 Vec3D) f64 {
return (v.x * v2.x) + (v.y * v2.y) + (v.z * v2.z)
}
fn (v Vec3D) scale(scalar f64) Vec3D {
return Vec3D{
x: v.x * scalar
y: v.y * scalar
z: v.z * scalar
}
}
fn (v Vec3D) norm_squared() f64 {
return v.dot(v)
}
fn (v Vec3D) norm() f64 {
return math.sqrt(v.norm_squared())
}
struct SimState {
mut:
position Vec3D
velocity Vec3D
accel Vec3D
}
// magnets lie at [
// math.cos(index * 2 * math.pi / 3) * magnet_spacing
// math.sin(index * 2 * math.pi / 3) * magnet_spacing
// -magnet_height
// ]
struct SimParams {
rope_length f64
bearing_mass f64
magnet_spacing f64
magnet_height f64
magnet_strength f64
gravity f64
}
fn (params SimParams) get_rope_vector(state SimState) Vec3D {
rope_origin := Vec3D{
x: 0
y: 0
z: params.rope_length
}
return state.position.add(rope_origin.scale(-1))
}
fn (mut state SimState) satisfy_rope_constraint(params SimParams) {
mut rope_vector := params.get_rope_vector(state)
rope_vector = rope_vector.scale(params.rope_length / rope_vector.norm())
state.position = Vec3D{
x: 0
y: 0
z: params.rope_length
}.add(rope_vector)
}
fn (params SimParams) get_grav_force(state SimState) Vec3D {
return Vec3D{
x: 0
y: 0
z: -params.bearing_mass * params.gravity
}
}
fn (params SimParams) get_magnet_position(theta f64) Vec3D {
return Vec3D{
x: math.cos(theta) * params.magnet_spacing
y: math.sin(theta) * params.magnet_spacing
z: -params.magnet_height
}
}
fn (params SimParams) get_magnet_force(theta f64, state SimState) Vec3D {
magnet_position := params.get_magnet_position(theta)
mut diff := magnet_position.add(state.position.scale(-1))
distance_squared := diff.norm_squared()
diff = diff.scale(1.0 / math.sqrt(distance_squared))
return diff.scale(params.magnet_strength / distance_squared)
}
fn (params SimParams) get_magnet_dist(theta f64, state SimState) f64 {
return params.get_magnet_position(theta).add(state.position.scale(-1)).norm()
}
fn (params SimParams) get_magnet1_force(state SimState) Vec3D {
return params.get_magnet_force(0.0 * math.pi / 3.0, state)
}
fn (params SimParams) get_magnet2_force(state SimState) Vec3D {
return params.get_magnet_force(2.0 * math.pi / 3.0, state)
}
fn (params SimParams) get_magnet3_force(state SimState) Vec3D {
return params.get_magnet_force(4.0 * math.pi / 3.0, state)
}
fn (params SimParams) get_tension_force(state SimState, f_passive Vec3D) Vec3D {
rope_vector := params.get_rope_vector(state)
rope_vector_norm := rope_vector.scale(1.0 / rope_vector.norm())
return rope_vector_norm.scale(-1.0 * rope_vector_norm.dot(f_passive))
}
fn (mut state SimState) increment(delta_t f64, params SimParams) {
// basically just add up all forces =>
// get an accelleration =>
// add to velocity =>
// ensure rope constraint is satisfied
// force due to gravity
f_gravity := params.get_grav_force(state)
// force due to each magnet
f_magnet1 := params.get_magnet1_force(state)
// force due to each magnet
f_magnet2 := params.get_magnet2_force(state)
// force due to each magnet
f_magnet3 := params.get_magnet3_force(state)
// passive forces
f_passive := f_gravity.add(f_magnet1.add(f_magnet2.add(f_magnet3)))
// force due to tension of the rope
f_tension := params.get_tension_force(state, f_passive)
// sum up all the fores
f_sum := f_tension.add(f_passive)
// get the acceleration
accel := f_sum.scale(1.0 / params.bearing_mass)
state.accel = accel
// update the velocity
state.velocity = state.velocity.add(accel.scale(delta_t))
// update the position
state.position = state.position.add(state.velocity.scale(delta_t))
// ensure the position satisfies rope constraint
state.satisfy_rope_constraint(params)
}
fn (state SimState) done() bool {
return state.velocity.norm() < 0.05 && state.accel.norm() < 0.01
}
struct PPMWriter {
mut:
file os.File
}
struct ImageSettings {
width int
height int
}
struct Pixel {
r byte
g byte
b byte
}
fn (mut writer PPMWriter) start_for_file(fname string, settings ImageSettings) {
writer.file = os.create(fname) or { panic("can't create file $fname") }
writer.file.writeln('P6 $settings.width $settings.height 255') or {}
}
fn (mut writer PPMWriter) next_pixel(p Pixel) {
writer.file.write([p.r, p.g, p.b]) or {}
}
fn (mut writer PPMWriter) finish() {
writer.file.close()
}
fn sim_runner(mut state SimState, params SimParams) Pixel {
// do the simulation!
for _ in 0 .. 1000 {
state.increment(0.0005, params)
if state.done() {
println('done!')
break
}
}
// find the closest magnet
m1_dist := params.get_magnet_dist(0, state)
m2_dist := params.get_magnet_dist(2.0 * math.pi / 3.0, state)
m3_dist := params.get_magnet_dist(4.0 * math.pi / 3.0, state)
if m1_dist < m2_dist && m1_dist < m3_dist {
return Pixel{
r: 255
g: 0
b: 0
}
} else if m2_dist < m1_dist && m2_dist < m3_dist {
return Pixel{
r: 0
g: 255
b: 0
}
} else {
return Pixel{
r: 0
g: 0
b: 255
}
}
}
struct SimResult {
id u64
p Pixel
}
struct SimRequest {
id u64
params SimParams
mut:
initial SimState
}
fn sim_worker(request_chan chan SimRequest, result_chan chan SimResult) {
// serve sim requests as they come in
for {
mut request := <-request_chan or { break }
result_chan <- SimResult{
id: request.id
p: sim_runner(mut request.initial, request.params)
}
}
}
struct ValidPixel {
Pixel
mut:
valid bool
}
fn image_worker(mut writer PPMWriter, result_chan chan SimResult, total_pixels u64) {
// as new pixels come in, write them to the image file
mut current_index := u64(0)
mut pixel_buf := []ValidPixel{len: int(total_pixels), init: ValidPixel{
valid: false
}}
for {
result := <-result_chan or { break }
pixel_buf[result.id].Pixel = result.p
pixel_buf[result.id].valid = true
for current_index < total_pixels && pixel_buf[current_index].valid {
writer.next_pixel(pixel_buf[current_index].Pixel)
current_index++
}
if current_index >= total_pixels {
break
}
}
}
fn main() {
params := SimParams{
rope_length: 0.25
bearing_mass: 0.03
magnet_spacing: 0.05
magnet_height: 0.03
magnet_strength: 10.0
gravity: 4.9
}
mut writer := PPMWriter{}
writer.start_for_file('test.ppm', ImageSettings{
width: width
height: height
})
defer {
writer.finish()
}
result_chan := chan SimResult{}
request_chan := chan SimRequest{}
// start a worker on each core
for _ in 0 .. parallel_workers {
go sim_worker(request_chan, result_chan)
}
go fn (request_chan chan SimRequest, params SimParams) {
mut index := u64(0)
println('')
for y in 0 .. height {
term.clear_previous_line()
println('Line: $y')
for x in 0 .. width {
// setup initial conditions
mut state := SimState{}
state.position = Vec3D{
x: 0.1 * ((f64(x) - 0.5 * f64(width - 1)) / f64(width - 1))
y: 0.1 * ((f64(y) - 0.5 * f64(height - 1)) / f64(height - 1))
z: 0.0
}
state.velocity = Vec3D{}
state.satisfy_rope_constraint(params)
request_chan <- SimRequest{
id: index
initial: state
params: params
}
index++
}
}
request_chan.close()
}(request_chan, params)
image_worker(mut writer, result_chan, width * height)
}