diff --git a/kwant/builder.py b/kwant/builder.py index 01fbc07817ff0afa9ae03464404dfc751e981910..8f189ede67d0aaa1726196256b2dfc56a7070ce3 100644 --- a/kwant/builder.py +++ b/kwant/builder.py @@ -2153,7 +2153,7 @@ class _VectorizedFinalizedBuilderMixin(_FinalizedBuilderMixin): except Exception as exc: _raise_user_error(exc, val) - ham = _normalize_term_value(ham, len(to_site_array)) + ham = system._normalize_matrix_blocks(ham, len(to_site_array)) return ham @@ -2392,33 +2392,11 @@ class _IdBySite: raise KeyError(site) -def _normalize_term_value(value, expected_n_values): - try: - value = np.asarray(value, dtype=complex) - except TypeError: - raise ValueError( - "Matrix elements declared with incompatible shapes." - ) from None - # Upgrade to vector of matrices - if len(value.shape) == 1: - value = value[:, np.newaxis, np.newaxis] - if len(value.shape) != 3: - msg = ( - "Vectorized value functions must return an array of" - "scalars or an array of matrices." - ) - raise ValueError(msg) - if value.shape[0] != expected_n_values: - raise ValueError("Value functions must return a single value per " - "onsite/hopping.") - return value - - def _sort_term(term, value): term = np.asarray(term) if not callable(value): - value = _normalize_term_value(value, len(term)) + value = system._normalize_matrix_blocks(value, len(term)) # Ensure that values still correspond to the correct # sites in 'term' once the latter has been sorted. value = value[term.argsort()] diff --git a/kwant/system.py b/kwant/system.py index f25761bfd36c6aff708d396ae4f1c0cdd1cc6873..6e77ea67d0bba7bca40ebd9b93a5ec4cbdaccdfb 100644 --- a/kwant/system.py +++ b/kwant/system.py @@ -742,6 +742,36 @@ def is_vectorized(syst): return isinstance(syst, (FiniteVectorizedSystem, InfiniteVectorizedSystem)) +def _normalize_matrix_blocks(matrix_blocks, expected_length): + """Normalize a sequence of matrices into a single 3D numpy array + + Parameters + ---------- + matrix_blocks : sequence of complex array-like + expected_length : int + """ + try: + matrix_blocks = np.asarray(matrix_blocks, dtype=complex) + except TypeError: + raise ValueError( + "Matrix elements declared with incompatible shapes." + ) from None + # Upgrade to vector of matrices + if len(matrix_blocks.shape) == 1: + matrix_blocks = matrix_blocks[:, np.newaxis, np.newaxis] + if len(matrix_blocks.shape) != 3: + msg = ( + "Vectorized value functions must return an array of" + "scalars or an array of matrices." + ) + raise ValueError(msg) + if matrix_blocks.shape[0] != expected_length: + raise ValueError("Value functions must return a single value per " + "onsite/hopping.") + return matrix_blocks + + + class PrecalculatedLead: def __init__(self, modes=None, selfenergy=None): """A general lead defined by its self energy.