1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38 package ffx.potential.commands;
39
40 import ffx.potential.bonded.Atom;
41 import ffx.potential.cli.PotentialCommand;
42 import ffx.utilities.FFXBinding;
43 import org.graalvm.polyglot.Context;
44 import org.graalvm.polyglot.Value;
45 import org.graalvm.python.embedding.GraalPyResources;
46 import picocli.CommandLine.Command;
47 import picocli.CommandLine.Parameters;
48
49 import java.nio.file.Path;
50 import java.nio.file.Paths;
51
52
53
54
55
56
57
58
59
60 @Command(description = " Compute the ANI-2x energy.", name = "ANI-2x")
61 public class ANI extends PotentialCommand {
62
63
64
65
66 @Parameters(arity = "1", paramLabel = "file",
67 description = "The atomic coordinate file in PDB or XYZ format.")
68 private String filename = null;
69
70 public ANI() {
71 super();
72 }
73
74 public ANI(FFXBinding binding) {
75 super(binding);
76 }
77
78 public ANI(String[] args) {
79 super(args);
80 }
81
82 @Override
83 public ANI run() {
84 if (!init()) {
85 return this;
86 }
87
88
89 activeAssembly = getActiveAssembly(filename);
90 if (activeAssembly == null) {
91 logger.info(helpString());
92 return this;
93 }
94
95
96 filename = activeAssembly.getFile().getAbsolutePath();
97
98 logger.info("\n Running Energy on " + filename);
99
100 String FFX_HOME = System.getProperty("basedir");
101 Path graalpy = Paths.get(FFX_HOME, "python-resources");
102 String graalpyString = System.getProperty("graalpy", graalpy.toString());
103 graalpy = Paths.get(graalpyString);
104 logger.info(" graalpy (-Dgraalpy=path.to.graalpy): " + graalpy);
105 String torchScript = "ANI2x.pt";
106 torchScript = System.getProperty("torchscript", torchScript);
107 logger.info(" torchscript (-Dtorchscript=path.to.torchscript): " + torchScript);
108
109 Atom[] atoms = activeAssembly.getAtomArray();
110 int nAtoms = atoms.length;
111 int[] species = new int[nAtoms];
112 double[][] coords = new double[nAtoms][3];
113 for (int i = 0; i < nAtoms; i++) {
114 Atom a = atoms[i];
115 species[i] = a.getAtomicNumber();
116 coords[i][0] = a.getX();
117 coords[i][1] = a.getY();
118 coords[i][2] = a.getZ();
119 }
120 double energy = 0.0;
121 double[] grad = new double[nAtoms * 3];
122
123 try (Context context = GraalPyResources.contextBuilder(graalpy)
124 .option("python.WarnExperimentalFeatures", "false").build()) {
125
126 Value polyglotBindings = context.getPolyglotBindings();
127 polyglotBindings.putMember("coords", coords);
128 polyglotBindings.putMember("species", species);
129 polyglotBindings.putMember("torchScript", torchScript);
130
131
132 String torch = """
133 import platform
134 import polyglot
135 import torch
136
137 # Load the Java arrays into the Torch tensors.
138 species = polyglot.import_value('species')
139 coords = polyglot.import_value('coords')
140 torchScript = polyglot.import_value('torchScript')
141 speciesTensor = torch.tensor([species], dtype=torch.int64)
142 coordinatesTensor = torch.tensor([coords], dtype=torch.double)
143
144 # Load and evaluate the ANI-2x forward method.
145 ani = torch.jit.load(torchScript)
146 gradient = ani(speciesTensor, coordinatesTensor)""";
147
148 torch = torch.stripIndent();
149 logger.info(torch);
150
151
152 Value result = context.eval("python", torch);
153 Value bindings = context.getBindings("python");
154 Value ret = bindings.getMember("gradient");
155 for (int i = 0; i < nAtoms * 3; i++) {
156 grad[i] = ret.getArrayElement(i).asDouble();
157 }
158 energy = ret.getArrayElement(nAtoms * 3).asDouble();
159 logger.info(" ANI-2x Energy (Hartree): " + energy);
160
161
162 } catch (Exception e) {
163 logger.info(" Exception:\n" + e);
164 }
165
166 return this;
167 }
168
169 }