<jupyter_start><jupyter_text>Schematic figureGenerate a schematic figure to demonstrate a few ways how discounting may relate to BMI.<jupyter_code>import numpy as np
import matplotlib.pyplot as plt
%config InlineBackend.figure_format = 'retina'

from matplotlib import rcParams
rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['Arial']
rcParams['font.size'] = 16
#plt.rcParams.update({'font.size': 16})

np.random.seed(1234)
fig, ax = plt.subplots(1, 3, figsize=(8, 3), sharey=True)

# fig, (ax0, ax1, ax2, cax) = plt.subplots(ncols=4, figsize=(16, 4), sharey=True, 
#                                          gridspec_kw={"width_ratios":[1,1,1, 0.05]})

bmi_lim = [15, 50]
text_ypos = 45
cm = plt.cm.get_cmap('viridis')
N = 500
marker_size = 10**2

# correlation ----------------------------------------------------------------
logk_mu, logk_sd = -6, 1.5
bmi_mu, bmi_std = 25, 3
ρ = 0.4
diag = ρ * logk_sd * bmi_std

mu = [logk_mu, bmi_mu]
cov = [[logk_sd**2, diag], [diag, bmi_std**2]]

logk, bmi = np.random.multivariate_normal(mu, cov, N).T
#ax[0].scatter(logk, bmi)
ax[0].scatter(logk, bmi, c=logk, cmap=cm, s=marker_size)
# fig.colorbar(sc, ax=ax[0])
ax[0].set(xlabel='$\log(k)$', ylabel='BMI [kg/m$^2$]', title='Simple correlation', 
          ylim=bmi_lim, 
          xlim=[-11, -1],
          xticks=np.linspace(-10,-2, 5))
ax[0].text(-10.25, text_ypos, 'BMI ~ log(k)', fontsize=9, usetex=False, family="monospace")

# age ----------------------------------------------------------------
age = np.linspace(18, 85, N)
logk = np.random.normal(logk_mu, logk_sd, N)
β0, β1, β3 = 25, 0.2, 1
bmi = β0 + β1 * age + β3 * logk

ax[1].scatter(age, bmi, c=logk, cmap=cm, s=marker_size)
ax[1].set(xlabel='Age [years]', title='Main effects model', ylim=bmi_lim)
# fig.colorbar(sc, ax=ax[1])
ax[1].text(20, text_ypos, r'BMI ~ age + log(k)', fontsize=9, usetex=False, family="monospace")


# moderation ----------------------------------------------------------------
age = np.linspace(18, 85, N)
#logk = np.random.normal(logk_mu, logk_sd, N)
β0, β1, β2 = 20, 0.3, 0.03
bmi = β0 + (β1 * age) + (β2 * age*logk)

sc = ax[2].scatter(age, bmi, c=logk, cmap=cm, s=marker_size)
# fig.colorbar(sc, ax=ax[2])

ax[2].set(xlabel='Age [years]', title='Rate of gain model', ylim=bmi_lim)
ax[2].text(20, text_ypos, r'BMI ~ age + log(k):age', fontsize=9, usetex=False, family="monospace")


# add in colorbar as extra floating axis --------------------

im = plt.gca().get_children()[0]
k=0.02
# left, bottom, width, height
cax = fig.add_axes([0.92,0.15,0.015,0.7]) 
fig.colorbar(im, cax=cax, label='$\log(k)$')



# plt.show()

# figure level formatting ----------------------------------------------------------------
plt.subplots_adjust(wspace=0.1)

# save

plt.savefig('schematic.pdf', bbox_inches='tight')
plt.savefig('schematic.png', bbox_inches='tight', dpi=300)<jupyter_output><empty_output><jupyter_text>Test<jupyter_code>fig, (ax0, ax1, ax2, cax) = plt.subplots(ncols=4, figsize=(16, 4), sharey=True, 
                                         gridspec_kw={"width_ratios":[1,1,1, 0.05]})<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.26191237607930923, 'fraction_numerical': 0.0463703229932843, 'mean_word_length': 2.91, 'pattern_counts': {'":': 2, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '189306', 'license': 'MIT', 'n_tokens_mistral': 1350, 'n_tokens_neox': 1247, 'n_words': 248, 'path': 'figs/make_schematic_figure.ipynb', 'repo_name': 'drbenvincent/Veillard-Vincent', 'size': '3127'}