Merge pull request #92 from choderalab/no_dgl_import

no global import dgl
choderalab · Oct 26, 2021 · 51ea6db · 51ea6db
2 parents c93b501 + ef5c8e3
commit 51ea6db
Show file tree

Hide file tree

Showing 16 changed files with 24 additions and 284 deletions.
diff --git a/espaloma/app/experiment.py b/espaloma/app/experiment.py
@@ -3,9 +3,6 @@
 # =============================================================================
 import abc
 import copy
-
-import numpy as np
-import dgl
 import torch
 
 import espaloma as esp

diff --git a/espaloma/data/utils.py b/espaloma/data/utils.py
@@ -3,7 +3,6 @@
 # =============================================================================
 import random
 
-import dgl
 import numpy as np
 import pandas as pd
 import torch
@@ -105,6 +104,7 @@ def split(ds, partition):
 
 def batch(ds, batch_size, seed=2666):
     """Batch graphs and values after shuffling."""
+    import dgl
     # get the numebr of data
     n_data_points = len(ds)
     n_batches = n_data_points // batch_size  # drop the rest
@@ -127,6 +127,7 @@ def batch(ds, batch_size, seed=2666):
 
 
 def collate_fn(graphs):
+    import dgl
     return esp.HomogeneousGraph(dgl.batch(graphs))
 
 

diff --git a/espaloma/graphs/graph.py b/espaloma/graphs/graph.py
@@ -2,8 +2,6 @@
 # IMPORTS
 # =============================================================================
 import abc
-
-import dgl
 import openff.toolkit
 
 import espaloma as esp
@@ -72,7 +70,7 @@ def __init__(self, mol=None, homograph=None, heterograph=None):
     def save(self, path):
         import os
         import json
-
+        import dgl
         os.mkdir(path)
         dgl.save_graphs(path + "/homograph.bin", [self.homograph])
         dgl.save_graphs(path + "/heterograph.bin", [self.heterograph])
@@ -82,6 +80,7 @@ def save(self, path):
     @classmethod
     def load(cls, path):
         import json
+        import dgl
 
         homograph = dgl.load_graphs(path + "/homograph.bin")[0][0]
         heterograph = dgl.load_graphs(path + "/heterograph.bin")[0][0]
@@ -117,6 +116,7 @@ def get_homograph_from_mol(mol):
 
     @staticmethod
     def get_heterograph_from_graph_and_mol(graph, mol):
+        import dgl
         assert isinstance(
             graph, dgl.DGLGraph
         ), "graph can only be dgl Graph object."

diff --git a/espaloma/graphs/utils/read_heterogeneous_graph.py b/espaloma/graphs/utils/read_heterogeneous_graph.py
@@ -4,7 +4,6 @@
 # =============================================================================
 # IMPORTS
 # =============================================================================
-import dgl
 import numpy as np
 import torch
 from espaloma.graphs.utils import offmol_indices
@@ -262,6 +261,7 @@ def from_homogeneous_and_mol(g, offmol):
             axis=1,
         )
 
+    import dgl
     hg = dgl.heterograph({key: list(value) for key, value in hg.items()})
 
     hg.nodes["n1"].data["h0"] = g.ndata["h0"]

diff --git a/espaloma/graphs/utils/read_homogeneous_graph.py b/espaloma/graphs/utils/read_homogeneous_graph.py
@@ -4,7 +4,6 @@
 # =============================================================================
 # IMPORTS
 # =============================================================================
-import dgl
 import torch
 
 # =============================================================================
@@ -118,6 +117,7 @@ def fp_rdkit(atom):
 # MODULE FUNCTIONS
 # =============================================================================
 def from_openff_toolkit_mol(mol, use_fp=True):
+    import dgl
     # initialize graph
     from rdkit import Chem
 
@@ -207,6 +207,7 @@ def from_oemol(mol, use_fp=True):
 
 
 def from_rdkit_mol(mol, use_fp=True):
+    import dgl
     from rdkit import Chem
 
     # initialize graph

diff --git a/espaloma/mm/energy.py b/espaloma/mm/energy.py
@@ -1,7 +1,6 @@
 # =============================================================================
 # IMPORTS
 # =============================================================================
-import dgl
 import torch
 
 import espaloma as esp
@@ -254,7 +253,7 @@ def energy_in_graph(
     """
     # TODO: this is all very restricted for now
     # we need to make this better
-
+    import dgl
     if "n2" in terms:
         # apply energy function
 
@@ -417,6 +416,7 @@ def forward(self, g):
 class CarryII(torch.nn.Module):
     def forward(self, g):
         import math
+        import dgl
 
         g.multi_update_all(
             {

diff --git a/espaloma/mm/geometry.py b/espaloma/mm/geometry.py
@@ -1,7 +1,6 @@
 # =============================================================================
 # IMPORTS
 # =============================================================================
-import dgl
 import torch
 
 # =============================================================================
@@ -181,7 +180,7 @@ def geometry_in_graph(g):
     This function modifies graphs in-place.
 
     """
-
+    import dgl
     # Copy coordinates to higher-order nodes.
     g.multi_update_all(
         {

diff --git a/espaloma/mm/nonbonded.py b/espaloma/mm/nonbonded.py
@@ -1,14 +1,11 @@
 # =============================================================================
 # IMPORTS
 # =============================================================================
-import dgl
 import torch
 
 # =============================================================================
 # CONSTANTS
 # =============================================================================
-from simtk import unit
-
 import espaloma as esp
 
 # CODATA 2018
@@ -43,7 +40,7 @@ def _arithmetic_mean(nodes):
 # COMBINATION RULES FOR NONBONDED
 # =============================================================================
 def lorentz_berthelot(g, suffix=""):
-
+    import dgl
     g.multi_update_all(
         {
             "n1_as_%s_in_%s"

diff --git a/espaloma/nn/layers/dgl_legacy.py b/espaloma/nn/layers/dgl_legacy.py
@@ -2,16 +2,10 @@
 
 """
 
-import math
-from copy import deepcopy
-
-import dgl
-
 # =============================================================================
 # IMPORTS
 # =============================================================================
 import torch
-from dgl.nn import pytorch as dgl_pytorch
 
 # =============================================================================
 # CONSTANT
@@ -35,7 +29,7 @@ def __init__(
         kwargs={},
     ):
         super(GN, self).__init__()
-
+        from dgl.nn import pytorch as dgl_pytorch
         if kwargs == {}:
             if model_name in DEFAULT_MODEL_KWARGS:
                 kwargs = DEFAULT_MODEL_KWARGS[model_name]
@@ -58,6 +52,7 @@ def forward(self, g, x):
 
 
 def gn(model_name="GraphConv", kwargs={}):
+    from dgl.nn import pytorch as dgl_pytorch
     if model_name == "GINConv":
         return lambda in_features, out_features: dgl_pytorch.conv.GINConv(
             apply_func=torch.nn.Linear(in_features, out_features),

diff --git a/espaloma/nn/readout/charge_equilibrium.py b/espaloma/nn/readout/charge_equilibrium.py
@@ -5,8 +5,6 @@
 # IMPORTS
 # =============================================================================
 import torch
-import dgl
-import math
 
 # =============================================================================
 # UTILITY FUNCTIONS
@@ -67,6 +65,7 @@ def __init__(self):
     def forward(self, g, total_charge=0.0):
         """ apply charge equilibrium to all molecules in batch """
         # calculate $s ^ {-1}$ and $ es ^ {-1}$
+        import dgl
         g.apply_nodes(
             lambda node: {"s_inv": node.data["s"] ** -1}, ntype="n1"
         )

diff --git a/espaloma/nn/readout/graph_level_readout.py b/espaloma/nn/readout/graph_level_readout.py
@@ -1,7 +1,6 @@
 # =============================================================================
 # IMPORTS
 # =============================================================================
-import dgl
 import torch
 import espaloma as esp
 
@@ -17,10 +16,13 @@ def __init__(
         config_local,
         config_global,
         out_name,
-        pool=dgl.function.sum,
+        pool=None,
     ):
 
         super(GraphLevelReadout, self).__init__()
+        import dgl
+        if pool is None:
+            pool = dgl.function.sum
         self.in_features = in_features
         self.config_local = config_local
         self.config_global = config_global
@@ -42,6 +44,7 @@ def __init__(
         self.out_name = out_name
 
     def forward(self, g):
+        import dgl
         g.apply_nodes(
             lambda node: {"h_global": self.d_local(None, node.data["h"])},
             ntype="n1",

diff --git a/espaloma/nn/readout/janossy.py b/espaloma/nn/readout/janossy.py
@@ -1,7 +1,6 @@
 # =============================================================================
 # IMPORTS
 # =============================================================================
-import dgl
 import torch
 
 import espaloma as esp
@@ -98,6 +97,7 @@ def forward(self, g):
         g : dgl.DGLHeteroGraph,
             input graph.
         """
+        import dgl
 
         # copy
         g.multi_update_all(
@@ -233,7 +233,8 @@ def forward(self, g):
         g : dgl.DGLHeteroGraph,
             input graph.
         """
-
+        import dgl
+
         # copy
         g.multi_update_all(
             {

diff --git a/espaloma/nn/sequential.py b/espaloma/nn/sequential.py
@@ -1,9 +1,7 @@
 """ Chain mutiple layers of GN together.
 """
-import dgl
 import torch
 
-
 class _Sequential(torch.nn.Module):
     """Sequentially staggered neural networks."""
 
@@ -115,7 +113,7 @@ def forward(self, g, x=None):
         g : `dgl.DGLHeteroGraph`
             output graph
         """
-
+        import dgl
         # get homogeneous subgraph
         g_ = dgl.to_homo(g.edge_type_subgraph(["n1_neighbors_n1"]))
 

diff --git a/espaloma/redux/energy.py b/espaloma/redux/energy.py