119 lines
3.8 KiB
Python
119 lines
3.8 KiB
Python
import numpy as np
|
|
|
|
class Simulation:
|
|
def __init__(self):
|
|
self.node_ids = {}
|
|
self.next_node_id = 0
|
|
|
|
self.branches = []
|
|
self.shunt_caps = []
|
|
|
|
self.branch_count = 0
|
|
self.cstate_count = 0
|
|
|
|
self.dirichlet = []
|
|
|
|
def add_node(self, label:str="") -> int:
|
|
if label in self.node_ids:
|
|
raise ValueError(f"Node {label} already exists")
|
|
idx = self.next_node_id
|
|
self.node_ids[label] = idx
|
|
self.next_node_id += 1
|
|
return idx
|
|
|
|
def _get_node_idx(self, label: str) -> int:
|
|
if label not in self.node_ids:
|
|
raise KeyError(f"Unknown node '{label}'")
|
|
return self.node_ids[label]
|
|
|
|
|
|
def add_branch(self, node1:str, node2:str, R:float=0.0, L:float=0.0, C:float=0.0, label:str="") -> int:
|
|
branch_idx = self.branch_count
|
|
self.branch_count += 1
|
|
|
|
if label == "":
|
|
label = f"b_{branch_idx}"
|
|
|
|
a, b = self._get_node_idx(node1), self._get_node_idx(node2)
|
|
|
|
c_idx = None
|
|
if C and C > 0.0:
|
|
c_idx = self.cstate_count
|
|
self.cstate_count += 1
|
|
self.branches.append((a, b, R, L, C, branch_idx, c_idx, label))
|
|
|
|
return branch_idx, c_idx
|
|
|
|
def add_shunt_C(self, node1: str, node2: str, C: float):
|
|
a, b = self._get_node_idx(node1), self._get_node_idx(node2)
|
|
self.shunt_caps.append((a, b, float(C)))
|
|
|
|
def build_matrices(self):
|
|
Nv = self.next_node_id
|
|
Ni = self.branch_count
|
|
Nc = self.cstate_count
|
|
N = Nv + Ni + Nc
|
|
|
|
E = np.zeros((N,N), dtype=float)
|
|
A = np.zeros((N,N), dtype=float)
|
|
|
|
for (a, b, R, L, C, branch_id, cap_id, _) in self.branches:
|
|
i = Nv + branch_id # Index of this branch's current
|
|
|
|
# KCL for the nodes that have this current. Current leaves a and enters b
|
|
A[a, i] += 1.0
|
|
A[b, i] -= 1.0
|
|
|
|
# KVL for branch: v(a)-v(b) - R*i - L di/dt - v_c = 0
|
|
A[i, a] += 1.0
|
|
A[i, b] -= 1.0
|
|
if R: A[i, i] += -R
|
|
if L: E[i, i] += L
|
|
if C and cap_id is not None:
|
|
vc = Nv + Ni + cap_id
|
|
A[i, vc] += 1.0 # subtract v_c in KVL
|
|
E[vc, vc] += C # C * d v_c/dt = i
|
|
A[vc, i] -= 1.0
|
|
|
|
# Shunt capacitors into E(v,v) block
|
|
for (a, b, C) in self.shunt_caps:
|
|
if a != b:
|
|
E[a, a] -= C; E[a, b] += C
|
|
E[b, a] += C; E[b, b] -= C
|
|
|
|
return E, A, N
|
|
|
|
|
|
def step_BE(self, E, A, x_n, h, dirichlet_now):
|
|
# Dirichlet_now: dict label->value at t_{n+1}
|
|
Nv = self.next_node_id
|
|
Ni = self.branch_count
|
|
Nc = self.cstate_count
|
|
N = Nv + Ni + Nc
|
|
|
|
# Fixed voltage indices and values
|
|
fixed_v_idx = np.array([self._get_node_idx(lbl) for lbl in dirichlet_now.keys()], dtype=int)
|
|
v_d = np.array([float(dirichlet_now[lbl]) for lbl in dirichlet_now.keys()], dtype=float)
|
|
|
|
# Free sets
|
|
all_v_idx = np.arange(Nv, dtype=int)
|
|
free_v_idx = np.array(sorted(set(all_v_idx) - set(fixed_v_idx)), dtype=int)
|
|
free_i_idx = np.arange(Nv, Nv+Ni+Nc, dtype=int)
|
|
free_idx = np.concatenate([free_v_idx, free_i_idx])
|
|
|
|
# Partitioned matrices
|
|
Ef_f = E[np.ix_(free_idx, free_idx)]
|
|
Af_f = A[np.ix_(free_idx, free_idx)]
|
|
Ef_d = E[np.ix_(free_idx, fixed_v_idx)]
|
|
Af_d = A[np.ix_(free_idx, fixed_v_idx)]
|
|
|
|
# Left matrix and RHS for BE
|
|
LHS = Ef_f - h * Af_f
|
|
# print(np.linalg.cond(LHS))
|
|
rhs = (E[np.ix_(free_idx, np.arange(N))] @ x_n) - (Ef_d - h * Af_d) @ v_d
|
|
|
|
x_free_next = np.linalg.solve(LHS, rhs)
|
|
x_next = np.zeros_like(x_n)
|
|
x_next[free_idx] = x_free_next
|
|
x_next[fixed_v_idx]= v_d
|
|
return x_next |