From 29895b6a63b525147b902b18607bca168c7c1e5e Mon Sep 17 00:00:00 2001
From: arnaudon <alexis.arnaudon@epfl.ch>
Date: Wed, 30 Oct 2024 09:35:13 +0100
Subject: [PATCH] more

---
 netsalt/modes.py       | 10 ++++++++--
 netsalt/non_abelian.py | 20 +++++++++++---------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/netsalt/modes.py b/netsalt/modes.py
index cb8220b..7eb8bcf 100644
--- a/netsalt/modes.py
+++ b/netsalt/modes.py
@@ -281,11 +281,17 @@ def mode_on_nodes(mode, graph):
         )
     # this 1.5 is fairly arbitrary, it is so get similar igenvalues, close to 0, regardless of the
     # choice of quality_threshold, which may pick up other small ones, not 0 if set to large values
-    n_eig = len([1 for a in np.abs(min_eigenvalue) < 1.5 * min(np.abs(min_eigenvalue)) if a])
+    mask = np.logical_or(
+        np.abs(min_eigenvalue) < 1.5 * min(np.abs(min_eigenvalue)),
+        np.abs(min_eigenvalue) < 1e-4,
+    )
+
+    n_eig = len([1 for a in mask if a])
+    print(min_eigenvalue, np.abs(min_eigenvalue))
     if n_eig > 1:
         L.info(f"We found {n_eig} vanishing eigenvalues, we will use the sum of their eigenvectors")
         print(f"We found {n_eig} vanishing eigenvalues, we will use the sum of their eigenvectors")
-    return node_solution[:, np.abs(min_eigenvalue) < 1.5 * min(np.abs(min_eigenvalue))].sum(axis=1)
+    return node_solution[:, mask].sum(axis=1)
 
 
 def flux_on_edges(mode, graph):
diff --git a/netsalt/non_abelian.py b/netsalt/non_abelian.py
index 00d303f..b6745bd 100644
--- a/netsalt/non_abelian.py
+++ b/netsalt/non_abelian.py
@@ -48,14 +48,15 @@ def Ad(chi_mat):
 
 def set_so3_wavenumber(graph, wavenumber):
     """Set so3 matrix wavenumber."""
-    chis = [graph[u][v].get("chi", None) for u, v in graph.edges]
+    chis = np.array([graph[u][v].get("chi", None) for u, v in graph.edges])
     if chis[0] is None:
         chi = hat_inv(np.array([0.0, 0.0, 1.0]))
         chis = np.array(len(graph.edges) * [chi])
     else:
         if len(np.shape(chis[0])) == 1:
             chis = np.array([hat_inv(chi) for chi in chis])
-    graph.graph["ks"] = chis * np.real(wavenumber) - np.eye(3) * np.imag(wavenumber)
+    chi_loss = np.array([proj_perp(chi) for chi in chis])
+    graph.graph["ks"] = chis * np.real(wavenumber) - chi_loss * np.imag(wavenumber)
     graph.graph["chis"] = chis
     graph.graph["wavenumber"] = wavenumber
 
@@ -69,14 +70,15 @@ def _ext(i):
     B = sparse.lil_matrix((len(graph.edges) * 2 * DIM, len(graph) * DIM), dtype=np.complex128)
     BT = sparse.lil_matrix((len(graph) * DIM, len(graph.edges) * 2 * DIM), dtype=np.complex128)
     for ei, (u, v) in enumerate(graph.edges):
+
+        k = graph.graph["ks"][ei]
+        chi = graph.graph["chis"][ei]
+        length = graph.graph["lengths"][ei]
+
         one = np.eye(DIM)
-        expl = Ad(graph.graph["lengths"][ei] * graph.graph["ks"][ei])
-        expl = np.array(expl.dot(proj_perp(graph.graph["chis"][ei])), dtype=np.complex128)
-        expl += (
-            abelian_scale
-            * np.exp(1.0j * graph.graph["lengths"][ei] * graph.graph["wavenumber"])
-            * proj_paral(graph.graph["chis"][ei])
-        )
+        expl = Ad(length * k)
+        expl = np.array(expl.dot(proj_perp(chi)), dtype=np.complex128)
+        expl += abelian_scale * np.exp(1.0j * length * graph.graph["wavenumber"]) * proj_paral(chi)
 
         B[_ext(2 * ei), _ext(u)] = -one
         B[_ext(2 * ei), _ext(v)] = expl