code: def hookes_hamiltonian_from_graph_fn(
      graph: jraph.GraphsTuple,
  ) -> jraph.GraphsTuple:
      """Computes Hamiltonian of a Hooke's potential system represented in a graph.
  
      While this function hardcodes the Hamiltonian for a Hooke's potential, a
      learned Hamiltonian Graph Network (https://arxiv.org/abs/1909.12790) could
      be implemented by replacing the hardcoded formulas by learnable MLPs that
      take as inputs all of the concatenated features to the edge_fn, node_fn,
      and global_fn, and outputs a single scalar value in the global_fn.
  
      Args:
        graph: `GraphsTuple` where the nodes contain:
            - "mass": [num_particles]
            - "position": [num_particles, num_dims]
            - "momentum": [num_particles, num_dims]
            and the edges contain:
            - "spring_constant": [num_interations]
  
      Returns:
        `GraphsTuple` with features:
            - edge features: "hookes_potential" [num_interactions]
            - node features: "kinetic_energy" [num_particles]
            - global features: "hamiltonian" [batch_size]
  
      """
  
      def update_edge_fn(edges, senders, receivers, globals_):
          del globals_
          distance = jnp.linalg.norm(
              senders["position"] - receivers["position"]
          )
          hookes_potential_per_edge = (
              0.5 * edges["spring_constant"] * distance**2
          )
          return frozendict(
              {"hookes_potential": hookes_potential_per_edge}
          )
  
      def update_node_fn(nodes, sent_edges, received_edges, globals_):
          del sent_edges, received_edges, globals_
          momentum_norm = jnp.linalg.norm(nodes["momentum"])
          kinetic_energy_per_node = momentum_norm**2 / (
              2 * nodes["mass"]
          )
          return frozendict({"kinetic_energy": kinetic_energy_per_node})
  
      def update_global_fn(nodes, edges, globals_):
          del globals_
          # At this point we will receive node and edge features aggregated (summed)
          # for all nodes and edges in each graph.
          hamiltonian_per_graph = (
              nodes["kinetic_energy"] + edges["hookes_potential"]
          )
          return frozendict({"hamiltonian": hamiltonian_per_graph})
  
      gn = jraph.GraphNetwork(
          update_edge_fn=update_edge_fn,
          update_node_fn=update_node_fn,
          update_global_fn=update_global_fn,
      )
  
      return gn(graph)
  
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 1 for calling a global function with example usage: hookes_hamiltonian_from_graph_fn(graph) and returns: gn, frozendict, frozendict, frozendict