Python and Viv · Numerics
LU factorization and solve
What to notice
- Swapping two rows: NumPy's fancy index
LU[[k, p]] = LU[[p, k]]is Viv'sLU[k, :], LU[p, :] = LU[p, :], LU[k, :]. The right side is computed (and copied) before either row is written. LU[i, k + 1:] -= LU[i, k] * LU[k, k + 1:]is the same statement in both: a compound assignment into a slice.k + np.argmax(np.abs(LU[k:, k]))ismax([i for i from k to n], key=fn(i) -> abs(LU[i, k]))in Viv: the first row with the largest entry, asargmaxpicks.- A dot product of two slices,
LU[i, :i] @ y[:i], issum(LU[i, :i] * y[:i]). b[perm](indexing with an array) is the comprehension[b[j] for j in perm]in Viv.permholds doubles, and an index that is a double is read as an integer.np.linalg.solve(A, b)returns the solution;lin_solve(A, b, x)writes it into an array the caller made.for i in reversed(range(n))isfor i to n reversed.
Python lu_solve.py · 48 lines
Viv lu_solve.viv · 59 lines
py2viv translated from lu_solve.py
What py2viv adds before the program
(nothing)
(nothing)
386 lines: a header and the guard functions its translation calls
1###
2Translated from lu_solve.py by py2viv 0.2.0 (numeric Python -> Viv).
3Python source sha256: 94e8d0881c0a0a3439981387e724fe1aa44ed7c2a8a1825956f47011a53fc4a2
4Each print(...) is one print(...) line; fmt(x) writes a float as Python's repr does.
5A `raise`, a failed assert and every guard `stop` with Python's exception line.
6###
7
8#| 0, 1, ..., n - 1.
9def py_iota(n) {
10 a = zeros(n)
11 for i to n {
12 a[i] = i
13 }
14 return a
15}
16
17#| Every entry of a real array, row-major, as a vector (a copy).
18def py_flat(a: Arr) {
19 return take(a, py_iota(numel(a)))
20}
21
22#| A copy of a real array (np.copy, a.copy(), np.array(a)).
23def py_copy(a: Arr) {
24 return a * 1.0
25}
26
27#| np.arange(start, stop, step), as NumPy fills it: ceil((stop - start) / step) entries,
28#| entry k = start + k * delta with delta = (start + step) - start.
29def py_arange(start, stop_, step) {
30 if step == 0 {
31 stop("ZeroDivisionError: division by zero")
32 }
33 if not (abs((stop_ - start) / step) < 1.0 / 0.0) {
34 stop("py2viv: np.arange of a length that is not finite (NumPy raises ValueError)")
35 }
36 n = max(ceil((stop_ - start) / step), 0)
37 a = zeros(n)
38 delta = (start + step) - start
39 for k to n {
40 a[k] = start + k * delta
41 }
42 return a
43}
44
45#| np.argmax(a): the flat index of the first largest entry.
46def py_argmax(a: Arr) {
47 f = py_flat(a)
48 k = 0
49 if f[0] != f[0] {
50 return 0
51 }
52 for i from 1 to len(f) {
53 if f[i] != f[i] {
54 return i
55 }
56 if f[i] > f[k] {
57 k = i
58 }
59 }
60 return k
61}
62
63#| np.linalg.solve(A, b) for a vector b. Viv's lin_solve(A, b, x) writes x and returns 0.
64def py_solve(A: Arr, b: Arr) {
65 if not (all(abs(A) < 1.0 / 0.0) and all(abs(b) < 1.0 / 0.0)) {
66 stop("py2viv: np.linalg.solve with a NaN or an infinity (LAPACK's answer is its own)")
67 }
68 x = zeros(len(b))
69 lin_solve(A, b, x)
70 if all(abs(A) < 1.0 / 0.0) and all(abs(b) < 1.0 / 0.0) and not all(abs(x) < 1.0 / 0.0) {
71 stop("numpy.linalg.LinAlgError: Singular matrix")
72 }
73 return x
74}
75
76#| An int result, as Python's int: a -0.0 becomes 0 (an int has no sign on zero), and the run
77#| stops at 2^53, past which a double no longer holds every integer (Python's int does).
78def py_int(x) {
79 if abs(x) >= 9007199254740992 {
80 stop("py2viv: an integer reached 2^53; Python's int is exact there and a double is not")
81 }
82 return x + 0.0
83}
84
85#| a / b with a float: ZeroDivisionError on a zero divisor.
86def py_div_f(a, b) {
87 if b == 0 {
88 stop("ZeroDivisionError: float division by zero")
89 }
90 return a / b
91}
92
93#| np.isclose(a, b, rtol, atol) of two numbers: |a - b| <= atol + rtol |b| when both are
94#| finite, a == b otherwise (a NaN is never close).
95def py_npisclose(a, b, rtol, atol) {
96 if abs(a) < 1.0 / 0.0 and abs(b) < 1.0 / 0.0 {
97 return abs(a - b) <= atol + rtol * abs(b)
98 }
99 return a == b
100}
101
102#| np.allclose(a, b, rtol, atol) of two arrays of one shape: np.isclose entry by entry.
103def py_allclose(a: Arr, b: Arr, rtol, atol) {
104 fa = py_flat(a)
105 fb = py_flat(b)
106 if len(fa) != len(fb) {
107 stop("py2viv: np.allclose of two arrays of different sizes (NumPy broadcasts them)")
108 }
109 for i to len(fa) {
110 if not py_npisclose(fa[i], fb[i], rtol, atol) {
111 return false
112 }
113 }
114 return true
115}
116
117#| A slice bound as Python reads it on an axis of length n: a negative one counts from the
118#| end; the result is clamped to 0 ... n.
119def py_lo(lo, n) {
120 b = lo
121 if b < 0 {
122 b = b + n
123 if b < 0 {
124 b = 0
125 }
126 }
127 if b > n {
128 b = n
129 }
130 return b
131}
132
133#| An index next to a slice, as NumPy reads it on an axis of length n: a negative one counts
134#| from the end (one still out of bounds stops the run, as NumPy raises IndexError).
135def py_ix(i, n) {
136 if i < 0 {
137 return i + n
138 }
139 return i
140}
141
142#| The shape NumPy broadcasts two shapes da and db to (aligned at the last axis, a length 1
143#| stretched); the run stops where NumPy raises ValueError.
144def py_bdims(da: Arr, db: Arr) {
145 n = len(da)
146 if len(db) > n {
147 n = len(db)
148 }
149 e = zeros(n)
150 for i to n {
151 x = 1
152 y = 1
153 if i >= n - len(da) {
154 x = da[i - n + len(da)]
155 }
156 if i >= n - len(db) {
157 y = db[i - n + len(db)]
158 }
159 if x == y or y == 1 {
160 e[i] = x
161 }
162 else {
163 if x == 1 {
164 e[i] = y
165 }
166 else {
167 stop("ValueError: operands could not be broadcast together with shapes")
168 }
169 }
170 }
171 return e
172}
173
174#| The flat indices, in row-major order of the shape e, of the entries of an array of shape d
175#| that broadcasting to e reads (a length 1 of d is read again along that axis of e).
176def py_bidx(d: Arr, e: Arr) {
177 n = len(e)
178 m = len(d)
179 st = zeros(n)
180 s = 1
181 for k to m {
182 j = m - 1 - k
183 if d[j] != 1 {
184 st[j + n - m] = s
185 }
186 s = s * d[j]
187 }
188 idx = zeros(1)
189 for i to n {
190 w = e[i]
191 nxt = zeros(len(idx) * w)
192 for p to len(idx) {
193 for c to w {
194 nxt[p * w + c] = idx[p] + c * st[i]
195 }
196 }
197 idx = nxt
198 }
199 return idx
200}
201
202#| a broadcast to the shape e (a shape py_bdims gave for a and another array): a itself when
203#| it has that shape, else a new array whose entries are copies of a's.
204def py_bto(a: Arr, e: Arr) {
205 d = shape(a)
206 if len(d) == len(e) {
207 if all(d == e) {
208 return a
209 }
210 }
211 return py_shaped(take(a, py_bidx(d, e)), e)
212}
213
214#| a - b of two real arrays, broadcast.
215def py_bsub(a: Arr, b: Arr) {
216 e = py_bdims(shape(a), shape(b))
217 return py_bto(a, e) - py_bto(b, e)
218}
219
220#| The result r of `x += v` (and -=, *=, /=), which NumPy writes into x: the run stops where
221#| broadcasting made r larger than x (NumPy raises ValueError).
222def py_inpl(r: Arr, x: Arr) {
223 if len(shape(r)) != len(shape(x)) {
224 stop("ValueError: non-broadcastable output operand with shape")
225 }
226 if any(shape(r) != shape(x)) {
227 stop("ValueError: non-broadcastable output operand with shape")
228 }
229 return r
230}
231
232#| 1 when a value of shape dv broadcasts to exactly the shape e (each of its lengths,
233#| aligned at the last axis, is e's or 1), else 0.
234def py_bfits(dv: Arr, e: Arr) {
235 n = len(e)
236 m = len(dv)
237 if m > n {
238 return 0
239 }
240 for k to m {
241 x = dv[m - 1 - k]
242 if x != 1 and x != e[n - 1 - k] {
243 return 0
244 }
245 }
246 return 1
247}
248
249#| One entry of an integer index array on an axis of length n: a negative one counts from the
250#| end; NumPy raises IndexError for one out of bounds.
251def py_ixchk(i, n) {
252 j = i
253 if j < 0 {
254 j = j + n
255 }
256 if j < 0 or j >= n {
257 stop("IndexError: index")
258 }
259 return j
260}
261
262#| a with an integer index array idx on axis k (the other axes whole): the result's shape is
263#| a's with axis k replaced by idx's shape; a negative index counts from the end.
264def py_gath(a: Arr, idx: Arr, k) {
265 s = py_ax(a, k)
266 before = s[0]
267 along = s[1]
268 after = s[2]
269 f = py_flat(idx)
270 for q to len(f) {
271 f[q] = py_ixchk(f[q], along)
272 }
273 src = zeros(before * len(f) * after)
274 p = 0
275 for o to before {
276 for q to len(f) {
277 for i to after {
278 src[p] = (o * along + f[q]) * after + i
279 p += 1
280 }
281 }
282 }
283 d = shape(a)
284 di = shape(idx)
285 e = []
286 for t to k {
287 e = append(e, d[t])
288 }
289 for t to len(di) {
290 e = append(e, di[t])
291 }
292 for t from k + 1 to len(d) {
293 e = append(e, d[t])
294 }
295 return py_shaped(take(a, src), e)
296}
297
298#| The entries a store through integer indices writes: v broadcast to the shape e of the
299#| indexed result, row-major; NumPy raises ValueError where v does not broadcast to e.
300def py_vals(v: Arr, e: Arr) {
301 if py_bfits(shape(v), e) == 0 {
302 stop("ValueError: shape mismatch: value array of shape")
303 }
304 return py_flat(py_bto(v, e))
305}
306
307#| a[..., idx, ...] = v with an integer index array idx on axis k (the other axes whole):
308#| the entries of v (broadcast) written in row-major order of the indexed result, so of
309#| repeated indices the last write wins.
310def py_setax(a: Arr, idx: Arr, k, v: Arr) {
311 s = py_ax(a, k)
312 before = s[0]
313 along = s[1]
314 after = s[2]
315 f = py_flat(idx)
316 for q to len(f) {
317 f[q] = py_ixchk(f[q], along)
318 }
319 dst = zeros(before * len(f) * after)
320 p = 0
321 for o to before {
322 for q to len(f) {
323 for i to after {
324 dst[p] = (o * along + f[q]) * after + i
325 p += 1
326 }
327 }
328 }
329 d = shape(a)
330 di = shape(idx)
331 e = []
332 for t to k {
333 e = append(e, d[t])
334 }
335 for t to len(di) {
336 e = append(e, di[t])
337 }
338 for t from k + 1 to len(d) {
339 e = append(e, d[t])
340 }
341 put(a, dst, py_vals(v, e))
342 return 0
343}
344
345#| np.allclose(a, b, rtol, atol) of two arrays, broadcast.
346def py_ballclose(a: Arr, b: Arr, rtol, atol) {
347 e = py_bdims(shape(a), shape(b))
348 return py_allclose(py_bto(a, e), py_bto(b, e), rtol, atol)
349}
350
351#| The sizes around axis k of a (k = -1: every axis, the array flat): [before, along, after].
352def py_ax(a: Arr, k) {
353 d = shape(a)
354 if k < 0 {
355 return [1, numel(a), 1]
356 }
357 before = 1
358 for i to k {
359 before = before * d[i]
360 }
361 after = 1
362 for i from k + 1 to len(d) {
363 after = after * d[i]
364 }
365 return [before, d[k], after]
366}
367
368#| The entries of f, row-major, as an array of shape e (a vector of 1 to 4 lengths).
369def py_shaped(f: Arr, e: Arr) {
370 if len(e) == 1 {
371 return f
372 }
373 if len(e) == 2 {
374 z = zeros(e[0], e[1])
375 }
376 else {
377 if len(e) == 3 {
378 z = zeros(e[0], e[1], e[2])
379 }
380 else {
381 z = zeros(e[0], e[1], e[2], e[3])
382 }
383 }
384 put(z, py_iota(numel(z)), f)
385 return z
386}
1 imports
4import numpy as np
2# nothing to import: lin_solve is a builtin
(nothing)
2 LU factorization with partial pivoting
7def lu_factor(A):
8 n = A.shape[0]
9 LU = A.copy()
10 perm = np.arange(n)
11 for k in range(n - 1):
12 p = k + int(np.argmax(np.abs(LU[k:, k])))
13 if p != k:
14 LU[[k, p]] = LU[[p, k]]
15 perm[[k, p]] = perm[[p, k]]
16 for i in range(k + 1, n):
17 LU[i, k] /= LU[k, k]
18 LU[i, k + 1:] -= LU[i, k] * LU[k, k + 1:]
19 return LU, perm
5def lu_factor(A) {
6 n = shape(A)[0]
7 LU = copy(A)
8 perm = [i for i to n]
9 for k to n - 1 {
10 p = max([i for i from k to n], key=fn(i) -> abs(LU[i, k]))
11 if p != k {
12 LU[k, :], LU[p, :] = LU[p, :], LU[k, :]
13 perm[k], perm[p] = perm[p], perm[k]
14 }
15 for i from k + 1 to n {
16 LU[i, k] /= LU[k, k]
17 LU[i, k + 1:] -= LU[i, k] * LU[k, k + 1:]
18 }
19 }
20 return LU, perm
21}
392def lu_factor(A: Arr) {
393 n = len(A)
394 LU = py_copy(A)
395 perm = py_arange(0, n, 1)
396 for k to py_int(n - 1) {
397 p = py_int(k + py_argmax(abs(LU[py_lo(k, len(LU)):, py_ix(k, shape(LU)[1])])))
398 if p != k {
399 py_setax(LU, [k, p], 0, py_gath(LU, [p, k], 0))
400 py_setax(perm, [k, p], 0, py_gath(perm, [p, k], 0))
401 }
402 for i from py_int(k + 1) to n {
403 LU[i, k] /= LU[k, k]
404 LU[py_ix(i, len(LU)), py_lo(py_int(k + 1), shape(LU)[1]):] = py_inpl(py_bsub(LU[py_ix(i, len(LU)), py_lo(py_int(k + 1), shape(LU)[1]):], LU[i, k] * LU[py_ix(k, len(LU)), py_lo(py_int(k + 1), shape(LU)[1]):]), LU[py_ix(i, len(LU)), py_lo(py_int(k + 1), shape(LU)[1]):])
405 }
406 }
407 return LU, perm
408}
3 forward and back substitution
22def lu_solve(LU, perm, b):
23 n = len(b)
24 y = b[perm].astype(float)
25 for i in range(n):
26 y[i] -= LU[i, :i] @ y[:i]
27 x = y.copy()
28 for i in reversed(range(n)):
29 x[i] = (x[i] - LU[i, i + 1:] @ x[i + 1:]) / LU[i, i]
30 return x
24def lu_solve(LU, perm, b) {
25 n = len(b)
26 y = [b[j] for j in perm]
27 for i to n {
28 y[i] -= sum(LU[i, :i] * y[:i])
29 }
30 x = copy(y)
31 for i to n reversed {
32 x[i] = (x[i] - sum(LU[i, i + 1:] * x[i + 1:])) / LU[i, i]
33 }
34 return x
35}
410def lu_solve(LU: Arr, perm: Arr, b: Arr) {
411 n = len(b)
412 y = py_copy(py_gath(b, perm, 0))
413 for i to n {
414 y[i] -= sum(LU[py_ix(i, len(LU)), :py_lo(i, shape(LU)[1])] * y[:py_lo(i, len(y))])
415 }
416 x = py_copy(y)
417 for i to n reversed {
418 x[i] = py_div_f(x[i] - sum(LU[py_ix(i, len(LU)), py_lo(py_int(i + 1), shape(LU)[1]):] * x[py_lo(py_int(i + 1), len(x)):]), LU[i, i])
419 }
420 return x
421}
4 solve a system
33A = np.array([[2.0, 1.0, 1.0],
34 [4.0, -6.0, 0.0],
35 [-2.0, 7.0, 2.0]])
36b = np.array([5.0, -2.0, 9.0])
37LU, perm = lu_factor(A)
38x = lu_solve(LU, perm, b)
39print("pivot order", perm.tolist())
40print("x =", x.tolist())
38A = [[2.0, 1.0, 1.0],
39 [4.0, -6.0, 0.0],
40 [-2.0, 7.0, 2.0]]
41b = [5.0, -2.0, 9.0]
42LU, perm = lu_factor(A)
43x = lu_solve(LU, perm, b)
44print("pivot order", perm)
45print("x =", fmt(x))
423A = [[2.0, 1.0, 1.0], [4.0, -6.0, 0.0], [-2.0, 7.0, 2.0]]
424b = [5.0, -2.0, 9.0]
425LU, perm = lu_factor(A)
426x = lu_solve(LU, perm, b)
427print("pivot order", fmt(perm, "d"))
428print("x =", fmt(x))
5 against the library solver
43x_lib = np.linalg.solve(A, b)
44print("library x =", x_lib.tolist())
48x_lib = zeros(len(b))
49lin_solve(A, b, x_lib)
50print("library x =", fmt(x_lib))
430x_lib = py_solve(A, b)
431print("library x =", fmt(x_lib))
6 checks
47assert np.allclose(A @ x, b, atol=1e-12)
48assert np.allclose(x, x_lib, atol=1e-12)
53test "solves" {
54 assert_allclose(A @ x, b, 1e-12)
55 assert_allclose(x, x_lib, 1e-12)
56}
433if not (py_ballclose(A @ x, b, 1e-05, 1e-12)) {
434 stop("AssertionError")
435}
436if not (py_ballclose(x, x_lib, 1e-05, 1e-12)) {
437 stop("AssertionError")
438}
What they print
identical — the two programs print the same text.
Dimmed lines are left out of the comparison. Each Viv test block prints one PASS line; the Python checks the same facts with assert, which prints nothing.
$ python3 lu_solve.py
pivot order [1, 0, 2]x = [1.0, 1.0, 2.0]library x = [1.0, 1.0, 2.0]
$ viv run lu_solve.viv
pivot order [1, 0, 2]x = [1.0, 1.0, 2.0]library x = [1.0, 1.0, 2.0]PASS solves
$ viv run lu_solve.py (py2viv’s translation)
The translation prints exactly what the Python prints.
pivot order [1, 0, 2]x = [1.0, 1.0, 2.0]library x = [1.0, 1.0, 2.0]