Browse code

use vectorized system interface more idiomatically in tests

This is a pure refactor.

Joseph Weston authored on 04/12/2019 18:25:34
Showing 2 changed files
... ...
@@ -2025,7 +2025,7 @@ class _VectorizedFinalizedBuilderMixin(_FinalizedBuilderMixin):
2025 2025
             "elements at once using 'hamiltonian_term'.",
2026 2026
             KwantDeprecationWarning
2027 2027
         )
2028
-        site_offsets = np.cumsum([0] + [len(s) for s in self.site_arrays])
2028
+        site_offsets, _, _ = self.site_ranges.transpose()
2029 2029
         if i == j:
2030 2030
             which_term = self._onsite_term_by_site_id[i]
2031 2031
             (w, _), (off, _) = self.subgraphs[self.terms[which_term].subgraph]
... ...
@@ -321,10 +321,10 @@ def random_hopping_integral(rng):
321 321
 
322 322
 
323 323
 def check_onsite(fsyst, sites, subset=False, check_values=True):
324
-    vectorized = isinstance(fsyst, (system.FiniteVectorizedSystem, system.InfiniteVectorizedSystem))
324
+    vectorized = system.is_vectorized(fsyst)
325 325
 
326 326
     if vectorized:
327
-        site_offsets = np.cumsum([0] + [len(s) for s in fsyst.site_arrays])
327
+        site_offsets, _, _ = fsyst.site_ranges.transpose()
328 328
 
329 329
     freq = {}
330 330
     for node in range(fsyst.graph.num_nodes):
... ...
@@ -353,10 +353,10 @@ def check_onsite(fsyst, sites, subset=False, check_values=True):
353 353
 
354 354
 
355 355
 def check_hoppings(fsyst, hops):
356
-    vectorized = isinstance(fsyst, (system.FiniteVectorizedSystem, system.InfiniteVectorizedSystem))
356
+    vectorized = system.is_vectorized(fsyst)
357 357
 
358 358
     if vectorized:
359
-        site_offsets = np.cumsum([0] + [len(s) for s in fsyst.site_arrays])
359
+        site_offsets, _, _ = fsyst.site_ranges.transpose()
360 360
 
361 361
     assert fsyst.graph.num_edges == 2 * len(hops)
362 362
     for edge_id, edge in enumerate(fsyst.graph):