Solve the equation `a x = b` for `x`, assuming a is a triangular matrix. Parameters ---------- a : (M, M) array_like A triangular matrix b : (M,) or (M, N) array_like Right-hand side matrix in `a x = b` lower : bool, optional Use only data contained
(a, b, lower=False)
| 1103 | |
| 1104 | |
| 1105 | def solve_triangular(a, b, lower=False): |
| 1106 | """ |
| 1107 | Solve the equation `a x = b` for `x`, assuming a is a triangular matrix. |
| 1108 | |
| 1109 | Parameters |
| 1110 | ---------- |
| 1111 | a : (M, M) array_like |
| 1112 | A triangular matrix |
| 1113 | b : (M,) or (M, N) array_like |
| 1114 | Right-hand side matrix in `a x = b` |
| 1115 | lower : bool, optional |
| 1116 | Use only data contained in the lower triangle of `a`. |
| 1117 | Default is to use upper triangle. |
| 1118 | |
| 1119 | Returns |
| 1120 | ------- |
| 1121 | x : (M,) or (M, N) array |
| 1122 | Solution to the system `a x = b`. Shape of return matches `b`. |
| 1123 | """ |
| 1124 | |
| 1125 | if a.ndim != 2: |
| 1126 | raise ValueError("a must be 2 dimensional") |
| 1127 | if b.ndim <= 2: |
| 1128 | if a.shape[1] != b.shape[0]: |
| 1129 | raise ValueError("a.shape[1] and b.shape[0] must be equal") |
| 1130 | if a.chunks[1] != b.chunks[0]: |
| 1131 | msg = ( |
| 1132 | "a.chunks[1] and b.chunks[0] must be equal. " |
| 1133 | "Use .rechunk method to change the size of chunks." |
| 1134 | ) |
| 1135 | raise ValueError(msg) |
| 1136 | else: |
| 1137 | raise ValueError("b must be 1 or 2 dimensional") |
| 1138 | |
| 1139 | vchunks = len(a.chunks[1]) |
| 1140 | hchunks = 1 if b.ndim == 1 else len(b.chunks[1]) |
| 1141 | token = tokenize(a, b, lower) |
| 1142 | name = "solve-triangular-" + token |
| 1143 | |
| 1144 | # for internal calculation |
| 1145 | # (name, i, j, k, l) corresponds to a_ij.dot(b_kl) |
| 1146 | name_mdot = "solve-tri-dot-" + token |
| 1147 | |
| 1148 | def _b_init(i, j): |
| 1149 | if b.ndim == 1: |
| 1150 | return b.name, i |
| 1151 | else: |
| 1152 | return b.name, i, j |
| 1153 | |
| 1154 | def _key(i, j): |
| 1155 | if b.ndim == 1: |
| 1156 | return name, i |
| 1157 | else: |
| 1158 | return name, i, j |
| 1159 | |
| 1160 | dsk = {} |
| 1161 | if lower: |
| 1162 | for i in range(vchunks): |
no test coverage detected