<jupyter_start><jupyter_text>Parallel simulation of a bubble raft**Important**: This code also works on multi-host TPU setup without any changes !! The key thing to do with a multi-host TPU setup is to ssh the file and run it on all the host at the same time. In order to do that please refer to this [notebook](https://github.com/probml/probml-notebooks/blob/main/notebooks/tpu_colab_tutorial.ipynb). This notebook is based on the first example from the [JAX MD cookbook](https://colab.research.google.com/github/google/jax-md/blob/master/notebooks/jax_md_cookbook.ipynb) i.e the simulating bubble raft example ![alt text](https://upload.wikimedia.org/wikipedia/commons/8/89/Bubblerraft2.jpg)   Installation<jupyter_code>!pip install -q git+https://www.github.com/google/jax-md
import jax

try:
    import jax.tools.colab_tpu

    jax.tools.colab_tpu.setup_tpu()
except KeyError:
    import os
jax.devices()
import os
import jax.numpy as np

from jax import jit
from jax import vmap, pmap
import jax.numpy as jnp
from jax import random
from jax import lax


from jax_md import space
from jax_md import simulate
from jax_md import energy

os.environ["XLA_USE_32BIT_LONG"] = "1"
jax.local_device_count()<jupyter_output><empty_output><jupyter_text>Hyperparameters<jupyter_code>f32 = np.float32
ensemble_size = 1000
nlocal = 8
N = 32
dt = 1e-1
simulation_steps = np.arange(1000)
key = random.PRNGKey(0)<jupyter_output><empty_output><jupyter_text>Defining the box and the energy function<jupyter_code>def box_size_at_number_density(particle_count, number_density):
    return f32((particle_count / number_density) ** 0.5)


box_size = box_size_at_number_density(particle_count=N, number_density=1)
displacement, shift = space.periodic(box_size)
energy_fun = energy.soft_sphere_pair(displacement)<jupyter_output><empty_output><jupyter_text>Defining the solution<jupyter_code>def simulation(key, temperature):
    pos_key, sim_key = random.split(key)

    R = random.uniform(pos_key, (N, 2), maxval=box_size)

    init_fn, apply_fn = simulate.brownian(energy_fun, shift, dt, temperature)
    state = init_fn(sim_key, R)

    do_step = lambda state, t: (apply_fn(state, t=t), t)
    state, _ = lax.scan(do_step, state, simulation_steps)
    return state.position<jupyter_output><empty_output><jupyter_text>Parallelsing the simulation<jupyter_code>vectorized_simulation = vmap(simulation, in_axes=(0, None))
parallel_vectorized_simulation = pmap(vectorized_simulation, in_axes=(0, None))
vectorized_energy = vmap(energy_fun)
parallel_vectorized_energy = pmap(vectorized_energy)<jupyter_output><empty_output><jupyter_text>Getting the random keys<jupyter_code>simulation_keys_lst = []
for i in range(nlocal):
    key, *simulation_keys = random.split(key, ensemble_size + 1)
    simulation_keys = jnp.stack(simulation_keys)
    simulation_keys_lst.append(simulation_keys)
simulation_keys = jnp.stack(simulation_keys_lst)<jupyter_output><empty_output><jupyter_text>Running the simulation<jupyter_code>bubble_positions = parallel_vectorized_simulation(simulation_keys, 1e-5)
bubble_energies = parallel_vectorized_energy(bubble_positions)
import numpy as onp
import matplotlib.pyplot as plt


def format_plot(x, y):
    plt.xlabel(x, fontsize=20)
    plt.ylabel(y, fontsize=20)


bubble_energies = jax.pmap(lambda x: jax.lax.all_gather(x, "i"), axis_name="i")(bubble_energies)[0]
counts, bins = onp.histogram(bubble_energies, bins=50)
plt.plot(bins[:-1] * 10**5, counts, "o")
format_plot("$E\\times 10 ^{-5}$", "$P(E)$")
plt.savefig("plot.png")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10814606741573034, 'fraction_numerical': 0.014044943820224719, 'mean_word_length': 6.496842105263158, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 29, 'https://': 4, 'lorem ipsum': 0, 'www.': 1, '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': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '581172', 'license': 'MIT', 'n_tokens_mistral': 1339, 'n_tokens_neox': 1275, 'n_words': 268, 'path': 'notebooks/misc/jax_md_parallel_tpu_example.ipynb', 'repo_name': 'karm-patel/pyprobml', 'size': '3560'}