From 3caa5b074b4de43790c4beb93c244df25b2a40af Mon Sep 17 00:00:00 2001
From: Geoff Hutchison <geoff.hutchison@gmail.com>
Date: Mon, 14 Oct 2024 23:03:37 -0400
Subject: [PATCH] Initial support for reading charges and forces from extended
 XYZ

Signed-off-by: Geoff Hutchison <geoff.hutchison@gmail.com>
---
 avogadro/io/xyzformat.cpp | 52 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/avogadro/io/xyzformat.cpp b/avogadro/io/xyzformat.cpp
index 9c3e391c42..f063722141 100644
--- a/avogadro/io/xyzformat.cpp
+++ b/avogadro/io/xyzformat.cpp
@@ -87,6 +87,36 @@ bool XyzFormat::read(std::istream& inStream, Core::Molecule& mol)
       mol.setUnitCell(cell);
     }
   }
+  // check to see if there's an extended XYZ Properties= line
+  // e.g. Properties=species:S:1:pos:R:3
+  // https://gitlab.com/ase/ase/-/merge_requests/62
+  start = buffer.find("Properties=");
+  unsigned int chargeColumn = 0;
+  unsigned int forceColumn = 0;
+  vector<double> charges;
+  if (start != std::string::npos) {
+    start = start + 11; // skip over "Properties="
+    unsigned int stop = buffer.find(' ', start);
+    unsigned int length = stop - start;
+    // we want to track columns after the position
+    // (esp. charge, spin, force, velocity, etc.)
+    std::string properties = buffer.substr(start, length);
+    vector<string> tokens(split(properties, ':'));
+    unsigned int column = 0;
+    for (size_t i = 0; i < tokens.size(); i += 3) {
+      // increment column based on the count of the property
+      if (i + 2 < tokens.size()) {
+        column += lexicalCast<unsigned int>(tokens[i + 2]);
+      }
+
+      // we can safely assume species and pos are present
+      if (tokens[i] == "charge") {
+        chargeColumn = column;
+      } else if (tokens[i] == "force" || tokens[i] == "forces") {
+        forceColumn = column;
+      } // TODO other properties (velocity, spin, selection, etc.)
+    }
+  }
 
   // Parse atoms
   for (size_t i = 0; i < numAtoms; ++i) {
@@ -114,6 +144,18 @@ bool XyzFormat::read(std::istream& inStream, Core::Molecule& mol)
 
     Atom newAtom = mol.addAtom(atomicNum);
     newAtom.setPosition3d(pos);
+
+    // check for charge and force columns
+    if (chargeColumn > 0 && chargeColumn < tokens.size()) {
+      charges.push_back(lexicalCast<double>(tokens[chargeColumn]));
+      // we set the charges after all atoms are added
+    }
+    if (forceColumn > 0 && forceColumn < tokens.size()) {
+      Vector3 force(lexicalCast<double>(tokens[forceColumn]),
+                    lexicalCast<double>(tokens[forceColumn + 1]),
+                    lexicalCast<double>(tokens[forceColumn + 2]));
+      newAtom.setForceVector(force);
+    }
   }
 
   // Check that all atoms were handled.
@@ -126,6 +168,15 @@ bool XyzFormat::read(std::istream& inStream, Core::Molecule& mol)
     return false;
   }
 
+  // set the charges
+  if (!charges.empty()) {
+    MatrixX chargesMatrix = MatrixX::Zero(mol.atomCount(), 1);
+    for (size_t i = 0; i < charges.size(); ++i) {
+      chargesMatrix(i, 0) = charges[i];
+    }
+    mol.setPartialCharges("From File", chargesMatrix);
+  }
+
   // Do we have an animation?
   size_t numAtoms2;
   // check if the next frame has the same number of atoms
@@ -251,6 +302,7 @@ std::vector<std::string> XyzFormat::fileExtensions() const
 {
   std::vector<std::string> ext;
   ext.emplace_back("xyz");
+  ext.emplace_back("exyz");
   ext.emplace_back("extxyz");
   return ext;
 }