template for user defined cython potential of gala

custom_oscillator_potential.py

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

template_potential.py

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])