(
J_Ginv_i: torch.Tensor,
J_Ginv_j: torch.Tensor,
ii: torch.Tensor,
jj: torch.Tensor,
res: torch.Tensor,
ep: float,
lm: float,
freen: int,
)
| 34 | |
| 35 | |
| 36 | def solve_system_py( |
| 37 | J_Ginv_i: torch.Tensor, |
| 38 | J_Ginv_j: torch.Tensor, |
| 39 | ii: torch.Tensor, |
| 40 | jj: torch.Tensor, |
| 41 | res: torch.Tensor, |
| 42 | ep: float, |
| 43 | lm: float, |
| 44 | freen: int, |
| 45 | ) -> torch.Tensor: |
| 46 | # Ensure all tensors are on CPU |
| 47 | device = res.device |
| 48 | J_Ginv_i = J_Ginv_i.cpu() |
| 49 | J_Ginv_j = J_Ginv_j.cpu() |
| 50 | ii = ii.cpu() |
| 51 | jj = jj.cpu() |
| 52 | res = res.clone().cpu() |
| 53 | |
| 54 | r = res.size(0) # Number of edges |
| 55 | n = max(ii.max().item(), jj.max().item()) + 1 # Number of nodes |
| 56 | |
| 57 | res_vec = res.view(-1).numpy().astype(np.float64) |
| 58 | |
| 59 | rows, cols, data = [], [], [] |
| 60 | ii_np = ii.numpy() |
| 61 | jj_np = jj.numpy() |
| 62 | J_Ginv_i_np = J_Ginv_i.numpy() |
| 63 | J_Ginv_j_np = J_Ginv_j.numpy() |
| 64 | |
| 65 | for x in range(r): |
| 66 | i = ii_np[x] |
| 67 | j = jj_np[x] |
| 68 | if i == j: |
| 69 | raise ValueError("Self-edges are not allowed") |
| 70 | |
| 71 | for k in range(7): |
| 72 | for l in range(7): |
| 73 | row_idx = x * 7 + k |
| 74 | col_idx_i = i * 7 + l |
| 75 | val_i = J_Ginv_i_np[x, k, l] |
| 76 | rows.append(row_idx) |
| 77 | cols.append(col_idx_i) |
| 78 | data.append(val_i) |
| 79 | |
| 80 | col_idx_j = j * 7 + l |
| 81 | val_j = J_Ginv_j_np[x, k, l] |
| 82 | rows.append(row_idx) |
| 83 | cols.append(col_idx_j) |
| 84 | data.append(val_j) |
| 85 | |
| 86 | J = coo_matrix((data, (rows, cols)), shape=(r * 7, n * 7)).tocsc() |
| 87 | |
| 88 | b_vec = -J.T @ res_vec |
| 89 | |
| 90 | A_mat = J.T @ J |
| 91 | |
| 92 | diag = A_mat.diagonal() |
| 93 | new_diag = diag * (1.0 + lm) + ep |
no test coverage detected