import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from plotstyle import *

fig, (ax1, ax2) = plt.subplots(1, 2)
fig.subplots_adjust(wspace=0.35, hspace=0.5,
                    **adjust_fs(fig, left=4, top=0.5, bottom=2.7))

n = 200
rng = np.random.RandomState(3981)
x = rng.randn(n)
y = rng.randn(n)

z = x*x+0.2*y
r =np.corrcoef(x,z)[0,1]
ax1.text(0, 4.0, 'r=%.1f' % r, ha='center')
ax1.text(0, 5.6, r'$y = x^2+\xi/5$', ha='center')
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_xlim(-3.0, 3.0)
ax1.set_ylim(-0.5, 6.0)
ax1.plot(x, z, **psAm)

z = 0.5*x*y
r =np.corrcoef(x,z)[0,1]
ax2.text(0, 1.5, 'r=%.1f' % r, ha='center')
ax2.text(0, 2.8, r'$y = x \cdot \xi/2$', ha='center')
ax2.set_xlabel('x')
ax2.set_ylabel('y')
ax2.set_xlim(-3.0, 3.0)
ax2.set_ylim(-3.0, 3.0)
ax2.plot(x, z, **psAm)

plt.savefig('nonlincorrelation.pdf')