Add setter for Basis.four_element_traces (#16)

* Add setter for Basis.four_element_traces() * Add test for setter
qutech · Mar 5, 2020 · 0403091 · 0403091
1 parent 5d52491
commit 0403091
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/filter_functions/basis.py b/filter_functions/basis.py
@@ -368,6 +368,10 @@ def four_element_traces(self) -> COO:
 
         return self._four_element_traces
 
+    @four_element_traces.setter
+    def four_element_traces(self, traces):
+        self._four_element_traces = traces
+
     def normalize(self) -> None:
         """Normalize the basis in-place"""
         if self.ndim == 2:

diff --git a/tests/test_basis.py b/tests/test_basis.py
@@ -111,10 +111,14 @@ def test_basis_properties(self):
 
             if base.d < 8:
                 # Test very resource intense
+                ref = np.einsum('iab,jbc,kcd,lda', *(base,)*4)
                 self.assertArrayAlmostEqual(base.four_element_traces.todense(),
-                                            np.einsum('iab,jbc,kcd,lda',
-                                                      *(base,)*4),
-                                            atol=1e-16)
+                                            ref, atol=1e-16)
+
+                # Test setter
+                base._four_element_traces = None
+                base.four_element_traces = ref
+                self.assertArrayEqual(base.four_element_traces, ref)
 
             base._print_checks()