Simplify arguments and reintroduce aggregation mode

src/schnetpack/atomistic.py

@@ -124,7 +124,7 @@ class Atomwise(OutputModule):
 
     """
 
-    def __init__(self, n_in=128, n_out=1, pool_mode='sum', n_layers=2, n_neurons=None,
+    def __init__(self, n_in=128, n_out=1, aggregation_mode='sum', n_layers=2, n_neurons=None,
                  activation=schnetpack.nn.activations.shifted_softplus, return_contributions=False,
                  requires_dr=False, create_graph=False, mean=None, stddev=None, atomref=None, max_z=100, outnet=None,
                  train_embeddings=False):
@@ -157,9 +157,9 @@ def __init__(self, n_in=128, n_out=1, pool_mode='sum', n_layers=2, n_neurons=Non
         # Make standardization separate
         self.standardize = schnetpack.nn.base.ScaleShift(mean, stddev)
 
-        if pool_mode == 'sum':
+        if aggregation_mode == 'sum':
             self.atom_pool = schnetpack.nn.base.Aggregate(axis=1, mean=False)
-        elif pool_mode == 'avg':
+        elif aggregation_mode == 'avg':
             self.atom_pool = schnetpack.nn.base.Aggregate(axis=1, mean=True)
 
     def forward(self, inputs):
@@ -211,11 +211,11 @@ class Energy(Atomwise):
 
             If requires_dr is true additionally returns forces
         """
-    def __init__(self, n_in, pool_mode='sum', n_layers=2, n_neurons=None,
+    def __init__(self, n_in, aggregation_mode='sum', n_layers=2, n_neurons=None,
                  activation=schnetpack.nn.activations.shifted_softplus,
                  return_contributions=False, create_graph=False,
                  return_force=False, mean=None, stddev=None, atomref=None, max_z=100, outnet=None):
-        super(Energy, self).__init__(n_in, 1, pool_mode, n_layers, n_neurons, activation,
+        super(Energy, self).__init__(n_in, 1, aggregation_mode, n_layers, n_neurons, activation,
                                      return_contributions, return_force, create_graph, mean, stddev,
                                      atomref, 100, outnet)
 
@@ -297,14 +297,14 @@ class ElementalAtomwise(Atomwise):
     Particularly useful for networks of Behler-Parrinello type.
     """
 
-    def __init__(self, n_in, n_out=1, pool_mode='sum', n_layers=3, requires_dr=False, create_graph=False,
+    def __init__(self, n_in, n_out=1, aggregation_mode='sum', n_layers=3, requires_dr=False, create_graph=False,
                  elements=frozenset((1, 6, 7, 8, 9)), n_hidden=50,
                  activation=schnetpack.nn.activations.shifted_softplus,
                  return_contributions=False, mean=None, stddev=None, atomref=None, max_z=100):
         outnet = schnetpack.nn.blocks.GatedNetwork(n_in, n_out, elements, n_hidden=n_hidden, n_layers=n_layers,
                                                    activation=activation)
 
-        super(ElementalAtomwise, self).__init__(n_in, n_out, pool_mode, n_layers, None, activation,
+        super(ElementalAtomwise, self).__init__(n_in, n_out, aggregation_mode, n_layers, None, activation,
                                                 return_contributions, requires_dr, create_graph, mean, stddev, atomref,
                                                 max_z, outnet)
 
@@ -331,14 +331,14 @@ class ElementalEnergy(Energy):
         max_z (int): ignored if atomref is not learned (default: 100)
     """
 
-    def __init__(self, n_in, n_out=1, pool_mode='sum', n_layers=3, return_force=False, create_graph=False,
+    def __init__(self, n_in, n_out=1, aggregation_mode='sum', n_layers=3, return_force=False, create_graph=False,
                  elements=frozenset((1, 6, 7, 8, 9)), n_hidden=50,
                  activation=schnetpack.nn.activations.shifted_softplus, return_contributions=False, mean=None,
                  stddev=None, atomref=None, max_z=100):
         outnet = schnetpack.nn.blocks.GatedNetwork(n_in, n_out, elements, n_hidden=n_hidden, n_layers=n_layers,
                                                    activation=activation)
 
-        super(ElementalEnergy, self).__init__(n_in, pool_mode, n_layers, None, activation, return_contributions,
+        super(ElementalEnergy, self).__init__(n_in, aggregation_mode, n_layers, None, activation, return_contributions,
                                               create_graph=create_graph, return_force=return_force, mean=mean,
                                               stddev=stddev, atomref=atomref, max_z=max_z, outnet=outnet)
 

src/scripts/schnetpack_ani1.py

@@ -25,24 +25,19 @@ def get_parser():
     ## command-specific
     cmd_parser = argparse.ArgumentParser(add_help=False)
     cmd_parser.add_argument('--cuda', help='Set flag to use GPU(s)', action='store_true')
-    cmd_parser.add_argument('--logger', help='Choose logger for training process (default: %(default)s)',
-                            choices=['csv', 'tensorboard'], default='csv')
     cmd_parser.add_argument('--parallel',
                             help='Run data-parallel on all available GPUs (specify with environment variable'
                                  + ' CUDA_VISIBLE_DEVICES)', action='store_true')
     cmd_parser.add_argument('--batch_size', type=int,
                             help='Mini-batch size for training and prediction (default: %(default)s)',
                             default=400)
-    cmd_parser.add_argument('--log_every_n_epochs', type=int,
-                            help='Log metrics every given number of epochs (default: %(default)s)',
-                            default=1)
-    cmd_parser.add_argument('--property', type=str, help='ANI-1 property to be predicted (default: %(default)s)',
-                            default="energy", choices=ANI1.properties)
 
     ## training
     train_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
     train_parser.add_argument('datapath', help='Path / destination of ANI1 dataset directory')
     train_parser.add_argument('modelpath', help='Destination for models and logs')
+    train_parser.add_argument('--property', type=str, help='ANI-1 property to be predicted (default: %(default)s)',
+                              default="energy", choices=ANI1.properties)
     train_parser.add_argument('--seed', type=int, default=None, help='Set random seed for torch and numpy.')
     train_parser.add_argument('--overwrite', help='Remove previous model directory.', action='store_true')
 
@@ -63,6 +58,12 @@ def get_parser():
     train_parser.add_argument('--lr_min', type=float, help='Minimal learning rate (default: %(default)s)',
                               default=1e-6)
 
+    train_parser.add_argument('--logger', help='Choose logger for training process (default: %(default)s)',
+                              choices=['csv', 'tensorboard'], default='csv')
+    train_parser.add_argument('--log_every_n_epochs', type=int,
+                              help='Log metrics every given number of epochs (default: %(default)s)',
+                              default=1)
+
     ## evaluation
     eval_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
     eval_parser.add_argument('datapath', help='Path of ANI1 dataset directory')
@@ -72,6 +73,8 @@ def get_parser():
 
     # model-specific parsers
     model_parser = argparse.ArgumentParser(add_help=False)
+    model_parser.add_argument('--aggregation_mode', type=str, default='sum', choices=['sum', 'avg'],
+                              help=' (default: %(default)s)')
 
     #######  SchNet  #######
     schnet_parser = argparse.ArgumentParser(add_help=False, parents=[model_parser])
@@ -172,9 +175,9 @@ def loss(batch, result):
     trainer.train(device)
 
 
-def evaluate(args, model, train_loader, val_loader, test_loader, device):
-    metrics = [spk.metrics.MeanAbsoluteError(args.property, 0),
-               spk.metrics.RootMeanSquaredError(args.property, 0)]
+def evaluate(args, model, property, train_loader, val_loader, test_loader, device):
+    metrics = [spk.metrics.MeanAbsoluteError(property, 0),
+               spk.metrics.RootMeanSquaredError(property, 0)]
 
     header = []
     results = []
@@ -221,7 +224,8 @@ def get_model(args, atomref=None, mean=None, stddev=None, train_loader=None, par
     if args.model == 'schnet':
         representation = spk.representation.SchNet(args.features, args.features, args.interactions,
                                                    args.cutoff, args.num_gaussians)
-        atomwise_output = spk.atomistic.Atomwise(args.features, mean=mean, stddev=stddev, atomref=atomref)
+        atomwise_output = spk.atomistic.Atomwise(args.features, mean=mean, stddev=stddev, atomref=atomref,
+                                                 aggregation_mode=args.aggregation_mode)
         model = spk.atomistic.AtomisticModel(representation, atomwise_output)
 
     elif args.model == 'wacsf':
@@ -247,6 +251,7 @@ def get_model(args, atomref=None, mean=None, stddev=None, train_loader=None, par
         # Build HDNN model
         atomwise_output = spk.atomistic.ElementalEnergy(representation.n_symfuncs, n_hidden=args.n_nodes,
                                                         n_layers=args.n_layers, mean=mean, stddev=stddev,
+                                                        aggregation_mode=args.aggregation_mode,
                                                         atomref=atomref, elements=elements)
         model = spk.atomistic.AtomisticModel(representation, atomwise_output)
 
@@ -285,9 +290,9 @@ def get_model(args, atomref=None, mean=None, stddev=None, train_loader=None, par
 
     # will download ANI1 if necessary, calculate_triples is required for wACSF angular functions
     logging.info('ANI1 will be loaded...')
-    ani1 = spk.datasets.ANI1(args.datapath, download=True, properties=[args.property],
+    ani1 = spk.datasets.ANI1(args.datapath, download=True, properties=[train_args.property],
                              collect_triples=args.model == 'wacsf')
-    atomref = ani1.get_reference(args.property)
+    atomref = ani1.get_reference(train_args.property)
 
     # splits the dataset in test, val, train sets
     split_path = os.path.join(args.modelpath, 'split.npz')
@@ -330,7 +335,7 @@ def get_model(args, atomref=None, mean=None, stddev=None, train_loader=None, par
         test_loader = spk.data.AtomsLoader(data_test, batch_size=args.batch_size,
                                            num_workers=4, pin_memory=True)
         with torch.no_grad():
-            evaluate(args, model, train_loader, val_loader, test_loader, device)
+            evaluate(args, model, train_args.property, train_loader, val_loader, test_loader, device)
         logging.info("... done!")
     else:
         print('Unknown mode:', args.mode)

src/scripts/schnetpack_matproj.py

@@ -25,25 +25,19 @@ def get_parser():
     ## command-specific
     cmd_parser = argparse.ArgumentParser(add_help=False)
     cmd_parser.add_argument('--cuda', help='Set flag to use GPU(s)', action='store_true')
-    cmd_parser.add_argument('--logger', help='Choose logger for training process (default: %(default)s)',
-                            choices=['csv', 'tensorboard'], default='csv')
     cmd_parser.add_argument('--parallel',
                             help='Run data-parallel on all available GPUs (specify with environment variable'
                                  + ' CUDA_VISIBLE_DEVICES)', action='store_true')
     cmd_parser.add_argument('--batch_size', type=int,
                             help='Mini-batch size for training and prediction (default: %(default)s)',
                             default=32)
-    cmd_parser.add_argument('--log_every_n_epochs', type=int,
-                            help='Log metrics every given number of epochs (default: %(default)s)',
-                            default=1)
-    cmd_parser.add_argument('--property', type=str,
-                            help='Materials Project property to be predicted (default: %(default)s)',
-                            default="formation_energy_per_atom", choices=MaterialsProject.properties)
-
     ## training
     train_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
     train_parser.add_argument('datapath', help='Path / destination of Materials Project dataset directory')
     train_parser.add_argument('modelpath', help='Destination for models and logs')
+    train_parser.add_argument('--property', type=str,
+                              help='Materials Project property to be predicted (default: %(default)s)',
+                              default="formation_energy_per_atom", choices=MaterialsProject.properties)
     train_parser.add_argument('--apikey', help='API key for Materials Project (see https://materialsproject.org/open)')
     train_parser.add_argument('--seed', type=int, default=None, help='Set random seed for torch and numpy.')
     train_parser.add_argument('--overwrite', help='Remove previous model directory.', action='store_true')
@@ -64,6 +58,12 @@ def get_parser():
     train_parser.add_argument('--lr_min', type=float, help='Minimal learning rate (default: %(default)s)',
                               default=1e-6)
 
+    train_parser.add_argument('--logger', help='Choose logger for training process (default: %(default)s)',
+                              choices=['csv', 'tensorboard'], default='csv')
+    train_parser.add_argument('--log_every_n_epochs', type=int,
+                              help='Log metrics every given number of epochs (default: %(default)s)',
+                              default=1)
+
     ## evaluation
     eval_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
     eval_parser.add_argument('datapath', help='Path of MaterialsProject dataset directory')
@@ -75,6 +75,8 @@ def get_parser():
 
     # model-specific parsers
     model_parser = argparse.ArgumentParser(add_help=False)
+    model_parser.add_argument('--aggregation_mode', type=str, default='avg', choices=['sum', 'avg'],
+                              help=' (default: %(default)s)')
 
     #######  SchNet  #######
     schnet_parser = argparse.ArgumentParser(add_help=False, parents=[model_parser])
@@ -148,11 +150,11 @@ def loss(batch, result):
     trainer.train(device)
 
 
-def evaluate(args, model, train_loader, val_loader, test_loader, device):
-    header = ['Subset', args.property + ' MAE', args.property + ' RMSE']
+def evaluate(args, model, property, train_loader, val_loader, test_loader, device):
+    header = ['Subset', property + ' MAE', property + ' RMSE']
 
-    metrics = [spk.metrics.MeanAbsoluteError(args.property, 0),
-               spk.metrics.RootMeanSquaredError(args.property, 0)
+    metrics = [spk.metrics.MeanAbsoluteError(property, 0),
+               spk.metrics.RootMeanSquaredError(property, 0)
                ]
 
     results = []
@@ -195,7 +197,8 @@ def get_model(args, atomref=None, mean=None, stddev=None, train_loader=None, par
     if args.model == 'schnet':
         representation = spk.representation.SchNet(args.features, args.features, args.interactions,
                                                    args.cutoff, args.num_gaussians, normalize_filter=True)
-        atomwise_output = spk.atomistic.Atomwise(args.features, pool_mode='avg', mean=mean, stddev=stddev,
+        atomwise_output = spk.atomistic.Atomwise(args.features, aggregation_mode=args.aggregation_mode,
+                                                 mean=mean, stddev=stddev,
                                                  atomref=atomref, train_embeddings=True)
         model = spk.atomistic.AtomisticModel(representation, atomwise_output)
     else:
@@ -248,7 +251,7 @@ def export_model(args):
 
     # will download MaterialsProject if necessary
     mp = spk.datasets.MaterialsProject(args.datapath, args.cutoff, apikey=args.apikey, download=True,
-                                       properties=[args.property])
+                                       properties=[train_args.property])
 
     # splits the dataset in test, val, train sets
     split_path = os.path.join(args.modelpath, 'split.npz')
@@ -285,7 +288,7 @@ def export_model(args):
         train(args, model, train_loader, val_loader, device)
     elif args.mode == 'eval':
         test_loader = spk.data.AtomsLoader(data_test, batch_size=args.batch_size,
-                                      num_workers=2, pin_memory=True)
-        evaluate(args, model, train_loader, val_loader, test_loader, device)
+                                           num_workers=2, pin_memory=True)
+        evaluate(args, model, train_args.property, train_loader, val_loader, test_loader, device)
     else:
         print('Unknown mode:', args.mode)

src/scripts/schnetpack_md17.py

@@ -25,17 +25,12 @@ def get_parser():
     ## command-specific
     cmd_parser = argparse.ArgumentParser(add_help=False)
     cmd_parser.add_argument('--cuda', help='Set flag to use GPU(s)', action='store_true')
-    cmd_parser.add_argument('--logger', help='Choose logger for training process (default: %(default)s)',
-                            choices=['csv', 'tensorboard'], default='csv')
     cmd_parser.add_argument('--parallel',
                             help='Run data-parallel on all available GPUs (specify with environment variable'
                                  + ' CUDA_VISIBLE_DEVICES)', action='store_true')
     cmd_parser.add_argument('--batch_size', type=int,
                             help='Mini-batch size for training and prediction (default: %(default)s)',
                             default=100)
-    cmd_parser.add_argument('--log_every_n_epochs', type=int,
-                            help='Log metrics every given number of epochs (default: %(default)s)',
-                            default=1)
 
     ## training
     train_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
@@ -66,6 +61,12 @@ def get_parser():
                               help='Energy-force trade-off. For rho=0, use forces only. (default: %(default)s)',
                               default=0.1)
 
+    train_parser.add_argument('--logger', help='Choose logger for training process (default: %(default)s)',
+                              choices=['csv', 'tensorboard'], default='csv')
+    train_parser.add_argument('--log_every_n_epochs', type=int,
+                              help='Log metrics every given number of epochs (default: %(default)s)',
+                              default=1)
+
     ## evaluation
     eval_parser = argparse.ArgumentParser(add_help=False, parents=[cmd_parser])
     eval_parser.add_argument('datapath', help='Path / destination of MD17 dataset directory')