Add representation sorting by row-norm

2 files changed, 29 insertions, 2 deletions

diff --git a/ml_exp/compound.py b/ml_exp/compound.py
index 03cfb8519..a85daefa4 100644
--- a/ml_exp/compound.py
+++ b/ml_exp/compound.py
@@ -78,12 +78,14 @@ class Compound:
 
     def gen_cm(self,
                size=23,
+               sort=False,
                flatten=True,
                as_eig=True,
                bohr_ru=False):
         """
         Generate the Coulomb Matrix for the compund.
         size: compound size.
+        sort: if the representation should be sorted row-norm-wise.
         flatten: if the representation should be 1D.
         as_eig: if the representation should be as the eigenvalues.
         bohr_ru: if radius units should be in bohr's radius units.
@@ -91,15 +93,17 @@ class Compound:
         self.cm = coulomb_matrix(self.coordinates,
                                  self.nc,
                                  size=size,
+                                 sort=sort,
+                                 flatten=flatten,
                                  as_eig=as_eig,
-                                 bohr_ru=bohr_ru,
-                                 flatten=flatten)
+                                 bohr_ru=bohr_ru)
 
     def gen_ljm(self,
                 diag_value=None,
                 sigma=1.0,
                 epsilon=1.0,
                 size=23,
+                sort=False,
                 flatten=True,
                 as_eig=True,
                 bohr_ru=False):
@@ -109,6 +113,7 @@ class Compound:
         sigma: sigma value.
         epsilon: epsilon value.
         size: compound size.
+        sort: if the representation should be sorted row-norm-wise.
         flatten: if the representation should be 1D.
         as_eig: if the representation should be as the eigenvalues.
         bohr_ru: if radius units should be in bohr's radius units.
@@ -119,6 +124,7 @@ class Compound:
                                         sigma=sigma,
                                         epsilon=epsilon,
                                         size=size,
+                                        sort=sort,
                                         flatten=flatten,
                                         as_eig=as_eig,
                                         bohr_ru=bohr_ru)
@@ -142,11 +148,13 @@ class Compound:
     def gen_am(self,
                use_forces=False,
                size=23,
+               sort=False,
                flatten=True):
         """
         Generate the Adjacency Matrix for the compund.
         use_forces: if the use of forces instead of k_cx should be used.
         size: compound size.
+        sort: if the representation should be sorted row-norm-wise.
         flatten: if the representation should be 1D.
         """
         self.am = adjacency_matrix(self.bonds_i,
@@ -154,6 +162,7 @@ class Compound:
                                    self.bonds_f,
                                    use_forces=use_forces,
                                    size=size,
+                                   sort=sort,
                                    flatten=flatten)
 
     def gen_ei(self,
diff --git a/ml_exp/representations.py b/ml_exp/representations.py
index fcaa0f33a..47164ca7c 100644
--- a/ml_exp/representations.py
+++ b/ml_exp/representations.py
@@ -28,6 +28,7 @@ from ml_exp.data import POSSIBLE_BONDS
 def coulomb_matrix(coords,
                    nc,
                    size=23,
+                   sort=False,
                    flatten=True,
                    as_eig=True,
                    bohr_ru=False):
@@ -36,6 +37,7 @@ def coulomb_matrix(coords,
     coords: compound coordinates.
     nc: nuclear charge data.
     size: compound size.
+    sort: if the representation should be sorted row-norm-wise.
     flatten: if the representation should be 1D.
     as_eig: if the representation should be as the eigenvalues.
     bohr_ru: if radius units should be in bohr's radius units.
@@ -77,6 +79,10 @@ size. Arrays are not of the right shape.')
 
         return np.pad(cm_eigs, (0, size - n), 'constant')
     else:
+        if sort:
+            si = np.argsort(np.linalg.norm(cm, axis=-1))[::-1]
+            cm = cm[si]
+
         if flatten:
             return np.pad(cm, ((0, size - n), (0, size - n)),
                           'constant').flatten()
@@ -90,6 +96,7 @@ def lennard_jones_matrix(coords,
                          sigma=1.0,
                          epsilon=1.0,
                          size=23,
+                         sort=False,
                          flatten=True,
                          as_eig=True,
                          bohr_ru=False):
@@ -101,6 +108,7 @@ def lennard_jones_matrix(coords,
     sigma: sigma value.
     epsilon: epsilon value.
     size: compound size.
+    sort: if the representation should be sorted row-norm-wise.
     flatten: if the representation should be 1D.
     as_eig: if the representation should be as the eigenvalues.
     bohr_ru: if radius units should be in bohr's radius units.
@@ -149,6 +157,10 @@ size. Arrays are not of the right shape.')
 
         return np.pad(lj_eigs, (0, size - n), 'constant')
     else:
+        if sort:
+            si = np.argsort(np.linalg.norm(lj, axis=-1))[::-1]
+            lj = lj[si]
+
         if flatten:
             return np.pad(lj, ((0, size - n), (0, size - n)),
                           'constant').flatten()
@@ -216,6 +228,7 @@ def adjacency_matrix(bonds_i,
                      bonds_f,
                      use_forces=False,
                      size=23,
+                     sort=False,
                      flatten=True):
     """
     Calculates the adjacency matrix given the bond list.
@@ -225,6 +238,7 @@ def adjacency_matrix(bonds_i,
     bonds_f: list of force values.
     use_forces: if the use of forces instead of k_cx should be used.
     size: compund size.
+    sort: if the representation should be sorted row-norm-wise.
     flatten: if the representation should be 1D.
     """
     if bonds_i is None:
@@ -249,6 +263,10 @@ the current compound.')
                     am[i, j] = bonds_k[i]
                     am[j, i] = am[i, j]
 
+    if sort:
+        si = np.argsort(np.linalg.norm(am, axis=-1))
+        am = am[si]
+
     if flatten:
         return np.pad(am, ((0, size - n), (0, size - n)), 'constant').flatten()
     else: