feat: add support for moving/function coloring

src/main.cpp

@@ -1,13 +1,15 @@
 #include "screen.hpp"
 #include "vec.hpp"
+#include <cmath>
 #include <cstdio>
 #include <swarm.hpp>
 
-#define FPS             60
+static constexpr int   kFPS = 60;
-#define DT		(1.0 / FPS)
+static constexpr float kDT  = 1.0 / kFPS;
 
 float f(vec<2> x) {
-    return (x.x - 49)*(x.x - 49) + (x.y - 13)*(x.y - 13);
+    return 50 * (std::pow(std::sin((x.x - 10)/2),2) + std::pow(std::sin(x.y/2),2))
+	+ std::pow(std::abs(x.x - 10), 1.2) + std::pow(std::abs(x.y), 1.2);
 }
 
 int
@@ -19,35 +21,61 @@ main(int argc, char **argv) {
     Swarm<2> swarm(f);
     Screen scr(80, 24);
 
-    swarm.add_particle({0, 0}, {7, 3});
+    TanhColorizer tanh_col;
+
+    float min_f = 0, max_f = 0;
+    scr.shader = [&](int w, int h, float x, float y) -> Screen::Symbol {
+	float z = f({x, y});
+	min_f = std::min(min_f, z);
+	max_f = std::max(max_f, z);
+	return Screen::Symbol {
+	    .sym = ' ',
+	    .color = tanh_col(z),
+	};
+    };
+
+    swarm.add_particle({0, 0}, {5, 3});
     swarm.add_particle({4, 1}, {1, 5});
-    swarm.add_particle({16, 10}, {9, 10});
+    swarm.add_particle({16, 10}, {4, 5});
     swarm.add_particle({7, 4}, {4, -3});
     swarm.add_particle({14, 1}, {1, 5});
     swarm.add_particle({1, 10}, {9, 10});
     swarm.add_particle({5, 4}, {5, -3});
     swarm.add_particle({4, 6}, {-1, 5});
-    swarm.add_particle({2, 6}, {9, 0});
+    swarm.add_particle({2, 6}, {7, 0});
     swarm.add_particle({7, 7}, {4, -3});
 
+    for(const auto &p : swarm.get_particles()) {
+	scr.points.push_back(&p.get_position());
+    }
+
     enter_noncanonical_mode();
+    // swap to alt buffer
     printf("\e[?1049h");
 
     bool pause = false;
     bool frame_step = false;
 
-    for(int i = 0; i < iter_count * FPS/4; ++i) {
+    // initialize colorizer scale/translate
+    scr.draw();
+    tanh_col.scale = std::max(std::abs(max_f), std::abs(min_f));
+    tanh_col.translate = (max_f - min_f) / 2.0;
+
+    auto update_and_draw = [&]() -> void {
+	scr.draw();
+    };
+
+    for(int i = 0; i < iter_count * kFPS/4;) {
 	if(!pause) {
 	    scr.clear();
-	    for(const auto &p : swarm.get_particles()) {
-		const auto & pos = p.get_position();
-		scr.at(pos.x, pos.y) = '#';
-	    }
-	    scr.draw();
 
-	    swarm.move(DT * 4);
+	    update_and_draw();
-	    if(i % (FPS/4) == (FPS/4)-1)
+
+	    swarm.move(kDT * 4);
+	    if(i % (kFPS/4) == (kFPS/4)-1)
 		swarm.step();
+
+	    ++i;
 	}
 
 	if(frame_step) {
@@ -67,17 +95,72 @@ main(int argc, char **argv) {
 	    pause = false;
 	    frame_step = true;
 	    break;
+
+	case 'd':
+	    scr.dx += scr.sy;
+	    update_and_draw();
+	    break;
+	case 'a':
+	    scr.dx -= scr.sy;
+	    update_and_draw();
+	    break;
+	case 's':
+	    scr.dy += scr.sx;
+	    update_and_draw();
+	    break;
+	case 'w':
+	    scr.dy -= scr.sx;
+	    update_and_draw();
+	    break;
+
+	case 'I':
+	    tanh_col.scale *= 1.2;
+	    update_and_draw();
+	    break;
+	case 'O':
+	    tanh_col.scale /= 1.2;
+            update_and_draw();
+            break;
+
+	case 'K':
+	    tanh_col.translate *= 1.2;
+	    update_and_draw();
+	    break;
+	case 'L':
+	    tanh_col.translate /= 1.2;
+            update_and_draw();
+            break;
+
+	case 'i':
+	    scr.dx += 0.5 * (1 - 0.75) * scr.sx * scr.w;
+	    scr.dy += 0.5 * (1 - 0.75) * scr.sy * scr.h;
+	    scr.sx *= 0.75;
+	    scr.sy *= 0.75;
+	    update_and_draw();
+	    break;
+	case 'o':
+	    scr.dx -= 0.5 * (1 / 0.75 - 1) * scr.sx * scr.w;
+	    scr.dy -= 0.5 * (1 / 0.75 - 1) * scr.sy * scr.h;
+	    scr.sx /= 0.75;
+	    scr.sy /= 0.75;
+	    update_and_draw();
+	    break;
+
+	case 'p':
+	    auto [y, x] = swarm.best();
+	    printf("Current best: f(%.3f, %.3f) = %.3f", x.x, x.y, y);
+	    break;
 	}
 
-	usleep(1000 * 1000 / FPS);
+	usleep(1000 * 1000 / kFPS);
     }
 
 cleanup:
+    // swap back to original buffer
     printf("\e[?1049l");
     enter_canonical_mode();
 
     auto [y, x] = swarm.best();
-
     printf("Best: f(" V2_FMT ") = %.3f", V2_ARG(x), y);
 
     return 0;