Now integrate() can return the full trajectory and plotted in the Python; improved readablity

This commit is contained in:
Yohai Meiron 2020-04-12 21:37:03 -04:00
parent a2219af116
commit a624814a0c
5 changed files with 144 additions and 106 deletions

14
Makefile Normal file
View file

@ -0,0 +1,14 @@
OPTIMIZATION ?= 3
CXXFLAGS += -O$(OPTIMIZATION)
LIB += -lgsl
EXECUTABLE ?= main
default:
$(CXX) $(CPPFLAGS) $(CXXFLAGS) $(INC) loadtxt.cpp main.cpp -o $(EXECUTABLE) $(LIB)
lib:
$(CXX) $(CPPFLAGS) $(CXXFLAGS) -fPIC $(INC) loadtxt.cpp main.cpp -shared -o lib$(EXECUTABLE).so $(LIB)
clean:
rm -f *.o *.so $(EXECUTABLE)

View file

@ -1,14 +1,12 @@
#include <algorithm> #include <algorithm>
#include <cstdlib> #include <cstdlib>
#include <fstream> #include <fstream>
#include <vector> #include "loadtxt.h"
//TODO if cols is an empty vector, get all columns from the file //TODO if cols is an empty vector, get all columns from the file
//TODO error checking //TODO error checking
class Loadtxt { Loadtxt::Loadtxt(std::string file_name, std::vector<int> cols)
public:
Loadtxt(std::string file_name, std::vector<int> cols)
{ {
std::sort(cols.begin(), cols.end()); std::sort(cols.begin(), cols.end());
n_cols = cols.size(); n_cols = cols.size();
@ -30,11 +28,13 @@ public:
buffer = (double*)realloc(buffer, n_cols * sizeof(double) * (++row)); buffer = (double*)realloc(buffer, n_cols * sizeof(double) * (++row));
n_rows = row; n_rows = row;
} }
~Loadtxt()
Loadtxt::~Loadtxt()
{ {
free(buffer); free(buffer);
} }
std::vector<std::vector<double>> get_cols()
std::vector<std::vector<double>> Loadtxt::get_cols()
{ {
std::vector<std::vector<double>> data(n_cols); std::vector<std::vector<double>> data(n_cols);
for (int col=0; col<n_cols; col++) data[col] = std::vector<double>(n_rows); for (int col=0; col<n_cols; col++) data[col] = std::vector<double>(n_rows);
@ -45,9 +45,8 @@ public:
} }
return data; return data;
} }
private:
const char *whitespaces = " \t"; void Loadtxt::line_to_buf(std::vector<int> cols, std::string line, double *buffer)
void line_to_buf(std::vector<int> cols, std::string line, double *buffer)
{ {
int n_cols = cols.size(); int n_cols = cols.size();
line = line.substr(line.find_first_not_of(whitespaces)); line = line.substr(line.find_first_not_of(whitespaces));
@ -65,9 +64,6 @@ private:
if (col++ == cols[idx]) buffer[idx++] = std::stod(line); if (col++ == cols[idx]) buffer[idx++] = std::stod(line);
if (idx < n_cols) throw; if (idx < n_cols) throw;
} }
double *buffer;
int n_rows, n_cols;
};
// Below is a deomonstration. The file has multiple columns but we are only // Below is a deomonstration. The file has multiple columns but we are only
// interested in the second and fourth. We pass the file name and the column // interested in the second and fourth. We pass the file name and the column

14
loadtxt.h Normal file
View file

@ -0,0 +1,14 @@
#pragma once
#include <string>
#include <vector>
class Loadtxt {
public:
Loadtxt(std::string file_name, std::vector<int> cols);
~Loadtxt();
std::vector<std::vector<double>> get_cols();
private:
const char *whitespaces = " \t";
void line_to_buf(std::vector<int> cols, std::string line, double *buffer);
double *buffer;
int n_rows, n_cols;
};

View file

@ -1,12 +1,16 @@
#include <iostream> #include <algorithm>
#include <string>
#include <vector>
#include <fstream> #include <fstream>
#include <numeric>
#include <gsl/gsl_errno.h> #include <gsl/gsl_errno.h>
#include <gsl/gsl_math.h> #include <gsl/gsl_math.h>
#include <gsl/gsl_odeiv2.h> #include <gsl/gsl_odeiv2.h>
#include <gsl/gsl_spline.h> #include <gsl/gsl_spline.h>
#include <iostream>
#include <numeric>
#include <string>
#include <stdexcept>
#include <vector>
#include "loadtxt.h"
extern "C" const int gsl_success() { return GSL_SUCCESS; } // It's zero, but just for clarity sake. extern "C" const int gsl_success() { return GSL_SUCCESS; } // It's zero, but just for clarity sake.
@ -22,6 +26,7 @@ public:
spline = gsl_spline_alloc(gsl_interp_cspline, x.size()); spline = gsl_spline_alloc(gsl_interp_cspline, x.size());
gsl_spline_init(spline, x.data(), y.data(), x.size()); gsl_spline_init(spline, x.data(), y.data(), x.size());
} }
Interp() {}
inline double operator()(double x) const inline double operator()(double x) const
{ {
return gsl_spline_eval(spline, x, acc); return gsl_spline_eval(spline, x, acc);
@ -54,32 +59,20 @@ class Galaxy {
public: public:
Galaxy(std::string file_name) Galaxy(std::string file_name)
{ {
std::vector<double> t_data, M_halo_data, b_halo_data; auto data = Loadtxt("file.dat", {1, 2, 3}).get_cols();
auto& t_data = data[0];
std::ifstream file(file_name); auto& halo_m_data = data[1];
std::string line; auto& halo_b_data = data[2];
while (std::getline(file, line)) { std::transform(t_data.begin(), t_data.end(), t_data.begin(), [](const double& x){ return x-2.145; });
auto pos = line.find('#'); std::transform(halo_b_data.begin(), halo_b_data.end(), halo_b_data.begin(), [](const double& x){ return x*0.7664209365408798; });
if (pos != std::string::npos) line = line.substr(0, pos); interp_halo_m = Interp(t_data, halo_m_data);
pos = line.find_first_not_of(" \t"); interp_halo_b = Interp(t_data, halo_b_data);
if (pos == std::string::npos) continue;
double data[3];
sscanf(line.c_str(), "%*s %lf %lf %lf", &data[0], &data[1], &data[2]);
t_data.push_back(data[0]);
M_halo_data.push_back(data[1]);
b_halo_data.push_back(data[2]); // note, this is not half-mass radius
}
interp_M_halo = new Interp(t_data, M_halo_data);
interp_b_halo = new Interp(t_data, b_halo_data);
} }
int func(double t, const double y[], double f[], void *params) int func(double t, const double y[], double f[], void *params)
{ {
double M_halo = (*interp_M_halo)(t); double halo_m = interp_halo_m(t);
double b_halo = (*interp_b_halo)(t); double halo_b = interp_halo_b(t);
/*printf("xxxxxxxxx %e, %e msun\n", t, M_halo); Plummer plummer(halo_m, halo_b);
printf("xxxxxxxxx %e, %e kpc\n", t, b_halo);
exit(0);*/
Plummer plummer(M_halo, b_halo);
f[0] = y[3]; // vx -> x' f[0] = y[3]; // vx -> x'
f[1] = y[4]; // vy -> y' f[1] = y[4]; // vy -> y'
f[2] = y[5]; // vz -> z' f[2] = y[5]; // vz -> z'
@ -88,27 +81,25 @@ public:
} }
private: private:
Interp *interp_M_halo; Interp interp_halo_m;
Interp *interp_b_halo; Interp interp_halo_b;
} galaxy("file.dat"); } galaxy("file.dat");
// Not very nice to have it as a global variable but GSL will have problem otherwise. // Not very nice to have it as a global variable but GSL will have problem otherwise.
int jac(double t, const double y[], double *dfdy, double dfdt[], void *params) int jac(double t, const double y[], double *dfdy, double dfdt[], void *params) { return GSL_SUCCESS; }
{
return GSL_SUCCESS;
}
inline int func(double t, const double y[], double f[], void *params)
int func(double t, const double y[], double f[], void *params)
{ {
return galaxy.func(t, y, f, params); return galaxy.func(t, y, f, params);
} }
extern "C" extern "C"
int integrate(const double y0[], const double t_max, double y[]) int integrate(const double y0[], const double t_max, const double step_size, double y[])
{ {
double t = 2.145; double t = 0;
constexpr double h = 1./4096.; constexpr double h = 1./4096.;
if (step_size/h - (int)(step_size/h) != 0) throw std::runtime_error("step_size must be a multiple of h");
constexpr double epsabs = 1e-7; constexpr double epsabs = 1e-7;
constexpr double epsrel = 0; constexpr double epsrel = 0;
const gsl_odeiv2_step_type *T = gsl_odeiv2_step_rk8pd; const gsl_odeiv2_step_type *T = gsl_odeiv2_step_rk8pd;
@ -118,12 +109,23 @@ int integrate(const double y0[], const double t_max, double y[])
gsl_odeiv2_system sys = {func, jac, 6, nullptr}; gsl_odeiv2_system sys = {func, jac, 6, nullptr};
gsl_odeiv2_driver *d = gsl_odeiv2_driver_alloc_y_new(&sys, T, h, epsabs, epsrel); gsl_odeiv2_driver *d = gsl_odeiv2_driver_alloc_y_new(&sys, T, h, epsabs, epsrel);
int step = 0;
const int step_max = t_max / step_size;
std::copy(y0, y0+6, y); std::copy(y0, y0+6, y);
int status = gsl_odeiv2_driver_apply(d, &t, t_max, y); for (int step=0; step<step_max; step++) {
return status; std::copy(y+step*6, y+(step+1)*6, y+(step+1)*6);
int status = gsl_odeiv2_driver_apply(d, &t, (step+1)*step_size, y+(step+1)*6);
if (status != GSL_SUCCESS) return status;
}
return GSL_SUCCESS;
} }
int main() int main()
{ {
std::cout << "bye" << std::endl; std::cout << "bye" << std::endl;
double y[12];
double y0[] = {80,0,0,0,80,0};
for (int i=0; i<30000; i++)
integrate(y0, 10, 10, y);
return 0;
} }

28
plot.py
View file

@ -4,21 +4,33 @@ import ctypes, subprocess
recompile = True recompile = True
if recompile: if recompile:
p = subprocess.Popen('g++ -shared -o libmain.so -fPIC main.cpp -lgsl'.split()) p = subprocess.Popen('g++ -shared -o libmain.so -fPIC loadtxt.cpp main.cpp -lgsl'.split())
p.wait() p.wait()
if p.returncode != 0: raise RuntimeError(p.returncode) if p.returncode != 0: raise RuntimeError(p.returncode)
libmain = ctypes.CDLL('./libmain.so') libmain = ctypes.CDLL('./libmain.so')
def integrate(y0, t_max): def integrate(y0, t_max, step_size=None):
y0 = (ctypes.c_double*6)(*y0) y0 = (ctypes.c_double*6)(*y0)
y = (ctypes.c_double*6)() if step_size is None: step_size = t_max
status = libmain.integrate(y0, ctypes.c_double(t_max), y) size = int(t_max // step_size) + 1
y = (ctypes.c_double*size*6)()
status = libmain.integrate(y0, ctypes.c_double(t_max), ctypes.c_double(step_size), y)
y = np.array(y).reshape(size,6)
return np.array(y), status return np.array(y), status
gsl_success = libmain.gsl_success() gsl_success = libmain.gsl_success()
#t_array = linspace(plot_tmin, plot_tmax, plot_points) from pylab import *
#r_array = empty_like(t_array) t_max = 10
ic = [80,0,0,0,80,0]
res = integrate(ic, t_max, step_size=1/4096)
x, y, z, vx, vy, vz = res[0].T
zzz = x[-1]
plot(x,y)
res = integrate([10,0,0,0,200,0], 2.245) res = integrate(ic, t_max)
print(res, gsl_success) x, y, z, vx, vy, vz = res[0].T
plot(x,y,'o')
print(zzz - x[-1])
# gca().set_aspect('equal')
show()