@@ -2387,7 +2387,7 @@ def _sort_hopping_term(term, value):

     return term, value

-def _make_onsite_terms(builder, sites, site_offsets, term_offset):

+def _make_onsite_terms(builder, sites, site_arrays, term_offset):

     # Construct onsite terms.

     #

     # onsite_subgraphs

@@ -2409,6 +2409,8 @@ def _make_onsite_terms(builder, sites, site_offsets, term_offset):

     #   Hamiltonian term associated with each site/hopping. For

     #   Hermitian conjugate hoppings "-term - 1" is stored instead.

+    site_offsets = np.cumsum([0] + [len(sa) for sa in site_arrays])

+

     onsite_subgraphs = []

     onsite_term_values = []

     onsite_term_parameters = []

@@ -2479,7 +2481,7 @@ def _make_onsite_terms(builder, sites, site_offsets, term_offset):

             onsite_term_errors, _onsite_term_by_site_id)

-def _make_hopping_terms(builder, graph, sites, site_offsets, cell_size, term_offset):

+def _make_hopping_terms(builder, graph, sites, site_arrays, cell_size, term_offset):

     # Construct the hopping terms

     #

     # The logic is the same as for the onsite terms, with the following

@@ -2490,6 +2492,8 @@ def _make_hopping_terms(builder, graph, sites, site_offsets, cell_size, term_off

     #   is a part of. For Hermitian conjugate hoppings "-term_number -1"

     #   is stored instead.

+    site_offsets = np.cumsum([0] + [len(sa) for sa in site_arrays])

+

     hopping_subgraphs = []

     hopping_term_values = []

     hopping_term_parameters = []

@@ -2628,17 +2632,14 @@ class FiniteVectorizedSystem(_VectorizedFinalizedBuilderMixin, system.FiniteVect

         del id_by_site  # cleanup due to large size

         site_arrays = _make_site_arrays(builder.H)

-        # We need this to efficiently find which array a given

-        # site belongs to

-        site_offsets = np.cumsum([0] + [len(arr) for arr in site_arrays])

         (onsite_subgraphs, onsite_terms, onsite_term_values,

          onsite_term_errors, _onsite_term_by_site_id) =\

-            _make_onsite_terms(builder, sites, site_offsets, term_offset=0)

+            _make_onsite_terms(builder, sites, site_arrays, term_offset=0)

         (hopping_subgraphs, hopping_terms, hopping_term_values,

          hopping_term_errors, _hopping_term_by_edge_id) =\

-            _make_hopping_terms(builder, graph, sites, site_offsets,

+            _make_hopping_terms(builder, graph, sites, site_arrays,

                                 len(sites), term_offset=len(onsite_terms))

         # Construct the combined onsite/hopping term datastructures

@@ -2941,15 +2942,13 @@ class InfiniteVectorizedSystem(_VectorizedFinalizedBuilderMixin, system.Infinite

             + _make_site_arrays(interface)

         )

-        site_offsets = np.cumsum([0] + [len(arr) for arr in site_arrays])

-

         (onsite_subgraphs, onsite_terms, onsite_term_values,

          onsite_term_errors, _onsite_term_by_site_id) =\

-            _make_onsite_terms(builder, sites, site_offsets, term_offset=0)

+            _make_onsite_terms(builder, sites, site_arrays, term_offset=0)

         (hopping_subgraphs, hopping_terms, hopping_term_values,

          hopping_term_errors, _hopping_term_by_edge_id) =\

-            _make_hopping_terms(builder, graph, sites, site_offsets,

+            _make_hopping_terms(builder, graph, sites, site_arrays,

                                 cell_size, term_offset=len(onsite_terms))

         # Construct the combined onsite/hopping term datastructures