├── .gitignore ├── dftransform.py ├── rungekutta.h ├── interpol.h ├── bestfitcmp.py ├── Makefile ├── README.rst ├── rungekutta.c ├── angmom.c ├── interpol.c ├── predict.py ├── testexo.c ├── exofind.py ├── exo.c ├── rvfit.py └── rvglobalfit.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.png 3 | *.pdf 4 | *.prof 5 | -------------------------------------------------------------------------------- /dftransform.py: -------------------------------------------------------------------------------- 1 | /home/user/Downloads/p04dk-v1.2.0-final/cpp/transform.py -------------------------------------------------------------------------------- /rungekutta.h: -------------------------------------------------------------------------------- 1 | void rungekutta(double * t, double * thetas, double * diff, 2 | int n, double h, double t0, double theta_0); 3 | -------------------------------------------------------------------------------- /interpol.h: -------------------------------------------------------------------------------- 1 | double interpol_diff_at(double x); 2 | double interpol_diff_at_normalized(double x); 3 | void interpol_diff_setup(double xn_in, double xm, double yn, double ym, double yyn, double yym); 4 | 5 | double interpol_lin_at(double x); 6 | double interpol_lin_at_normalized(double x); 7 | void interpol_lin_setup(double xn_in, double xm, double yn, double ym); 8 | 9 | double interpol_sq_at(double x); 10 | double interpol_sq_at_normalized(double x); 11 | void interpol_sq_setup(double xn_in, double xm, double xo, double yn, double ym, double yo); 12 | 13 | 14 | 15 | -------------------------------------------------------------------------------- /bestfitcmp.py: -------------------------------------------------------------------------------- 1 | # 2 | 3 | import numpy 4 | import json 5 | import sys 6 | 7 | n = len(numpy.loadtxt(sys.argv[1], ndmin = 2)) 8 | data = [numpy.loadtxt(d, ndmin = 2) for d in sys.argv[2:]] 9 | fits = [(d[:,1].min(), d.shape[1] - 2) for d in data] 10 | 11 | def aicc(bestfit, nparams): 12 | """ bestfit is -2 * loglikelihood of best fit; nparams is number of model parameters """ 13 | k = nparams 14 | return (2 * k + bestfit + 2 * k * (k + 1) / (n - k - 1.)) / 2. 15 | 16 | for (bestfit, nparams), filename in zip(fits, sys.argv[2:]): 17 | print filename, nparams, bestfit, aicc(bestfit, nparams) 18 | 19 | 20 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | 2 | CC := gcc 3 | CFLAGS := -O3 -fPIC -std=c99 -Wall -Wextra -pedantic -lm 4 | CFLAGS += -DINTERPOL=1 5 | LFLAGS += -L$(MULTINEST) -lnest3 -latllapack -lpthread 6 | 7 | all: exo.so exo_noprior.so testexo 8 | 9 | exo.so: exo.o interpol.o rungekutta.o 10 | ${CC} -DUSE_PRIORS -DNOISE ${CFLAGS} ${LFLAGS} $^ -o $@ -shared 11 | exo_noprior.so: exo.o interpol.o rungekutta.o 12 | ${CC} -DNOISE ${CFLAGS} ${LFLAGS} $^ -o $@ -shared 13 | 14 | %.o: %.c 15 | ${CC} ${CFLAGS} $^ -o $@ -c 16 | 17 | testexo: testexo.c exo.o interpol.o rungekutta.o 18 | ${CC} ${CFLAGS} ${LFLAGS} $^ -pg -g -o $@ -lgsl 19 | 20 | 21 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | ExoFind 2 | ========================================== 3 | 4 | A radial velocity fitting code based on ExoFit. 5 | 6 | This is incomplete work, but the ideas were: 7 | 8 | * Make ExoFit callable from Python. Uses fast C code underneath. 9 | * Connect ExoFit with MultiNest instead of MCMC 10 | * Use innovative data-based prior transformations. The FFT of the residuals 11 | can indicate periods to prefer when searching the parameter space. 12 | This biasing is cancelled out again in the likelihood (to make Bayesian 13 | inference consistent), but should make exploring the parameter space simpler. 14 | * Use multi-dimensional exploration to detect planets, expanding and contracting 15 | the search space as needed. Achieved using differential evolution algorithm 16 | connected with the Akaike Information Criterion. 17 | * Plotting RV curves. 18 | 19 | No guarantees. 20 | -------------------------------------------------------------------------------- /rungekutta.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | extern double f(double t, double theta); 6 | 7 | 8 | /* 9 | #define RK4_DEBUG 10 | */ 11 | void rungekutta(double * t, double * thetas, double * diff, 12 | int n, double h, double t0, double theta_0 13 | ) { 14 | double kn1, kn2, kn3, kn4; 15 | double theta_i; 16 | double t_i; 17 | 18 | int i; 19 | 20 | theta_i = thetas[0] = theta_0; 21 | #ifdef RK4_DEBUG 22 | printf("%d\t%f\t%f\n", 0, theta_i, t0); 23 | #endif 24 | 25 | for(i = 1; i < n; i++) { 26 | t_i = t[i - 1] = (i - 1) * h + t0; 27 | /* function evaluated at previous point */ 28 | kn1 = f(t_i , theta_i); 29 | kn2 = f(t_i + h/2., theta_i + h * kn1/2.); 30 | kn3 = f(t_i + h/2., theta_i + h * kn2/2.); 31 | kn4 = f(t_i + h , theta_i + h * kn3); 32 | 33 | diff[i - 1] = kn1; 34 | 35 | theta_i += h/6. * (kn1 + 2*kn2 + 2*kn3 + kn4); 36 | thetas[i] = theta_i; 37 | #ifdef RK4_DEBUG 38 | printf("%d\t%f\t%f\t%f (above)\n", i, theta_i, i * h + t0, kn1); 39 | #endif 40 | 41 | } 42 | i = n - 1; 43 | t_i = t[i] = i * h + t0; 44 | diff[i] = f(t_i, theta_i); 45 | #ifdef RK4_DEBUG 46 | printf("%d\t%f\t%f\t%f\n", i, theta_i, t_i, diff[i]); 47 | printf("end of RK4\n"); 48 | #endif 49 | } 50 | 51 | 52 | -------------------------------------------------------------------------------- /angmom.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "interpol.h" 6 | #include "rungekutta.h" 7 | 8 | double p[] = { 9 | #include "p.dat" 10 | }; 11 | 12 | #ifndef NEVAL 13 | #define NEVAL 10 14 | #endif 15 | 16 | double P; 17 | double e; 18 | double chi; 19 | double A; 20 | 21 | double f(double t, double theta) { 22 | /*printf("A = %f, e = %f, theta = %f, P = %f, chi = %f\n", A, e, theta, P, chi);*/ 23 | return A * (1 - e * cos(theta + P * chi)); 24 | } 25 | 26 | /*double function(double x) { 27 | return sin(x) * (pow(x, 2)/10 + 1); 28 | } 29 | double derivative(double x) { 30 | return (x * sin(x)) / 5 + (pow(x,2) / 10 + 1) * cos(x); 31 | }*/ 32 | 33 | int main(void) { 34 | 35 | double minx; 36 | double maxx; 37 | double v_x[NPOINTS]; 38 | double v_y[NPOINTS]; 39 | double v_yy[NPOINTS]; 40 | 41 | double x; 42 | double y; 43 | 44 | int i; 45 | int j; 46 | int k; 47 | double h; 48 | 49 | double theta_0 = 0.1; 50 | minx = theta_0; 51 | e = 0.5; 52 | chi = 0.1; 53 | P = 1; 54 | h = P / NPOINTS; 55 | maxx = minx + 2 * P; 56 | 57 | for(k = 0; k < 100000; k++) { 58 | e *= 0.9; 59 | chi *= 0.9; 60 | P *= 1.01; 61 | A = 2 * 3.14 / (P * pow(1 - pow(e, 2), 3./2)); 62 | rungekutta(v_x, v_y, v_yy, NPOINTS, (maxx - minx)/NPOINTS, minx, theta_0); 63 | for(i = 0; i < NPOINTS; i++) { 64 | printf("%f\tnan\t%f\n", v_x[i], v_y[i]); 65 | } 66 | 67 | for(i = 0; i < NEVAL; i++) { 68 | x = p[i] / 20 + 1; 69 | if (x < minx || x > maxx) { 70 | fprintf(stderr, "skipped %f\n", x); 71 | continue; 72 | } 73 | j = floor(NPOINTS * (x - minx)/(maxx - minx)); 74 | 75 | #if INTERPOL == 1 76 | interpol_lin_setup(v_x[j], v_x[j+1], v_y[j], v_y[j+1]); 77 | y = interpol_lin_at(x); 78 | #elif INTERPOL == 2 79 | if (j+2 >= NPOINTS) 80 | j--; 81 | interpol_sq_setup(v_x[j], v_x[j+1], v_x[j+2], v_y[j], v_y[j+1], v_y[j+2]); 82 | y = interpol_sq_at(x); 83 | #elif INTERPOL == 3 84 | interpol_diff_setup(v_x[j], v_x[j+1], v_y[j], v_y[j+1], v_yy[j], v_yy[j+1]); 85 | y = interpol_diff_at(x); 86 | #else 87 | #error "INTERPOL not in [1 (lin), 2 (square) or 3 (diff)]" 88 | #endif 89 | 90 | printf("%f\t%f\n", x, y); 91 | } 92 | } 93 | 94 | return 0; 95 | } 96 | 97 | -------------------------------------------------------------------------------- /interpol.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include "interpol.h" 5 | 6 | /* 7 | * INTERPOLATION 8 | * linear, square, and with derivatives 9 | */ 10 | 11 | static double a, b, c, d; 12 | static double deltax; 13 | static double xn; 14 | 15 | /* derivative interpolation */ 16 | 17 | double interpol_diff_at(double x) { 18 | return a + 19 | b * (x - xn) / deltax + 20 | c * pow((x - xn) / deltax, 2) + 21 | d * pow((x - xn) / deltax, 3); 22 | } 23 | 24 | double interpol_diff_at_normalized(double x) { 25 | return a + b * x + c * pow(x, 2) + d * pow(x, 3); 26 | } 27 | 28 | void interpol_diff_setup(double xn_in, double xm, double yn, double ym, double yyn, double yym) { 29 | double deltay = ym - yn; 30 | double deltayy = yym - yyn; 31 | xn = xn_in; 32 | deltax = xm - xn; 33 | 34 | a = yn; 35 | b = yyn * deltax; 36 | d = (yym + yyn) * deltax - 2 * deltay; 37 | c = (deltayy * deltax - 3 * d) / 2; 38 | } 39 | 40 | /* linear interpolation */ 41 | 42 | double interpol_lin_at(double x) { 43 | return a + 44 | b * (x - xn) / deltax; 45 | } 46 | 47 | double interpol_lin_at_normalized(double x) { 48 | return a + b * x; 49 | } 50 | 51 | 52 | void interpol_lin_setup(double xn_in, double xm, double yn, double ym) { 53 | double deltay = ym - yn; 54 | xn = xn_in; 55 | deltax = xm - xn; 56 | 57 | a = yn; 58 | b = deltay; 59 | } 60 | 61 | /* square interpolation */ 62 | 63 | double interpol_sq_at(double x) { 64 | return a + 65 | b * x + 66 | c * pow(x, 2); 67 | } 68 | 69 | double interpol_sq_at_normalized(double x) { 70 | return a + 71 | b * (x * deltax + xn) + 72 | c * pow((x * deltax + xn), 2); 73 | } 74 | 75 | void interpol_sq_setup(double xn_in, double xm, double xo, double yn, double ym, double yo) { 76 | /*double deltay = ym - yn;*/ 77 | double xn_in2 = xn = xn_in; 78 | double xo2 = pow(xo,2); 79 | double xm2 = pow(xm,2); 80 | double xn2 = pow(xn,2); 81 | 82 | double anom = xn*(xo2*ym-xm2*yo)+xn2*(xm*yo-xo*ym)+(xm2*xo-xm*xo2)*yn; 83 | double div = xn*(xo2-xm2)-xm*xo2+xm2*xo+xn2*(xm-xo); 84 | 85 | double bnom = xn2*(yo-ym)-xm2*yo+(xm2-xo2)*yn+xo2*ym; 86 | 87 | double cnom = xn*(yo-ym)-xm*yo+(xm-xo)*yn+xo*ym; 88 | 89 | (void)xn_in2; 90 | 91 | deltax = xo - xn; 92 | 93 | a = anom / div; 94 | b = -bnom / div; 95 | c = cnom / div; 96 | } 97 | 98 | /* END OF INTERPOLATION */ 99 | 100 | -------------------------------------------------------------------------------- /predict.py: -------------------------------------------------------------------------------- 1 | import numpy 2 | from numpy import pi, log, exp, nan, sin, cos, arccos 3 | from ctypes import * 4 | lib = cdll.LoadLibrary('./exo_noprior.so') 5 | lib.LogLike.restype = c_double 6 | lib.ModLogPrior.restype = c_double 7 | 8 | n_params = None 9 | params_low = None 10 | params_high = None 11 | cube = None 12 | 13 | kavg = 0 14 | kmax = 1000 15 | sturn = 0.1 16 | smax = 10 17 | pmax = 100 18 | fs = 0.01 19 | pmin = 0.01 20 | 21 | # number of dimensions our problem has 22 | parameters = ['V'] 23 | # plus possibly a noise term 24 | parameters += ['s'] 25 | 26 | def update_n_params(new_n_params): 27 | global n_params 28 | global cube 29 | global params_low 30 | global params_high 31 | n_params = new_n_params 32 | 33 | # number of dimensions our problem has 34 | global parameters 35 | parameters = ['V'] 36 | # plus possibly a noise term 37 | parameters += ['s'] 38 | 39 | n_planets = (n_params - 2) / 5 40 | 41 | for i in range(n_planets): 42 | parameters += ['%s%d' % (var, i) for var in ['P', 'K', 'chi', 'e', 'omega']] 43 | 44 | params_low = (c_double * n_params)() 45 | params_high = (c_double * n_params)() 46 | cube = (c_double * n_params)() 47 | 48 | for i, p in enumerate(parameters): 49 | if p.startswith('V'): 50 | params_low[i], params_high[i] = kavg - 2 * kmax, kavg + 2 * kmax, 51 | if p.startswith('P'): 52 | params_low[i], params_high[i] = pmin, pmax 53 | if p.startswith('K'): 54 | params_low[i], params_high[i] = 1., kmax 55 | if p.startswith('chi'): 56 | params_low[i], params_high[i] = 0., 1. 57 | if p.startswith('e'): 58 | params_low[i], params_high[i] = 0., 1. 59 | if p.startswith('omega'): 60 | params_low[i], params_high[i] = 0., 2 * pi 61 | if p.startswith('s'): 62 | params_low[i], params_high[i] = sturn, smax 63 | lib.set_param_limits(params_low, params_high) 64 | 65 | def prepare(params): 66 | global n_params 67 | global cube 68 | global params_low 69 | global params_high 70 | if len(params) != n_params: 71 | update_n_params(len(params)) 72 | assert not numpy.any(numpy.isinf(params)) 73 | assert not numpy.any(numpy.isnan(params)) 74 | for i,a in enumerate(params): 75 | cube[i] = a 76 | 77 | def predict(c, x = None): 78 | if x is None: 79 | n = 1000 80 | x = (c_double * 1000)(*numpy.linspace(data[:,0].min(), data[:,0].max(), 1000)) 81 | else: 82 | n = len(x) 83 | x = (c_double * len(x))(*x) 84 | y = (c_double * len(x))() 85 | prepare(c) 86 | global cube 87 | global n_params 88 | lib.predict(cube, n_params, 0, n, x, y) 89 | return numpy.asarray([x[i] for i in range(n)]), numpy.asarray([y[i] for i in range(n)]) 90 | 91 | if __name__ == '__main__': 92 | x = numpy.linspace(0, 40, 1000) 93 | # parameters: 94 | # V: Velocity of the entire stellar system 95 | # s: noise term (not doing anything for predict) 96 | # parameters for each added planet: 97 | # P: Period (days) 98 | # K: Amplitude (m/s) 99 | # chi: anomaly parameter 100 | # e: eccentricity of orbit 101 | # omega: starting point (longitude of periastron) 102 | import matplotlib.pyplot as plt 103 | x, y = predict([0, 0.1, 10, 2, 0, 0, 0], x) 104 | plt.plot(x, y, ':') 105 | x, y = predict([0, 0.1, 10, 2, 0.2, 0.3, 0.1], x) 106 | plt.plot(x, y) 107 | plt.ylabel('RV Line Shift [m/s]') 108 | plt.xlabel('time [d]') 109 | plt.savefig('predict.pdf', bbox_inches='tight') 110 | plt.close() 111 | 112 | -------------------------------------------------------------------------------- /testexo.c: -------------------------------------------------------------------------------- 1 | 2 | #include 3 | #include 4 | #include 5 | 6 | extern double LogLike(double *params, int ndim, int npars); 7 | 8 | /* 9 | int main(int argc, char ** argv) { 10 | int i; 11 | int j = 0; 12 | int n; 13 | double params[100]; 14 | if (argc != 2) { 15 | fprintf(stderr, "SYNOPSIY: %s < infile \n"); 16 | return -127; 17 | } 18 | n = atoi(argv[1]); 19 | for (j = 0; !feof(stdin); j++) { 20 | for (i = 0; i < n + 3; i++) { 21 | if (fscanf(stdin, "%le", ¶ms[i]) != 1) { 22 | printf("couldn't read token %i.\n" % i); 23 | } 24 | } 25 | 26 | LogLike(params, n, 0); 27 | } 28 | printf("handled %i lines\n" % j); 29 | 30 | return 0; 31 | } 32 | */ 33 | 34 | double low[100] = {-2000, 0.2, 1, 0, 0, 0, 0.2, 1, 0, 0, 0}; 35 | double high[100] = {2000, 10000, 20000, 1, 1, 1, 10000, 20000, 1, 1, 1}; 36 | double newdata[1000] = { 37 | 1400.152000, -25.300000, 2.100000, 38 | 1413.877490, -18.200000, 2.800000, 39 | 1434.465730, -22.500000, 1.400000, 40 | 1455.053960, -26.200000, 2.000000, 41 | 1489.367690, -52.700000, 2.100000, 42 | 1496.230430, -50.800000, 2.400000, 43 | 1503.093180, -51.200000, 1.700000, 44 | 1564.857880, -45.300000, 1.700000, 45 | 1571.720630, -46.800000, 2.000000, 46 | 1585.446120, -50.700000, 1.600000, 47 | 1592.308860, -44.700000, 1.900000, 48 | 1612.897100, -14.700000, 1.900000, 49 | 1619.759840, -3.070000, 1.900000, 50 | 1654.073570, 45.300000, 2.400000, 51 | 1688.387290, 50.300000, 1.600000, 52 | 1695.250040, 54.900000, 1.700000, 53 | 1715.838270, 75.300000, 2.900000, 54 | 1722.701020, 78.000000, 2.700000, 55 | 1736.426510, 78.000000, 1.800000, 56 | 1743.289250, 81.100000, 2.600000, 57 | 1770.740240, 65.500000, 2.300000, 58 | 1777.602980, 57.800000, 2.000000, 59 | 1784.465730, 55.200000, 2.700000, 60 | 1798.191220, 53.400000, 2.400000, 61 | 1811.916710, 56.000000, 1.900000, 62 | 1818.779450, 58.800000, 2.200000, 63 | 1832.504940, 58.200000, 1.700000, 64 | 1901.132390, 25.900000, 1.600000, 65 | 1907.995140, 29.600000, 2.000000, 66 | 1921.720630, 33.700000, 2.200000, 67 | 1928.583370, 40.100000, 3.700000, 68 | 1935.446120, 36.800000, 2.900000, 69 | 1962.897100, 22.400000, 2.200000, 70 | 1969.759840, 16.300000, 2.200000 71 | }; 72 | 73 | void set_param_limits(double * new_params_low, double * new_params_high); 74 | void set_data(double * new_data, int new_n_data); 75 | void setup_data() { 76 | set_param_limits(low, high); 77 | set_data(newdata, 34); 78 | } 79 | 80 | extern void check(int ndim); 81 | 82 | #define IFDEBUG if(0) 83 | 84 | int main(int argc, char ** argv) { 85 | const gsl_rng_type * T; 86 | gsl_rng * r; 87 | int n = 1000000; 88 | int i; 89 | int j; 90 | double params[100]; 91 | double v; 92 | int n_params; 93 | if (argc != 2) { 94 | fprintf(stderr, "SYNOPSIY: %s < infile \n", argv[0]); 95 | return -127; 96 | } 97 | n_params = atoi(argv[1]); 98 | 99 | gsl_rng_env_setup(); 100 | T = gsl_rng_default; 101 | r = gsl_rng_alloc (T); 102 | setup_data(); 103 | 104 | for (i = 0; i < n; i++) 105 | { 106 | IFDEBUG printf(" iteration %d\n", i); 107 | for (j = 0; j < n_params; j++) { 108 | params[j] = gsl_rng_uniform (r); 109 | IFDEBUG printf(" %f\t", params[j]); 110 | } 111 | IFDEBUG check(n_params); 112 | IFDEBUG printf(" calc ... \n"); 113 | v = LogLike(params, n_params, 0); 114 | IFDEBUG printf(" calc ... %e\n", v); 115 | } 116 | 117 | gsl_rng_free (r); 118 | } 119 | 120 | -------------------------------------------------------------------------------- /exofind.py: -------------------------------------------------------------------------------- 1 | import pymultinest 2 | import math 3 | import os, sys 4 | import numpy 5 | import json 6 | 7 | from ctypes import cdll 8 | #cdll.LoadLibrary('libnest3.so') 9 | lib = cdll.LoadLibrary('exo.so') 10 | from ctypes import * 11 | 12 | # put data into double[n][3] array 13 | filename = sys.argv[1] 14 | data = numpy.loadtxt(filename) 15 | cdata = ((c_double * 3) * data.shape[0])() 16 | for i in range(data.shape[0]): 17 | for j in range(3): 18 | cdata[i][j] = data[i,j] 19 | lib.set_data(cdata, data.shape[0]) 20 | if 'MODE' in os.environ: 21 | fastmode = 'fast' in os.environ['MODE'] 22 | plotmode = 'noplot' not in os.environ['MODE'] 23 | else: 24 | fastmode = False 25 | plotmode = True 26 | 27 | # simplify parameter space: 28 | # strongest impact is at most the variation of the data 29 | kmax = (data[:,1].max() - data[:,1].min()) 30 | # sampling theorem -- highest frequency = longest period = 1 / dataduration / 2. 31 | pmax = (data[:,0].max() - data[:,0].min()) / 2 32 | 33 | # exclude very bad fits (worse than 5 sigma on all datapoints) 34 | log_bad_fit = (-0.5 * 3**2.) * len(data) 35 | log_very_bad_fit = (-0.5 * 5**2.) * len(data) 36 | 37 | def plot_marg(parameters, values, i, s, p): 38 | import matplotlib.pyplot as plt 39 | plt.xlabel(parameters[i]) 40 | 41 | m = s['marginals'][i] 42 | plt.xlim(m['5sigma']) 43 | 44 | oldax = plt.gca() 45 | x,w,patches = oldax.hist(values[:,i], bins=20, edgecolor='grey', color='grey', histtype='stepfilled', alpha=0.2) 46 | oldax.set_ylim(0, x.max()) 47 | 48 | newax = plt.gcf().add_axes(oldax.get_position(), sharex=oldax, frameon=False) 49 | p.plot_marginal(i, ls='-', color='blue', linewidth=3) 50 | newax.set_ylim(0, 1) 51 | 52 | ylim = newax.get_ylim() 53 | y = ylim[0] + 0.05*(ylim[1] - ylim[0]) 54 | center = m['median'] 55 | low1, high1 = m['1sigma'] 56 | newax.errorbar(x=center, y=y, 57 | xerr=numpy.transpose([[center - low1, high1 - center]]), 58 | color='blue', linewidth=2, marker='s') 59 | oldax.set_yticks([]) 60 | newax.set_yticks([]) 61 | newax.set_ylabel("Probability") 62 | ylim = oldax.get_ylim() 63 | newax.set_xlim(m['5sigma']) 64 | oldax.set_xlim(m['5sigma']) 65 | #plt.close() 66 | 67 | 68 | def plot(a, parameters): 69 | import matplotlib.pyplot as plt 70 | p = pymultinest.PlotMarginal(a) 71 | 72 | n_params = len(parameters) 73 | values = a.get_equal_weighted_posterior() 74 | s = a.get_stats() 75 | assert n_params == len(s['marginals']) 76 | prefix = a.outputfiles_basename 77 | 78 | plt.figure(figsize=(5*n_params, 5*n_params)) 79 | for i in range(n_params): 80 | plt.subplot(n_params, n_params, n_params * i + i + 1) 81 | plot_marg(parameters, values, i, s, p) 82 | for j in range(i): 83 | plt.subplot(n_params, n_params, n_params * j + i + 1) 84 | p.plot_conditional(i, j, bins=20, cmap = plt.cm.gray_r) 85 | plt.xlabel(parameters[i]) 86 | plt.ylabel(parameters[j]) 87 | #plt.savefig('cond_%s_%s.pdf' % (params[i], params[j]), bbox_tight=True) 88 | #plt.close() 89 | 90 | plt.savefig(prefix + 'cond.pdf') 91 | plt.close() 92 | plt.figure(figsize=(4*4, 2*math.ceil(n_params / 4.))) 93 | for i in range(n_params): 94 | plt.subplot(math.ceil(n_params / 4.), 4, i + 1) 95 | plot_marg(parameters, values, i, s, p) 96 | plt.savefig(prefix + 'marg.pdf', bbox_inches='tight') 97 | plt.close() 98 | 99 | def mkdir(path): 100 | if not os.path.exists(path): os.mkdir(path) 101 | 102 | def run(n_planets, previous_periods = [], log_zero = -1e90): 103 | # number of dimensions our problem has 104 | parameters = ['V'] 105 | # plus possibly a noise term 106 | parameters += ['s'] 107 | for i in range(n_planets): 108 | parameters += ['%s%d' % (var, i) for var in ['P', 'K', 'chi', 'e', 'omega']] 109 | n_params = len(parameters) 110 | params_low = (c_double * n_params)() 111 | params_high = (c_double * n_params)() 112 | simplification = 0. 113 | wrapped_params = [] 114 | for i, p in enumerate(parameters): 115 | if p.startswith('V'): 116 | params_low[i], params_high[i] = -2000., +2000. 117 | wrapped_params.append(0) 118 | if p.startswith('P'): 119 | if not fastmode or len(previous_periods) == 0: 120 | params_low[i], params_high[i] = 0.2, pmax 121 | else: 122 | params_low[i], params_high[i] = previous_periods.pop() 123 | if params_low[i] < 0.2: params_low[i] = 0.2 124 | if params_high[i] > pmax: params_high[i] = pmax 125 | print 'fixing period (parameter %s) to %f..%f' % (p, params_low[i], params_high[i]) 126 | simplification += math.log((math.log(params_high[i]) - math.log(params_low[i])) / (math.log(pmax) - math.log(0.2))) 127 | wrapped_params.append(0) 128 | if p.startswith('K'): 129 | params_low[i], params_high[i] = 1., kmax 130 | wrapped_params.append(0) 131 | if p.startswith('chi'): 132 | params_low[i], params_high[i] = 0., 1. 133 | wrapped_params.append(1) 134 | if p.startswith('e'): 135 | params_low[i], params_high[i] = 0., 1. 136 | wrapped_params.append(0) 137 | if p.startswith('omega'): 138 | params_low[i], params_high[i] = 0., 2 * math.pi 139 | wrapped_params.append(1) 140 | if p.startswith('s'): 141 | params_low[i], params_high[i] = 0.1, 100 142 | wrapped_params.append(0) 143 | lib.set_param_limits(params_low, params_high) 144 | #print 'n_params', n_params 145 | #lib.check(n_params) 146 | 147 | lib.LogLike.restype = c_double 148 | 149 | mkdir(filename + '.out') 150 | basename = filename + '.out/%d/' % n_planets 151 | mkdir(basename) 152 | if simplification != 0: 153 | print 'simplification: %f' % simplification 154 | pymultinest.run(lib.LogLike, None, n_params, resume = True, verbose = True, 155 | outputfiles_basename=basename, sampling_efficiency = 'model', 156 | n_live_points = 2000, evidence_tolerance = math.log(5), 157 | wrapped_params = wrapped_params, const_efficiency_mode = fastmode, 158 | null_log_evidence = log_zero, 159 | max_iter = 1000000, 160 | ) # max_modes=1000, 161 | 162 | # lets analyse the results 163 | a = pymultinest.Analyzer(n_params = n_params, outputfiles_basename=basename) 164 | s = a.get_stats() 165 | 166 | json.dump(parameters, file('%sparams.json' % a.outputfiles_basename, 'w'), indent=2) 167 | json.dump(s, file('%s.json' % a.outputfiles_basename, 'w'), indent=2) 168 | if plotmode: 169 | print 'plotting ...' 170 | plot(a, parameters) 171 | print 'plotting done.' 172 | print 'evidence: %.1f +- %.1f' % (s['global evidence'] + simplification, s['global evidence error']) 173 | return ( s['global evidence'] + simplification, s['global evidence error'] ) 174 | 175 | def get_previous_periods(n_planets): 176 | if n_planets == 0: 177 | return [] 178 | basename = filename + '.out/%d/' % n_planets 179 | parameters = json.load(file('%sparams.json' % basename)) 180 | a = pymultinest.Analyzer(n_params = len(parameters), outputfiles_basename=basename) 181 | s = a.get_stats() 182 | return map(lambda (p,m): (m['5sigma'][0] - (m['5sigma'][1] - m['5sigma'][0]), m['5sigma'][1] + (m['5sigma'][1] - m['5sigma'][0])), 183 | filter(lambda (p,m): p.startswith('P'), 184 | zip(parameters, s['marginals']))) 185 | def get_previous_best_fit(n_planets): 186 | basename = filename + '.out/%d/' % n_planets 187 | s = numpy.loadtxt("%s.txt" % basename) 188 | return (-0.5*s[:,1]).max() 189 | 190 | 191 | def evidence(diff): 192 | if diff > math.log(100): 193 | return 'decisive' 194 | if diff > math.log(30): 195 | return 'very strong' 196 | if diff > math.log(10): 197 | return 'strong' 198 | return 'indifferent' 199 | 200 | 201 | if len(sys.argv) == 3: 202 | n = int(sys.argv[2]) 203 | run(n) 204 | sys.exit(0) 205 | 206 | # number of planets 207 | n = 0 # start value, we'll increase it later 208 | 209 | ev, everr = run(0) 210 | while True: 211 | previous_periods = get_previous_periods(n) 212 | previous_best_fit = get_previous_best_fit(n) 213 | log_zero = log_very_bad_fit 214 | print 'using log_zero %.1f' % log_zero 215 | n = n + 1 216 | evnext, evnexterr = run(n, previous_periods, log_zero = log_zero) 217 | if evnext - evnexterr > ev + everr + math.log(10): 218 | print "%d planets preferred over %d" % (n, n - 1) 219 | print "evidence difference: %s" % evidence((evnext - evnexterr - (ev + everr))) 220 | elif evnext + evnexterr > ev - everr + math.log(10): 221 | print "%d planned could be preferred over %d if accuracy was better" % (n, n - 1) 222 | else: 223 | print "stick with %d planets, %d are not preferred" % (n - 1, n) 224 | break 225 | ev, everr = evnext, evnexterr 226 | 227 | # write out marginal posteriors in a nice way 228 | # calculate pdf on a_s * sin i = K*T*(1-e^2)**0.5 / 2 / pi 229 | # calculate pdf on m_p * sin i = K*mstar^(2/3)*T^(1/3)*(1-e^2)**0.5 / 2 / pi / G 230 | # calculate pdf on a = mstar^(1/3) * T^(2/3) * G 231 | 232 | # ? calculate pdf of m_p / mstar using an assumption on sin i 233 | 234 | # write out nice html / latex table of discovered planet system 235 | 236 | # go through all params with associated prob. p(params) 237 | # go through all times when observation could be made 238 | # go through all observations that could be made at that time with the given params (proportional to its probability): p(y|params, time) 239 | # search time that maximizes 240 | # \sum{ log( p(y|params, time) ) * p(y|params, time) * p(params) } 241 | # - \sum{ log( p(y|time) ) * p(y|time) dy } 242 | 243 | 244 | 245 | 246 | 247 | -------------------------------------------------------------------------------- /exo.c: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include "rungekutta.h" 4 | #include "interpol.h" 5 | 6 | #define MAX_NPOINTS 1024 7 | #define MAX_PLANETS 20 8 | #ifndef M_PI 9 | #define M_PI 3.14159265358979323846 10 | #endif 11 | 12 | #define NOISE 13 | 14 | #define gsl_matrix_get(data, i, j) data[i*3 + j] 15 | /* 16 | #define USE_RK4 17 | */ 18 | struct planetorbit { 19 | double P; 20 | double e; 21 | double chi; 22 | double omega; 23 | double K; 24 | 25 | double B; 26 | double v_x[MAX_NPOINTS]; 27 | double v_y[MAX_NPOINTS]; 28 | double v_yy[MAX_NPOINTS]; 29 | double minx; 30 | double maxx; 31 | double theta_0; 32 | }; 33 | 34 | static struct planetorbit * current; 35 | 36 | struct planetorbit planets[MAX_PLANETS]; 37 | 38 | double * params_low; 39 | double * params_high; 40 | double * data; 41 | unsigned int n_data; 42 | 43 | /* for RK4 */ 44 | double f(double t, double theta) { 45 | (void)t; 46 | return current->B * (1 - current->e * cos(theta + current->P * current->chi)); 47 | } 48 | 49 | double UniformPrior(double c, double xmin, double xmax){ 50 | return xmin + c * (xmax - xmin); 51 | } 52 | double LogPrior(double c, double xmin, double xmax){ 53 | double log_min = log(xmin); 54 | double log_max = log(xmax); 55 | return exp(log_min + c * (log_max - log_min)); 56 | } 57 | double ModLogPrior(double c, double xturn, double xmax) { 58 | return xturn * (exp( (c + 1e-10) * log(xmax/xturn + 1)) - 1); 59 | } 60 | 61 | void set_param_limits(double * new_params_low, double * new_params_high) { 62 | params_low = new_params_low; 63 | params_high = new_params_high; 64 | } 65 | void set_data(double * new_data, int new_n_data) { 66 | data = new_data; 67 | n_data = new_n_data; 68 | } 69 | 70 | /* 71 | * estimate theta from t0 = gsl_matrix_get(m->data, 0, 0), P, e, chi 72 | */ 73 | static void theta_setup() { 74 | current->minx = 0; 75 | current->theta_0 = 0; 76 | current->maxx = current->minx + current->P; 77 | current->B = 2 * M_PI / (current->P * pow(1 - pow(current->e, 2), 3./2)); 78 | rungekutta(current->v_x, current->v_y, current->v_yy, 79 | n_data, (current->maxx - current->minx)/n_data, 80 | current->minx, current->theta_0); 81 | } 82 | 83 | static double theta_eval(double x) { 84 | int j; 85 | double y; 86 | int n_points = n_data; 87 | 88 | if (x < current->minx || x > current->maxx) 89 | x = fmod(x, current->P); 90 | 91 | j = floor(n_points * (x - current->minx) / (current->maxx - current->minx)); 92 | 93 | if (j + 1 > n_points - 1) 94 | j = n_points - 2; 95 | 96 | #if INTERPOL == 1 97 | interpol_lin_setup( 98 | current->v_x[j], current->v_x[j+1], 99 | current->v_y[j], current->v_y[j+1]); 100 | y = interpol_lin_at(x); 101 | #elif INTERPOL == 2 102 | if (j+2 >= n_points) 103 | j--; 104 | interpol_sq_setup( 105 | current->v_x[j], current->v_x[j+1], current->v_x[j+2], 106 | current->v_y[j], current->v_y[j+1], current->v_y[j+2]); 107 | y = interpol_sq_at(x); 108 | #elif INTERPOL == 3 109 | assert(j < n_points); 110 | assert(j+1 < n_points); 111 | interpol_diff_setup( 112 | current->v_x[j], current->v_x[j+1], 113 | current->v_y[j], current->v_y[j+1], 114 | current->v_yy[j], current->v_yy[j+1]); 115 | y = interpol_diff_at(x); 116 | #else 117 | #error "INTERPOL not in [1 (lin), 2 (square) or 3 (diff)]" 118 | #endif 119 | return y; 120 | } 121 | 122 | void check(int ndim) { 123 | unsigned int i; 124 | printf(" PARAMETERS \n"); 125 | for (i = 0; (signed)i < ndim; i++) { 126 | printf(" param %i/%i low, high: %f %f\n", i, ndim, params_low[i], params_high[i]); 127 | } 128 | printf(" DATA \n"); 129 | for (i = 0; i < n_data; i++) { 130 | printf(" data point %i/%i: %f %f %f\n", i+1, n_data, gsl_matrix_get(data,i,0), gsl_matrix_get(data,i,1), gsl_matrix_get(data,i,2)); 131 | } 132 | } 133 | 134 | const double D2S = 60*60*24; 135 | 136 | // this function calculates eccentric anomaly 137 | double ecc_anomaly(double time, double prd, double ecc, double peri_pass) 138 | { 139 | time = time * D2S; 140 | double T = prd * D2S; 141 | double n = 2 * M_PI / T; 142 | double tau = peri_pass * T; 143 | double M = n * (time + tau); 144 | double E0; // start value 145 | double Mstar;// check equation 6.6.9 Danby 146 | double sigma; 147 | int int_M2PI; 148 | // to get an interger value for M/(2*PI) 149 | // so that for M we have a value between 0 and 2PI 150 | int_M2PI = M / (2 * M_PI); 151 | Mstar = M - int_M2PI * 2 * M_PI; 152 | 153 | // define a SIGN function 154 | sigma = fabs(sin(Mstar)); 155 | 156 | E0 = Mstar + sigma * 0.85 * ecc; 157 | // the value for k=0.85 is arbitrary 158 | // the only condition is 0<= k <=1 check, again Danby 159 | 160 | double TINY = 1e-6; 161 | int count = 0; 162 | for(count = 0; count < 100; count++) { 163 | double eSinE = ecc * sin(E0); // a dummy 164 | double f = E0 - eSinE - Mstar; 165 | if (fabs(f) < TINY) 166 | break; 167 | double eCosE = ecc * cos(E0); 168 | double f1 = 1 - eCosE; 169 | double f2 = eSinE; 170 | double f3 = eCosE; 171 | double dE0 = -f / f1; 172 | dE0 = -f / (f1 + 0.5 * dE0 * f2); 173 | dE0 = -f / (f1 + 0.5 * dE0 * f2 + dE0 * dE0 * f3 / 6); 174 | E0 = E0 + dE0; 175 | } 176 | 177 | return E0; 178 | } 179 | 180 | double true_anomaly(double time, double prd, double ecc, double peri_pass) 181 | { 182 | double E = ecc_anomaly(time,prd,ecc,peri_pass); 183 | double f = acos((cos(E) - ecc) / (1 - ecc * cos(E))); 184 | //acos gives the principal values ie [0:PI] 185 | //when E goes above PI we need another condition 186 | if (E > M_PI) 187 | f = 2 * M_PI - f; 188 | return f; 189 | } 190 | 191 | double calc(unsigned int n_planets, double ti) { 192 | double fi = 0; 193 | unsigned int j; 194 | 195 | for(j = 0; j < n_planets; j++) { 196 | current = &planets[j]; 197 | #ifdef USE_RK4 198 | fi += current->K * (cos(theta_eval(ti + current->chi * current->P) 199 | + current->omega) + current->e * cos(current->omega)); 200 | #else 201 | fi += current->K * (sin(true_anomaly(ti, current->P, current->e, current->chi) 202 | + current->omega + M_PI) + current->e * sin(current->omega + M_PI)); 203 | #endif 204 | } 205 | 206 | return fi; 207 | } 208 | 209 | /* (9) 210 | * A not needed. just sum over the datapoints. 211 | */ 212 | unsigned int set_params(double *params, int ndim) { 213 | unsigned int i; 214 | unsigned int j; 215 | 216 | double s; 217 | 218 | unsigned int n_planets; 219 | i = 0; 220 | n_planets = (ndim - 1) / 5; 221 | 222 | /* V prior: uniform improper */ 223 | /* (17) priors */ 224 | #ifdef USE_PRIORS 225 | params[i] = UniformPrior(params[i], params_low[i], params_high[i]); 226 | #endif 227 | /*V = params[0]; */ 228 | i++; 229 | 230 | if (ndim > (signed) n_planets * 5 + 1) { 231 | #ifdef USE_PRIORS 232 | params[i] = ModLogPrior(params[i], params_low[i], params_high[i]); 233 | #endif 234 | s = params[i++]; 235 | } else { 236 | s = 0; 237 | } 238 | 239 | for (j = 0; j < n_planets; j++) { 240 | current = &planets[j]; 241 | 242 | /* P prior: jeffreys */ 243 | #ifdef USE_PRIORS 244 | params[i] = LogPrior(params[i], params_low[i], params_high[i]); 245 | #endif 246 | current->P = params[i++]; /* days */ 247 | 248 | /* K prior: mod jeffreys */ 249 | #ifdef USE_PRIORS 250 | params[i] = ModLogPrior(params[i], params_low[i], params_high[i]); 251 | #endif 252 | current->K = params[i++]; 253 | 254 | /* Chi prior */ 255 | #ifdef USE_PRIORS 256 | params[i] = UniformPrior(params[i], params_low[i], params_high[i]); 257 | #endif 258 | current->chi = params[i++]; 259 | 260 | /* e */ 261 | #ifdef USE_PRIORS 262 | params[i] = UniformPrior(params[i], params_low[i], params_high[i]); 263 | #endif 264 | current->e = params[i++]; 265 | /* omega */ 266 | #ifdef USE_PRIORS 267 | params[i] = UniformPrior(params[i], params_low[i], params_high[i]); 268 | #endif 269 | current->omega = params[i++]; 270 | 271 | #ifdef USE_RK4 272 | theta_setup(); 273 | #endif 274 | } 275 | return n_planets; 276 | } 277 | 278 | int planets_are_sorted(unsigned int n_planets, int ndim) { 279 | unsigned int i; 280 | unsigned int j = 1; 281 | double lastp = planets[0].P; 282 | if (n_planets == 0) 283 | return 0; 284 | if (ndim > (signed) n_planets * 5 + 1) { 285 | j = 2; 286 | } 287 | for (i = 1; i < n_planets; i++) { 288 | /* avoid the planets that have been forced to a range */ 289 | if (params_low[i * 5 + j] / params_high[i * 5 + j] < 100) { 290 | continue; 291 | } 292 | 293 | /* avoid double solutions, by making sure the periods 294 | * are sorted in increasing order */ 295 | if (planets[i].P < lastp) 296 | return 1; 297 | lastp = planets[i].P; 298 | } 299 | return 0; 300 | } 301 | 302 | double LogLike(double *params, int ndim, int npars) { 303 | unsigned int n_planets = set_params(params, ndim); 304 | unsigned int i; 305 | 306 | double ti; 307 | double vi; 308 | double fi; 309 | double varsum; 310 | double vari; 311 | double prob = 0; 312 | double V = params[0]; 313 | #ifdef NOISE 314 | double s = params[1]; 315 | #endif 316 | 317 | /*if (planets_are_sorted(n_planets, ndim) != 0) 318 | return -1e300;*/ 319 | 320 | for (i = 0; i < n_data; i++) { 321 | ti = gsl_matrix_get(data,i,0); 322 | vi = gsl_matrix_get(data,i,1); 323 | vari = pow(gsl_matrix_get(data,i,2), 2); 324 | varsum = vari 325 | #ifdef NOISE 326 | + pow(s, 2) 327 | #endif 328 | ; 329 | 330 | /* (5) f_i = */ 331 | fi = V + calc(n_planets, ti); 332 | 333 | /* (9) 2*p = */ 334 | prob -= pow(vi - fi, 2) / varsum; 335 | 336 | /* (10) 2*A = */ 337 | prob -= log(varsum); 338 | } 339 | 340 | prob += log(2 * M_PI) * n_data; 341 | /*for (i = 0; (signed)i < ndim; i++) { 342 | printf(" %.3f \t", params[i]); 343 | } 344 | printf(" --> %f\n", prob);*/ 345 | 346 | return prob; 347 | } 348 | 349 | void predict(double *params, int ndim, int npars, int n_data, double * x, double * y) { 350 | unsigned int n_planets = set_params(params, ndim); 351 | unsigned int i; 352 | double V = params[0]; 353 | 354 | for (i = 0; i < n_data; i++) { 355 | y[i] = V + calc(n_planets, x[i]); 356 | } 357 | } 358 | 359 | -------------------------------------------------------------------------------- /rvfit.py: -------------------------------------------------------------------------------- 1 | import pymultinest 2 | import math 3 | import os, sys 4 | import numpy 5 | import json 6 | from random import Random 7 | from time import time 8 | import inspyred 9 | prng = Random() 10 | prng.seed(123) 11 | #prng.seed(time()) 12 | 13 | from ctypes import * 14 | #cdll.LoadLibrary('libnest3.so') 15 | lib = cdll.LoadLibrary('exo.so') 16 | lib.LogLike.restype = c_double 17 | 18 | # put data into double[n][3] array 19 | filename = sys.argv[1] 20 | data = numpy.loadtxt(filename) 21 | cdata = ((c_double * 3) * data.shape[0])() 22 | for i in range(data.shape[0]): 23 | for j in range(3): 24 | cdata[i][j] = data[i,j] 25 | lib.set_data(cdata, data.shape[0]) 26 | 27 | # simplify parameter space: 28 | # strongest impact is at most the variation of the data 29 | kmax = (data[:,1].max() - data[:,1].min()) 30 | print 'kmax: %f' % (kmax) 31 | pmax = (data[:,0].max() - data[:,0].min()) * 2 32 | # sampling theorem -- highest frequency --> longest period = 1 / dataduration / 2. 33 | fs = 1 / (data[1:,0] - data[:-1,0]).min() 34 | pmin = 1 / (fs / 2.) 35 | print 'P limits: %f..%f' % (pmin, pmax) 36 | #pmax = 15000 37 | 38 | # exclude very bad fits (worse than 5 sigma on all datapoints) 39 | log_bad_fit = (-0.5 * 3**2.) * len(data) 40 | log_very_bad_fit = (-0.5 * 5**2.) * len(data) 41 | 42 | 43 | def mkdir(path): 44 | if not os.path.exists(path): os.mkdir(path) 45 | 46 | n_params = None 47 | params_low = None 48 | params_high = None 49 | cube = None 50 | 51 | # number of dimensions our problem has 52 | parameters = ['V'] 53 | # plus possibly a noise term 54 | parameters += ['s'] 55 | 56 | def set_param_limits(): 57 | global n_params 58 | global params_low 59 | global params_high 60 | global cube 61 | 62 | n_params = len(parameters) 63 | params_low = (c_double * n_params)() 64 | params_high = (c_double * n_params)() 65 | cube = (c_double * n_params)() 66 | 67 | for i, p in enumerate(parameters): 68 | if p.startswith('V'): 69 | params_low[i], params_high[i] = -2000., +2000. 70 | if p.startswith('P'): 71 | if True or len(previous_periods) == 0: 72 | params_low[i], params_high[i] = pmin, pmax 73 | else: 74 | params_low[i], params_high[i] = previous_periods.pop() 75 | if params_low[i] < 0.2: params_low[i] = 0.2 76 | if params_high[i] > pmax: params_high[i] = pmax 77 | print 'fixing period (parameter %s) to %f..%f' % (p, params_low[i], params_high[i]) 78 | if p.startswith('K'): 79 | params_low[i], params_high[i] = 1., kmax 80 | if p.startswith('chi'): 81 | params_low[i], params_high[i] = 0., 1. 82 | if p.startswith('e'): 83 | params_low[i], params_high[i] = 0., 1. 84 | if p.startswith('omega'): 85 | params_low[i], params_high[i] = 0., 2 * math.pi 86 | if p.startswith('s'): 87 | params_low[i], params_high[i] = 0.1, 100 88 | lib.set_param_limits(params_low, params_high) 89 | 90 | set_param_limits() 91 | # later add more params 92 | n_planets = 0 93 | #for i in range(n_planets): 94 | # parameters += ['%s%d' % (var, i) for var in ['P', 'K', 'chi', 'e', 'omega']] 95 | 96 | def likelihood(*args): 97 | if len(args) == 1: 98 | args = args[0] 99 | if numpy.any(numpy.isnan(args)): 100 | return 1e300 101 | for i,a in enumerate(args): 102 | cube[i] = a 103 | v = lib.LogLike(cube, n_params, 0) 104 | #print args, v 105 | return -float(v) 106 | def transform(*args): 107 | for i,a in enumerate(args): 108 | cube[i] = a 109 | lib.set_params(cube, len(args)) 110 | return [float(cube[i]) for i in range(len(args))] 111 | 112 | fn = {} 113 | fn[2] = lambda c0, c1: likelihood(c0, c1) 114 | fn[7] = lambda c0, c1, c2, c3, c4, c5, c6: likelihood(c0, c1, c2, c3, c4, c5, c6) 115 | fn[12] = lambda c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11: likelihood(c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11) 116 | fn[17] = lambda c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16: likelihood(c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16) 117 | fn[22] = lambda c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21: likelihood(c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21) 118 | fn[25] = lambda c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26: likelihood(c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26) 119 | 120 | # do minuit search 121 | def run_minuit(start, errors): 122 | import minuit 123 | m = minuit.Minuit(fn[n_params]) 124 | minuitparams = ['c%d' % i for i in range(len(start))] 125 | for s, e, paramname in zip(start, errors, minuitparams): 126 | m.limits[paramname] = (0,1) 127 | m.errors[paramname] = e 128 | m.values[paramname] = s 129 | 130 | m.maxcalls = 100000 131 | # minimization succeeds when edm is less than 0.001*tol*up 132 | #m.tol = 50 133 | m.printMode = 0 134 | m.up = 0.5 135 | m.hesse() 136 | #m.migrad() 137 | for pminuit, paramname, val in zip(minuitparams, parameters, transform(*map(m.values.get, minuitparams))): 138 | print "\t%s: %f [+-%f]" % (paramname, val, m.errors[pminuit]) 139 | #print 'values', m.values 140 | #print 'errors', m.errors 141 | optparams = map(m.values.get, minuitparams) 142 | errors = map(m.errors.get, minuitparams) 143 | return {'opt': optparams, 'errors': errors, 'minuit':m} 144 | 145 | def run_fmin(start, errors): 146 | import scipy.optimize 147 | n = len(start) 148 | cons = [lambda x: x[i] for i in range(n)] + [lambda x: 1 - x[i] for i in range(n)] 149 | start = scipy.optimize.fmin_cobyla(func=likelihood, x0=start, cons=cons, disp=1, rhobeg = numpy.mean(errors)) 150 | start = scipy.optimize.fmin(func=likelihood, x0=start, disp=1) 151 | return start 152 | 153 | import matplotlib.pyplot as plt 154 | plt.ion() 155 | plt.errorbar(x=data[:,0], y=data[:,1], yerr=data[:,2], label='data', ls=' ') 156 | plt.show() 157 | 158 | def run_de(seeds): 159 | class GlobalOpt(inspyred.benchmarks.Benchmark): 160 | def __init__(self): 161 | inspyred.benchmarks.Benchmark.__init__(self, n_params) 162 | self.bounder = inspyred.ec.Bounder(0, 1) 163 | self.maximize = True 164 | 165 | def generator(self, random, args): 166 | return [random.uniform(0,1) for _ in range(n_params)] 167 | 168 | def evaluator(self, candidates, args): 169 | fitness = [] 170 | for c in candidates: 171 | fitness.append(-likelihood(*c)) 172 | return fitness 173 | 174 | problem = GlobalOpt() 175 | 176 | #ea.terminator = inspyred.ec.terminators.diversity_termination 177 | def stats_obs(population, num_generations, num_evaluations, args): 178 | stats = inspyred.ec.analysis.fitness_statistics(population) 179 | worst_fit = '{0:>10}'.format(stats['worst'])[:10] 180 | best_fit = '{0:>10}'.format(stats['best'])[:10] 181 | avg_fit = '{0:>10}'.format(stats['mean'])[:10] 182 | med_fit = '{0:>10}'.format(stats['median'])[:10] 183 | std_fit = '{0:>10}'.format(stats['std'])[:10] 184 | 185 | sys.stdout.write("{0:>10} {1:>10} {2:>10} {3:>10} {4:>10} {5:>10} {6:>10}\r".format( 186 | num_generations, num_evaluations, worst_fit, best_fit, med_fit, avg_fit, std_fit)) 187 | sys.stdout.flush() 188 | 189 | 190 | if n_params <= 2: 191 | timelimit = 2 192 | else: 193 | timelimit = 120 + 10 * n_params 194 | 195 | """print 'SA:' 196 | ea = inspyred.ec.SA(prng) 197 | ea.terminator = inspyred.ec.terminators.evaluation_termination 198 | ea.observer = stats_obs 199 | final_pop = ea.evolve(generator=problem.generator, 200 | evaluator=problem.evaluator, 201 | seeds = seeds, 202 | bounder=problem.bounder, 203 | maximize=problem.maximize, 204 | max_evaluations=30000) 205 | 206 | seeds = [c.candidate for c in final_pop] 207 | print """ 208 | ea = inspyred.ec.DEA(prng) 209 | ea.observer = stats_obs # inspyred.ec.observers.stats_observer 210 | #ea.observer = [inspyred.ec.observers.file_observer, inspyred.ec.observers.stats_observer] 211 | #ea.observer = inspyred.ec.observers.plot_observer 212 | ea.terminator = inspyred.ec.terminators.time_termination 213 | #ea.selector = inspyred.ec.selectors.fitness_proportionate_selection 214 | ea.selector = inspyred.ec.selectors.truncation_selection 215 | 216 | print 'EA:' 217 | final_pop = ea.evolve(generator=problem.generator, 218 | evaluator=problem.evaluator, 219 | pop_size=200, 220 | seeds=seeds, 221 | bounder=problem.bounder, 222 | maximize=problem.maximize, 223 | mutation_rate = 0.1, 224 | gaussian_stdev = 0.04, 225 | min_diversity = 0.1, 226 | max_time = timelimit, 227 | num_selected = 20, 228 | ) 229 | print 230 | #plt.close() 231 | return final_pop 232 | 233 | def plot(args, label = ''): 234 | x = (c_double * 1000)(*numpy.linspace(data[:,0].min(), data[:,0].max(), 1000)) 235 | y = (c_double * 1000)() 236 | for i,a in enumerate(args): 237 | cube[i] = a 238 | lib.predict(cube, len(cube), 0, 1000, x, y) 239 | plt.plot(x, y, '-', label=label) 240 | plt.legend(loc = 'upper left') 241 | #plt.savefig(label + '.pdf') 242 | plt.savefig(label + '.png') 243 | 244 | def aicc(res): 245 | n = 2 * len(data) 246 | k = len(res['opt']) 247 | return (k - 2 * res['optvalue'] + 2 * k * (k + 1) / (n - k - 1.)) / 2. 248 | 249 | 250 | # optimize globally first 251 | pop = run_de(seeds = [[prng.uniform(0,1) for _ in range(n_params)] for _ in range(10000)]) 252 | print 'optimization found highest loglikelihood %.1f' % -pop[-1].fitness, pop[-1].candidate, numpy.std([c.candidate for c in pop[-10:]], axis=0) 253 | plot(pop[-1].candidate, '%d planets' % n_planets) 254 | 255 | start = pop[-1].candidate 256 | errors = numpy.std([c.candidate for c in pop[-10:]], axis=0) 257 | start = run_fmin(start, errors) 258 | res = run_minuit(start = start, errors = errors) 259 | res['optvalue'] = -likelihood(*res['opt']) 260 | del res['minuit'] 261 | allresults = {0:res} 262 | 263 | while True: 264 | # add a planet 265 | parameters += ['%s%d' % (var, n_planets) for var in ['P', 'K', 'chi', 'e', 'omega']] 266 | n_planets += 1 267 | print '---- %d planets -------------------------' % n_planets 268 | 269 | set_param_limits() 270 | 271 | print 'creating seed population...' 272 | # create seed population from last population: 273 | seeds = [] 274 | for c in pop[-10:]: 275 | seeds += [c.candidate + [prng.uniform(0,1) for _ in range(5)] for i in range(100000)] 276 | 277 | print 'running DE...' 278 | pop = run_de(seeds = [[prng.uniform(0,1) for _ in range(n_params)] for _ in range(10000)]) 279 | print 'optimization found highest loglikelihood %.1f' % -pop[-1].fitness, pop[-1].candidate, numpy.std([c.candidate for c in pop[-10:]], axis=0) 280 | plot(pop[-1].candidate, '%d planets' % n_planets) 281 | 282 | res = run_minuit(start = pop[-1].candidate, errors = numpy.std([c.candidate for c in pop[-10:]], axis=0) ) 283 | del res['minuit'] 284 | res['optvalue'] = -likelihood(*res['opt']) 285 | print res 286 | allresults[n_planets] = res 287 | json.dump([allresults[i] for i in range(n_planets + 1)], 288 | file('result.json', 'w'), indent=4) 289 | 290 | 291 | print 'aicc of %d planets: %.1f' % (n_planets, aicc(allresults[n_planets])) 292 | for k,v in allresults.iteritems(): 293 | print ' %d: %f' % (k,v['optvalue']) 294 | lndiff = aicc(allresults[n_planets]) - aicc(allresults[n_planets - 1]) 295 | print 'ln diff: %.1f' % lndiff 296 | if lndiff > 0: 297 | print 'only taking %d planets' % (n_planets - 1) 298 | break 299 | #if abs(lndiff) < numpy.log(100): 300 | # break 301 | 302 | 303 | 304 | 305 | -------------------------------------------------------------------------------- /rvglobalfit.py: -------------------------------------------------------------------------------- 1 | import numpy 2 | import math 3 | import json 4 | import sys 5 | from random import Random 6 | import copy 7 | import evolve 8 | import inspyred 9 | import time 10 | import numpy, scipy, scipy.interpolate, scipy.stats 11 | from numpy import pi, log, exp, nan, sin, cos, arccos 12 | import matplotlib.pyplot as plt 13 | import dftransform 14 | import shutil 15 | plt.rcParams['pdf.compression'] = 0 16 | prng = Random() 17 | prng.seed(time.time()) 18 | #prng.seed(123) 19 | #prng.seed(1234) 20 | 21 | from ctypes import * 22 | lib = cdll.LoadLibrary('exo_noprior.so') 23 | lib.LogLike.restype = c_double 24 | lib.ModLogPrior.restype = c_double 25 | 26 | # put data into double[n][3] array 27 | evolve.parse_args() 28 | filename = evolve.args.data 29 | prefix = filename + '_out_' 30 | 31 | data = numpy.loadtxt(filename) 32 | cdata = ((c_double * 3) * data.shape[0])() 33 | for i in range(data.shape[0]): 34 | for j in range(3): 35 | cdata[i][j] = data[i,j] 36 | lib.set_data(cdata, data.shape[0]) 37 | 38 | # simplify parameter space: 39 | # strongest impact is at most the variation of the data 40 | kmax = (data[:,1].max() - data[:,1].min()) 41 | kavg = data[:,1].mean() 42 | print 'kmax: %f' % (kmax) 43 | pmax = (data[:,0].max() - data[:,0].min()) * 10 44 | # sampling theorem -- highest frequency --> longest period = 1 / dataduration / 2. 45 | fs = 1 / (data[1:,0] - data[:-1,0]).min() 46 | pmin = 1 / (fs / 2.) 47 | print 'P limits: %f..%f' % (pmin, pmax) 48 | sturn = data[:,2].min() 49 | smax = (data[:,1].max() - data[:,1].min()) 50 | print 's limits: %f..%f' % (sturn, smax) 51 | #pmax = 15000 52 | 53 | # exclude very bad fits (worse than 5 sigma on all datapoints) 54 | log_bad_fit = (-0.5 * 3**2.) * len(data) 55 | log_good_fit = (-0.5 * 3**2.) * len(data) 56 | log_very_bad_fit = (-0.5 * 5**2.) * len(data) 57 | 58 | n_params = None 59 | params_low = None 60 | params_high = None 61 | cube = None 62 | 63 | n_data = len(data) 64 | 65 | def update_n_params(new_n_params): 66 | global n_params 67 | global cube 68 | global params_low 69 | global params_high 70 | n_params = new_n_params 71 | 72 | # number of dimensions our problem has 73 | global parameters 74 | parameters = ['V'] 75 | # plus possibly a noise term 76 | parameters += ['s'] 77 | 78 | n_planets = (n_params - 2) / 5 79 | 80 | for i in range(n_planets): 81 | parameters += ['%s%d' % (var, i) for var in ['P', 'K', 'chi', 'e', 'omega']] 82 | 83 | params_low = (c_double * n_params)() 84 | params_high = (c_double * n_params)() 85 | cube = (c_double * n_params)() 86 | 87 | for i, p in enumerate(parameters): 88 | if p.startswith('V'): 89 | params_low[i], params_high[i] = kavg - 2 * kmax, kavg + 2 * kmax, 90 | if p.startswith('P'): 91 | params_low[i], params_high[i] = pmin, pmax 92 | if p.startswith('K'): 93 | params_low[i], params_high[i] = 1., kmax 94 | if p.startswith('chi'): 95 | params_low[i], params_high[i] = 0., 1. 96 | if p.startswith('e'): 97 | params_low[i], params_high[i] = 0., 1. 98 | if p.startswith('omega'): 99 | params_low[i], params_high[i] = 0., 2 * math.pi 100 | if p.startswith('s'): 101 | params_low[i], params_high[i] = sturn, smax 102 | lib.set_param_limits(params_low, params_high) 103 | 104 | def prepare(params): 105 | global n_params 106 | global cube 107 | global params_low 108 | global params_high 109 | if len(params) != n_params: 110 | update_n_params(len(params)) 111 | assert not numpy.any(numpy.isinf(params)) 112 | assert not numpy.any(numpy.isnan(params)) 113 | for i,a in enumerate(params): 114 | cube[i] = a 115 | 116 | def predict(c, x = None): 117 | if x is None: 118 | n = 1000 119 | x = (c_double * 1000)(*numpy.linspace(data[:,0].min(), data[:,0].max(), 1000)) 120 | else: 121 | n = len(x) 122 | x = (c_double * len(x))(*x) 123 | y = (c_double * len(x))() 124 | prepare(c) 125 | global cube 126 | global n_params 127 | lib.predict(cube, n_params, 0, n, x, y) 128 | return numpy.asarray([x[i] for i in range(n)]), numpy.asarray([y[i] for i in range(n)]) 129 | 130 | def plot_predict_resid(newcube, predicted, resid): 131 | plt.figure("predict", figsize=(7,10)) 132 | plt.subplot(2,1,1) 133 | lastplanet = newcube[-5:] 134 | if len(lastplanet) == 5: 135 | lastplanet[0] = 1. / lastplanet[0] 136 | lastplanet[1] = lastplanet[1]**2 137 | lastplanet[2] = (numpy.fmod(lastplanet[2] + 0.5, 1) - 0.5) *2*pi 138 | lastplanet[4] = (numpy.fmod(lastplanet[4] + 0.5, 1) - 0.5) 139 | 140 | #plt.title(", ".join(["%.2f" % c for c in lastplanet])) 141 | plt.title("freq=%.5f pow=%.1f ang1=%.3f e=%.2f ang2=%.3f" % tuple(lastplanet), fontsize='small') 142 | elif len(lastplanet) == 2: 143 | plt.title("V=%.3f s=%.3f" % tuple(lastplanet)) 144 | else: 145 | plt.title("V=%.3f" % tuple(lastplanet)) 146 | #plt.errorbar(x=data[:,0], y=data[:,1], yerr=data[:,2], marker='x', ls=' ') 147 | plt.plot(data[:,0], predicted, 's ') 148 | plt.subplot(2,1,2) 149 | plt.plot(data[:,0], resid, 's ') 150 | 151 | def plot_fft_results(freq, pow, cum, angles, x, y, freqi, powi, angi): 152 | plt.figure(figsize=(10,15)) 153 | plt.subplot(5,1,1) 154 | plt.plot(freq, pow, 'x') 155 | plt.gca().set_xscale('log') 156 | plt.xlabel('frequency') 157 | plt.ylabel('power') 158 | plt.hlines(powi, plt.xlim()[0], plt.xlim()[1], alpha=0.3) 159 | plt.subplot(5,1,2) 160 | plt.plot(freq, cum, 'x') 161 | plt.gca().set_xscale('log') 162 | plt.ylabel('cumulative probability') 163 | plt.xlabel('frequency') 164 | plt.vlines(freqi, 0, 1, alpha=0.3) 165 | plt.subplot(5,1,3) 166 | plt.plot(cum, pow**0.5, 'x') 167 | plt.xlabel('cumulative probability') 168 | plt.ylabel('amplitude') 169 | plt.hlines(powi**0.5, plt.xlim()[0], plt.xlim()[1], alpha=0.3) 170 | plt.subplot(5,1,4) 171 | plt.plot(cum, numpy.fmod(angles+10*numpy.pi, 2*numpy.pi), 'x') 172 | plt.xlabel('cumulative probability') 173 | plt.ylabel('phase') 174 | plt.hlines(numpy.fmod(angi+10*numpy.pi, 2*numpy.pi), plt.xlim()[0], plt.xlim()[1], alpha=0.3) 175 | plt.subplot(5,1,5) 176 | plt.plot(x, y, 'x') 177 | plt.xlabel('time') 178 | plt.ylabel('value') 179 | xnew = numpy.linspace(x.min(), x.max(), 200) 180 | ynew = powi**0.5 * numpy.cos(xnew * freqi * pi * 2 + angi) 181 | plt.plot(xnew, ynew, '-', alpha=0.5, lw=3, 182 | label="freq=%.5f\npow=%.3f\nang=%.3f" % (freqi, powi, angi)) 183 | plt.legend(loc='best', prop=dict(size=10)) 184 | 185 | Vrv = scipy.stats.norm(numpy.mean(data[:,1]), numpy.std(data[:,1])) 186 | 187 | def fft_prior(cube, plot = False): 188 | global params_low 189 | global params_high 190 | global n_params 191 | if len(cube) != n_params: 192 | update_n_params(len(cube)) 193 | assert len(cube) == n_params, [cube, n_params] 194 | # K0 is gaussian with average + stdev of data 195 | bias = 0. 196 | print 'FFT transforming', cube 197 | 198 | newcube = [] 199 | lowprob = scipy.stats.norm().cdf(-5) 200 | V = Vrv.ppf(lowprob + (cube[0]) / (1 + 2 * lowprob)) 201 | newcube.append(V) 202 | bias += log(Vrv.pdf(V)) 203 | # would have been uniform prior, so don't need anything too special here 204 | 205 | ndim = len(cube) 206 | n_planets = (ndim - 1) / 5 207 | i0 = 1 208 | if ndim > n_planets * 5 + 1: 209 | # using s 210 | noise = float(lib.ModLogPrior(c_double(cube[1]), c_double(params_low[1]), c_double(params_high[1]) )) 211 | newcube.append(noise) 212 | # no bias here, same as before 213 | i0 = i0 + 1 214 | 215 | # now P, K, chi, e, omega 216 | # K * (sin(true_anomaly(ti + chi * P) + omega) + e * sin(omega)) 217 | # K * (sin(true_anomaly(ti, P, e, chi) + omega + PI) + e * sin(omega + PI)) 218 | # omega = phase 219 | # P = period 220 | # ti = 221 | # chi = 222 | # e = eccentricity: uniform prior 223 | for i in range(n_planets): 224 | #for k in range(2, len(cube)): 225 | # cube[k] = 0.5 226 | 227 | j0 = i * 5 + i0 228 | # calculate psd of residuals 229 | #print 'calculating psd of residuals' 230 | x, preddata = predict(newcube, data[:,0]) 231 | resdata = data[:,1] - preddata 232 | if plot: 233 | #print 'plotting prediction...' 234 | plt.figure("predict_%02d" % i) 235 | plot_predict_resid(newcube, preddata, resdata) 236 | plt.subplot(2,1,1) 237 | plt.errorbar(x=data[:,0], y=data[:,1], yerr=data[:,2], label='data', marker='x', ls=' ', color='black') 238 | plt.savefig("predict_%02d_.pdf" % i) 239 | shutil.move("predict_%02d_.pdf" % i, "predict_%02d.pdf" % i) 240 | plt.close() 241 | #print 'plotting prediction ... done' 242 | 243 | #print 'dft prepare...' 244 | dftransform.prepare(data[:,0], resdata) 245 | dftransform.updateNyquist() 246 | dftransform.updateZeropoint() 247 | #print 'dft calc...' 248 | out = dftransform.calc() 249 | #print 'dft calc... done' 250 | freq, pow, angles = out[1:,0], out[1:,1], out[1:,2] 251 | 252 | low = pow.sum() / len(pow) / 20. 253 | quitelow = pow.max() / 4. 254 | pdf = numpy.where(pow < quitelow, low, pow) 255 | pdf = pdf / freq # weighting: prefer lower frequencies 256 | cum = pdf.cumsum() 257 | cum /= cum[-1] 258 | 259 | freqi = numpy.interp(x=cube[j0], xp=cum, fp=freq) 260 | powi = numpy.interp(x=cube[j0], xp=cum, fp=pow) 261 | angi = numpy.interp(x=cube[j0], xp=cum, fp=angles) + 3*pi/2. 262 | if plot: 263 | #print 'plotting results ...' 264 | plot_fft_results(freq, pow, cum, angles, data[:,0], resdata, freqi, powi, angi) 265 | plt.savefig("fft_%02d_.pdf" % i) 266 | shutil.move("fft_%02d_.pdf" % i, "fft_%02d.pdf" % i) 267 | plt.close() 268 | #print 'plotting results ... done' 269 | 270 | # powi**0.5 * numpy.sin(xnew * freqi * pi * 2 + angi) 271 | 272 | # generate distributions 273 | # draw from distributions using cube[] ppf 274 | # also compute probability there, vs. usual draw --> bias 275 | 276 | period = 1. / (freqi) 277 | bias -= scipy.stats.uniform.logpdf(log(period), log(pmin), log(pmax)) 278 | 279 | amps = powi**0.5 280 | amprv = scipy.stats.norm(amps, amps / 5.) # so that at lowprob, amprv.ppf gives 0 281 | amplitude = amprv.ppf(lowprob + (cube[j0 + 1]) / (1 + 2 * lowprob)) 282 | bias += log( amprv.pdf(amplitude) ) 283 | bias -= scipy.stats.uniform.logpdf(log(amplitude), log(1.), log(kmax)) 284 | 285 | # make angle 2 sigma correspond to 2 pi 286 | chirv = scipy.stats.norm(angi / (2*pi), 1./2.) 287 | chi = numpy.fmod(chirv.ppf(lowprob + (cube[j0 + 2]) / (1 + 2 * lowprob)) + 3., 1) 288 | bias += log( chirv.pdf(chi) ) 289 | #print 'new angle: ', angi, (numpy.fmod(chi + 0.5, 1) - 0.5) * 2*pi 290 | if plot: 291 | #print 'plotting results ...' 292 | plot_fft_results(freq, pow, cum, angles, data[:,0], resdata, 293 | 1 / period, amplitude**2, 294 | chi * 2*pi) 295 | #print 'plotting results ... saving' 296 | plt.savefig("fftrand_%02d_.pdf" % i) 297 | shutil.move("fftrand_%02d_.pdf" % i, "fftrand_%02d.pdf" % i) 298 | plt.close() 299 | #print 'plotting results ... done' 300 | # would have been uniform prior, so don't need anything too special here 301 | 302 | # draw e, no bias there 303 | e = cube[j0 + 3]**2 304 | 305 | # calculate omega from chi and angle 306 | cosE = cos(chi / (1 + e)) 307 | f0 = arccos( (cosE - e) / ( 1 - e * cosE) ) 308 | omega0 = numpy.fmod(f0 - (chi) + 4*pi, 2*pi) 309 | omegarv = scipy.stats.norm(omega0, 1./5. * (2*numpy.pi)) 310 | omega = numpy.fmod(omegarv.ppf(lowprob + (cube[j0 + 4]) / (1 + 2 * lowprob)) + 5., 1) 311 | bias += log( omegarv.pdf(omega) ) 312 | #chi = numpy.fmod(chi + 0.25, 1) 313 | omega = numpy.fmod(omega + 3*pi/2., 2*pi) 314 | 315 | newcube += [period, amplitude, chi, e, omega] 316 | print ' Period%2d : %2.5f' % (i, period) 317 | print ' Amplitude%2d: %2.5f' % (i, amplitude) 318 | # next round 319 | print ' Velocity: %2.2f' % V 320 | print ' Noise : %.3f' % noise 321 | #if plot: 322 | # sys.stdout.write("press return: ") 323 | # sys.stdout.flush() 324 | # sys.stdin.readline() 325 | 326 | assert not numpy.any(numpy.isinf(newcube)), newcube 327 | assert not numpy.any(numpy.isnan(newcube)), newcube 328 | # so finally we have parameters and a parameter space deformation correction bias 329 | return bias, newcube 330 | 331 | class Planets(evolve.AbstractComponents): 332 | """ rv planet """ 333 | def __init__(self, n): 334 | self.n = n 335 | self.component_length = 5 # ['P', 'K', 'chi', 'e', 'omega'] 336 | self.head = 2 # V s 337 | self.length = self.head + self.component_length * n 338 | self.low_bounds = [0] * self.length 339 | self.high_bounds = [1] * self.length 340 | def decode(self, encoding, plot=False): 341 | bias, fftparams = fft_prior(encoding, plot=plot) 342 | return fftparams 343 | def visualize(self, encodings, limits): 344 | plt.figure("viz", figsize=(7,10)) 345 | for i, encoding in enumerate(encodings[:len(encodings)/3+2]): 346 | #for i, encoding in enumerate(encodings[0:1]): 347 | alpha = 1 if i == 0 else 0.3 348 | fftparams = self.decode(encoding, plot = (i == 0)) 349 | plt.figure("viz") 350 | x, predicted = predict(fftparams, data[:,0]) 351 | resid = data[:,1] - predicted 352 | plt.subplot(3,1,1) 353 | plt.plot(data[:,0], predicted, 's ', alpha=alpha) 354 | plt.subplot(3,1,2) 355 | plt.plot(predicted, 's ', alpha=alpha) 356 | plt.subplot(3,1,3) 357 | plt.plot(resid, 's ', alpha=alpha) 358 | 359 | plt.subplot(3,1,1) 360 | plt.errorbar(x=data[:,0], y=data[:,1], yerr=data[:,2], label='data', marker='x', ls=' ', color='black') 361 | plt.subplot(3,1,2) 362 | plt.errorbar(x=range(len(data[:,1])), y=data[:,1], yerr=data[:,2], label='data', marker='x', ls=' ', color='black') 363 | plt.savefig(prefix + "_viz_.pdf") 364 | shutil.move(prefix + "_viz_.pdf", prefix + "_viz.pdf") 365 | plt.close() 366 | 367 | 368 | class RVLikelihoodModel(object): 369 | def __init__(self, encoder, limits, events=[]): 370 | self.events = events 371 | self.encoder = encoder 372 | self.limits = limits 373 | 374 | def likelihood(self, candidate): 375 | fftparams = self.encoder.decode(candidate) 376 | prepare(fftparams) 377 | v = lib.LogLike(cube, len(cube), 0) # - bias 378 | #print params, [cube[i] for i in range(n_params)], v 379 | return v 380 | 381 | def generate_events(self, candidate, random, k): 382 | raise NotImplementedError 383 | 384 | # mutation that adds or removes a planet 385 | def rotation_variator(random, candidates, args): 386 | mutants = [] 387 | #print 'dimension_variator' 388 | r = args.get('rotate_rate', 0.05) 389 | 390 | for c in candidates: 391 | n_planets = (len(c) - 2) / 5 392 | for i in range(n_planets): 393 | j = 2 + i * 5 + 2 394 | k = 2 + i * 5 + 4 395 | if c[j] > 0.95 and random.uniform(0,1) < r: 396 | print 'applying rotation', c[j] 397 | c[j] = 1 - c[j] 398 | elif c[j] < 0.05 and random.uniform(0,1) < r: 399 | print 'applying rotation', c[j] 400 | c[j] = 1 - c[j] 401 | if c[k] > 0.95 and random.uniform(0,1) < r: 402 | print 'applying rotation', c[k] 403 | c[k] = 1 - c[k] 404 | elif c[k] < 0.05 and random.uniform(0,1) < r: 405 | print 'applying rotation', c[k] 406 | c[k] = 1 - c[k] 407 | mutants.append(c) 408 | return mutants 409 | 410 | if __name__ == '__main__': 411 | limits = [[data[:,0].min(), data[:,0].max()]] 412 | 413 | encoder = Planets(evolve.args.n_components) 414 | 415 | pm = RVLikelihoodModel(encoder=encoder, limits=limits) 416 | ea, evolve_args = evolve.setup(pm) 417 | # inject rotation variator 418 | ea.variator += [rotation_variator] 419 | 420 | gc_debug = False 421 | if gc_debug: 422 | from collections import defaultdict 423 | from gc import get_objects 424 | before = defaultdict(int) 425 | after = defaultdict(int) 426 | for i in get_objects(): 427 | before[type(i)] += 1 428 | import cherrypy 429 | import dowser 430 | 431 | """cherrypy.tree.mount(dowser.Root()) 432 | cherrypy.config.update({ 433 | 'environment': 'embedded', 434 | 'server.socket_port': 8081 435 | }) 436 | cherrypy.engine.start()""" 437 | evolve.evolve(ea, evolve_args) 438 | if gc_debug: 439 | objs = get_objects() 440 | for i in objs: 441 | after[type(i)]+=1 442 | items = [(k,after[k]-before[k]) for k in after if after[k]-before[k]] 443 | for j, (k,v) in enumerate(sorted(items, key = lambda (k,v):v, reverse=True)): 444 | print ' ', k, v, [str(i) for i in evolve.random.sample(objs, 1000) if type(i) == k][:10] 445 | if j > 5: break 446 | del objs 447 | del items 448 | 449 | from guppy import hpy 450 | h = hpy() 451 | print h.heap() 452 | del h 453 | #cherrypy.engine.block() 454 | 455 | --------------------------------------------------------------------------------