From 5423af633c77937a849f1f603f15eacfaf3fdc91 Mon Sep 17 00:00:00 2001
From: Paul Saxe <paul@thesaxes.net>
Date: Sat, 29 Jun 2024 10:26:41 -0400
Subject: [PATCH] Bugfix: bonding incorrect in	some cases.     * The	use of
 PDB with	Packmol	could in some cases lead to multiple bonds in the     
  wrong place. This release fixes that problem.

---
 HISTORY.rst             |  4 ++++
 packmol_step/packmol.py | 33 ++++++++++++++++++++++++++++-----
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/HISTORY.rst b/HISTORY.rst
index d2f7802..a9ecf5d 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -1,6 +1,10 @@
 =======
 History
 =======
+2024.6.29 -- Bugfix: bonding incorrect in some cases.
+    * The use of PDB with Packmol could in some cases lead to multiple bonds in the
+      wrong place. This release fixes that problem.
+      
 2024.6.21.2 -- Another internal release for Docker.
 
 2024.6.21.1 -- Internal release for Docker
diff --git a/packmol_step/packmol.py b/packmol_step/packmol.py
index 9d940db..4685901 100644
--- a/packmol_step/packmol.py
+++ b/packmol_step/packmol.py
@@ -273,13 +273,20 @@ def run(self):
         self.logger.debug(pprint.pformat(result))
 
         # Get the bond orders and extra parameters like ff atom types
-        bond_orders = []
         extra_data = {}
         total_q = 0.0
+        offset = 0
+        i_indices = []
+        j_indices = []
+        bond_orders = []
         for molecule in molecules:
             n = molecule["number"]
-            orders = molecule["configuration"].bonds.get_column_data("bondorder")
-            bond_orders.extend(orders * n)
+            for _ in range(n):
+                for i, j, bond_order in molecule["bonds"]:
+                    i_indices.append(i + offset)
+                    j_indices.append(j + offset)
+                    bond_orders.append(bond_order)
+                offset += molecule["configuration"].n_atoms
 
             total_q += n * molecule["configuration"].charge
 
@@ -303,8 +310,12 @@ def run(self):
         configuration.coordinate_system = "Cartesian"
         configuration.from_pdb_text(result["packmol.pdb"]["data"])
 
-        # And set the bond orders and extra data we saved earlier.
-        configuration.bonds.get_column("bondorder")[:] = bond_orders
+        ids = configuration.atoms.ids
+        i_atoms = [ids[x] for x in i_indices]
+        j_atoms = [ids[x] for x in j_indices]
+        configuration.bonds.append(i=i_atoms, j=j_atoms, bondorder=bond_orders)
+
+        # And set the extra data we saved earlier.
         for key, values in extra_data.items():
             if key not in configuration.atoms:
                 if "atom_types_" in key:
@@ -438,6 +449,9 @@ def round_copies(n_copies, molecules):
                 if ff is not None:
                     if ff_key not in tmp_configuration.atoms or assign_ff_always:
                         ff.assign_forcefield(tmp_configuration)
+                    else:
+                        if any(typ is None for typ in tmp_configuration.atoms[ff_key]):
+                            ff.assign_forcefield(tmp_configuration)
 
             tmp_mass = tmp_configuration.mass * ureg.g / ureg.mol
             tmp_mass.ito("kg")
@@ -455,6 +469,14 @@ def round_copies(n_copies, molecules):
                 n_fluid_atoms += count * tmp_configuration.n_atoms
                 fluid_mass += count * tmp_mass
 
+            # Keep track of the bonding information to add at end_bond
+            bonds = []
+            index = {_id: i for i, _id in enumerate(tmp_configuration.atoms.ids)}
+            for row in tmp_configuration.bonds.bonds():
+                bonds.append((index[row["i"]], index[row["j"]], row["bondorder"]))
+            tmp_configuration.bonds.clear()
+            tmp_configuration.db.commit()
+
             molecules.append(
                 {
                     "configuration": tmp_configuration,
@@ -463,6 +485,7 @@ def round_copies(n_copies, molecules):
                     "mass": tmp_mass,
                     "n_atoms": tmp_configuration.n_atoms,
                     "definition": definition,
+                    "bonds": bonds,
                 }
             )