allow to pump edges independently

examples/line_PRA/create_graph.py

@@ -30,4 +30,6 @@
     nopump_edges = len(graph.edges()) - pump_edges
     pump = np.append(np.ones(pump_edges), np.zeros(nopump_edges))
     pump[0] = 0
-    yaml.dump(pump.astype(int).tolist(), open("pump.yaml", "w"))
+    # convert new version
+    pump = [[0.0, float(p)] for p in pump]
+    yaml.dump(pump, open("pump.yaml", "w"))

netsalt/modes.py

@@ -148,9 +148,15 @@ def find_modes(graph, qualities):
     return modes_df
 
 
-def _convert_edges(vector):
+def _convert_edges(vector, D0=None):
     """Convert single edge values to double edges."""
     edge_vector = np.zeros(2 * len(vector), dtype=np.complex128)
+    if D0 is not None and len(np.shape(vector)) == 2:
+        vector = vector[:, 0] + D0 * vector[:, 1]
+
+    if D0 is None and len(np.shape(vector)) == 2:
+        vector = [1 if v[0] + v[1] > 0 else 0 for v in vector]
+
     edge_vector[::2] = vector
     edge_vector[1::2] = vector
     return edge_vector
@@ -161,10 +167,10 @@ def _get_dielectric_constant_matrix(params):
     return sc.sparse.diags(_convert_edges(params["dielectric_constant"]))
 
 
-def _get_mask_matrices(params):
+def _get_mask_matrices(params, D0=None):
     """Return sparse diagonal matrices of pump and inner edge masks."""
     in_mask = sc.sparse.diags(_convert_edges(np.array(params["inner"])))
-    pump_mask = sc.sparse.diags(_convert_edges(params["pump"])).dot(in_mask)
+    pump_mask = sc.sparse.diags(_convert_edges(params["pump"], D0)).dot(in_mask)
     return in_mask, pump_mask
 
 
@@ -219,10 +225,10 @@ def compute_overlapping_single_edges(passive_mode, graph):
     return np.real(pump_norm / inner_norm)
 
 
-def compute_overlapping_factor(passive_mode, graph):
+def compute_overlapping_factor(passive_mode, graph, D0):
     """Compute the overlapping factor of a mode with the pump."""
     dielectric_constant = _get_dielectric_constant_matrix(graph.graph["params"])
-    in_mask, pump_mask = _get_mask_matrices(graph.graph["params"])
+    in_mask, pump_mask = _get_mask_matrices(graph.graph["params"], D0)
     inner_dielectric_constants = dielectric_constant.dot(in_mask)
 
     node_solution = mode_on_nodes(passive_mode, graph)
@@ -241,10 +247,13 @@ def compute_overlapping_factor(passive_mode, graph):
 def pump_linear(mode_0, graph, D0_0, D0_1):
     """Find the linear approximation of the new wavenumber."""
     graph.graph["params"]["D0"] = D0_0
-    overlapping_factor = compute_overlapping_factor(mode_0, graph)
+    overlapping_factor_0 = compute_overlapping_factor(mode_0, graph, D0_0)
+    overlapping_factor_1 = compute_overlapping_factor(mode_0, graph, D0_1)
     freq = to_complex(mode_0)
-    gamma_overlap = gamma(freq, graph.graph["params"]) * overlapping_factor
-    return from_complex(freq * np.sqrt((1.0 + gamma_overlap * D0_0) / (1.0 + gamma_overlap * D0_1)))
+    _gamma = gamma(freq, graph.graph["params"])
+    gamma_overlap_0 = _gamma * overlapping_factor_0
+    gamma_overlap_1 = _gamma * overlapping_factor_1
+    return from_complex(freq * np.sqrt((1.0 + gamma_overlap_0) / (1.0 + gamma_overlap_1)))
 
 
 def mode_on_nodes(mode, graph):
@@ -426,7 +435,7 @@ def _compute_mode_competition_element(lengths, params, data, with_gamma=True):
 
     matrix_element = 0
     for ei, length in enumerate(lengths):
-        if params["pump"][ei] > 0.0 and params["inner"][ei]:
+        if sum(params["pump"][ei]) > 0.0 and params["inner"][ei]:
             k_mu = k_mus[ei]
             k_nu = k_nus[ei]
 
@@ -865,5 +874,5 @@ def lasing_threshold_linear(mode, graph, D0):
     return 1.0 / (
         q_value(mode)
         * -np.imag(gamma(to_complex(mode), graph.graph["params"]))
-        * np.real(compute_overlapping_factor(mode, graph))
+        * np.real(compute_overlapping_factor(mode, graph, D0))
     )

netsalt/physics.py

@@ -92,9 +92,8 @@ def dispersion_relation_pump(freq, params=None):
     if "pump" not in params or "D0" not in params:
         return freq * np.sqrt(params["dielectric_constant"])
 
-    return freq * np.sqrt(
-        params["dielectric_constant"] + gamma(freq, params) * params["D0"] * params["pump"]
-    )
+    pump = params["pump"][:, 0] + params["D0"] * params["pump"][:, 1]
+    return freq * np.sqrt(params["dielectric_constant"] + gamma(freq, params) * pump)
 
 
 def set_dielectric_constant(graph, params, custom_values=None):

netsalt/plotting.py

@@ -123,10 +123,6 @@ def plot_stem_spectra(
     ax2.set_xlabel(r"$\lambda$")
     ax2.set_ylabel("Gain spectrum (a.u.)")
 
-    ax3 = ax.twiny()
-    lams = 2 * np.pi / ks
-    ax3.set_xlim(lams[0], lams[-1])
-
     if save_option:
         _savefig(graph, fig, folder, filename)
 
@@ -274,7 +270,7 @@ def plot_pump_profile(graph, pump, figsize=(5, 4), ax=None, node_size=1.0, c="0.
         ax = plt.gca()
 
     positions = [graph.nodes[u]["position"] for u in graph]
-    pumped_edges = [e for e, _pump in zip(graph.edges, pump) if _pump > 0.0]
+    pumped_edges = [e for e, _pump in zip(graph.edges, pump) if sum(_pump) > 0.0]
     nx.draw_networkx_edges(
         graph,
         pos=positions,

netsalt/quantum_graph.py

@@ -57,7 +57,7 @@ def _not_equal(data1, data2, force=False):
     if force:
         return True
     if isinstance(data1, np.ndarray):
-        return all(data1 != data2)
+        return not np.array_equal(data1, data2)
     return data1 != data2
 
 
@@ -165,14 +165,14 @@ def set_total_length(graph, total_length=None, max_extent=None, inner=True, with
 def _set_pump_on_graph(graph):
     """Set the pump values on the graph from params."""
     if "pump" not in graph.graph["params"]:
-        graph.graph["params"]["pump"] = np.ones(len(graph.edges))
+        graph.graph["params"]["pump"] = np.ones((len(graph.edges), 2))
     for ei, e in enumerate(graph.edges):
         graph[e[0]][e[1]]["pump"] = graph.graph["params"]["pump"][ei]
 
 
 def _set_pump_on_params(graph, params):
     """Set the pump values on the graph from params."""
-    params["pump"] = np.ones(len(graph.edges))
+    params["pump"] = np.ones((len(graph.edges), 2))
     for ei, e in enumerate(graph.edges):
         params["pump"][ei] = graph[e[0]][e[1]]["pump"]
 

tests/data/run_simple/create_graph.py

@@ -30,4 +30,6 @@
     nopump_edges = len(graph.edges()) - pump_edges
     pump = np.append(np.ones(pump_edges), np.zeros(nopump_edges))
     pump[0] = 0
-    yaml.dump(pump.astype(int).tolist(), open("pump.yaml", "w"))
+    # convert new version
+    pump = [[0.0, float(p)] for p in pump]
+    yaml.dump(pump, open("pump.yaml", "w"))

tests/data/run_simple/out/pump_profile.yaml

@@ -1,10 +1,20 @@
-- 0
-- 1
-- 1
-- 1
-- 1
-- 0
-- 0
-- 0
-- 0
-- 0
+- - 0.0
+  - 0.0
+- - 0.0
+  - 1.0
+- - 0.0
+  - 1.0
+- - 0.0
+  - 1.0
+- - 0.0
+  - 1.0
+- - 0.0
+  - 0.0
+- - 0.0
+  - 0.0
+- - 0.0
+  - 0.0
+- - 0.0
+  - 0.0
+- - 0.0
+  - 0.0

tests/test_functional.py

@@ -55,7 +55,9 @@ def create_graph():
     nopump_edges = len(graph.edges()) - pump_edges
     pump = np.append(np.ones(pump_edges), np.zeros(nopump_edges))
     pump[0] = 0
-    yaml.dump(pump.astype(int).tolist(), open("pump.yaml", "w"))
+    # convert new version
+    pump = [[0.0, float(p)] for p in pump]
+    yaml.dump(pump, open("pump.yaml", "w"))
 
 
 @pytest.fixture