cart-and-pole/main.c

/*
 * Copyright (c) Camden Dixie O'Brien
 * SPDX-License-Identifier: AGPL-3.0-only
 */

#include <SDL2/SDL.h>
#include <assert.h>
#include <math.h>

#define WIN_WIDTH 800
#define WIN_HEIGHT 800
#define MARGIN_X 100
#define MARGIN_Y 100

#define SCENE_WIDTH (WIN_WIDTH - 2 * MARGIN_X)
#define SCENE_HEIGHT (WIN_HEIGHT - 2 * MARGIN_Y)

#define CART_WIDTH 40
#define CART_HEIGHT 20
#define MASS_RADIUS 20
#define SETPOINT_LN_LEN 40

#define A_KP 80.0
#define A_KD 10.0

#define S_KP 0.08
#define S_KD 0.08

#define MIN(x, y) ((x) < (y) ? (x) : (y))
#define MAX(x, y) ((x) > (y) ? (x) : (y))
#define CLAMP(x, min, max) MAX(MIN(x, max), min)

typedef struct {
	double s, ps; // Cart displacement
	double a, pa; // Pendulum angle (from vertical)
} state_t;

static const double mc = 2; // Cart mass
static const double mp = 4; // Pendulum mass
static const double l = 0.5; // Pendulum length
static const double g = 9.8; // Gravity

static void update(state_t *s, double dt, double k)
{
	const double co = cos(s->a);
	const double si = sin(s->a);

	const double s_dot
	    = (s->ps * l - s->pa * co) / (l * (mc + mp * pow(si, 2)));
	const double ps_dot = k;
	const double a_dot = (s->pa * (mc + mp) - s->ps * mp * l * co)
	    / (mp * pow(l, 2) * (mc + mp * pow(si, 2)));
	const double pa_dot
	    = (co * si
	       * (pow(s->ps, 2) * mp * pow(l, 2)
	          - 2 * s->ps * s->pa * mp * l * co + pow(s->pa, 2) * (mc + mp)))
	        / (pow(l, 2) * pow(mc + mp * pow(si, 2), 2))
	    - (s->ps * s->pa * si) / (l * (mc + mp * pow(si, 2)))
	    + mp * g * l * si;

	s->s += dt * s_dot;
	s->ps += dt * ps_dot;
	s->a += dt * a_dot;
	s->pa += dt * pa_dot;
}

static double control(const state_t *s, double pos, double dt)
{
	static double prev_s = 0.0, prev_a = 0.0;

	const double err_s = pos - s->s;
	const double deriv_s = (err_s - prev_s) / dt;
	const double target_a = CLAMP(S_KP * err_s + S_KD * deriv_s, -0.3, 0.3);
	prev_s = err_s;

	const double err_a = s->a - target_a;
	const double deriv_a = (err_a - prev_a) / dt;
	const double out = CLAMP(A_KP * err_a + A_KD * deriv_a, -10.0, 10.0);
	prev_a = err_a;

	return out;
}

static void render(SDL_Renderer *r, const state_t *s, double pos)
{
	static const double sf = SCENE_WIDTH;
	static const double sc_l = sf * l;
	static const double sc_vdisp = (SCENE_HEIGHT - sc_l) / 2;
	static const SDL_Point sc_anchor = {
		.x = WIN_WIDTH / 2,
		.y = MARGIN_Y + (int)(sc_vdisp + sc_l),
	};

	// Cart centre
	const SDL_Point sc_a = {
		.x = sc_anchor.x + (int)(sf * s->s),
		.y = sc_anchor.y,
	};
	// Pendulum centre
	const SDL_Point sc_b = {
		.x = sc_a.x + (int)(sc_l * sin(s->a)),
		.y = sc_a.y - (int)(sc_l * cos(s->a)),
	};

	// Cart
	const SDL_Rect c = {
		.x = sc_a.x - CART_WIDTH / 2,
		.y = sc_a.y - CART_HEIGHT / 2,
		.w = CART_WIDTH,
		.h = CART_HEIGHT,
	};
	SDL_RenderFillRect(r, &c);

	// Rod
	SDL_RenderDrawLine(r, sc_a.x, sc_a.y, sc_b.x, sc_b.y);

	// Mass
	const SDL_Rect m = {
		.x = sc_b.x - MASS_RADIUS,
		.y = sc_b.y - MASS_RADIUS,
		.w = 2 * MASS_RADIUS,
		.h = 2 * MASS_RADIUS,
	};
	SDL_RenderFillRect(r, &m);

	// Position setpoint
	const int sc_setpoint = sc_anchor.x + (int)(sf * pos);
	SDL_RenderDrawLine(
	    r, sc_setpoint, WIN_HEIGHT - MARGIN_Y, sc_setpoint,
	    WIN_HEIGHT - MARGIN_Y - SETPOINT_LN_LEN);
}

int main(void)
{
	int err = SDL_Init(SDL_INIT_VIDEO);
	assert(err == 0);

	SDL_Window *window = SDL_CreateWindow(
	    "Cart and Rod Simulation", SDL_WINDOWPOS_UNDEFINED,
	    SDL_WINDOWPOS_UNDEFINED, WIN_WIDTH, WIN_HEIGHT, 0);
	assert(NULL != window);

	SDL_Renderer *renderer = SDL_CreateRenderer(
	    window, -1, SDL_RENDERER_ACCELERATED | SDL_RENDERER_PRESENTVSYNC);
	assert(NULL != renderer);

	SDL_DisplayMode mode;
	assert(SDL_GetWindowDisplayMode(window, &mode) == 0);
	const double dt = 1.0 / mode.refresh_rate;

	state_t state = { .a = 0.1 };
	double pos = 0.0;
	while (1) {
		SDL_Event e;
		while (SDL_PollEvent(&e)) {
			if (e.type == SDL_QUIT)
				goto quit;
			if (e.type == SDL_KEYDOWN) {
				switch (e.key.keysym.sym) {
				case SDLK_LEFT:
					pos -= 0.1;
					break;
				case SDLK_RIGHT:
					pos += 0.1;
					break;
				case SDLK_j:
					state.a -= 0.1;
					break;
				case SDLK_k:
					state.a += 0.1;
					break;
				}
			}
		}

		const double k = control(&state, pos, dt);
		update(&state, dt, k);

		SDL_SetRenderDrawColor(renderer, 255, 255, 255, 255);
		SDL_RenderClear(renderer);
		SDL_SetRenderDrawColor(renderer, 0, 0, 0, 255);
		render(renderer, &state, pos);

		SDL_RenderPresent(renderer);
	}

quit:
	SDL_DestroyRenderer(renderer);
	SDL_DestroyWindow(window);
	SDL_Quit();

	return 0;
}