code / 2pend

   1 // A double pendulum simulator by Bruce Hill
   2 // Released under the MIT license, see LICENSE for details.
   3 #include <fcntl.h>
   4 #include <math.h>
   5 #include <stdio.h>
   6 #include <stdlib.h>
   7 #include <string.h>
   8 #include <sys/ioctl.h>
   9 #include <termios.h>
  10 #include <unistd.h>
  11 
  12 typedef struct {
  13     double a, p;
  14 } pendulum_t;
  15 
  16 typedef struct {
  17     pendulum_t p1, p2;
  18 } state_t;
  19 
  20 static const double M1 = 1.0, M2 = 1.0, G = 9.80, L1 = .5, L2 = .5;
  21 struct winsize wsize = {0};
  22 
  23 const int color_ramp[] = {16,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,
  24     247,248,249,250,251,252,253,254,255,231,231,231,231,231};
  25 
  26 static FILE *tty_out = NULL;
  27 int ttyin_fd = 0;
  28 
  29 static state_t get_derivative(state_t state)
  30 {
  31     double a1 = state.p1.a, a2 = state.p2.a, p1 = state.p1.p, p2 = state.p2.p;
  32     double da1, da2, dp1, dp2;
  33     da1 = (6/(M1*L1*L1)) * (2*p1 - 3*cos(a1-a2)*p2)/(16-9*cos(a1-a2)*cos(a1-a2));
  34     da2 = (6/(M2*L2*L2)) * (8*p2 - 3*cos(a1-a2)*p1)/(16-9*cos(a1-a2)*cos(a1-a2));
  35     dp1 = -.5 * M1*L1*L1 * (da1*da2*sin(a1-a2) + 3*G/L1*sin(a1));
  36     dp2 = -.5 * M2*L2*L2 * (-da1*da2*sin(a1-a2) + G/L1*sin(a2));
  37     state_t derivative = {{da1, dp1}, {da2, dp2}};
  38     return derivative;
  39 }
  40 
  41 static state_t step(state_t state, state_t derivative, double dt)
  42 {
  43     state_t next = {
  44         {state.p1.a+dt*derivative.p1.a, state.p1.p+dt*derivative.p1.p},
  45         {state.p2.a+dt*derivative.p2.a, state.p2.p+dt*derivative.p2.p}
  46     };
  47     return next;
  48 }
  49 
  50 static inline void draw_pix(double x, double y, int W, int H, double *prev)
  51 {
  52     double b = sqrt(2*((x-floor(x+.5))*(x-floor(x+.5)) + (y-floor(y+.5))*(y-floor(y+.5))));
  53     int ix = (int)x, iy = (int)y;
  54     if (b > prev[iy*W+ix] && b > .1) {
  55         prev[iy*W+ix] = b;
  56         int color = color_ramp[(int)(sizeof(color_ramp)/sizeof(int) * b)];
  57         fprintf(tty_out, "\033[%d;%dH\033[48;5;%dm \033[0m", iy+1, ix+1, color);
  58     }
  59 }
  60 
  61 static void draw(state_t state)
  62 {
  63     // TODO: properly walk along cells
  64     int W = wsize.ws_col, H = wsize.ws_row;
  65     double *brightness = calloc(W*H, sizeof(double));
  66     double len = W/2 > H ? H/2 : (W/2)/2;
  67     double x = W/2, y = H/2;
  68     double step = 0.25;
  69     for (double p = 0; p <= len*L1; p += step) {
  70         draw_pix(x, y, W, H, brightness);
  71         draw_pix(x+1.0, y, W, H, brightness);
  72         x += step*cos(state.p1.a + M_PI/2)*2;
  73         y += step*sin(state.p1.a + M_PI/2);
  74     }
  75     for (double p = 0; p <= len*L2; p += step) {
  76         draw_pix(x, y, W, H, brightness);
  77         draw_pix(x+1.0, y, W, H, brightness);
  78         x += step*cos(state.p2.a + M_PI/2)*2;
  79         y += step*sin(state.p2.a + M_PI/2);
  80     }
  81     free(brightness);
  82 }
  83 
  84 int main(int argc, char **argv)
  85 {
  86     int ret = 0;
  87     ttyin_fd = open("/dev/tty", O_RDONLY | O_NONBLOCK);
  88     tty_out = fopen("/dev/tty", "w");
  89 
  90     struct termios orig_termios, termios;
  91     tcgetattr(fileno(tty_out), &orig_termios);
  92     memcpy(&termios, &orig_termios, sizeof(termios));
  93     cfmakeraw(&termios);
  94     termios.c_cc[VMIN] = 10;
  95     termios.c_cc[VTIME] = 1;
  96     if (tcsetattr(fileno(tty_out), TCSAFLUSH, &termios) == -1)
  97         return 1;
  98 
  99     ioctl(STDOUT_FILENO, TIOCGWINSZ, &wsize);
 100     fputs("\033[?25l\033[?1049h", tty_out);
 101 
 102     state_t state;
 103 
 104   randomize:
 105     state.p1.a = M_PI*7/8 + M_PI/4*(double)rand()/RAND_MAX;
 106     state.p2.a = M_PI*7/8 + M_PI/4*(double)rand()/RAND_MAX;
 107     state.p1.p = 0;
 108     state.p2.p = 0;
 109 
 110     while (1) {
 111         const int steps = 5;
 112         double dt = (1./60.) / (double)steps;
 113         for (int i = 0; i < steps; i++) {
 114             // RK4 integration
 115             state_t k1 = get_derivative(state);
 116             state_t k2 = get_derivative(step(state, k1, dt/2.));
 117             state_t k3 = get_derivative(step(state, k2, dt/2.));
 118             state_t k4 = get_derivative(step(state, k3, dt));
 119             state.p1.a += dt/6.*(k1.p1.a + 2.*k2.p1.a + 2.*k3.p1.a + k4.p1.a);
 120             state.p2.a += dt/6.*(k1.p2.a + 2.*k2.p2.a + 2.*k3.p2.a + k4.p2.a);
 121             state.p1.p += dt/6.*(k1.p1.p + 2.*k2.p1.p + 2.*k3.p1.p + k4.p1.p);
 122             state.p2.p += dt/6.*(k1.p2.p + 2.*k2.p2.p + 2.*k3.p2.p + k4.p2.p);
 123         }
 124         fputs("\033[2J", tty_out);
 125         draw(state);
 126         fflush(tty_out);
 127         usleep(60000);
 128         char ch;
 129         if (read(ttyin_fd, &ch, 1) == 1) {
 130             if (ch == '\3') {
 131                 ret = 1;
 132                 goto done;
 133             } else if (ch == 'r') {
 134                 goto randomize;
 135             } else if (ch == '\x1b') {
 136                 if (read(ttyin_fd, &ch, 1) != 1)
 137                     goto done;
 138             } else if (ch == 'q') {
 139                 goto done;
 140             }
 141         }
 142     }
 143   done:
 144     tcsetattr(fileno(tty_out), TCSAFLUSH, &orig_termios);
 145     fputs("\033[?25h\033[?1049l", tty_out);
 146     fclose(tty_out);
 147     close(ttyin_fd);
 148     return ret;
 149 }