forked from sympy/sympy
-
Notifications
You must be signed in to change notification settings - Fork 1
/
autowrap_ufuncify.py
executable file
·82 lines (59 loc) · 2.35 KB
/
autowrap_ufuncify.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
#!/usr/bin/env python
"""
Setup ufuncs for the legendre polynomials
-----------------------------------------
This example demonstrates how you can use the ufuncify utility in SymPy
to create fast, customized universal functions for use with numpy
arrays. An autowrapped sympy expression can be significantly faster than
what you would get by applying a sequence of the ufuncs shipped with
numpy. [0]
You need to have numpy installed to run this example, as well as a
working fortran compiler.
[0]:
http://ojensen.wordpress.com/2010/08/10/fast-ufunc-ish-hydrogen-solutions/
"""
import sys
from sympy.external import import_module
np = import_module('numpy')
if not np:
sys.exit("Cannot import numpy. Exiting.")
import sympy.mpmath as mpmath
from sympy.utilities.autowrap import ufuncify
from sympy.utilities.lambdify import implemented_function
from sympy import symbols, legendre, Plot, pprint
def main():
print(__doc__)
x = symbols('x')
# a numpy array we can apply the ufuncs to
grid = np.linspace(-1, 1, 1000)
# set mpmath precision to 20 significant numbers for verification
mpmath.mp.dps = 20
print("Compiling legendre ufuncs and checking results:")
# Let's also plot the ufunc's we generate
plot1 = Plot(visible=False)
for n in range(6):
# Setup the SymPy expression to ufuncify
expr = legendre(n, x)
print("The polynomial of degree %i is" % n)
pprint(expr)
# This is where the magic happens:
binary_poly = ufuncify(x, expr)
# It's now ready for use with numpy arrays
polyvector = binary_poly(grid)
# let's check the values against mpmath's legendre function
maxdiff = 0
for j in range(len(grid)):
precise_val = mpmath.legendre(n, grid[j])
diff = abs(polyvector[j] - precise_val)
if diff > maxdiff:
maxdiff = diff
print("The largest error in applied ufunc was %e" % maxdiff)
assert maxdiff < 1e-14
# We can also attach the autowrapped legendre polynomial to a sympy
# function and plot values as they are calculated by the binary function
g = implemented_function('g', binary_poly)
plot1[n] = g(x), [200]
print("Here's a plot with values calculated by the wrapped binary functions")
plot1.show()
if __name__ == '__main__':
main()