example: add a pendulum simulation (#9992)

examples/pendulum_sim/sim.v

@@ -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)
+}