Making integration test pull an actual dataset and do a few steps of training and evaluation.

PiperOrigin-RevId: 331524333

Author: alvarosg

learning_to_simulate/README.md

@@ -0,0 +1,33 @@
+# Learning to Simulate Complex Physics with Graph Networks
+
+This is a model implementation for the ICML 2020 submission (also available in
+arXiv [arxiv.org/abs/2002.09405](https://arxiv.org/abs/2002.09405). If you use
+the code here please cite this paper:
+
+    @article{sanchezgonzalez2020learning,
+        title={Learning to Simulate Complex Physics with Graph Networks},
+        author={Alvaro Sanchez-Gonzalez and
+                Jonathan Godwin and
+                Tobias Pfaff and
+                Rex Ying and
+                Jure Leskovec and
+                Peter W. Battaglia},
+        url={https://arxiv.org/abs/2002.09405},
+        year={2020},
+        eprint={2002.09405},
+        archivePrefix={arXiv},
+        primaryClass={cs.LG}
+    }
+
+## Contents
+
+*   `model.py`: implementation of the graph network use as the learnable part of
+    the model.
+*   `model_demo.py`: example connecting the model to input dummy data.
+
+## Running demo
+
+(From one directory above)
+
+    pip install -r learning_to_simulate/requirements.txt
+    python -m learning_to_simulate.model_demo

learning_to_simulate/connectivity_utils.py

@@ -0,0 +1,104 @@
+# Lint as: python3
+# pylint: disable=g-bad-file-header
+# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Tools to compute the connectivity of the graph."""
+
+import numpy as np
+from sklearn import neighbors
+import tensorflow.compat.v1 as tf
+
+
+def _compute_connectivity(positions, radius):
+  """Get the indices of connected edges with radius connectivity.
+
+  Args:
+    positions: Positions of nodes in the graph. Shape:
+      [num_nodes_in_graph, num_dims].
+    radius: Radius of connectivity.
+
+  Returns:
+    senders indices [num_edges_in_graph]
+    receiver indices [num_edges_in_graph]
+
+  """
+  tree = neighbors.KDTree(positions)
+  receivers_list = tree.query_radius(positions, r=radius)
+  num_nodes = len(positions)
+  senders = np.repeat(range(num_nodes), [len(a) for a in receivers_list])
+  receivers = np.concatenate(receivers_list, axis=0)
+  return senders, receivers
+
+
+def _compute_connectivity_for_batch(positions, n_node, radius):
+  """`compute_connectivity` for a batch of graphs.
+
+  Args:
+    positions: Positions of nodes in the batch of graphs. Shape:
+      [num_nodes_in_batch, num_dims].
+    n_node: Number of nodes for each graph in the batch. Shape:
+      [num_graphs in batch].
+    radius: Radius of connectivity.
+
+  Returns:
+    senders indices [num_edges_in_batch]
+    receiver indices [num_edges_in_batch]
+    number of edges per graph [num_graphs_in_batch]
+
+  """
+
+  # TODO(alvarosg): Consider if we want to support batches here or not.
+  # Separate the positions corresponding to particles in different graphs.
+  positions_per_graph_list = np.split(positions, np.cumsum(n_node[:-1]), axis=0)
+  receivers_list = []
+  senders_list = []
+  n_edge_list = []
+  num_nodes_in_previous_graphs = 0
+
+  # Compute connectivity for each graph in the batch.
+  for positions_graph_i in positions_per_graph_list:
+    senders_graph_i, receivers_graph_i = _compute_connectivity(
+        positions_graph_i, radius)
+
+    num_edges_graph_i = len(senders_graph_i)
+    n_edge_list.append(num_edges_graph_i)
+
+    # Because the inputs will be concatenated, we need to add offsets to the
+    # sender and receiver indices according to the number of nodes in previous
+    # graphs in the same batch.
+    receivers_list.append(receivers_graph_i + num_nodes_in_previous_graphs)
+    senders_list.append(senders_graph_i + num_nodes_in_previous_graphs)
+
+    num_nodes_graph_i = len(positions_graph_i)
+    num_nodes_in_previous_graphs += num_nodes_graph_i
+
+  # Concatenate all of the results.
+  senders = np.concatenate(senders_list, axis=0).astype(np.int32)
+  receivers = np.concatenate(receivers_list, axis=0).astype(np.int32)
+  n_edge = np.stack(n_edge_list).astype(np.int32)
+
+  return senders, receivers, n_edge
+
+
+def compute_connectivity_for_batch_pyfunc(positions, n_node, radius):
+  """`_compute_connectivity_for_batch` wrapped in a pyfunc."""
+  senders, receivers, n_edge = tf.py_function(
+      _compute_connectivity_for_batch,
+      [positions, n_node, radius], [tf.int32, tf.int32, tf.int32])
+  senders.set_shape([None])
+  receivers.set_shape([None])
+  n_edge.set_shape(n_node.get_shape())
+  return senders, receivers, n_edge
+

learning_to_simulate/download_dataset.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+# Copyright 2020 Deepmind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Usage:
+#     bash download_dataset.sh ${DATASET_NAME} ${OUTPUT_DIR}
+# Example:
+#     bash download_dataset.sh WaterDrop /tmp/
+
+set -e
+
+DATASET_NAME="${1}"
+OUTPUT_DIR="${2}/${DATASET_NAME}"
+
+BASE_URL="https://storage.googleapis.com/learning-to-simulate-complex-physics/Datasets/${DATASET_NAME}/"
+
+mkdir -p ${OUTPUT_DIR}
+for file in metadata.json train.tfrecord valid.tfrecord test.tfrecord
+do
+wget -O "${OUTPUT_DIR}/${file}" "${BASE_URL}${file}"
+done

learning_to_simulate/graph_network.py

@@ -0,0 +1,173 @@
+# Lint as: python3
+# pylint: disable=g-bad-file-header
+# Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""Graph network implementation accompanying ICML 2020 submission.
+
+   "Learning to Simulate Complex Physics with Graph Networks"
+
+   Alvaro Sanchez-Gonzalez*, Jonathan Godwin*, Tobias Pfaff*, Rex Ying,
+   Jure Leskovec, Peter W. Battaglia
+
+   https://arxiv.org/abs/2002.09405
+
+The Sonnet `EncodeProcessDecode` module provided here implements the learnable
+parts of the model.
+It assumes an encoder preprocessor has already built a graph with
+connectivity and features as described in the paper, with features normalized
+to zero-mean unit-variance.
+
+Dependencies include Tensorflow 1.x, Sonnet 1.x and the Graph Nets 1.1 library.
+"""
+
+import graph_nets as gn
+import sonnet as snt
+import tensorflow as tf
+
+
+def build_mlp(
+    hidden_size: int, num_hidden_layers: int, output_size: int) -> snt.Module:
+  """Builds an MLP."""
+  return snt.nets.MLP(
+      output_sizes=[hidden_size] * num_hidden_layers + [output_size])
+
+
+class EncodeProcessDecode(snt.AbstractModule):
+  """Encode-Process-Decode function approximator for learnable simulator."""
+
+  def __init__(
+      self,
+      latent_size: int,
+      mlp_hidden_size: int,
+      mlp_num_hidden_layers: int,
+      num_message_passing_steps: int,
+      output_size: int,
+      name: str = "EncodeProcessDecode"):
+    """Inits the model.
+
+    Args:
+      latent_size: Size of the node and edge latent representations.
+      mlp_hidden_size: Hidden layer size for all MLPs.
+      mlp_num_hidden_layers: Number of hidden layers in all MLPs.
+      num_message_passing_steps: Number of message passing steps.
+      output_size: Output size of the decode node representations as required
+        by the downstream update function.
+      name: Name of the model.
+    """
+
+    super().__init__(name=name)
+
+    self._latent_size = latent_size
+    self._mlp_hidden_size = mlp_hidden_size
+    self._mlp_num_hidden_layers = mlp_num_hidden_layers
+    self._num_message_passing_steps = num_message_passing_steps
+    self._output_size = output_size
+
+    with self._enter_variable_scope():
+      self._networks_builder()
+
+  def _build(self, input_graph: gn.graphs.GraphsTuple) -> tf.Tensor:
+    """Forward pass of the learnable dynamics model."""
+
+    # Encode the input_graph.
+    latent_graph_0 = self._encode(input_graph)
+
+    # Do `m` message passing steps in the latent graphs.
+    latent_graph_m = self._process(latent_graph_0)
+
+    # Decode from the last latent graph.
+    return self._decode(latent_graph_m)
+
+  def _networks_builder(self):
+    """Builds the networks."""
+
+    def build_mlp_with_layer_norm():
+      mlp = build_mlp(
+          hidden_size=self._mlp_hidden_size,
+          num_hidden_layers=self._mlp_num_hidden_layers,
+          output_size=self._latent_size)
+      return snt.Sequential([mlp, snt.LayerNorm()])
+
+    # The encoder graph network independently encodes edge and node features.
+    encoder_kwargs = dict(
+        edge_model_fn=build_mlp_with_layer_norm,
+        node_model_fn=build_mlp_with_layer_norm)
+    self._encoder_network = gn.modules.GraphIndependent(**encoder_kwargs)
+
+    # Create `num_message_passing_steps` graph networks with unshared parameters
+    # that update the node and edge latent features.
+    # Note that we can use `modules.InteractionNetwork` because
+    # it also outputs the messages as updated edge latent features.
+    self._processor_networks = []
+    for _ in range(self._num_message_passing_steps):
+      self._processor_networks.append(
+          gn.modules.InteractionNetwork(
+              edge_model_fn=build_mlp_with_layer_norm,
+              node_model_fn=build_mlp_with_layer_norm))
+
+    # The decoder MLP decodes node latent features into the output size.
+    self._decoder_network = build_mlp(
+        hidden_size=self._mlp_hidden_size,
+        num_hidden_layers=self._mlp_num_hidden_layers,
+        output_size=self._output_size)
+
+  def _encode(
+      self, input_graph: gn.graphs.GraphsTuple) -> gn.graphs.GraphsTuple:
+    """Encodes the input graph features into a latent graph."""
+
+    # Copy the globals to all of the nodes, if applicable.
+    if input_graph.globals is not None:
+      broadcasted_globals = gn.blocks.broadcast_globals_to_nodes(input_graph)
+      input_graph = input_graph.replace(
+          nodes=tf.concat([input_graph.nodes, broadcasted_globals], axis=-1),
+          globals=None)
+
+    # Encode the node and edge features.
+    latent_graph_0 = self._encoder_network(input_graph)
+    return latent_graph_0
+
+  def _process(
+      self, latent_graph_0: gn.graphs.GraphsTuple) -> gn.graphs.GraphsTuple:
+    """Processes the latent graph with several steps of message passing."""
+
+    # Do `m` message passing steps in the latent graphs.
+    # (In the shared parameters case, just reuse the same `processor_network`)
+    latent_graph_prev_k = latent_graph_0
+    for processor_network_k in self._processor_networks:
+      latent_graph_k = self._process_step(
+          processor_network_k, latent_graph_prev_k)
+      latent_graph_prev_k = latent_graph_k
+
+    latent_graph_m = latent_graph_k
+    return latent_graph_m
+
+  def _process_step(
+      self, processor_network_k: snt.Module,
+      latent_graph_prev_k: gn.graphs.GraphsTuple) -> gn.graphs.GraphsTuple:
+    """Single step of message passing with node/edge residual connections."""
+
+    # One step of message passing.
+    latent_graph_k = processor_network_k(latent_graph_prev_k)
+
+    # Add residuals.
+    latent_graph_k = latent_graph_k.replace(
+        nodes=latent_graph_k.nodes+latent_graph_prev_k.nodes,
+        edges=latent_graph_k.edges+latent_graph_prev_k.edges)
+    return latent_graph_k
+
+  def _decode(self, latent_graph: gn.graphs.GraphsTuple) -> tf.Tensor:
+    """Decodes from the latent graph."""
+    return self._decoder_network(latent_graph.nodes)