more

examples/non_abelian/test.py

@@ -4,31 +4,47 @@
 from netsalt.quantum_graph import create_quantum_graph, laplacian_quality, mode_quality
 from netsalt.modes import mode_on_nodes
 
+from netsalt.modes import scan_frequencies
+from netsalt.plotting import plot_scan
 from netsalt.physics import dispersion_relation_linear, set_dispersion_relation
 import networkx as nx
 
 
-from netsalt.non_abelian import construct_so3_laplacian, so3_mode_on_nodes
+from netsalt.non_abelian import construct_so3_laplacian, so3_mode_on_nodes, scan_frequencies_so3
 
 
-def make_graph():
+def make_graph(n):
     graph = nx.cycle_graph(n)
     graph.add_edge(0, 8)
     graph.add_edge(0, 20)
     graph.add_edge(10, 15)
     x = np.linspace(0, 2 * np.pi * (1 - 1.0 / (len(graph) - 1)), len(graph))
     pos = np.array([np.cos(x), np.sin(x)]).T
     pos = list(pos)
+    graph.add_edge(1, n)
+    graph.add_edge(15, n + 1)
+    pos += [[1.2, 0]]
+    pos += [[-1.2, 0]]
 
     return graph, pos
 
 
 if __name__ == "__main__":
 
-    params = {"open_model": "open", "c": 1.0}
+    params = {
+        "open_model": "open",
+        "c": 1.0,
+        "k_min": 10.0,
+        "k_max": 12.0,
+        "k_n": 200,
+        "alpha_min": 0.0,
+        "alpha_max": 0.4,
+        "alpha_n": 50,
+        "n_workers": 7,
+    }
     n = 30
-    graph, pos = make_graph()
-    graph_u1, pos = make_graph()
+    graph, pos = make_graph(n)
+    graph_u1, pos = make_graph(n)
 
     nx.draw(graph, pos=pos)
 
@@ -37,12 +53,23 @@ def make_graph():
 
     set_dispersion_relation(graph_u1, dispersion_relation_linear)
 
-    ks = np.linspace(1.0, 3, 200)
+    qualities_u1 = scan_frequencies(graph_u1)
+    plot_scan(graph_u1, qualities_u1)
+    plt.suptitle('u1')
+
+    qualities = scan_frequencies_so3(graph)
+    plot_scan(graph, qualities)
+    plt.suptitle("so3")
+    plt.show()
+
+
+def lkj():
+    ks = np.linspace(10.0, 15, 500)
     qs = []
     qs_u1 = []
     for k in tqdm(ks):
-        kim = 0  # 0.0966j
-        L = construct_so3_laplacian(k - kim, graph, abelian_scale=1.0)
+        kim = 0.05
+        L = construct_so3_laplacian(k + 1j * kim, graph)
         qs.append(laplacian_quality(L))
         qs_u1.append(mode_quality([k, kim], graph_u1))
 
@@ -53,7 +80,6 @@ def make_graph():
     plt.yscale("log")
     plt.show()
 
-
     k = ks[np.argmin(qs)]
     L = construct_so3_laplacian(k, graph)
     mode = so3_mode_on_nodes(L)

netsalt/non_abelian.py

@@ -1,8 +1,12 @@
 """Module for non-abelian quantum graphs."""
 import numpy as np
 import scipy as sc
+from tqdm import tqdm
 
+import multiprocessing
 from scipy import sparse, linalg
+from netsalt.quantum_graph import laplacian_quality
+from netsalt.utils import get_scan_grid, to_complex
 
 
 def hat_inv(xi_vec):
@@ -46,16 +50,15 @@ def Ad(chi_mat):
 
 def set_so3_wavenumber(graph, wavenumber, chis=None):
     if chis is None:
-        x = np.array([0.1, 0.5, 1.0])
+        x = np.array([0.0, 0.0, 1.0])
         chi_vec = wavenumber * x / np.linalg.norm(x)
         chi_mat = hat_inv(chi_vec)
         graph.graph["ks"] = len(graph.edges) * [chi_mat]
 
-        x2 = np.array([1.0, 0.5, 0.0])
+        x2 = np.array([1.0, 0.0, 0.0])
         chi_vec2 = wavenumber * x2 / np.linalg.norm(x2)
         chi_mat2 = hat_inv(chi_vec2)
         graph.graph["ks"][:10] = 10 * [chi_mat2]
-        graph.graph["ks"][20:30] = 10 * [chi_mat2]
     else:
         graph.graph["ks"] = chis
 
@@ -72,12 +75,13 @@ def _ext(i):
 
         one = np.eye(dim)
         expl = Ad(graph.graph["lengths"][ei] * graph.graph["ks"][ei])
-        #expl = np.array(expl.dot(proj_perp(graph.graph["ks"][ei])), dtype=np.complex128)
+        expl = np.array(expl.dot(proj_perp(graph.graph["ks"][ei])), dtype=np.complex128)
         expl += (
             abelian_scale
             * proj_paral(graph.graph["ks"][ei])
             * np.exp(1.0j * graph.graph["lengths"][ei] * norm(graph.graph["ks"][ei]))
         )
+
         out = True if (len(graph[u]) == 1 or len(graph[v]) == 1) else False
 
         Bout[_ext(2 * ei), _ext(u)] = -one
@@ -113,8 +117,7 @@ def _ext(i):
         winv = linalg.inv(w)
 
         if with_k:
-            # winv = (chi.dot(proj_perp(chi)) + 1.0j * norm(chi) * proj_paral(chi)).dot(winv)
-            winv = (chi + 1.0j * norm(chi) * proj_paral(chi)).dot(winv)
+            winv = (chi.dot(proj_perp(chi)) + 1.0j * norm(chi) * proj_paral(chi)).dot(winv)
 
         Winv[_ext(2 * ei), _ext(2 * ei)] = winv
         Winv[_ext(2 * ei + 1), _ext(2 * ei + 1)] = winv
@@ -143,3 +146,41 @@ def so3_mode_on_nodes(laplacian, quality_thresh=1e2):
         )
 
     return node_solution[:, 0]
+
+
+class WorkerScan:
+    """Worker to scan complex frequency."""
+
+    def __init__(self, graph, quality_method="eigenvalue"):
+        self.graph = graph
+        self.quality_method = quality_method
+        np.random.seed(42)
+
+    def __call__(self, freq):
+        L = construct_so3_laplacian(to_complex(freq), self.graph, abelian_scale=1.0)
+        return laplacian_quality(L)
+
+
+def scan_frequencies_so3(graph, quality_method="eigenvalue"):
+    """Scan a range of complex frequencies and return mode qualities."""
+    ks, alphas = get_scan_grid(graph)
+    freqs = [[k, a] for k in ks for a in alphas]
+
+    worker_scan = WorkerScan(graph, quality_method=quality_method)
+    chunksize = max(1, int(0.1 * len(freqs) / graph.graph["params"]["n_workers"]))
+    with multiprocessing.Pool(graph.graph["params"]["n_workers"]) as pool:
+        qualities_list = list(
+            tqdm(
+                pool.imap(worker_scan, freqs, chunksize=chunksize),
+                total=len(freqs),
+            )
+        )
+
+    id_k = [k_i for k_i in range(len(ks)) for a_i in range(len(alphas))]
+    id_a = [a_i for k_i in range(len(ks)) for a_i in range(len(alphas))]
+    qualities = sc.sparse.coo_matrix(
+        (qualities_list, (id_k, id_a)),
+        shape=(graph.graph["params"]["k_n"], graph.graph["params"]["alpha_n"]),
+    ).toarray()
+
+    return qualities