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 picocli.CommandLine.Command;
44 import picocli.CommandLine.Parameters;
45
46 import java.io.BufferedReader;
47 import java.io.IOException;
48 import java.nio.charset.StandardCharsets;
49 import java.nio.file.Files;
50 import java.nio.file.Path;
51 import java.nio.file.Paths;
52 import java.util.ArrayList;
53 import java.util.List;
54
55 import static java.lang.Double.parseDouble;
56
57
58
59
60
61 @Command(description = " Compute the ANI-2x energy via an external GraalPy subprocess.",
62 name = "ANIExternal")
63 public class ANIExternal extends PotentialCommand {
64
65 private static final String PYTHON_SCRIPT = """
66 import os
67 import sys
68 import torch
69
70 def flatten(value):
71 if isinstance(value, (list, tuple)):
72 flattened = []
73 for item in value:
74 flattened.extend(flatten(item))
75 return flattened
76 return [value]
77
78 input_path = sys.argv[1]
79 output_path = sys.argv[2]
80
81 with open(input_path, encoding="utf-8") as handle:
82 torch_script = handle.readline().rstrip("\\n")
83 n_atoms = int(handle.readline().strip())
84 species = [int(value) for value in handle.readline().split()]
85 coords = []
86 for _ in range(n_atoms):
87 coords.append([float(value) for value in handle.readline().split()])
88
89 species_tensor = torch.tensor([species], dtype=torch.int64)
90 coordinates_tensor = torch.tensor([coords], dtype=torch.double)
91 ani = torch.jit.load(torch_script)
92 gradient = ani(species_tensor, coordinates_tensor)
93 if hasattr(gradient, "detach"):
94 gradient = gradient.detach().cpu()
95 if hasattr(gradient, "tolist"):
96 gradient = gradient.tolist()
97 flat = flatten(gradient)
98 energy = flat[-1]
99 gradients = flat[:-1]
100
101 with open(output_path, "w", encoding="utf-8") as handle:
102 handle.write(f"{energy}\\n")
103 handle.write(" ".join(str(value) for value in gradients))
104 handle.write("\\n")
105 handle.flush()
106
107 # Torch may leave background activity that upsets GraalPy shutdown on exit.
108 os._exit(0)
109 """;
110
111
112
113
114 @Parameters(arity = "1", paramLabel = "file",
115 description = "The atomic coordinate file in PDB or XYZ format.")
116 private String filename = null;
117
118 public ANIExternal() {
119 super();
120 }
121
122 public ANIExternal(FFXBinding binding) {
123 super(binding);
124 }
125
126 public ANIExternal(String[] args) {
127 super(args);
128 }
129
130 @Override
131 public ANIExternal run() {
132 if (!init()) {
133 return this;
134 }
135
136 activeAssembly = getActiveAssembly(filename);
137 if (activeAssembly == null) {
138 logger.info(helpString());
139 return this;
140 }
141
142 filename = activeAssembly.getFile().getAbsolutePath();
143 logger.info("\n Running External ANI Energy on " + filename);
144
145 String ffxHome = System.getProperty("basedir");
146 Path defaultGraalPyRoot = Paths.get(ffxHome, "python-resources");
147 String graalpyProperty = System.getProperty("graalpy", defaultGraalPyRoot.toString());
148 Path graalpyPath = Paths.get(graalpyProperty);
149 Path graalpyExecutable = resolvePythonExecutable(graalpyPath);
150 logger.info(" graalpy (-Dgraalpy=path.to.graalpy): " + graalpyPath);
151 logger.info(" executable: " + graalpyExecutable);
152
153 String torchScriptProperty = System.getProperty("torchscript", "ANI2x.pt");
154 Path torchScript = resolveTorchScript(ffxHome, torchScriptProperty);
155 logger.info(" torchscript (-Dtorchscript=path.to.torchscript): " + torchScript);
156
157 if (!Files.isExecutable(graalpyExecutable)) {
158 logger.severe(" External GraalPy executable was not found: " + graalpyExecutable);
159 return this;
160 }
161 if (!Files.isRegularFile(torchScript)) {
162 logger.severe(" TorchScript model was not found: " + torchScript);
163 return this;
164 }
165
166 Atom[] atoms = activeAssembly.getAtomArray();
167 int nAtoms = atoms.length;
168 int[] species = new int[nAtoms];
169 double[][] coords = new double[nAtoms][3];
170 for (int i = 0; i < nAtoms; i++) {
171 Atom atom = atoms[i];
172 species[i] = atom.getAtomicNumber();
173 coords[i][0] = atom.getX();
174 coords[i][1] = atom.getY();
175 coords[i][2] = atom.getZ();
176 }
177
178 Path inputFile = null;
179 Path outputFile = null;
180 try {
181 Path tempDir = baseDir == null ? null : baseDir.toPath();
182 inputFile = tempDir == null
183 ? Files.createTempFile("ani-external-input", ".txt")
184 : Files.createTempFile(tempDir, "ani-external-input", ".txt");
185 outputFile = tempDir == null
186 ? Files.createTempFile("ani-external-output", ".txt")
187 : Files.createTempFile(tempDir, "ani-external-output", ".txt");
188
189 writeInputFile(inputFile, torchScript, species, coords);
190 runSubprocess(graalpyExecutable, inputFile, outputFile, Paths.get(ffxHome, "python"));
191 readAndLogResults(outputFile, nAtoms);
192 } catch (Exception e) {
193 logger.severe(" External ANI subprocess failed: " + e);
194 } finally {
195 deleteIfExists(inputFile);
196 deleteIfExists(outputFile);
197 }
198
199 return this;
200 }
201
202 private Path resolvePythonExecutable(Path graalpyPath) {
203 if (Files.isRegularFile(graalpyPath)) {
204 return graalpyPath;
205 }
206
207 boolean isWindows = System.getProperty("os.name").toLowerCase().startsWith("windows");
208 String exeDir = isWindows ? "Scripts" : "bin";
209 String[] candidateNames = isWindows
210 ? new String[]{"python.exe", "graalpy.exe"}
211 : new String[]{"python", "graalpy"};
212
213 List<Path> candidates = new ArrayList<>();
214 for (String candidateName : candidateNames) {
215 candidates.add(graalpyPath.resolve("venv").resolve(exeDir).resolve(candidateName));
216 candidates.add(graalpyPath.resolve(exeDir).resolve(candidateName));
217 }
218
219 for (Path candidate : candidates) {
220 if (Files.isExecutable(candidate)) {
221 return candidate;
222 }
223 }
224
225 return candidates.getFirst();
226 }
227
228 private Path resolveTorchScript(String ffxHome, String torchScriptProperty) {
229 Path torchScript = Paths.get(torchScriptProperty);
230 if (torchScript.isAbsolute() || Files.exists(torchScript)) {
231 return torchScript.toAbsolutePath();
232 }
233
234 Path pythonDirTorchScript = Paths.get(ffxHome, "python", torchScriptProperty);
235 if (Files.exists(pythonDirTorchScript)) {
236 return pythonDirTorchScript.toAbsolutePath();
237 }
238
239 return torchScript.toAbsolutePath();
240 }
241
242 private void writeInputFile(Path inputFile, Path torchScript, int[] species, double[][] coords)
243 throws IOException {
244 List<String> lines = new ArrayList<>();
245 lines.add(torchScript.toAbsolutePath().toString());
246 lines.add(Integer.toString(species.length));
247
248 StringBuilder speciesLine = new StringBuilder();
249 for (int i = 0; i < species.length; i++) {
250 if (i > 0) {
251 speciesLine.append(' ');
252 }
253 speciesLine.append(species[i]);
254 }
255 lines.add(speciesLine.toString());
256
257 for (double[] coord : coords) {
258 lines.add(coord[0] + " " + coord[1] + " " + coord[2]);
259 }
260
261 Files.write(inputFile, lines, StandardCharsets.UTF_8);
262 }
263
264 private void runSubprocess(Path graalpyExecutable, Path inputFile, Path outputFile, Path workingDir)
265 throws IOException, InterruptedException {
266 ProcessBuilder processBuilder = new ProcessBuilder(
267 graalpyExecutable.toAbsolutePath().toString(),
268 "-c",
269 PYTHON_SCRIPT,
270 inputFile.toAbsolutePath().toString(),
271 outputFile.toAbsolutePath().toString());
272 processBuilder.directory(workingDir.toFile());
273 processBuilder.redirectErrorStream(true);
274
275 Process process = processBuilder.start();
276 try (BufferedReader reader = process.inputReader(StandardCharsets.UTF_8)) {
277 String line;
278 while ((line = reader.readLine()) != null) {
279 logger.info(line);
280 }
281 }
282
283 int exitCode = process.waitFor();
284 if (exitCode != 0) {
285 throw new IOException(" External ANI subprocess exited with code " + exitCode);
286 }
287 }
288
289 private void readAndLogResults(Path outputFile, int nAtoms) throws IOException {
290 List<String> lines = Files.readAllLines(outputFile, StandardCharsets.UTF_8);
291 if (lines.isEmpty()) {
292 throw new IOException(" External ANI subprocess produced no output.");
293 }
294
295 double energy = parseDouble(lines.getFirst().trim());
296 if (lines.size() > 1 && !lines.get(1).isBlank()) {
297 String[] gradientTokens = lines.get(1).trim().split("\\s+");
298 if (gradientTokens.length != nAtoms * 3) {
299 logger.warning(
300 " Expected " + (nAtoms * 3) + " gradient elements but received " + gradientTokens.length);
301 }
302 }
303 logger.info(" ANI-2x Energy (Hartree): " + energy);
304 }
305
306 private void deleteIfExists(Path path) {
307 if (path == null) {
308 return;
309 }
310 try {
311 Files.deleteIfExists(path);
312 } catch (IOException e) {
313 logger.fine(" Could not delete temporary file " + path + ": " + e.getMessage());
314 }
315 }
316 }