planetary motion case study.

Author: charlesm93

knitr/planetary_motion/init/init1.r

@@ -0,0 +1,7 @@
+k <- 0.0370383855031032
+p0 <- 
+c(-0.265409729119215, 0.995573007641702)
+q0 <- 
+c(1.04456730637102, 0.221941902167825)
+star <- 
+c(-0.0377712773150063, -0.0703029465489535)

knitr/planetary_motion/model/planetary_motion.stan

@@ -0,0 +1,79 @@
+
+functions {
+  real[] ode (real t,
+              real[] y,
+              real[] theta,
+              real[] x_r, int[] x_i) {
+    vector[2] q = to_vector({y[1], y[2]});
+    real r_cube = pow(dot_self(q), 1.5);
+    vector[2] p = to_vector({y[3], y[4]});
+    real m = x_r[1];
+    real k = theta[1];
+
+    vector[4] dydt;
+    dydt[1:2] = p ./ m;
+    dydt[3:4] = - k * q ./ r_cube;
+
+    return to_array_1d(dydt);
+  }
+}
+
+data {
+  int n;
+  real q_obs[n, 2];
+}
+
+transformed data {
+  real t0 = 0;
+  int n_coord = 2;
+  real q0[n_coord] = {1.0, 0.0};
+  real p0[n_coord] = {0.0, 1.0};
+  real y0[n_coord * 2] = append_array(q0, p0);
+
+  real m = 1.0;
+
+  real t[n];
+  for (i in 1:n) t[i] = i * 1.0 / 10;
+
+  int x_i[0];
+  
+  real<lower = 0> sigma_x = 0.01;
+  real<lower = 0> sigma_y = 0.01;
+  
+  // ODE tuning parameters
+  real rel_tol = 1e-6;  // 1e-12;
+  real abs_tol = 1e-6;  // 1e-12;
+  int max_steps = 1000;  // 1e6
+  
+}
+
+parameters {
+  real<lower = 0> k;
+  // real<lower = 0> sigma_x;
+  // real<lower = 0> sigma_y;
+}
+
+transformed parameters {
+  real y[n, n_coord * 2]
+    = integrate_ode_bdf(ode, y0, t0, t, {k}, {m}, x_i,
+                         rel_tol, abs_tol, max_steps);
+}
+
+model {
+  // sigma_x ~ normal(0, 1);
+  // sigma_y ~ normal(0, 1);
+  k ~ normal(0, 1);
+  // k ~ normal(1, 0.1);
+
+  q_obs[, 1] ~ normal(y[, 1], sigma_x);
+  q_obs[, 2] ~ normal(y[, 2], sigma_y);
+}
+
+generated quantities {
+  // real q_pred[n, 2];
+  real qx_pred[n];
+  real qy_pred[n];
+
+  qx_pred = normal_rng(y[, 1], sigma_x);
+  qy_pred = normal_rng(y[, 2], sigma_y);
+}

knitr/planetary_motion/model/planetary_motion_sim.stan

@@ -0,0 +1,68 @@
+
+// It's assumed that the planet orbits a star at position (0, 0),
+// and that the mass of the planet is negligable next to that of
+// the star (meaning the star doesn't move).
+
+functions {
+  real[] ode (real t,
+              real[] y,
+              real[] theta,
+              real[] x_r, int[] x_i) {
+    vector[2] q = to_vector({y[1], y[2]});
+    real r_cube = pow(dot_self(q), 1.5);
+    vector[2] p = to_vector({y[3], y[4]});
+    real m = x_r[1];
+    real k = x_r[2];
+
+    vector[4] dydt;
+    dydt[1:2] = p ./ m;
+    dydt[3:4] = - k * q ./ r_cube;
+
+    return to_array_1d(dydt);
+  }
+}
+
+data {
+  int n;
+  real<lower = 0> sigma_x;
+  real<lower = 0> sigma_y;
+}
+
+transformed data {
+  // intial state at t = 0
+  real t0 = 0;
+  int n_coord = 2;
+  real q0[n_coord] = {1.0, 0.0};
+  real p0[n_coord] = {0.0, 1.0};
+  real y0[n_coord * 2] = append_array(q0, p0);
+
+  // model parameters
+  real m = 1.0;
+  real k = 1.0;
+
+  // exact motion
+  real t[n];
+  for (i in 1:n) t[i] = (i * 1.0) / 10;
+  real x_r[2] = {m, k};
+
+  real theta[0];
+  int x_i[0];
+
+  real y[n, n_coord * 2]
+    = integrate_ode_rk45(ode, y0, t0, t, theta, x_r, x_i);
+}
+
+parameters {
+  real phi;
+}
+
+model {
+  phi ~ normal(0, 1);
+}
+
+generated quantities {
+  real q_obs[n, 2];
+
+  q_obs[, 1] = normal_rng(y[, 1], sigma_x);
+  q_obs[, 2] = normal_rng(y[, 2], sigma_y);
+}

knitr/planetary_motion/model/planetary_motion_star2.stan

@@ -0,0 +1,94 @@
+
+// fixes a calculation error in planetary_motion_star.stan.
+
+functions {
+  real[] ode (real t,
+              real[] y,
+              real[] theta,
+              real[] x_r, int[] x_i) {
+    vector[2] q = to_vector({y[1], y[2]});
+    vector[2] s = to_vector({theta[2], theta[3]});
+
+    real r_cube = pow(dot_self(q - s), 1.5);
+    vector[2] p = to_vector({y[3], y[4]});
+    real m = x_r[1];
+    real k = theta[1];
+
+    vector[4] dydt;
+    dydt[1:2] = p ./ m;
+    dydt[3:4] = - k * (q - s) ./ r_cube;
+
+    return to_array_1d(dydt);
+  }
+}
+
+data {
+  int n;
+  real q_obs[n, 2];
+  real time[n];
+
+  real<lower = 0> sigma;
+}
+
+transformed data {
+  real t0 = 0;
+  int n_coord = 2;
+  // real q0[n_coord] = {1.0, 0.0};
+  // real p0[n_coord] = {0.0, 1.0};
+  // real y0[n_coord * 2] = append_array(q0, p0);
+
+  real m = 1.0;
+
+  // real t[n];
+  // for (i in 1:n) t[i] = i * 1.0 / 10;
+
+  int x_i[0];
+
+  real<lower = 0> sigma_x = sigma;
+  real<lower = 0> sigma_y = sigma;
+
+  // ODE tuning parameters
+  real rel_tol = 1e-6;
+  real abs_tol = 1e-6;
+  int max_steps = 1000;
+  
+}
+
+parameters {
+  real<lower = 0> k;
+  real q0[n_coord];
+  real p0[n_coord];
+  real star[n_coord];
+}
+
+transformed parameters {
+  real y0[n_coord * 2] = append_array(q0, p0);
+  real theta[n_coord + 1] = append_array({k}, star);
+
+  real y[n, n_coord * 2]
+    = integrate_ode_rk45(ode, y0, t0, time, theta, {m}, x_i,
+                         rel_tol, abs_tol, max_steps);
+}
+
+model {
+  k ~ normal(1, 0.001);  // prior derive based on solar system 
+                         // (still pretty uninformative)
+
+  p0[1] ~ normal(0, 1);
+  p0[2] ~ lognormal(0, 1);  // impose p0 to be positive.
+  q0 ~ normal(0, 1);
+  star ~ normal(0, 0.5);
+  // star ~ normal(0, 1);
+
+  q_obs[, 1] ~ normal(y[, 1], sigma_x);
+  q_obs[, 2] ~ normal(y[, 2], sigma_y);
+}
+
+generated quantities {
+  // real q_pred[n, 2];
+  real qx_pred[n];
+  real qy_pred[n];
+
+  qx_pred = normal_rng(y[, 1], sigma_x);
+  qy_pred = normal_rng(y[, 2], sigma_y);
+}