bumped version, added some comments

experiments/nbody/data/generate_dataset.py

@@ -1,8 +1,8 @@
 """
-Generate chargedcd  and gravity datasets.
+Generate charged and gravity datasets.
 
-charged: python3 -u generate_dataset.py --simulation=charged --num-train 10000 --seed 43
-gravity: python3 -u generate_dataset.py --simulation=gravity --num-train 10000 --seed 43 --n-balls=100
+charged: python3 generate_dataset.py --simulation=charged --num-train=10000 --seed=43
+gravity: python3 generate_dataset.py --simulation=gravity --num-train=10000 --seed=43 --n-balls=100
 """
 import argparse
 import time

experiments/nbody/utils.py

@@ -1,20 +1,23 @@
 from typing import Callable, List, Optional, Tuple
 
+import e3nn_jax as e3nn
 import jax.numpy as jnp
 import jax.tree_util as tree
 import numpy as np
 import torch
-import e3nn_jax as e3nn
 from jraph import GraphsTuple, segment_mean
 from torch.utils.data import DataLoader
 from torch_geometric.nn import knn_graph
 
 from segnn_jax import SteerableGraphsTuple
+
 from .datasets import ChargedDataset, GravityDataset
 
 
 def O3Transform(
-    node_features_irreps: e3nn.Irreps, edge_features_irreps: e3nn.Irreps, lmax_attributes: int
+    node_features_irreps: e3nn.Irreps,
+    edge_features_irreps: e3nn.Irreps,
+    lmax_attributes: int,
 ) -> Callable:
     """
     Build a transformation function that includes (nbody) O3 attributes to a graph.
@@ -25,7 +28,7 @@ def _o3_transform(
         st_graph: SteerableGraphsTuple,
         loc: jnp.ndarray,
         vel: jnp.ndarray,
-        charges: jnp.ndarray
+        charges: jnp.ndarray,
     ) -> SteerableGraphsTuple:
 
         graph = st_graph.graph
@@ -40,7 +43,7 @@ def _o3_transform(
 
         vel_abs = jnp.sqrt(jnp.power(vel, 2).sum(1, keepdims=True))
         mean_loc = loc.mean(1, keepdims=True)
-        
+
         nodes = e3nn.IrrepsArray(
             node_features_irreps,
             jnp.concatenate((loc - mean_loc, vel, vel_abs), axis=-1),

main.py

@@ -2,6 +2,7 @@
 import time
 from functools import partial
 from typing import Tuple, Union
+
 import e3nn_jax as e3nn
 import haiku as hk
 import jax
@@ -20,7 +21,7 @@ def predict(
     state: hk.State,
     graph: SteerableGraphsTuple,
     mean_shift: Union[jnp.array, float] = 0,
-    mad_shift: Union[jnp.array, float] = 1
+    mad_shift: Union[jnp.array, float] = 1,
 ) -> Tuple[jnp.ndarray, hk.State]:
     pred, state = segnn.apply(params, state, graph)
     return (jnp.multiply(pred, mad_shift) + mean_shift), state
@@ -34,7 +35,7 @@ def mae(
     target: jnp.ndarray,
     mean_shift: Union[jnp.array, float] = 0,
     mad_shift: Union[jnp.array, float] = 1,
-    mask_last: bool = False
+    mask_last: bool = False,
 ) -> Tuple[float, hk.State]:
     pred, state = predict(params, state, graph, mean_shift, mad_shift)
     assert target.shape == pred.shape
@@ -53,7 +54,7 @@ def mse(
     target: jnp.ndarray,
     mean_shift: Union[jnp.array, float] = 0,
     mad_shift: Union[jnp.array, float] = 1,
-    mask_last: bool = False
+    mask_last: bool = False,
 ) -> Tuple[float, hk.State]:
     pred, state = predict(params, state, graph, mean_shift, mad_shift)
     assert target.shape == pred.shape

segnn_jax/__init__.py

@@ -13,4 +13,4 @@
     "SteerableGraphsTuple",
 ]
 
-__version__ = "0.2"
+__version__ = "0.3"

segnn_jax/blocks.py

@@ -80,7 +80,7 @@ def O3TensorProductGate(
         x (IrrepsArray): Left tensor
         y (IrrepsArray): Right tensor
         output_irreps: Output representation
-        biases: If set ot true will add biases
+        biases: Add biases
         scalar_activation: Activation function for scalars
         gate_activation: Activation function for higher order
         name: Name of the linear layer params

segnn_jax/segnn.py

@@ -111,6 +111,7 @@ def _message(
             msg = O3TensorProductGate(
                 msg, edge_attribute, output_irreps, name=f"message_{i}_{layer_num}"
             )
+        # NOTE: original implementation only applied batch norm to messages
         if norm == "batch":
             msg = e3nn.BatchNorm(irreps=output_irreps)(msg)
         return msg
@@ -136,6 +137,7 @@ def _update(
         )
         # residual connection
         nodes += update
+        # message norm
         if norm in ["batch", "instance"]:
             nodes = e3nn.BatchNorm(irreps=output_irreps, instance=(norm == "instance"))(
                 nodes
@@ -172,7 +174,7 @@ def __init__(
         blocks_per_layer: int = 2,
         embed_msg_features: bool = False,
     ):
-        super(SEGNN, self).__init__() # noqa
+        super(SEGNN, self).__init__()  # noqa # pylint: disable=R1725
 
         if isinstance(hidden_irreps, e3nn.Irreps):
             self._hidden_irreps_units = num_layers * [hidden_irreps]