complex current works.

PSD/complex_current_and_nodes.py

@@ -13,7 +13,7 @@ class NodeType(object):
 
 
 class SwingEquationNode(NodeType):
-    def __init__(self, number_of_variables, v_set = 1., infeed = 1., damping = 0.1, H = 1.):
+    def __init__(self, number_of_variables, v_set = 1., infeed = 1., damping = 0.1, H = 1.5):
         super(SwingEquationNode, self).__init__(number_of_variables)
         self.v_set = v_set
         self.v_set_squared = v_set ** 2
@@ -21,13 +21,21 @@ class SwingEquationNode(NodeType):
         self.damping = damping
         self.H_inv = 1./H
 
-    def node_dynamics(self, y, dy, i, t):
-        v = np.complex(y[0], y[1])
-        dv = 1.j * y[2] * v - (v.conjugate() * v - self.v_set_squared) * v
-        dy[0] = dv.real
-        dy[1] = dy.imag
-        dy[2] = self.H_inv * (self.infeed - self.damping * y[2] - (v * i.conjugate()).real)
-        return np.array
+    def node_dynamics(self, v, omega, i, t, dv, domega, index):
+        # The current is assumed to be negative when it is flowing away from the node
+        dv[index] = 1.j * omega[index] * v[index] - (v[index].real * v[index].real + v[index].imag * v[index].imag - self.v_set_squared) * v[index]
+        domega[index] = self.H_inv * (self.infeed - self.damping * omega[index] - (v[index] * i[index].conjugate()).real)
+
+    def node_dynamics_string(self, j="{index}"):
+        return """
+    dv[{index}] = 1.j * omega[{index}] * v[{index}] - (v[{index}].real * v[{index}].real + v[{index}].imag * v[{index}].imag - {v_set_squared}) * v[{index}]
+    domega[{index}] = {H_inv:.64f} * ({infeed} - {damping} * omega[{index}] - (v[{index}] * i[{index}].conjugate()).real)"""\
+            .format(
+            v_set_squared=self.v_set_squared,
+            H_inv=self.H_inv,
+            infeed=self.infeed,
+            damping=self.damping,
+            index=j)
 
 
 @njit(complex128[:](complex128[:], int32[:], int32[:], complex128[:]))
@@ -57,34 +65,91 @@ def define_network_rhs(node_list, Y):
     total_length = 2*length
 
     coupling_sp = sps.csr_matrix(Y)
-    data = coupling_sp.data
-    indptr = coupling_sp.indptr
-    indices = coupling_sp.indices
-
-    v_indices = [1, 2, 4, 5] # ...
 
     def network_rhs(y, t):
-        dydt = np.empty(total_length)
+        dydt = np.empty(total_length + length)
+
+        v = y[:total_length].view(np.complex128)
+        omega = y[total_length:]
+        dv = dydt[:total_length].view(np.complex128)
+        domega = dydt[total_length:]
+
         i = coupling_sp.dot(v)
 
+        for j, node in enumerate(node_list):
+            node.node_dynamics(v, omega, i, t, dv, domega, j)
         return dydt
 
     return network_rhs
 
+
+def define_network_rhs_codegen(node_list, Y):
+    assert len(node_list) == len(Y)
+    length = len(Y)
+    total_length = 2*length
+
+    coupling_sp = sps.csr_matrix(Y)
+
+    data = coupling_sp.data
+    indptr = coupling_sp.indptr
+    l_indptr = len(indptr)
+    indices = coupling_sp.indices
+
+    def_network_rhs_string = """
+@njit(float64[:](float64[:], float64))
+def network_rhs_numba(y, t):
+    dydt = np.empty(total_length + length)
+
+    v = y[:total_length].view(np.complex128)
+    omega = y[total_length:]
+    dv = dydt[:total_length].view(np.complex128)
+    domega = dydt[total_length:]
+
+    i = np.zeros(l_indptr - 1, dtype=np.complex128)
+    index = 0
+    for row, number_of_entries in enumerate(indptr[1:]):
+        while index < number_of_entries:
+            i[row] += data[index] * v[indices[index]]
+            index += 1
+
+"""
+
+    for j, node in enumerate(node_list):
+        def_network_rhs_string += node.node_dynamics_string(j)
+
+    def_network_rhs_string += """
+    return dydt
+"""
+
+    context = globals()
+    context.update(locals())
+
+    exec def_network_rhs_string in context
+
+    return network_rhs_numba
+
 if __name__ == "__main__":
 
     node_list = list()
-    node_list.append(SwingEquationNode(3))
-    node_list.append(SwingEquationNode(3))
+    node_list.append(SwingEquationNode(3, infeed=1.))
+    node_list.append(SwingEquationNode(3, infeed=-1.))
 
-    Y = 8.j * np.ones((2, 2), dtype=np.complex128)
+    Y = -8.j * np.ones((2, 2), dtype=np.complex128)
     Y[0, 0] *= -1.
     Y[1, 1] *= -1.
 
     rhs = define_network_rhs(node_list, Y)
+    nrhs = define_network_rhs_codegen(node_list, Y)
+
+    for i in range(10):
+        ic = np.random.rand(6)
+        print(rhs(ic, 0.) - nrhs(ic, 0.))
+
     root_rhs = lambda x: rhs(x, 0.)
 
-    ic = np.ones(4, dtype=np.float64) * 0.5
+    ic = np.ones(6, dtype=np.float64) * 0.5
+
+    print(root_rhs(ic))
 
     from scipy.optimize import root
 
@@ -92,9 +157,9 @@ if __name__ == "__main__":
     if result.success:
         print(rhs(result.x, 0.))
         ic = result.x
+        ic[5] += 1.
     else:
         print("failed")
-        exit()
 
     times = np.linspace(0., 100., 1000)
 
@@ -103,6 +168,12 @@ if __name__ == "__main__":
     states = odeint(rhs, ic, times)
 
     import matplotlib.pyplot as plt
+    from mpl_toolkits.mplot3d import Axes3D
+
     plt.figure()
-    plt.plot(times, states[:,2:])
+    plt.plot(times, states[:, 4:])
+    fig = plt.figure()
+    ax = fig.gca(projection='3d')
+    ax.plot(states[:, 0], states[:, 1], times)
+    ax.plot(states[:, 2], states[:, 3], times)
     plt.show()
\ No newline at end of file