Add bos to compound, and bugfix to bos

1 files changed, 14 insertions, 1 deletions

diff --git a/ml_exp/compound.py b/ml_exp/compound.py
index 595078d55..8b6af0ae9 100644
--- a/ml_exp/compound.py
+++ b/ml_exp/compound.py
@@ -23,7 +23,7 @@ SOFTWARE.
 import numpy as np
 from ml_exp.data import NUCLEAR_CHARGE
 from ml_exp.representations import coulomb_matrix, lennard_jones_matrix,\
-    first_neighbor_matrix, adjacency_matrix
+    first_neighbor_matrix, adjacency_matrix, bag_of_stuff
 
 
 class Compound:
@@ -115,6 +115,19 @@ class Compound:
                                    forces,
                                    size=size)
 
+    def gen_bos(self,
+                size=23,
+                stuff='bonds'):
+        """
+        Generate the Bag of Stuff for the compound.
+        size: compound size.
+        stuff: elements of the bag, by default the known bag of bonds.
+        """
+        self.bos = bag_of_stuff(self.cm,
+                                self.atoms,
+                                size=size,
+                                stuff=stuff)
+
     def read_xyz(self,
                  filename):
         """