import apebench
import jax.numpy as jnp

for scenario in [
    "diff_disp",
    "diff_burgers",
    "diff_kdv",
    "diff_ks_cons",
    "diff_ks",
]:
    scene = apebench.scenarios.scenario_dict[scenario]()

    ref_trj = scene.get_ref_sample_data()

    jnp.save(f"ref_sample_rollouts/{scene.get_scenario_name()}.npy", ref_trj)