proteneer

# processed from https://raw.githubusercontent.com/proteneer/timemachine/8ff69ee3c7bf248a7bdc5caf9fcae2e37fdb86ac/jax/harmonic_dij.py

import functools
import time
from tqdm import tqdm
import numpy as onp
from jax.config import config; config.update("jax_enable_x64", True)
import jax
import jax.numpy as np