From 9c4bf187a929336982d5ec4b70ad8fac82785cba Mon Sep 17 00:00:00 2001
From: Johanna <johanna@zijderveld.de>
Date: Wed, 27 Mar 2024 23:43:16 +0100
Subject: [PATCH] add minimal example and function for hopping dict to flat and
 back

---
 examples/dict_to_flat.py | 41 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 examples/dict_to_flat.py

diff --git a/examples/dict_to_flat.py b/examples/dict_to_flat.py
new file mode 100644
index 0000000..a58efe1
--- /dev/null
+++ b/examples/dict_to_flat.py
@@ -0,0 +1,41 @@
+# %%
+import numpy as np
+from codes.kwant_helper import utils
+from codes import kwant_examples
+
+# %%
+# Example hopping dictionary to use:
+graphene_builder, int_builder = kwant_examples.graphene_extended_hubbard()
+tb_model = utils.builder2tb_model(graphene_builder)
+
+
+# %%
+def hop_dict_to_flat(hop_dict):
+    sorted_vals = np.array(list(hop_dict.values()))[
+        np.lexsort(np.array(list(hop_dict.keys())).T)
+    ]
+    flat = sorted_vals[..., *np.triu_indices(sorted_vals.shape[-1])].flatten()
+    return flat
+
+
+def flat_to_hop_dict(flat, shape, hop_dict_keys):
+
+    matrix = np.zeros(shape, dtype=complex)
+    matrix[..., *np.triu_indices(shape[-1])] = flat.reshape(*shape[:-2], -1)
+    indices = np.arange(shape[-1])
+    diagonal = matrix[..., indices, indices]
+    matrix += np.moveaxis(matrix[-1::-1], -1, -2).conj()
+    matrix[..., indices, indices] -= diagonal
+
+    hop_dict_keys = np.array(list(hop_dict_keys))
+    sorted_keys = hop_dict_keys[np.lexsort(hop_dict_keys.T)]
+    hop_dict = dict(zip(map(tuple, sorted_keys), matrix))
+    return hop_dict
+
+
+# %%
+flat = hop_dict_to_flat(tb_model)
+shape = (len(tb_model.keys()), *list(tb_model.values())[0].shape)
+hop_dict = flat_to_hop_dict(flat, shape, tb_model.keys())
+
+# %%
-- 
GitLab