<jupyter_start><jupyter_text>Barycenter of time-series<jupyter_code>def l2_distance(a, b):
    return jnp.mean(jnp.sum(jnp.square(a - b), axis=-1))

def barycenter(timeseries, metric=l2_distance):
    # Function to calculate the loss
    @jax.jit
    def loss_fn(curve):
        return sum(metric(curve, ts) for ts in timeseries)

    init = jnp.zeros(timeseries[1].shape)
    result = minimize(loss_fn, init, method='bfgs').x
    return result
t  = jnp.linspace(0, 2*jnp.pi, 32)
timeseries_1 = jnp.sin(t)
timeseries_2 = jnp.cos(t)
series = [timeseries_1, timeseries_2]

print('Calculating L2 Barycenter')
l2_barycenter   = barycenter(series, metric=l2_distance)
print('Calculating SDTW Barycenter')
sdtw_barycenter = barycenter(series, metric=SoftDTW(gamma=0.1))
fig, axs = plt.subplots(1, 2, figsize=(15, 7), sharey=True)

for ax in axs:
    ax.plot(t, timeseries_1, label='timeseries_1')
    ax.plot(t, timeseries_2, label='timeseries_2')

axs[0].plot(t, l2_barycenter, label='l2_barycenter')
axs[0].set_title('L2 Barycenter')
axs[1].plot(t, sdtw_barycenter, label='SoftDTW barycenter (γ=0.1)')
axs[1].set_title('Soft-DTW Barycenter (γ=0.1)') 

plt.legend()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1495798319327731, 'fraction_numerical': 0.031092436974789917, 'mean_word_length': 5.616666666666666, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '232717', 'license': 'MIT', 'n_tokens_mistral': 507, 'n_tokens_neox': 489, 'n_words': 89, 'path': 'Barycenters.ipynb', 'repo_name': 'khdlr/softdtw_jax', 'size': '1190'}