Repository: dwavesystems/dimod

Issue:
Support metadata in COO (like vartype)
**Application**
Currently, (de-)serialization of BQMs is partial and with loss of information - vartype is not stored/loaded.

It would be nice to have a complete (de-)serialization of BQM to COO format, like we already have for JSON and BSON. E.g.:
```
bqm = BinaryQuadraticModel.from_coo(fp)
bqm.to_coo(fp)
```

**Proposed Solution**
Support `vartype` via `p` or `c` in COO.


Discussion:

Patch:
diff --git a/dimod/binary_quadratic_model.py b/dimod/binary_quadratic_model.py
index 31450ad8..995c36cd 100644
--- a/dimod/binary_quadratic_model.py
+++ b/dimod/binary_quadratic_model.py
@@ -1391,7 +1391,7 @@ class BinaryQuadraticModel(Sized, Container, Iterable):
 # conversions
 ##################################################################################################
 
-    def to_coo(self, fp=None):
+    def to_coo(self, fp=None, vartype_header=False):
         """Serialize the binary quadratic model to a COOrdinate_ format encoding.
 
         .. _COOrdinate: https://en.wikipedia.org/wiki/Sparse_matrix#Coordinate_list_(COO)
@@ -1403,6 +1403,10 @@ class BinaryQuadraticModel(Sized, Container, Iterable):
                 (i, j, bias), where :math:`i=j` for linear biases. If not provided,
                 returns a string.
 
+            vartype_header (bool, optional, default=False):
+                If true, the binary quadratic model's variable type as prepended to the
+                string or file as a header.
+
         .. _file object: https://docs.python.org/3/glossary.html#term-file-object
 
         .. note:: Variables must use index lables (numeric lables). Binary quadratic
@@ -1417,6 +1421,15 @@ class BinaryQuadraticModel(Sized, Container, Iterable):
                 0 1 0.50000
                 1 1 -1.50000
 
+            The Coordinate format with a header
+
+            .. code-block:: none
+
+                # vartype=SPIN
+                0 0 0.50000
+                0 1 0.50000
+                1 1 -1.50000
+
             This is an example of writing a binary quadratic model to a COOrdinate-format
             file.
 
@@ -1436,12 +1449,12 @@ class BinaryQuadraticModel(Sized, Container, Iterable):
         import dimod.io.coo as coo
 
         if fp is None:
-            return coo.dumps(self)
+            return coo.dumps(self, vartype_header)
         else:
-            coo.dump(self, fp)
+            coo.dump(self, fp, vartype_header)
 
     @classmethod
-    def from_coo(cls, obj, vartype):
+    def from_coo(cls, obj, vartype=None):
         """Deserialize a binary quadratic model from a COOrdinate_ format encoding.
 
         .. _COOrdinate: https://en.wikipedia.org/wiki/Sparse_matrix#Coordinate_list_(COO)
@@ -1453,12 +1466,15 @@ class BinaryQuadraticModel(Sized, Container, Iterable):
                 is stored as a list of 3-tuples, (i, j, bias), where :math:`i=j`
                 for linear biases.
 
-            vartype (:class:`.Vartype`/str/set):
+            vartype (:class:`.Vartype`/str/set, optional):
                 Variable type for the binary quadratic model. Accepted input values:
 
                 * :class:`.Vartype.SPIN`, ``'SPIN'``, ``{-1, 1}``
                 * :class:`.Vartype.BINARY`, ``'BINARY'``, ``{0, 1}``
 
+                If not provided, the vartype must be specified with a header in the
+                file.
+
         .. _file object: https://docs.python.org/3/glossary.html#term-file-object
 
         .. note:: Variables must use index lables (numeric lables). Binary quadratic
@@ -1466,10 +1482,19 @@ class BinaryQuadraticModel(Sized, Container, Iterable):
             zero.
 
         Examples:
-            This is an example of a binary quadratic model encoded in COOrdinate format.
+            An example of a binary quadratic model encoded in COOrdinate format.
+
+            .. code-block:: none
+
+                0 0 0.50000
+                0 1 0.50000
+                1 1 -1.50000
+
+            The Coordinate format with a header
 
             .. code-block:: none
 
+                # vartype=SPIN
                 0 0 0.50000
                 0 1 0.50000
                 1 1 -1.50000
@@ -1490,9 +1515,9 @@ class BinaryQuadraticModel(Sized, Container, Iterable):
         import dimod.io.coo as coo
 
         if isinstance(obj, str):
-            return coo.loads(obj, cls.empty(vartype))
+            return coo.loads(obj, cls=cls, vartype=vartype)
 
-        return coo.load(obj, cls.empty(vartype))
+        return coo.load(obj, cls=cls, vartype=vartype)
 
     def to_serializable(self, use_bytes=False):
         """Convert the binary quadratic model to a serializable object.
diff --git a/dimod/io/coo.py b/dimod/io/coo.py
index 870ffc13..74060bd3 100644
--- a/dimod/io/coo.py
+++ b/dimod/io/coo.py
@@ -18,7 +18,7 @@ import re
 
 from dimod.binary_quadratic_model import BinaryQuadraticModel
 
-_LINE_REGEX = r'^\s*(\d+)\s+(\d+)\s+([+-]?([0-9]*[.])?[0-9]+)\s*$'
+_LINE_REGEX = r'^\s*(\d+)\s+(\d+)\s+([+-]?(?:[0-9]*[.]):?[0-9]+)\s*$'
 """
 Each line should look like
 
@@ -27,58 +27,70 @@ Each line should look like
 where 0, 1 are the variable lables and 2.0 is the bias.
 """
 
+_VARTYPE_HEADER_REGEX = r'^[ \t\f]*#.*?vartype[:=][ \t]*([-_.a-zA-Z0-9]+)'
+"""
+The header should be in the first line and look like
+
+# vartype=SPIN
+
+"""
 
-def dumps(bqm):
+
+def dumps(bqm, vartype_header=False):
     """Dump a binary quadratic model to a string in COOrdinate format."""
-    return '\n'.join(_iter_triplets(bqm))
+    return '\n'.join(_iter_triplets(bqm, vartype_header))
 
 
-def dump(bqm, fp):
+def dump(bqm, fp, vartype_header=False):
     """Dump a binary quadratic model to a string in COOrdinate format."""
-    for triplet in _iter_triplets(bqm):
+    for triplet in _iter_triplets(bqm, vartype_header):
         fp.write('%s\n' % triplet)
 
 
-def loads(s, bqm):
+def loads(s, cls=BinaryQuadraticModel, vartype=None):
     """Load a COOrdinate formatted binary quadratic model from a string."""
-    pattern = re.compile(_LINE_REGEX)
-
-    for line in s.split('\n'):
-        match = pattern.search(line)
+    return load(s.split('\n'), cls=cls, vartype=vartype)
 
-        if match is not None:
-            u, v, bias = int(match.group(1)), int(match.group(2)), float(match.group(3))
-            if u == v:
-                bqm.add_variable(u, bias)
-            else:
-                bqm.add_interaction(u, v, bias)
 
-    return bqm
-
-
-def load(fp, bqm):
+def load(fp, cls=BinaryQuadraticModel, vartype=None):
     """Load a COOrdinate formatted binary quadratic model from a file."""
     pattern = re.compile(_LINE_REGEX)
+    vartype_pattern = re.compile(_VARTYPE_HEADER_REGEX)
 
+    triplets = []
     for line in fp:
-        match = pattern.search(line)
+        triplets.extend(pattern.findall(line))
+
+        vt = vartype_pattern.findall(line)
+        if vt:
+            if vartype is None:
+                vartype = vt[0]
+            elif vartype is not vt[0]:
+                raise ValueError("vartypes from headers and/or inputs do not match")
 
-        if match is not None:
-            u, v, bias = int(match.group(1)), int(match.group(2)), float(match.group(3))
-            if u == v:
-                bqm.add_variable(u, bias)
-            else:
-                bqm.add_interaction(u, v, bias)
+    if vartype is None:
+        raise ValueError("vartype must be provided either as a header or as an argument")
+
+    bqm = cls.empty(vartype)
+
+    for u, v, bias in triplets:
+        if u == v:
+            bqm.add_variable(int(u), float(bias))
+        else:
+            bqm.add_interaction(int(u), int(v), float(bias))
 
     return bqm
 
 
-def _iter_triplets(bqm):
+def _iter_triplets(bqm, vartype_header):
     if not isinstance(bqm, BinaryQuadraticModel):
         raise TypeError("expected input to be a BinaryQuadraticModel")
     if not all(isinstance(v, int) and v >= 0 for v in bqm.linear):
         raise ValueError("only positive index-labeled binary quadratic models can be dumped to COOrdinate format")
 
+    if vartype_header:
+        yield '# vartype=%s' % bqm.vartype.name
+
     # developer note: we could (for some threshold sparseness) sort the neighborhoods,
     # but this is simple and probably sufficient
 


Test patch:
diff --git a/tests/test_io_coo.py b/tests/test_io_coo.py
index dc38ca36..bc119f6f 100644
--- a/tests/test_io_coo.py
+++ b/tests/test_io_coo.py
@@ -40,17 +40,29 @@ class TestCOO(unittest.TestCase):
         contents = "0 0 1.000000\n0 1 2.000000\n2 3 0.400000"
         self.assertEqual(s, contents)
 
+    def test_dumps_sortable_SPIN_with_header(self):
+        bqm = dimod.BinaryQuadraticModel.from_ising({0: 1.}, {(0, 1): 2, (2, 3): .4})
+        s = coo.dumps(bqm, vartype_header=True)
+        contents = "# vartype=SPIN\n0 0 1.000000\n0 1 2.000000\n2 3 0.400000"
+        self.assertEqual(s, contents)
+
+    def test_dumps_sortable_BINARY_with_header(self):
+        bqm = dimod.BinaryQuadraticModel.from_qubo({(0, 0): 1., (0, 1): 2, (2, 3): .4})
+        s = coo.dumps(bqm, vartype_header=True)
+        contents = "# vartype=BINARY\n0 0 1.000000\n0 1 2.000000\n2 3 0.400000"
+        self.assertEqual(s, contents)
+
     def test_load(self):
         filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data', 'coo_qubo.qubo')
 
         with open(filepath, 'r') as fp:
-            bqm = coo.load(fp, dimod.BinaryQuadraticModel.empty(dimod.BINARY))
+            bqm = coo.load(fp, dimod.BinaryQuadraticModel, dimod.BINARY)
 
         self.assertEqual(bqm, dimod.BinaryQuadraticModel.from_qubo({(0, 0): -1, (1, 1): -1, (2, 2): -1, (3, 3): -1}))
 
     def test_loads(self):
         contents = "0 0 1.000000\n0 1 2.000000\n2 3 0.400000"
-        bqm = coo.loads(contents, dimod.BinaryQuadraticModel.empty(dimod.SPIN))
+        bqm = coo.loads(contents, dimod.BinaryQuadraticModel, dimod.SPIN)
         self.assertEqual(bqm, dimod.BinaryQuadraticModel.from_ising({0: 1.}, {(0, 1): 2, (2, 3): .4}))
 
     def test_functional_file_empty_BINARY(self):
@@ -64,7 +76,7 @@ class TestCOO(unittest.TestCase):
             coo.dump(bqm, fp=file)
 
         with open(filename, 'r') as file:
-            new_bqm = coo.load(file, dimod.BinaryQuadraticModel.empty(dimod.BINARY))
+            new_bqm = coo.load(file, dimod.BinaryQuadraticModel, dimod.BINARY)
 
         shutil.rmtree(tmpdir)
 
@@ -81,7 +93,7 @@ class TestCOO(unittest.TestCase):
             coo.dump(bqm, fp=file)
 
         with open(filename, 'r') as file:
-            new_bqm = coo.load(file, dimod.BinaryQuadraticModel.empty(dimod.SPIN))
+            new_bqm = coo.load(file, dimod.BinaryQuadraticModel, dimod.SPIN)
 
         shutil.rmtree(tmpdir)
 
@@ -98,7 +110,7 @@ class TestCOO(unittest.TestCase):
             coo.dump(bqm, fp=file)
 
         with open(filename, 'r') as file:
-            new_bqm = coo.load(file, dimod.BinaryQuadraticModel.empty(dimod.BINARY))
+            new_bqm = coo.load(file, dimod.BinaryQuadraticModel, dimod.BINARY)
 
         shutil.rmtree(tmpdir)
 
@@ -115,7 +127,7 @@ class TestCOO(unittest.TestCase):
             coo.dump(bqm, fp=file)
 
         with open(filename, 'r') as file:
-            new_bqm = coo.load(file, dimod.BinaryQuadraticModel.empty(dimod.SPIN))
+            new_bqm = coo.load(file, dimod.BinaryQuadraticModel, dimod.SPIN)
 
         shutil.rmtree(tmpdir)
 
@@ -126,7 +138,7 @@ class TestCOO(unittest.TestCase):
         bqm = dimod.BinaryQuadraticModel.empty(dimod.BINARY)
 
         s = coo.dumps(bqm)
-        new_bqm = coo.loads(s, dimod.BinaryQuadraticModel.empty(dimod.BINARY))
+        new_bqm = coo.loads(s, dimod.BinaryQuadraticModel, dimod.BINARY)
 
         self.assertEqual(bqm, new_bqm)
 
@@ -135,7 +147,7 @@ class TestCOO(unittest.TestCase):
         bqm = dimod.BinaryQuadraticModel.empty(dimod.SPIN)
 
         s = coo.dumps(bqm)
-        new_bqm = coo.loads(s, dimod.BinaryQuadraticModel.empty(dimod.SPIN))
+        new_bqm = coo.loads(s, dimod.BinaryQuadraticModel, dimod.SPIN)
 
         self.assertEqual(bqm, new_bqm)
 
@@ -144,7 +156,7 @@ class TestCOO(unittest.TestCase):
         bqm = dimod.BinaryQuadraticModel({0: 1.}, {(0, 1): 2, (2, 3): .4}, 0.0, dimod.BINARY)
 
         s = coo.dumps(bqm)
-        new_bqm = coo.loads(s, dimod.BinaryQuadraticModel.empty(dimod.BINARY))
+        new_bqm = coo.loads(s, dimod.BinaryQuadraticModel, dimod.BINARY)
 
         self.assertEqual(bqm, new_bqm)
 
@@ -153,6 +165,28 @@ class TestCOO(unittest.TestCase):
         bqm = dimod.BinaryQuadraticModel({0: 1.}, {(0, 1): 2, (2, 3): .4}, 0.0, dimod.SPIN)
 
         s = coo.dumps(bqm)
-        new_bqm = coo.loads(s, dimod.BinaryQuadraticModel.empty(dimod.SPIN))
+        new_bqm = coo.loads(s, dimod.BinaryQuadraticModel, dimod.SPIN)
 
         self.assertEqual(bqm, new_bqm)
+
+    def test_functional_SPIN_vartypeheader(self):
+        bqm = dimod.BinaryQuadraticModel({0: 1.}, {(0, 1): 2, (2, 3): .4}, 0.0, dimod.SPIN)
+
+        s = coo.dumps(bqm, vartype_header=True)
+        new_bqm = coo.loads(s, dimod.BinaryQuadraticModel)
+
+        self.assertEqual(bqm, new_bqm)
+
+    def test_no_vartype(self):
+        bqm = dimod.BinaryQuadraticModel({0: 1.}, {(0, 1): 2, (2, 3): .4}, 0.0, dimod.SPIN)
+
+        s = coo.dumps(bqm)
+        with self.assertRaises(ValueError):
+            coo.loads(s, dimod.BinaryQuadraticModel)
+
+    def test_conflicting_vartype(self):
+        bqm = dimod.BinaryQuadraticModel({0: 1.}, {(0, 1): 2, (2, 3): .4}, 0.0, dimod.SPIN)
+
+        s = coo.dumps(bqm, vartype_header=True)
+        with self.assertRaises(ValueError):
+            coo.loads(s, dimod.BinaryQuadraticModel, dimod.BINARY)