code: def main(_):
      # Get a state function and jit it.
      # We could switch to any other Hamiltonian and any other integrator here.
      # e.g. the non-symplectic `euler_integrator`.
      step_fn = functools.partial(
          single_integration_step,
          hamiltonian_from_graph_fn=hookes_hamiltonian_from_graph_fn,
          integrator_fn=verlet_integrator,
      )
      step_fn = jax.jit(step_fn)
  
      # Get a graph with the initial state.
      num_particles = 10
      graph = build_hookes_particle_state_graph(num_particles)
  
      # Iterate for multiple timesteps.
      num_steps = 200
      time_step = 0.002
      positions_sequence = []
      total_energies = []
      steps = []
      for step_i in range(num_steps):
          energy, graph = step_fn(graph, time_step)
          total_energies.append(energy)
          positions_sequence.append(graph.nodes["position"])
          steps.append(step_i + 1)
  
      # Plot results (positions and energy as a function of time).
      unused_fig, axes = plt.subplots(1, 2, figsize=(15, 5))
      positions_sequence_array = np.stack(positions_sequence, axis=0)
      axes[0].plot(
          positions_sequence_array[..., 0],
          positions_sequence_array[..., -1],
      )
      axes[0].set_xlabel("Particle position x")
      axes[0].set_ylabel("Particle position y")
  
      axes[1].plot(steps, total_energies)
      axes[1].set_ylim(0, max(total_energies) * 1.2)
      axes[1].set_xlabel("Simulation step")
      axes[1].set_ylabel("Total energy")
      plt.show()
  
imports: import functools
  from typing import Tuple, Callable
  from absl import app
  from frozendict import frozendict
  import jax
  import jax.numpy
  import jraph
  import matplotlib.pyplot
  import numpy
label_desc: Use image node_id 14 for calling a global function with example usage: main(_) without return types