How to use symengine - 10 common examples

from collections import OrderedDict, Counter, defaultdict
from itertools import combinations_with_replacement as cwr
import numpy as np
import pandas as pd
from numexpr import evaluate
try:
    from symengine import var, exp, cos, sin, Integer, Float
except ImportError:
    from sympy import symbols as var
    from sympy import exp, cos, sin, Integer, Float
from exa import Series
from exatomic.algorithms.overlap import _cartesian_shell_pairs, _iter_atom_shells
from exatomic.algorithms.numerical import fac, _tri_indices, _triangle, _enum_spherical


_x, _y, _z = var("_x _y _z")
_r = (_x ** 2 + _y ** 2 + _z ** 2) ** 0.5

lorder = ['s', 'p', 'd', 'f', 'g',
          'h', 'i', 'k', 'l', 'm']
lmap = OrderedDict()
rlmap = OrderedDict()
spher_ml_count = OrderedDict()
cart_ml_count = OrderedDict()
spher_lml_count = OrderedDict()
cart_lml_count = OrderedDict()
enum_cartesian = OrderedDict()
for i, L in enumerate(lorder):
    lmap[L] = i
    rlmap[i] = L
    spher_ml_count[L] = 2 * i + 1
    cart_ml_count[L] = (i + 1) * (i + 2) // 2

if time == sympy_t:
		return sympy_current_y(index)
	else:
		return sympy_past_y(time, index, sympy_anchors(time))
sympy_current_y = sympy.Function("current_y",real=True)
sympy_past_y = sympy.Function("past_y",real=True)
sympy_anchors = sympy.Function("anchors",real=True)

symengine_t = symengine.Symbol("t",real=True)
def symengine_y(index,time=symengine_t):
	if time == symengine_t:
		return symengine_current_y(index)
	else:
		return symengine_past_y(time, index, symengine_anchors(time))
symengine_current_y = symengine.Function("current_y",real=True)
symengine_past_y = symengine.Function("past_y",real=True)
symengine_anchors = symengine.Function("anchors",real=True)


symengine_manually = [
		symengine_t,
		symengine_y,
		symengine.cos,
	]

sympy_manually = [
		sympy_t,
		sympy_y,
		sympy.cos,
	]

jitcdde_provisions = [

# -*- coding: utf-8 -*-

"""
Integration test of jitcode_restricted_lyap and jitcode_transversal_lyap by comparing them to each other.
"""

from itertools import combinations
from jitcode import jitcode_restricted_lyap, y, jitcode_transversal_lyap
import numpy as np
from scipy.stats import sem
from symengine import Symbol

a = -0.025794
b =  0.01
c =  0.02
k = Symbol("k")

scenarios = [
	{
	"f":
		[
			y(0)*(a-y(0))*(y(0)-1.0) - y(3) + k*(y(1)-y(0)),
			y(1)*(a-y(1))*(y(1)-1.0) - y(4) + k*(y(2)-y(1)),
			y(2)*(a-y(2))*(y(2)-1.0) - y(5) + k*(y(0)-y(2)),
			b*y(0) - c*y(3),
			b*y(1) - c*y(4),
			b*y(2) - c*y(5),
		],
	"vectors":
		[
			[1.,1.,1.,0.,0.,0.],
			[0.,0.,0.,1.,1.,1.]

"""
Tests whether things works independent of where symbols are imported from.
"""

import jitcode
import jitcode.sympy_symbols
import sympy
import symengine

symengine_manually = [
		symengine.Symbol("t",real=True),
		symengine.Function("y",real=True),
		symengine.cos,
	]

sympy_manually = [
		sympy.Symbol("t",real=True),
		sympy.Function("y",real=True),
		sympy.cos,
	]

jitcode_provisions = [
		jitcode.t,
		jitcode.y,
		symengine.cos,
	]

T = 5

sech = lambda x: 2/(exp(x)+exp(-x))
eps = 1e-5
abs = lambda x: sqrt(x**2+eps**2)

ε = [ 0.03966, 0.03184, 0.02847 ]
ν = [ 1, 2.033, 3.066 ]
μ = [ 0.16115668456085775, 0.14093420256851111, 0.11465065353644151 ]
ybar_0 = 0
τ = 1.7735
ζ = [ 0.017940997406325931, 0.015689701773967984, 0.012763648066925721 ]

anchors_past = Symbol("anchors_past")
difference = Symbol("difference")
factor_μ = Symbol("factor_μ")
factor_ζ = Symbol("factor_ζ")

ydot = [ y(i) for i in range(3,6) ]
y_tot      = sum( current_y(i)                for i in range(  3) )
ydot_tot   = sum( current_y(i)                for i in range(3,6) )
y_past     = sum( past_y (t-τ,i,anchors_past) for i in range(  3) )
ydot_past  = sum( past_y (t-τ,i,anchors_past) for i in range(3,6) )
yddot_past = sum( past_dy(t-τ,i,anchors_past) for i in range(3,6) )

helpers = {
		( anchors_past, anchors(t-τ) ),
		( difference, ybar_0-y_past ),
		( factor_μ, sech(difference)**2 * (yddot_past + 2*ydot_past**2*tanh(difference)) ),
		( factor_ζ, 2*abs(y_tot)*ydot_tot ),
	}

class TestGeneratorChunking(TestGenerator):
	def generator(self):
		self.DDE.compile_C(chunk_size=1, extra_compile_args=compile_args)


f_dict = { y(i):entry for i,entry in enumerate(f) }

class TestDictionary(TestIntegration):
	@classmethod
	def setUpClass(self):
		self.DDE = jitcdde(f_dict)
		self.DDE.set_integration_parameters(**test_parameters)


delayed_y, y3m10, coupling_term = symengine.symbols("delayed_y y3m10 coupling_term")
f_alt_helpers = [
		(delayed_y, y(0,t-delay)),
		(coupling_term, 0.25 * (delayed_y - y(3))),
		(y3m10, y(3)-10)
	]

f_alt = [
		omega[0] * (-y(1) - y(2)),
		omega[0] * (y(0) + 0.165 * y(1)),
		omega[0] * (0.2 + y(2) * (y(0) - 10.0)),
		omega[1] * (-y(4) - y(5)) + coupling_term,
		omega[1] * (y3m10 + 10 + 0.165 * y(4)),
		omega[1] * (0.2 + y(5) * y3m10)
	]

class TestHelpers(TestIntegration):

def test_check_undefined_variable(self):
		x = symengine.symbols("x")
		DDE = jitcdde([x])
		with self.assertRaises(ValueError):
			DDE.check()

# as dictionary
f_dict = { y(i):entry for i,entry in reversed(list(enumerate(f))) }
with_dictionary = {"f_sym": f_dict}

# with generator

def f_generator():
	for entry in f:
		yield entry

with_generator = { "f_sym":f_generator, "n":n }

# with helpers

f1, f2, f3, f4 = symbols("f1, f2, f3, f4")
coupling, first_y, first_y_sq = symbols("coupling, first_y, first_y_sq")
a_alt, b1_alt, b2_alt, c_alt, k_alt = symbols("a_alt, b1_alt, b2_alt, c_alt, k_alt")
f_alt = [ f1, f2, f3, f4 ]

f_alt_helpers = [
	( a_alt , a  ),
	( b1_alt, b1 ),
	( b2_alt, b2 ),
	( c_alt,  c  ),
	( k_alt,  k  ),
	( first_y, y(0) ),
	( first_y_sq, first_y**2 ),
	( coupling, k_alt*(y(2)-first_y) ),
	( f1, a_alt*first_y_sq - a_alt*first_y + coupling - first_y**3 + first_y_sq - y(1) ),
	( f2, b1_alt*first_y - c_alt*y(1)),
	( f3, y(2) * ( a_alt-y(2) ) * ( y(2)-1.0 ) - y(3) - coupling ),

( k_alt,  k  ),
	( first_y, y(0) ),
	( first_y_sq, first_y**2 ),
	( coupling, k_alt*(y(2)-first_y) ),
	( f1, a_alt*first_y_sq - a_alt*first_y + coupling - first_y**3 + first_y_sq - y(1) ),
	( f2, b1_alt*first_y - c_alt*y(1)),
	( f3, y(2) * ( a_alt-y(2) ) * ( y(2)-1.0 ) - y(3) - coupling ),
	( f4, b2_alt*y(2) - c_alt*y(3) ),
]
shuffle(f_alt_helpers)

with_helpers = {"f_sym": f_alt, "helpers": f_alt_helpers}

# with parameters

a_par, c_par, k_par = symbols("a_par c_par k_par")
f_params_helpers = [ ( coupling, k_par*(y(2)-y(0)) ) ]
f_params = [
	y(0) * ( a_par-y(0) ) * ( y(0)-1.0 ) - y(1) + coupling,
	b1*y(0) - c_par*y(1),
	y(2) * ( a_par-y(2) ) * ( y(2)-1.0 ) - y(3) - coupling,
	b2*y(2) - c_par*y(3)
	]
params_args = (a,c,k)
n_params = 3

with_params = {
		"f_sym": f_params,
		"helpers": f_params_helpers,
		"control_pars": [a_par, c_par, k_par]
	}

f_callback = [
		omega[0] * (-y(1) - y(2)),
		omega[0] * (y(0) + 0.165 * y(1)),
		omega[0] * (0.2 + y(2) * call_minus_ten(y(0))),
		call_fourth_component(y(0,t-delay)),
		omega[1] * (y(3) + 0.165 * y(4)),
		omega[1] * (0.2 + y(5) * call_minus_ten(y(3))),
	]

class TestCallback(TestIntegration):
	@classmethod
	def setUpClass(self):
		self.DDE = jitcdde(f_callback,callback_functions=callbacks)
		self.DDE.set_integration_parameters(**test_parameters)

a,b,c,k,tau = symengine.symbols("a b c k tau")
parameters = [0.165, 0.2, 10.0, 0.25, 4.5]
f_params = [
		omega[0] * (-y(1) - y(2)),
		omega[0] * (y(0) + a * y(1)),
		omega[0] * (b + y(2) * (y(0) - c)),
		omega[1] * (-y(4) - y(5)) + k * (y(0,t-tau) - y(3)),
		omega[1] * (y(3) + a * y(4)),
		omega[1] * (b + y(5) * (y(3) - c))
	]

class TestParameters(TestIntegration):
	@classmethod
	def setUpClass(self):
		self.DDE = jitcdde(f_params, control_pars=[a,b,c,k,tau],max_delay=parameters[-1])
		self.DDE.set_integration_parameters(**test_parameters)