Solves a tridiagonal system Ax = b. The arguments A_upper, A_digonal, A_lower correspond to the three diagonals of A. Letting U = A_upper, D=A_digonal and L = A_lower, and assuming for simplicity that there are no batch dimensions, then the matrix A is assumed to be of size (k, k), with
(b, A_upper, A_diagonal, A_lower)
| 22 | return out |
| 23 | |
| 24 | def tridiagonal_solve(b, A_upper, A_diagonal, A_lower): |
| 25 | """Solves a tridiagonal system Ax = b. |
| 26 | |
| 27 | The arguments A_upper, A_digonal, A_lower correspond to the three diagonals of A. Letting U = A_upper, D=A_digonal |
| 28 | and L = A_lower, and assuming for simplicity that there are no batch dimensions, then the matrix A is assumed to be |
| 29 | of size (k, k), with entries: |
| 30 | |
| 31 | D[0] U[0] |
| 32 | L[0] D[1] U[1] |
| 33 | L[1] D[2] U[2] 0 |
| 34 | L[2] D[3] U[3] |
| 35 | . . . |
| 36 | . . . |
| 37 | . . . |
| 38 | L[k - 3] D[k - 2] U[k - 2] |
| 39 | 0 L[k - 2] D[k - 1] U[k - 1] |
| 40 | L[k - 1] D[k] |
| 41 | |
| 42 | Arguments: |
| 43 | b: A tensor of shape (..., k), where '...' is zero or more batch dimensions |
| 44 | A_upper: A tensor of shape (..., k - 1). |
| 45 | A_diagonal: A tensor of shape (..., k). |
| 46 | A_lower: A tensor of shape (..., k - 1). |
| 47 | |
| 48 | Returns: |
| 49 | A tensor of shape (..., k), corresponding to the x solving Ax = b |
| 50 | |
| 51 | Warning: |
| 52 | This implementation isn't super fast. You probably want to cache the result, if possible. |
| 53 | """ |
| 54 | |
| 55 | # This implementation is very much written for clarity rather than speed. |
| 56 | |
| 57 | A_upper, _ = torch.broadcast_tensors(A_upper[:, None, :], b[..., :-1]) |
| 58 | A_lower, _ = torch.broadcast_tensors(A_lower[:, None, :], b[..., :-1]) |
| 59 | A_diagonal, b = torch.broadcast_tensors(A_diagonal[:, None, :], b) |
| 60 | |
| 61 | channels = b.size(-1) |
| 62 | |
| 63 | new_b = np.empty(channels, dtype=object) |
| 64 | new_A_diagonal = np.empty(channels, dtype=object) |
| 65 | outs = np.empty(channels, dtype=object) |
| 66 | |
| 67 | new_b[0] = b[..., 0] |
| 68 | new_A_diagonal[0] = A_diagonal[..., 0] |
| 69 | for i in range(1, channels): |
| 70 | w = A_lower[..., i - 1] / new_A_diagonal[i - 1] |
| 71 | new_A_diagonal[i] = A_diagonal[..., i] - w * A_upper[..., i - 1] |
| 72 | new_b[i] = b[..., i] - w * new_b[i - 1] |
| 73 | |
| 74 | outs[channels - 1] = new_b[channels - 1] / new_A_diagonal[channels - 1] |
| 75 | for i in range(channels - 2, -1, -1): |
| 76 | outs[i] = (new_b[i] - A_upper[..., i] * outs[i + 1]) / new_A_diagonal[i] |
| 77 | |
| 78 | return torch.stack(outs.tolist(), dim=-1) |
| 79 | |
| 80 | |
| 81 | def FDT(strand): |
no test coverage detected