root finding nice

PSD/complex_current_and_nodes.py

@@ -111,6 +111,21 @@ def define_network_rhs(node_list, Y):
     return network_rhs
 
 
+def define_root_rhs(total_length, rhs):
+    def root_rhs(y):
+        dydt = rhs(y, 0.)
+        v = y[:total_length].view(np.complex128)
+        dv = dydt[:total_length].view(np.complex128)
+        domega = dydt[total_length:]
+        if v[0].imag >= v[0].real:
+            omega = -1. * dv[0].real/v[0].imag
+        else:
+            omega = dv[0].imag/v[0].real
+        a = 1.j * omega * v - dv
+        return np.concatenate((domega, a.view(np.float64)))
+    return root_rhs
+
+
 def define_network_rhs_codegen(node_list, Y):
     assert len(node_list) == len(Y)
     length = len(Y)
@@ -160,7 +175,7 @@ if __name__ == "__main__":
 
     node_list = list()
     node_list.append(SwingEquationNode(3, infeed=1.))
-    node_list.append(SwingEquationNode(3, infeed=-1))
+    node_list.append(SwingEquationNode(3, infeed=-1.1))
 
     Y = -8.j * np.ones((2, 2), dtype=np.complex128)
     Y[0, 0] *= -1.
@@ -168,12 +183,14 @@ if __name__ == "__main__":
 
     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.))
 
-    for i in range(10):
-        ic = np.random.rand(6)
-        print(rhs(ic, 0.) - nrhs(ic, 0.))
+    root_rhs_1 = lambda x : rhs(x, 0.)
 
-    root_rhs = lambda x: rhs(x, 0.)
+    root_rhs = define_root_rhs(2*len(Y), rhs)
 
     ic = np.ones(6, dtype=np.float64) * 0.5
 
@@ -185,7 +202,8 @@ if __name__ == "__main__":
     if result.success:
         print(rhs(result.x, 0.))
         ic = result.x
-        ic[5] += 1.
+        print(root_rhs(ic))
+        # ic[5] += 1.
     else:
         print("failed")
 
@@ -193,9 +211,6 @@ if __name__ == "__main__":
 
     from scipy.integrate import odeint
 
-    node_list[0].infeed += 0.1
-    rhs = define_network_rhs(node_list, Y)
-
     states = odeint(rhs, ic, times)
 
     import matplotlib.pyplot as plt