From 6229a3e7c331752a6c95f7a40e42b56e48a81273 Mon Sep 17 00:00:00 2001 From: Mike Henry <11765982+mikemhenry@users.noreply.github.com> Date: Mon, 21 Aug 2023 08:20:40 -0700 Subject: [PATCH] Branch to update to new dgl version (#179) * Branch to update to new dgl version * Also only linux builds * See if this makes the solver happy * see if the shim solves the GPU problem * see if the python 3.10 build works better * unpin pytorch since the conda-forge package should have the right pins * until we get CPU builds, I think I need to fake a cuda toolkit version * Remove shim * see if CPU builds JustWork TM * update to testing channel that has osx support * update API to new dgl 1.1.0 * figure out which version of python are giving us issues * bump ci * see if new logic works * Update espaloma.yaml * bump ci --- .github/workflows/CI.yaml | 5 +++-- devtools/conda-envs/espaloma.yaml | 5 +---- espaloma/app/experiment.py | 2 +- espaloma/data/dataset.py | 4 ++-- espaloma/mm/energy.py | 16 ++++++++-------- espaloma/mm/geometry.py | 6 +++--- espaloma/mm/nonbonded.py | 8 ++++---- espaloma/mm/tests/test_energy_ii.py | 14 +++++++------- espaloma/nn/readout/charge_equilibrium.py | 12 ++++++------ espaloma/nn/readout/graph_level_readout.py | 2 +- espaloma/nn/readout/janossy.py | 8 ++++---- 11 files changed, 40 insertions(+), 42 deletions(-) diff --git a/.github/workflows/CI.yaml b/.github/workflows/CI.yaml index 524603fc..0b08c0cf 100644 --- a/.github/workflows/CI.yaml +++ b/.github/workflows/CI.yaml @@ -29,10 +29,11 @@ jobs: strategy: fail-fast: false matrix: - os: ['macos', 'ubuntu'] + os: ['ubuntu','macos'] python-version: - - "3.9" + - "3.11" - "3.10" + - "3.9" env: OPENMM: ${{ matrix.cfg.openmm }} diff --git a/devtools/conda-envs/espaloma.yaml b/devtools/conda-envs/espaloma.yaml index 89560532..4865f11c 100644 --- a/devtools/conda-envs/espaloma.yaml +++ b/devtools/conda-envs/espaloma.yaml @@ -1,7 +1,6 @@ name: espaloma-test channels: - conda-forge - - dglteam - openeye dependencies: # Base dependencies @@ -20,9 +19,7 @@ dependencies: - openmmforcefields >=0.11.2 - tqdm - pydantic <2 # We need our deps to fix this - # Pytorch - - pytorch>=1.8.0 - - dgl <1 + - dgl =1.1.0 # Testing - pytest - pytest-cov diff --git a/espaloma/app/experiment.py b/espaloma/app/experiment.py index 06f1e7c0..debe234c 100644 --- a/espaloma/app/experiment.py +++ b/espaloma/app/experiment.py @@ -201,7 +201,7 @@ def test(self): # from time to time # make it just one giant graph # g = list(self.data) - # g = dgl.batch_hetero(g) + # g = dgl.batch(g) # g = g.to(self.device) if self.states is None: diff --git a/espaloma/data/dataset.py b/espaloma/data/dataset.py index 4e413c86..fc09be0f 100644 --- a/espaloma/data/dataset.py +++ b/espaloma/data/dataset.py @@ -293,13 +293,13 @@ def batch(graphs): import dgl if all(isinstance(graph, esp.graphs.graph.Graph) for graph in graphs): - return dgl.batch_hetero([graph.heterograph for graph in graphs]) + return dgl.batch([graph.heterograph for graph in graphs]) elif all(isinstance(graph, dgl.DGLGraph) for graph in graphs): return dgl.batch(graphs) elif all(isinstance(graph, dgl.DGLHeteroGraph) for graph in graphs): - return dgl.batch_hetero(graphs) + return dgl.batch(graphs) else: raise RuntimeError( diff --git a/espaloma/mm/energy.py b/espaloma/mm/energy.py index 0a9e687f..6608746d 100644 --- a/espaloma/mm/energy.py +++ b/espaloma/mm/energy.py @@ -361,7 +361,7 @@ def energy_in_graph( { "%s_in_g" % term: ( - dgl.function.copy_src(src="u%s" % suffix, out="m_%s" % term), + dgl.function.copy_u(u="u%s" % suffix, out="m_%s" % term), dgl.function.sum( msg="m_%s" % term, out="u_%s%s" % (term, suffix) ), @@ -462,31 +462,31 @@ def forward(self, g): g.multi_update_all( { "n2_as_0_in_n3": ( - dgl.function.copy_src("u", "m_u_0"), + dgl.function.copy_u("u", "m_u_0"), dgl.function.sum("m_u_0", "u_left"), ), "n2_as_1_in_n3": ( - dgl.function.copy_src("u", "m_u_1"), + dgl.function.copy_u("u", "m_u_1"), dgl.function.sum("m_u_1", "u_right"), ), "n2_as_0_in_n4": ( - dgl.function.copy_src("u", "m_u_0"), + dgl.function.copy_u("u", "m_u_0"), dgl.function.sum("m_u_0", "u_bond_left"), ), "n2_as_1_in_n4": ( - dgl.function.copy_src("u", "m_u_1"), + dgl.function.copy_u("u", "m_u_1"), dgl.function.sum("m_u_1", "u_bond_center"), ), "n2_as_2_in_n4": ( - dgl.function.copy_src("u", "m_u_2"), + dgl.function.copy_u("u", "m_u_2"), dgl.function.sum("m_u_2", "u_bond_right"), ), "n3_as_0_in_n4": ( - dgl.function.copy_src("u", "m3_u_0"), + dgl.function.copy_u("u", "m3_u_0"), dgl.function.sum("m3_u_0", "u_angle_left"), ), "n3_as_1_in_n4": ( - dgl.function.copy_src("u", "m3_u_1"), + dgl.function.copy_u("u", "m3_u_1"), dgl.function.sum("m3_u_1", "u_angle_right"), ), }, diff --git a/espaloma/mm/geometry.py b/espaloma/mm/geometry.py index e576c012..3462fd58 100644 --- a/espaloma/mm/geometry.py +++ b/espaloma/mm/geometry.py @@ -188,7 +188,7 @@ def geometry_in_graph(g): **{ "n1_as_%s_in_n%s" % (pos_idx, big_idx): ( - dgl.function.copy_src(src="xyz", out="m_xyz%s" % pos_idx), + dgl.function.copy_u(u="xyz", out="m_xyz%s" % pos_idx), dgl.function.sum( msg="m_xyz%s" % pos_idx, out="xyz%s" % pos_idx ), @@ -199,7 +199,7 @@ def geometry_in_graph(g): **{ "n1_as_%s_in_%s" % (pos_idx, term): ( - dgl.function.copy_src(src="xyz", out="m_xyz%s" % pos_idx), + dgl.function.copy_u(u="xyz", out="m_xyz%s" % pos_idx), dgl.function.sum( msg="m_xyz%s" % pos_idx, out="xyz%s" % pos_idx ), @@ -210,7 +210,7 @@ def geometry_in_graph(g): **{ "n1_as_%s_in_%s" % (pos_idx, term): ( - dgl.function.copy_src(src="xyz", out="m_xyz%s" % pos_idx), + dgl.function.copy_u(u="xyz", out="m_xyz%s" % pos_idx), dgl.function.sum( msg="m_xyz%s" % pos_idx, out="xyz%s" % pos_idx ), diff --git a/espaloma/mm/nonbonded.py b/espaloma/mm/nonbonded.py index d98fd4c6..f5211b55 100644 --- a/espaloma/mm/nonbonded.py +++ b/espaloma/mm/nonbonded.py @@ -48,8 +48,8 @@ def lorentz_berthelot(g, suffix=""): { "n1_as_%s_in_%s" % (pos_idx, term): ( - dgl.function.copy_src( - src="epsilon%s" % suffix, out="m_epsilon" + dgl.function.copy_u( + u="epsilon%s" % suffix, out="m_epsilon" ), geometric_mean(msg="m_epsilon", out="epsilon%s" % suffix), ) @@ -63,7 +63,7 @@ def lorentz_berthelot(g, suffix=""): { "n1_as_%s_in_%s" % (pos_idx, term): ( - dgl.function.copy_src(src="sigma%s" % suffix, out="m_sigma"), + dgl.function.copy_u(u="sigma%s" % suffix, out="m_sigma"), arithmetic_mean(msg="m_sigma", out="sigma%s" % suffix), ) for pos_idx in [0, 1] @@ -94,7 +94,7 @@ def multiply_charges(g, suffix=""): { "n1_as_%s_in_%s" % (pos_idx, term): ( - dgl.function.copy_src(src="q%s" % suffix, out="m_q"), + dgl.function.copy_u(u="q%s" % suffix, out="m_q"), dgl.function.sum(msg="m_q", out="_q") # lambda node: {"q%s" % suffix: node.mailbox["m_q"].prod(dim=1)} ) diff --git a/espaloma/mm/tests/test_energy_ii.py b/espaloma/mm/tests/test_energy_ii.py index 724506e9..5207949f 100644 --- a/espaloma/mm/tests/test_energy_ii.py +++ b/espaloma/mm/tests/test_energy_ii.py @@ -78,31 +78,31 @@ def forward(self, g): g.multi_update_all( { "n2_as_0_in_n3": ( - dgl.function.copy_src("u", "m_u_0"), + dgl.function.copy_u("u", "m_u_0"), dgl.function.sum("m_u_0", "u_left"), ), "n2_as_1_in_n3": ( - dgl.function.copy_src("u", "m_u_1"), + dgl.function.copy_u("u", "m_u_1"), dgl.function.sum("m_u_1", "u_right"), ), "n2_as_0_in_n4": ( - dgl.function.copy_src("u", "m_u_0"), + dgl.function.copy_u("u", "m_u_0"), dgl.function.sum("m_u_0", "u_bond_left"), ), "n2_as_1_in_n4": ( - dgl.function.copy_src("u", "m_u_1"), + dgl.function.copy_u("u", "m_u_1"), dgl.function.sum("m_u_1", "u_bond_center"), ), "n2_as_2_in_n4": ( - dgl.function.copy_src("u", "m_u_2"), + dgl.function.copy_u("u", "m_u_2"), dgl.function.sum("m_u_2", "u_bond_right"), ), "n3_as_0_in_n4": ( - dgl.function.copy_src("u", "m3_u_0"), + dgl.function.copy_u("u", "m3_u_0"), dgl.function.sum("m3_u_0", "u_angle_left"), ), "n3_as_1_in_n4": ( - dgl.function.copy_src("u", "m3_u_1"), + dgl.function.copy_u("u", "m3_u_1"), dgl.function.sum("m3_u_1", "u_angle_right"), ), }, diff --git a/espaloma/nn/readout/charge_equilibrium.py b/espaloma/nn/readout/charge_equilibrium.py index 558c1d12..65e60ea1 100644 --- a/espaloma/nn/readout/charge_equilibrium.py +++ b/espaloma/nn/readout/charge_equilibrium.py @@ -80,7 +80,7 @@ def forward(self, g, total_charge=0.0): if "q_ref" in g.nodes["n1"].data: # get total charge g.update_all( - dgl.function.copy_src(src="q_ref", out="m_q"), + dgl.function.copy_u(u="q_ref", out="m_q"), dgl.function.sum(msg="m_q", out="sum_q"), etype="n1_in_g", ) @@ -95,32 +95,32 @@ def forward(self, g, total_charge=0.0): ) g.update_all( - dgl.function.copy_src(src="sum_q", out="m_sum_q"), + dgl.function.copy_u(u="sum_q", out="m_sum_q"), dgl.function.sum(msg="m_sum_q", out="sum_q"), etype="g_has_n1", ) # get the sum of $s^{-1}$ and $m_s^{-1}$ g.update_all( - dgl.function.copy_src(src="s_inv", out="m_s_inv"), + dgl.function.copy_u(u="s_inv", out="m_s_inv"), dgl.function.sum(msg="m_s_inv", out="sum_s_inv"), etype="n1_in_g", ) g.update_all( - dgl.function.copy_src(src="e_s_inv", out="m_e_s_inv"), + dgl.function.copy_u(u="e_s_inv", out="m_e_s_inv"), dgl.function.sum(msg="m_e_s_inv", out="sum_e_s_inv"), etype="n1_in_g", ) g.update_all( - dgl.function.copy_src(src="sum_s_inv", out="m_sum_s_inv"), + dgl.function.copy_u(u="sum_s_inv", out="m_sum_s_inv"), dgl.function.sum(msg="m_sum_s_inv", out="sum_s_inv"), etype="g_has_n1", ) g.update_all( - dgl.function.copy_src(src="sum_e_s_inv", out="m_sum_e_s_inv"), + dgl.function.copy_u(u="sum_e_s_inv", out="m_sum_e_s_inv"), dgl.function.sum(msg="m_sum_e_s_inv", out="sum_e_s_inv"), etype="g_has_n1", ) diff --git a/espaloma/nn/readout/graph_level_readout.py b/espaloma/nn/readout/graph_level_readout.py index c4a53617..d0da31c8 100644 --- a/espaloma/nn/readout/graph_level_readout.py +++ b/espaloma/nn/readout/graph_level_readout.py @@ -53,7 +53,7 @@ def forward(self, g): ) g.update_all( - dgl.function.copy_src("h_global", "m"), + dgl.function.copy_u("h_global", "m"), self.pool("m", "h_global"), etype="n1_in_g", ) diff --git a/espaloma/nn/readout/janossy.py b/espaloma/nn/readout/janossy.py index 8abb38c5..329629b1 100644 --- a/espaloma/nn/readout/janossy.py +++ b/espaloma/nn/readout/janossy.py @@ -104,7 +104,7 @@ def forward(self, g): { "n1_as_%s_in_n%s" % (relationship_idx, big_idx): ( - dgl.function.copy_src("h", "m%s" % relationship_idx), + dgl.function.copy_u("h", "m%s" % relationship_idx), dgl.function.mean( "m%s" % relationship_idx, "h%s" % relationship_idx ), @@ -240,7 +240,7 @@ def forward(self, g): { "n1_as_%s_in_%s" % (relationship_idx, big_idx): ( - dgl.function.copy_src("h", "m%s" % relationship_idx), + dgl.function.copy_u("h", "m%s" % relationship_idx), dgl.function.mean( "m%s" % relationship_idx, "h%s" % relationship_idx ), @@ -358,7 +358,7 @@ def forward(self, g): { "n1_as_%s_in_%s" % (relationship_idx, big_idx): ( - dgl.function.copy_src("h", "m%s" % relationship_idx), + dgl.function.copy_u("h", "m%s" % relationship_idx), dgl.function.mean( "m%s" % relationship_idx, "h%s" % relationship_idx ), @@ -475,7 +475,7 @@ def forward(self, g): { "n1_as_%s_in_%s" % (relationship_idx, big_idx): ( - dgl.function.copy_src("h", "m%s" % relationship_idx), + dgl.function.copy_u("h", "m%s" % relationship_idx), dgl.function.mean( "m%s" % relationship_idx, "h%s" % relationship_idx ),