Allow coo

.ci_support/environment.yml

@@ -4,3 +4,4 @@ channels:
 dependencies:
   - numpy
   - cython
+  - scipy

README.md

@@ -25,7 +25,7 @@ neighbors = structure.get_neighbors()
 first_shell_tensor = neighbors.get_shell_matrix()[0]
 
 mc = MC(len(structure))
-mc.set_heisenberg_coeff(J*first_shell_tensor.toarray())
+mc.set_heisenberg_coeff(J * first_shell_tensor)
 
 mc.run(temperature=300, number_of_iterations=1000)
 ```

mamonca/mc.pyx

@@ -24,7 +24,7 @@ cdef class MC:
     >>> first_shell_tensor = neighbors.get_shell_matrix()[0]
     >>> 
     >>> mc = MC(len(structure))
-    >>> mc.set_heisenberg_coeff(J*first_shell_tensor.toarray())
+    >>> mc.set_heisenberg_coeff(J*first_shell_tensor)
     >>> 
     >>> mc.run(temperature=300, number_of_iterations=1000)
 
@@ -71,33 +71,49 @@ cdef class MC:
     def set_heisenberg_coeff(self, coeff, i=None, j=None, deg=1, index=0):
         """
             Args:
-                coeff (float/list/ndarray): Heisenberg coefficient. If a single number is given,
-                    the same parameter is applied to all the pairs defined in me and neigh.
-                    Instead of giving me and neigh, you can also give a n_atom x n_atom tensor.
-                i (list/ndarray): list of indices i (s. definition in the comment)
-                j (list/ndarray): list of indices j (s. definition in the comment)
+                coeff (float/list/ndarray/scipy.sparse): Heisenberg coefficient.
+                    If a single number is given, the same parameter is applied
+                    to all the pairs defined in me and neigh. Instead of
+                    giving me and neigh, you can also give a n_atom x n_atom
+                    tensor.
+                i (list/ndarray): list of indices i (s. def in the comment)
+                j (list/ndarray): list of indices j (s. def in the comment)
                 deg (int): polynomial degree
-                index (int): potential index for thermodynamic integration (0 or 1; choose 0 if
-                    not thermodynamic integration)
+                index (int): potential index for thermodynamic integration
+                    (0 or 1; choose 0 if not thermodynamic integration)
             Comment:
-                Heisenberg term is given by: -sum_ij coeff_ij*(m_i*m_j)^deg. Beware of the
-                negative sign.
+                Heisenberg term is given by: -sum_ij coeff_ij*(m_i*m_j)^deg.
+                Beware of the negative sign.
         """
         if i is None and j is None:
             n = self.c_mc.get_number_of_atoms()
-            if np.array(coeff).shape!=(n, n):
-                raise ValueError(
-                    'If i and j are not specified, coeff has to be a 2d tensor with the length '
-                    + 'equal to the number of atoms in each direction.'
-                )
-            i,j = np.where(coeff!=0)
-            coeff = coeff[coeff!=0]
-        coeff = np.array([coeff]).flatten()
-        i = np.array(i).flatten()
-        j = np.array(j).flatten()
-        if len(coeff)==1:
+            if isinstance(coeff, (np.ndarray, list)):
+                if np.array(coeff).shape!=(n, n):
+                    raise ValueError(
+                        "If i and j are not specified, coeff has to be a 2d"
+                        " tensor with the length equal to the number of atoms in"
+                        " each direction."
+                    )
+                i,j = np.where(coeff != 0)
+                coeff = coeff[coeff != 0]
+            else:
+                if hasattr(coeff, "tocoo"):
+                    coeff = coeff.tocoo()
+                try:
+                    i = coeff.row
+                    j = coeff.col
+                    coeff = coeff.data
+                except AttributeError:
+                    raise ValueError(
+                        "Input can be a sparse matrix or a numpy array of"
+                        f" shape ({n}, {n}) or you must specify i, j and coeff"
+                    )
+        coeff = np.asarray([coeff]).flatten()
+        i = np.asarray(i).flatten()
+        j = np.asarray(j).flatten()
+        if len(coeff) == 1:
             coeff = np.tile(coeff, len(i))
-        if len(coeff)!=len(i) or len(i)!=len(j):
+        if not len(coeff) == len(i) == len(j):
             raise ValueError('Length of vectors not the same')
         self.c_mc.set_heisenberg_coeff(coeff, i, j, deg, index)
 

tests/test_heisenberg.py

@@ -2,18 +2,35 @@
 import numpy as np
 import unittest
 import os
+from scipy.sparse import coo_matrix
 
 
 class TestFull(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
         cls.file_location = os.path.dirname(os.path.abspath(__file__))
-        ij = np.loadtxt(os.path.join(cls.file_location, "neighbors.txt"))
-        cls.mc = MC(np.max(ij) + 1)
-        cls.mc.set_heisenberg_coeff(0.1, *ij)
+        cls.ij = np.loadtxt(
+            os.path.join(cls.file_location, "neighbors.txt")
+        ).astype(int)
+        cls.n = np.max(cls.ij) + 1
+        cls.mc = MC(cls.n)
+        cls.mc.set_heisenberg_coeff(0.1, *cls.ij)
         cls.mc.run(temperature=300, number_of_iterations=1000)
 
+    def test_coo_mat(self):
+        data = 0.1 * np.ones(self.ij.shape[1])
+        mat = coo_matrix(
+            (data, (self.ij[0], self.ij[1])),
+            shape=(self.n, self.n),
+        )
+        for i in range(2):
+            mc = MC(self.n)
+            mc.set_heisenberg_coeff(mat)
+            mc.run(100, number_of_iterations=1000)
+            self.assertLess(mc.get_energy(), 0)
+            mat = mat + mat # transforms to csr
+
     def test_acceptance_ratio(self):
         self.assertGreater(self.mc.get_acceptance_ratio(), 0)
         self.assertLess(self.mc.get_acceptance_ratio(), 1)