#!/usr/bin/env python
# coding: utf-8

# # Chebfun in Python

# This is a demo following [the original paper of Battles & Trefethen][1].
# 
# [1]: http://people.maths.ox.ac.uk/trefethen/publication/PDF/2004_107.pdf

# ## Initialisation

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.figure_formats = {'svg', 'png'}")
import numpy as np
import matplotlib.pyplot as plt


# It mostly suffices to use the class `Chebfun`, but we also import the whole `pychebfun` module for convenience.

# In[ ]:


from pychebfun import Chebfun


# In[ ]:


import pychebfun


# In[ ]:


x = Chebfun.identity()


# This is only for this IPython notebook

# In[ ]:


from IPython.core.display import HTML


# We define a convenience plotting function

# In[ ]:


def cplot(fun):
    pychebfun.plot(fun, with_interpolation_points=True)


# ## Chebfuns and barycentric interpolation

# ### Creation of a chebfun

# A typical way of initialising a Chebfun is to use the function `chebfun`. It will work for any function.

# In[ ]:


cube = pychebfun.chebfun(lambda x:x**3)


# A more pythonic way to achieve the same result would be using `Chebfun.from_function`, it will also work for any function.

# In[ ]:


cube = Chebfun.from_function(lambda x:x**3)


# Another possibility is to use the variable `x` defined by `x = Chebfun.identity()`:
# 

# In[ ]:


x = Chebfun.identity()


# Note however, that this will only work if the function of x has been registered in chebfun. Here it works because chebfun knows how to compute arbitray powers:

# In[ ]:


cube = x**3


# Other examples could include:

# In[ ]:


print(np.cos(x))
print(np.tan(x))
print(np.exp(np.sin(x)))


# ### Display

# The _size_ of `f` is the number of interpolation point.

# In[ ]:


cube.size()


# Displaying a Chebfun gives the interpolation values at the Chebyshev interpolation points:

# In[ ]:


print(repr(cube))


# In[ ]:


Chebfun.from_function(lambda x: x**7 - 3*x + 5).size()


# In[ ]:


f_sin5 = Chebfun.from_function(lambda x: np.sin(5*np.pi*x))
f_sin5.size()


# ### Convergence

# For some functions, convergence is not achievable, and one must set a limit to the dichotomy algorithm.

# In[ ]:


f_abs = Chebfun.from_function(abs, N=1000)


# ### Evaluations

# `Chebfun` objects behave as ordinary function, which can be evaluated at any point.

# In[ ]:


cube(5)


# In[ ]:


cube(-.5)


# It is also possible (and more efficient) to evaluate a chebfun at several point at once.
# For instance, we evaluate the Chebfun for $\sin(5x)$ at several points at once:

# In[ ]:


f_sin5(np.linspace(0, .2, 5))


# ### Operations

# One can for instance add two chebfuns:

# In[ ]:


chebsum = cube+f_sin5
print(chebsum.size())
cplot(chebsum)


# Or multiply them:

# In[ ]:


chebmult = cube*f_sin5
print(chebmult.size())
cplot(chebmult)


# In[ ]:


f_sin = Chebfun.from_function(lambda x:np.sin(x))
print(f_sin.size())
print((f_sin*f_sin).size())


# It is also possible to add and multiply chebfuns with scalars:

# In[ ]:


plt.subplot(121)
cplot(3*f_sin)
plt.subplot(122)
cplot(3+f_sin)


# In[ ]:


f_cubic = Chebfun.from_function(lambda x: x**3 - x)
cplot(f_cubic)


# In[ ]:


cplot(f_sin5)


# ### 2D Chebfun

# In the paper, they suggest to create two Chebfuns as follows:

# In[ ]:


np.sin(16*x), np.sin(18*x)


# It is certainly possible, but we can't create a 2D Chebfun from them, at the moment.
# Instead we have to do it like this:

# In[ ]:


def vector_func(x):
    return np.vstack([np.sin(16*x), np.sin(18*x)]).T
cheb_vec = Chebfun.from_function(vector_func)
cplot(cheb_vec)


# ## Elementary functions

# Again, we explore the two ways to create a chebfun.
# First, the general way, which works for any smooth functions:

# In[ ]:


f_expsin = Chebfun.from_function(lambda x: np.exp(np.sin(x)))


# Then, the “operator overloading” way.
# We take advantage of the fact that `exp` and `sin` are defined on chebfuns.
# Not all functions are, though.
# This is very similar to ufuncs in numpy.

# In[ ]:


g_expsin = np.exp(np.sin(x))


# Of course, both chebfuns are equivalent, as we demonstrate here.

# In[ ]:


print(f_expsin.size())
print(g_expsin.size())


# In[ ]:


plt.subplot(121)
cplot(f_expsin)
plt.subplot(122)
cplot(g_expsin)


# In[ ]:


(f_expsin - g_expsin).norm()


# ## Approximation Theory

# ### Gibbs phenomenon

# By limiting the accuracy of the chebfun, one observes the celebrated [Gibbs phenomenon](http://en.wikipedia.org/wiki/Gibbs_phenomenon).

# In[ ]:


f_sign = Chebfun.from_function(lambda x: np.sign(x), N=25)
cplot(f_sign)


# Pychebfun implements the method `compare` which plots a graph of the error.

# In[ ]:


def absfive(x):
    return np.abs(x)**5
#errors = array([(Chebfun(absfive, N) - exact).norm() for N in range(10, 20)])
#loglog(errors)
pychebfun.compare(Chebfun.from_function(absfive), absfive)


# ### Interpolation of random data

# A chebfun passing through random values at 50 Chebyshev points.
# The method to initialise a Chebfun from data _at Chebyshev points_ is `from_data`:

# In[ ]:


rng = np.random.default_rng(0)


# In[ ]:


rand_chebfun = Chebfun.from_data(rng.random(51))
cplot(rand_chebfun)
rand_chebfun.size()


# In[ ]:


f_walk = Chebfun.from_data(rng.normal(size=(100,2)))
cplot(f_walk)


# ### Extrapolation

# Clearly, as chebfuns are based on polynomial interpolation, there is no guarantee that extrapolation outside the interval $[-1,1]$ will give any sensible result at all.
# This is an illustration of that phenomenon.

# In[ ]:


xx = np.linspace(-1.4, 1.4, 100)
error = np.log10(np.abs(f_sin5(xx)-np.sin(5*np.pi*xx))+1e-16)
plt.subplot(211)
plt.plot(xx, f_sin5(xx), marker='.')
plt.ylabel('interpolant')
plt.subplot(212)
plt.plot(xx,error, marker='.')
plt.ylabel('log_10(error)')


# ## Chebyshev expansions and FFT

# The Chebyshev coefficients of a chebfun are quickly computed using a [fast Fourier transform](https://en.wikipedia.org/wiki/Fast_fourier_transform).
# The corresponding method is `coefficients`

# In[ ]:


Chebfun.from_function(np.exp(x)).coefficients()


# In[ ]:


Chebfun.from_coeff(np.array([3,2,1.]))


# One way to create the basis Chebyshev polynomial of order 10.

# In[ ]:


T_10 = Chebfun.from_coeff(np.array([0.]*10+[1]))
cplot(T_10)


# Incidentally, the same result is achieved using `Chebfun.basis`:

# In[ ]:


cplot(Chebfun.basis(10))


# As an illustration of how fast the fast Fourier transform is, we create a chebfun with 100000 points.

# In[ ]:


r = rng.normal(size=100000)
f_randn = Chebfun.from_coeff(r)


# In[ ]:


np.sqrt(np.sum(np.square(r)))


# In[ ]:


f_randn.sum()
#f_randn.norm() # doesn't work yet


# In[ ]:


T_20 = Chebfun.basis(20)
T_20.norm()


# ## Integration and Differentiation

# ### Clenshaw Curtis Quadrature

# The method `sum` computes the integral of the chebfun using the Chebyshev expansion and integral of the Chebyshev basis polynomial (it is known as the [Clenshaw–Curtis quadrature](http://en.wikipedia.org/wiki/Clenshaw%E2%80%93Curtis_quadrature))

# In[ ]:


k_odd = 5
k_even = 2
HTML(r'For odd $k$ we check that $\int T_k = %s$ ($k=%s$), otherwise, $\int T_k = \frac{2}{1-k^2} = %s \approx %s$ for $k = %s$.' % 
(Chebfun.basis(k_odd).sum(), 
 k_odd,
2./(1-k_even**2), 
Chebfun.basis(k_even).sum(), 
k_even))


# ### Examples

# Computing some integrals.

# In[ ]:


f1 = np.sin(np.pi*x)**2
print(f1.size(), f1.sum())
HTML(r'$\int \sin(\pi x)^2 dx = %s$' % f1.sum())


# In[ ]:


f2 = Chebfun.from_function(lambda x: 1/(5+3*np.cos(np.pi*x)))
print(f2.size(), f2.sum())


# In[ ]:


f3 = Chebfun.from_function(lambda x: np.abs(x)**9 * np.log(np.abs(x)+1e-100))
print(f3.size(), f3.sum())


# In[ ]:


HTML(r'Computing the norm of $x^2$ gives: %s $\approx \sqrt{\int_{-1}^1 x^4 dx} = \sqrt{\frac{2}{5}} = %s$' % ((x**2).norm(), np.sqrt(2./5)))


# The `dot` method computes the Hilbert scalar product, so `f.dot(g)` corresponds to
# \\[\int_{-1}^1 f(x) g(x) \mathrm{d} x\\]

# In[ ]:


x.dot(x)


# The `integrate` method computes a primitive $F$ of a chebfun $f$, which is zero at zero i.e.,
# \\[F(x) = \int_{0}^x f(y) \mathrm{d}y\\]

# In[ ]:


f = Chebfun.from_function(lambda t: 2/np.sqrt(np.pi) * np.exp(-t**2))
erf2_ = f.integrate()
erf2 = erf2_ - erf2_(0)
from scipy.special import erf
randx = rng.random()
print(erf2(randx) - erf(randx))


# This allows to define continuous versions of the `prod` and `cumprod` functions for chebfuns:

# In[ ]:


def prod(f):
    return np.exp(np.log(f).sum())
def cumprod(f):
    return np.exp(np.log(f).integrate())


# In[ ]:


prod(np.exp(x))


# In[ ]:


prod(np.exp(np.exp(x)))


# In[ ]:


cplot(cumprod(np.exp(x)))


# ### Differentiation

# One can also differentiate chebfuns to arbitrary orders with the method `differentiate`

# In[ ]:


f_sin5 = np.sin(5*x)
fd_sin5 = f_sin5.differentiate(4)
print(fd_sin5.norm()/f_sin5.norm())


# In[ ]:


f = np.sin(np.exp(x**2)).differentiate()
print(f(1))


# In[ ]:


g = Chebfun.from_function(lambda x: 1/(2 + x**2))
h = g.differentiate()
print(h(1))


# In[ ]:


#f.differentiate().integrate() == f - f(-1)
#f.integrate().differentiate() == f


# ## Operations based on rootfinding

# As chebfun are polynomial it is possible to find all the roots in the interval $[-1,1]$ using the method `roots`:

# In[ ]:


print((x - np.cos(x)).roots())


# In[ ]:


print((x - np.cos(4*x)).roots())


# Zeros of the Bessel function $J_0$ on the interval $[0,20]$:

# In[ ]:


import scipy.special
def J_0(x):
    return scipy.special.jn(0,x)
#f = chebfun(lambda x: J_0(10*(1+x)))
f = Chebfun.from_function(J_0, domain=[0,20])
cplot(f)
roots = f.roots()
print('roots:', roots)
plt.plot(roots, np.zeros_like(roots), color='red', marker='.', linestyle='', markersize=10)
plt.grid()


# ### Extrema

# The methods `min` and `max` are not implemented but it is not difficult to do so, by finding the roots of the derivative.
# This may come in a future version of `pychebfun`.

# In[ ]:


f = x - x**2
#print('min', f.min())
#print('max', f.max())
#print('argmax', f.argmax())
#print('norm inf', f.norm('inf'))
#print('norm 1', f.norm(1))


# Total variation

# In[ ]:


def total_variation(f):
    return f.differentiate().norm(1)
#total_variation(x)
#total_variation(sin(5*pi*x))


# Here is an example of computation of extrema of a function.

# In[ ]:


f = np.sin(6*x) + np.sin(30*np.exp(x))
r = f.roots()
fd = f.differentiate()
fdd = fd.differentiate()
e = fd.roots()
ma = e[fdd(e) <= 0]
mi = e[fdd(e) > 0]
def plot_all():
    pychebfun.plot(f, with_interpolation_points=False)
    plt.plot(r, np.zeros_like(r), linestyle='', marker='o', color='gray', markersize=8)
    plt.plot(ma, f(ma), linestyle='', marker='o', color='green', markersize=8)
    plt.plot(mi, f(mi), linestyle='', marker='o', color='red', markersize=8)
    plt.grid()
plot_all()


# ## Applications in numerical analysis

# ### Quadrature

# We apply quadrature to the function
# $$ f(x) = \tan(x+\frac{1}{4}) + \cos(10x^2 + \exp(\exp(x)))$$

# In[ ]:


def s(x):
    return np.tan(x+1./4) + np.cos(10*x**2 + np.exp(np.exp(x)))
tancos = Chebfun.from_function(s)
cplot(tancos)
plt.title(str(tancos))


# Using the built-in `quad` quadrature integrator of `scipy.integrate`:

# In[ ]:


import scipy.integrate
scipy.integrate.quad(s, -1, 1, epsabs=1e-14)[0]


# Using Chebyshev integration, we obtain exactly the same value:

# In[ ]:


tancos.sum()


# ### ODE Solving

# One solves the ODE
# \begin{align}
# u'(x) &= \mathrm{e}^{-2.75 x u(x)} \\
# u(-1)&= 0
# \end{align}

# The problem is rewritten first as an integral problem, namely
# $$u(x) = T(u)$$
# with
# $$T(u) := \int_{-1}^{x} \mathrm{e}^{-2.75 yu(y)} dy$$
# This suggest using the following [fixed-point iteration](https://en.wikipedia.org/wiki/Fixed-point_iteration):

# In[ ]:


uold = Chebfun(0.)
du = 1.
for i in range(100):
    integrand = np.exp(-2.75*x*uold)
    uint = integrand.integrate()
    u = uint - uint(-1)
    du = u-uold
    uold = u
    if du.norm() < 1e-13:
        break
else:
    print('no convergence')
print('Converged in {} steps'.format(i))
print('u(1) = {:.14f}'.format(float(u(1))))
pychebfun.plot(u)
plt.grid()
plt.title(str(u))


# ### Boundary value problems

# One solves the boundary problem 
# $$u''(x) = \mathrm{e}^{4x}$$
# with boundary conditions
# \\[u(-1) = 0 \qquad u(1) = 0 \\]

# In[ ]:


f = np.exp(4*x)
u_ = f.integrate().integrate()
u = u_ - (u_(1)*(x+1) + u_(-1)*(1-x))/2.
cplot(u)
plt.grid()