template for user defined cython potential of gala
code = r'''
from __future__ import division, print_function
import numpy as np
from collections import OrderedDict
from gala.potential.potential.cpotential import CPotentialBase
from gala.potential.potential.cpotential cimport CPotentialWrapper
from gala.potential.potential.cpotential cimport densityfunc, energyfunc, gradientfunc, hessianfunc
from libc.math cimport sqrt, pow, log, exp, sin, cos, tan
__all__ = ['UserPotential']
"""
pars:
- pars[0] : G (Gravitational constant)
- pars[1:] : Potential parameters
"""
cdef double user_value(double t, double *pars, double *q, int n_dim) nogil:
return 0
cdef double user_density(double t, double *pars, double *q, int n_dim) nogil:
return 0
cdef void user_gradient(double t, double *pars, double *q, int n_dim, double *grad) nogil:
cdef int i
for i in range(n_dim):
grad[i] = 0
cdef void user_hessian(double t, double *pars, double *q, int n_dim, double *hess) nogil:
cdef int i
for i in range(n_dim*n_dim):
hess[i] = 0
cdef class UserWrapper(CPotentialWrapper):
def __init__(self, G, parameters, q0):
ndim = 3
self.init([G] + list(parameters), np.ascontiguousarray(q0), n_dim=ndim)
self.cpotential.value[0] = <energyfunc>(user_value)
self.cpotential.density[0] = <densityfunc>(user_density)
self.cpotential.gradient[0] = <gradientfunc>(user_gradient)
self.cpotential.hessian[0] = <hessianfunc>(user_hessian)
class UserPotential(CPotentialBase):
def __init__(self, m, units=None, origin=None):
param, param_types = OrderedDict(), OrderedDict()
#define parameters and their types
param['m'] = m
param_types['m'] = 'mass'
ndim = 3
super(UserPotential, self).__init__(parameters=param,
parameter_physical_types=param_types,
ndim=ndim,
units=units,
origin=origin,
#Wrapper=UserWrapper,
)
'''
def compile_potential(code, force=False):
import os, sys, gala
from handy.cython import cythonmagic
gala_path = os.path.dirname(gala.__file__)
args = dict(
include_dirs = [os.path.join(gala_path, 'potential')],
sources = [os.path.join(gala_path, 'potential/potential/src/cpotential.c')],
extra_compile_args = ['--std=gnu99'],
cimport_dirs = sys.path,
numpy = True,
)
return cythonmagic(code, force=force, **args)
usermod = compile_potential(code)
pot = usermod.UserPotential(1)
pot.energy([1, 1, 1]), pot.density([1, 1, 1])
code = r'''
from __future__ import division, print_function
import numpy as np
from collections import OrderedDict
from gala.potential.potential.cpotential import CPotentialBase
from gala.potential.potential.cpotential cimport CPotentialWrapper
from gala.potential.potential.cpotential cimport densityfunc, energyfunc, gradientfunc, hessianfunc
from libc.math cimport sqrt, pow, fabs, exp, log, log10, sin, cos, tan, nan
__all__ = ['UserPotential']
"""
pars:
- pars[0] : G (Gravitational constant)
- pars[1:] : Potential parameters
"""
cdef double user_value(double t, double *pars, double *q, int n_dim) nogil:
cdef:
double *omega=pars+1
double res=0.
int i
for i in range(n_dim):
res += 0.5 * omega[i] * q[i] * q[i]
return res
cdef double user_density(double t, double *pars, double *q, int n_dim) nogil:
return 0.
cdef void user_gradient(double t, double *pars, double *q, int n_dim, double *grad) nogil:
cdef:
double *omega=pars+1
int i
for i in range(n_dim):
grad[i] = omega[i] * q[i]
cdef void user_hessian(double t, double *pars, double *q, int n_dim, double *hess) nogil:
cdef:
double *omega=pars+1
int i
for i in range(n_dim):
hess[i * n_dim + i] = omega[i]
cdef class UserWrapper(CPotentialWrapper):
def __init__(self, G, parameters, q0):
ndim = len(parameters)
self.init([G] + list(parameters), np.ascontiguousarray(q0), ndim)
self.cpotential.value[0] = <energyfunc>(user_value)
self.cpotential.density[0] = <densityfunc>(user_density)
self.cpotential.gradient[0] = <gradientfunc>(user_gradient)
self.cpotential.hessian[0] = <hessianfunc>(user_hessian)
class UserPotential(CPotentialBase):
def __init__(self, omega, units=None, origin=None):
param, param_types = OrderedDict(), None
ndim = len(omega)
for i in range(ndim):
param['omega%d' % i] = omega[i]
super(UserPotential, self).__init__(parameters=param,
parameter_physical_types=param_types,
ndim=ndim,
units=units,
origin=origin,
#Wrapper=UserWrapper,
)
#self.c_instance.cpotential.n_dim = ndim
'''
def compile_potential(code, force=False):
import os, sys, gala
from handy.cython import cythonmagic
gala_path = os.path.dirname(gala.__file__)
args = dict(
include_dirs = [os.path.join(gala_path, 'potential')],
sources = [os.path.join(gala_path, 'potential/potential/src/cpotential.c')],
extra_compile_args = ['--std=gnu99'],
cimport_dirs = sys.path,
numpy = True,
)
return cythonmagic(code, force=force, **args)
usermod = compile_potential(code, False)
OscillatorPotential = usermod.UserPotential