code: def build_hookes_particle_state_graph(
      num_particles: int,
  ) -> jraph.GraphsTuple:
      """Generates a graph representing a Hooke's system in a random state."""
  
      mass = np.random.uniform(0, 5, [num_particles])
      velocity = get_random_uniform_norm2d_vectors(
          0, 0.1, num_particles
      )
      position = get_random_uniform_norm2d_vectors(0, 1, num_particles)
      momentum = velocity * np.expand_dims(mass, axis=-1)
      # Remove average momentum, so center of mass does not move.
      momentum = momentum - momentum.mean(0, keepdims=True)
  
      # Connect all particles to all particles.
      particle_indices = np.arange(num_particles)
      senders, receivers = np.meshgrid(
          particle_indices, particle_indices
      )
      senders, receivers = senders.flatten(), receivers.flatten()
  
      # Generate a symmetric random matrix of spring constants.
      # Generate random elements stringly in the lower triangular part.
      spring_constants = np.random.uniform(
          1e-2, 1e-1, [num_particles, num_particles]
      )
      spring_constants = (
          np.tril(spring_constants) + np.tril(spring_constants, -1).T
      )
      spring_constants = spring_constants.flatten()
  
      # Remove interactions of particles to themselves.
      mask = senders != receivers
      senders, receivers = senders[mask], receivers[mask]
      spring_constants = spring_constants[mask]
      num_interactions = receivers.shape[0]
  
      return jraph.GraphsTuple(
          n_node=np.asarray([num_particles]),
          n_edge=np.asarray([num_interactions]),
          nodes={
              "mass": mass,  # Scalar mass for each particle.
              "position": position,  # 2d position for each particle.
              "momentum": momentum,  # 2d momentum for each particle.
          },
          edges={
              # Scalar spring constant for each interaction
              "spring_constant": spring_constants,
          },
          globals={},
          senders=senders,
          receivers=receivers,
      )
  
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 2 for calling a global function with example usage: build_hookes_particle_state_graph(num_particles) and returns: jraph