Repository: dwavesystems/dimod
Pull request #1341 (2023-06-23 18:53:37)

Magic square generator
Adding a (normal) magic square generator to the generators package. 

[Magic squares](https://en.wikipedia.org/wiki/Magic_square) are quadratic constrained integer problems of size $n^2$ with no objective. 


Patch:
diff --git a/dimod/generators/__init__.py b/dimod/generators/__init__.py
index dc93c19a..f5f61d8d 100644
--- a/dimod/generators/__init__.py
+++ b/dimod/generators/__init__.py
@@ -24,3 +24,4 @@ from dimod.generators.knapsack import *
 from dimod.generators.binpacking import *
 from dimod.generators.multi_knapsack import *
 from dimod.generators.satisfiability import *
+from dimod.generators.magic_square import *
\ No newline at end of file
diff --git a/dimod/generators/magic_square.py b/dimod/generators/magic_square.py
new file mode 100644
index 00000000..26e12f4d
--- /dev/null
+++ b/dimod/generators/magic_square.py
@@ -0,0 +1,110 @@
+# Copyright 2023 D-Wave Systems Inc.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+from itertools import product
+
+from dimod.constrained import ConstrainedQuadraticModel
+from dimod.binary.binary_quadratic_model import quicksum
+from dimod.quadratic import Integer
+
+__all__ = ["magic_square"]
+
+
+def magic_square(size: int, power: int = 1) -> ConstrainedQuadraticModel:
+    """Generate a magic square of a particular size.
+    
+    Args:
+        size (int):
+            The sidelength of the magic square, the magic square will
+            have size * size integer variables.
+        power (int, optional):
+            If set to 1 the problem is a normal magic square, if set
+            to 2 it is a magic square of squares. Must be 1 or 2.
+            Defaults to 1.
+
+    Returns:
+        dimod.ConstrainedQuadraticModel:
+            A coonstrained quadratic model which represents a magic square.
+            The variables are labeled by "var_i_j" where i and j are two
+            non-negative integers which represent their column and row
+            respectively. If power is greater than 1 then the value of the
+            variable represents the base, so the magic square is satisfied
+            by var_i_j**power .
+    """
+    if power not in [1, 2]:
+        raise ValueError(f"power must be either 1 or 2, recieved {power}")
+    
+    variables = {}
+    
+    for i, j in product(range(size), range(size)):
+        variables[i, j] = Integer(f"var_{i}_{j}", lower_bound=1)
+    
+    # the sum across every row, column, diagonal
+    constraint_sum = Integer("sum", lower_bound=1)
+    
+    cqm = ConstrainedQuadraticModel()
+    
+    for i in range(size):
+        
+        if power == 1:
+            # Every row must equal the same value
+            cqm.add_constraint_from_comparison(
+                quicksum(variables[i, j] for j in range(size)) -
+                constraint_sum == 0, label=f"row_{i}")
+
+            # Every column must equal the same value
+            cqm.add_constraint_from_comparison(
+                    quicksum(variables[j, i] for j in range(size)) -
+                    constraint_sum == 0, label=f"col_{i}")
+        else:
+            # Every row must equal the same value
+            cqm.add_constraint_from_comparison(
+                quicksum(variables[i, j]**power for j in range(size)) -
+                constraint_sum == 0, label=f"row_{i}")
+
+            # Every column must equal the same value
+            cqm.add_constraint_from_comparison(
+                    quicksum(variables[j, i]**power for j in range(size)) -
+                    constraint_sum == 0, label=f"col_{i}")
+
+    if power == 1:
+        # the diagonals must equal the same value
+        cqm.add_constraint_from_comparison(
+            quicksum(variables[i, i] for i in range(size)) -
+            constraint_sum == 0, label="diagonal")
+        
+        cqm.add_constraint_from_comparison(
+            quicksum(variables[i, size-1-i] for i in range(size)) -
+            constraint_sum == 0, label="antidiagonal")
+    else:
+        # the diagonals must equal the same value
+        cqm.add_constraint_from_comparison(
+            quicksum(variables[i, i]**power for i in range(size)) -
+            constraint_sum == 0, label="diagonal")
+        
+        cqm.add_constraint_from_comparison(
+            quicksum(variables[i, size-1-i]**power for i in range(size)) -
+            constraint_sum == 0, label="antidiagonal")
+    
+    # every integer must be unique
+    cqm.add_constraint_from_comparison(
+        quicksum(
+            variables[i, j]**2 +
+            variables[k, l]**2 -
+            2*variables[i, j]*variables[k, l]
+            for i, j, k, l in product(range(size), repeat=4)
+            if ((k > i and l == j) or (l > j)))
+        >= (size**4 - size**2)/2, label="uniqueness")
+    
+    return cqm
diff --git a/releasenotes/notes/cqm-magic-square-generator-36a9de45609ce0be.yaml b/releasenotes/notes/cqm-magic-square-generator-36a9de45609ce0be.yaml
new file mode 100644
index 00000000..ee420604
--- /dev/null
+++ b/releasenotes/notes/cqm-magic-square-generator-36a9de45609ce0be.yaml
@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    Added `dimod.generators.magic_square` which allows for the generation of a `ConstrainedQuadaticModel` whos feasible 
+    solutions are normal magic squares. The function also generated magic squares of squares when passed `power=2` as a key-word argument.


Test patch:
diff --git a/tests/test_generators.py b/tests/test_generators.py
index 48192467..78daef33 100644
--- a/tests/test_generators.py
+++ b/tests/test_generators.py
@@ -746,6 +746,7 @@ class TestGates(unittest.TestCase):
              best.sample['p3']]
         self.assertEqual(p, [1, 0, 0, 1])
 
+
 class TestInteger(unittest.TestCase):
     def test_exceptions(self):
         with self.assertRaises(ValueError):
@@ -1030,3 +1031,91 @@ class TestSatisfiability(unittest.TestCase):
     def test_labels(self):
         self.assertEqual(dimod.generators.random_2in4sat(10, 1).variables, range(10))
         self.assertEqual(dimod.generators.random_2in4sat('abdef', 1).variables, 'abdef')
+
+
+class TestMagicSquares(unittest.TestCase):
+    def test_invalid_power(self):
+        for power in [-1, 0, 3, 5, 10]:
+            with self.assertRaises(ValueError):
+                dimod.generators.magic_square(1, power)
+
+    def test_size(self):
+        for size in [2**i for i in range(5)]:
+            magic_square = dimod.generators.magic_square(size)
+            
+            self.assertEqual(len(magic_square.constraints), 2*size + 3)
+            self.assertEqual(magic_square.num_variables(), size**2 + 1)
+    
+    def test_constraints(self):
+        
+        magic_square = dimod.generators.magic_square(5)
+        
+        constraint_names = magic_square.constraints.keys()
+        
+        for name in constraint_names:
+            if name != "uniqueness":
+                self.assertEqual(magic_square.constraints[name].rhs, 0)
+                
+                constraint_lhs = magic_square.constraints[name].lhs
+                
+                for name, term in constraint_lhs.linear.items():
+                    if name != "sum":
+                        self.assertEqual(term, 1)
+                    else:
+                        self.assertEqual(term, -1)
+                
+                for name, term in constraint_lhs.quadratic.items():
+                    self.assertEqual(term, 1)
+                
+            else:
+                self.assertEqual(magic_square.constraints[name].rhs, 300)
+                
+                constraint_lhs = magic_square.constraints[name].lhs
+                
+                for name, term in constraint_lhs.linear.items():
+                    self.assertEqual(term, 0)
+                    
+                self_squares = [(i, i) for i in magic_square.variables]
+                
+                for name, term in constraint_lhs.quadratic.items():
+                    if name not in self_squares:
+                        self.assertEqual(term, -2)
+                    else:
+                        self.assertEqual(term, 24)
+    
+    def test_constraints_squares(self):
+        
+        magic_square = dimod.generators.magic_square(5, power=2)
+        
+        constraint_names = magic_square.constraints.keys()
+        
+        for name in constraint_names:
+            if name != "uniqueness":
+                self.assertEqual(magic_square.constraints[name].rhs, 0)
+                
+                constraint_lhs = magic_square.constraints[name].lhs
+                
+                for name, term in constraint_lhs.linear.items():
+                    if name != "sum":
+                        self.assertEqual(term, 0)
+                    else:
+                        self.assertEqual(term, -1)
+                  
+                for name, term in constraint_lhs.quadratic.items():
+                    self.assertEqual(term, 1)
+                
+            else:
+                self.assertEqual(magic_square.constraints[name].rhs, 300)
+                
+                constraint_lhs = magic_square.constraints[name].lhs
+
+                for name, term in constraint_lhs.linear.items():
+                    self.assertEqual(term, 0)
+                    
+                self_squares = [(i, i) for i in magic_square.variables]
+                
+                for name, term in constraint_lhs.quadratic.items():
+                    if name not in self_squares:
+                        self.assertEqual(term, -2)
+                    else:
+                        self.assertEqual(term, 24)