Coverage for gamdpy/interactions/potential_functions/make_potential_function_from_sympy.py: 33%

1import numpy as np 

2import numba 

3import math 

4from numba import cuda 

5 

6def make_potential_function_from_sympy(ufunc, param_names) -> callable: 

7 """ Make a potential function from a sympy expression 

8 

9 Known to result in slow code. Use for testing and prototyping. 

10 

11 Parameters 

12 ---------- 

13 

14 ufunc : sympy expression 

15 The potential energy expression in Sympy's symbolic form 

16 It has to be a radial function and the pair distance symbol shuould be r 

17 

18 param_names : list 

19 List of parameters 

20 

21 Returns 

22 ------- 

23 

24 potential_function : callable 

25 A function that calculates the potential energy, force multiplier, and second derivative of the potential energy 

26 u, s, umm = potential_function(dist, params) 

27 

28 """ 

29 import sympy 

30 from sympy.abc import r 

31 from sympy.utilities.lambdify import lambdify 

32 

33 # Define functions in Sympy 

34 dufunc = sympy.simplify(sympy.diff(ufunc, r)) 

35 sfunc = sympy.simplify(-sympy.diff(ufunc, r) / r) 

36 ummfunc = sympy.simplify(sympy.diff(dufunc, r)) 

37 

38 # LC: There is a problem with lambdify and lists: 

39 # https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-reflection-for-list-and-set-types 

40 # Making param_names a tuple seems to fix the warning 

41 u_lam = numba.njit(lambdify([r, param_names], ufunc, 'numpy')) # Jitted python functions 

42 s_lam = numba.njit(lambdify([r, param_names], sfunc, 'numpy')) 

43 umm_lam = numba.njit(lambdify([r, param_names], ummfunc, 'numpy')) 

44 

45 #@numba.njit 

46 def potential_function(r, params): # pragma: no cover 

47 u = np.float32(u_lam(r, params)) 

48 s = np.float32(s_lam(r, params)) 

49 umm = np.float32(umm_lam(r, params)) 

50 return u, s, umm 

51 

52 return potential_function 

53