This site is version 0.0.1. Most of the code on this site is drafted by an agent and needs review. I apologize for the AI-speak, we are working on rewriting it!
viv-lang.org

Python and Viv · Numerics

Jacobi iteration

What to notice

Open the Viv in the workbench The atlas of the Viv py2viv translates this Python; the translation prints the same output.
Python jacobi.py · 34 lines
Viv jacobi.viv · 45 lines
py2viv translated from jacobi.py
What py2viv adds before the program
(nothing)
(nothing)

191 lines: a header and the guard functions its translation calls

1###
2Translated from jacobi.py by py2viv 0.2.0 (numeric Python -> Viv).
3Python source sha256: acb1ded27e05fd38a02f33a53766a794ebdf9afa55cc7336baf18531aee4fd11
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#| An int result, as Python's int: a -0.0 becomes 0 (an int has no sign on zero), and the run
23#| stops at 2^53, past which a double no longer holds every integer (Python's int does).
24def py_int(x) {
25 if abs(x) >= 9007199254740992 {
26 stop("py2viv: an integer reached 2^53; Python's int is exact there and a double is not")
27 }
28 return x + 0.0
29}
30
31#| np.max(a): NaN if any entry is NaN, else the largest (+0 over -0).
32def py_npmax(a: Arr) {
33 if any(a != a) {
34 return 0.0 / 0.0
35 }
36 f = py_flat(a)
37 m = f[0]
38 for i from 1 to len(f) {
39 m = py_zmax2(m, f[i])
40 }
41 return m
42}
43
44#| np.isclose(a, b, rtol, atol) of two numbers: |a - b| <= atol + rtol |b| when both are
45#| finite, a == b otherwise (a NaN is never close).
46def py_npisclose(a, b, rtol, atol) {
47 if abs(a) < 1.0 / 0.0 and abs(b) < 1.0 / 0.0 {
48 return abs(a - b) <= atol + rtol * abs(b)
49 }
50 return a == b
51}
52
53#| np.allclose(a, b, rtol, atol) of two arrays of one shape: np.isclose entry by entry.
54def py_allclose(a: Arr, b: Arr, rtol, atol) {
55 fa = py_flat(a)
56 fb = py_flat(b)
57 if len(fa) != len(fb) {
58 stop("py2viv: np.allclose of two arrays of different sizes (NumPy broadcasts them)")
59 }
60 for i to len(fa) {
61 if not py_npisclose(fa[i], fb[i], rtol, atol) {
62 return false
63 }
64 }
65 return true
66}
67
68#| The larger of two numbers that are not NaN, +0 over -0 (NumPy's rule; Viv's max of two
69#| zeros differs between engines).
70def py_zmax2(a, b) {
71 if a > b {
72 return a
73 }
74 if b > a {
75 return b
76 }
77 if a == 0 {
78 return a + b
79 }
80 return a
81}
82
83#| The shape NumPy broadcasts two shapes da and db to (aligned at the last axis, a length 1
84#| stretched); the run stops where NumPy raises ValueError.
85def py_bdims(da: Arr, db: Arr) {
86 n = len(da)
87 if len(db) > n {
88 n = len(db)
89 }
90 e = zeros(n)
91 for i to n {
92 x = 1
93 y = 1
94 if i >= n - len(da) {
95 x = da[i - n + len(da)]
96 }
97 if i >= n - len(db) {
98 y = db[i - n + len(db)]
99 }
100 if x == y or y == 1 {
101 e[i] = x
102 }
103 else {
104 if x == 1 {
105 e[i] = y
106 }
107 else {
108 stop("ValueError: operands could not be broadcast together with shapes")
109 }
110 }
111 }
112 return e
113}
114
115#| The flat indices, in row-major order of the shape e, of the entries of an array of shape d
116#| that broadcasting to e reads (a length 1 of d is read again along that axis of e).
117def py_bidx(d: Arr, e: Arr) {
118 n = len(e)
119 m = len(d)
120 st = zeros(n)
121 s = 1
122 for k to m {
123 j = m - 1 - k
124 if d[j] != 1 {
125 st[j + n - m] = s
126 }
127 s = s * d[j]
128 }
129 idx = zeros(1)
130 for i to n {
131 w = e[i]
132 nxt = zeros(len(idx) * w)
133 for p to len(idx) {
134 for c to w {
135 nxt[p * w + c] = idx[p] + c * st[i]
136 }
137 }
138 idx = nxt
139 }
140 return idx
141}
142
143#| a broadcast to the shape e (a shape py_bdims gave for a and another array): a itself when
144#| it has that shape, else a new array whose entries are copies of a's.
145def py_bto(a: Arr, e: Arr) {
146 d = shape(a)
147 if len(d) == len(e) {
148 if all(d == e) {
149 return a
150 }
151 }
152 return py_shaped(take(a, py_bidx(d, e)), e)
153}
154
155#| a - b of two real arrays, broadcast.
156def py_bsub(a: Arr, b: Arr) {
157 e = py_bdims(shape(a), shape(b))
158 return py_bto(a, e) - py_bto(b, e)
159}
160
161#| a / b of two real arrays, broadcast (a zero divisor gives inf or NaN, as NumPy).
162def py_bdiv(a: Arr, b: Arr) {
163 e = py_bdims(shape(a), shape(b))
164 return py_bto(a, e) / py_bto(b, e)
165}
166
167#| np.allclose(a, b, rtol, atol) of two arrays, broadcast.
168def py_ballclose(a: Arr, b: Arr, rtol, atol) {
169 e = py_bdims(shape(a), shape(b))
170 return py_allclose(py_bto(a, e), py_bto(b, e), rtol, atol)
171}
172
173#| The entries of f, row-major, as an array of shape e (a vector of 1 to 4 lengths).
174def py_shaped(f: Arr, e: Arr) {
175 if len(e) == 1 {
176 return f
177 }
178 if len(e) == 2 {
179 z = zeros(e[0], e[1])
180 }
181 else {
182 if len(e) == 3 {
183 z = zeros(e[0], e[1], e[2])
184 }
185 else {
186 z = zeros(e[0], e[1], e[2], e[3])
187 }
188 }
189 put(z, py_iota(numel(z)), f)
190 return z
191}
1 imports
4import numpy as np
2# nothing to import: matrices, @ and the reductions are built in
(nothing)
2 a diagonally dominant system
7A = np.array([[10.0, -1.0, 2.0, 0.0],
8 [-1.0, 11.0, -1.0, 3.0],
9 [2.0, -1.0, 10.0, -1.0],
10 [0.0, 3.0, -1.0, 8.0]])
11b = np.array([6.0, 25.0, -11.0, 15.0])
5A = [[10.0, -1.0, 2.0, 0.0],
6 [-1.0, 11.0, -1.0, 3.0],
7 [2.0, -1.0, 10.0, -1.0],
8 [0.0, 3.0, -1.0, 8.0]]
9b = [6.0, 25.0, -11.0, 15.0]
197A = [[10.0, -1.0, 2.0, 0.0], [-1.0, 11.0, -1.0, 3.0], [2.0, -1.0, 10.0, -1.0], [0.0, 3.0, -1.0, 8.0]]
198b = [6.0, 25.0, -11.0, 15.0]
3 Jacobi iteration: x_new = D^-1 (b - R x)
14def jacobi(A, b, tol, max_iter):
15 d = np.diag(A)
16 R = A - np.diag(d)
17 x = np.zeros(len(b))
18 for k in range(max_iter):
19 x_new = (b - R @ x) / d
20 if np.max(np.abs(x_new - x)) < tol:
21 return x_new, k + 1
22 x = x_new
23 return x, max_iter
12def jacobi(A, b, tol, max_iter) {
13 n = len(b)
14 d = zeros(n)
15 R = copy(A)
16 for i to n {
17 d[i] = A[i, i]
18 R[i, i] = 0.0
19 }
20 x = zeros(n)
21 for k to max_iter {
22 x_new = (b - R @ x) / d
23 if max(abs(x_new - x)) < tol {
24 return x_new, k + 1
25 }
26 x = x_new
27 }
28 return x, max_iter
29}
200def jacobi(A: Arr, b: Arr, tol, max_iter) {
201 d = diag(A, k=0)
202 R = py_bsub(A, diag(d, k=0))
203 x = zeros(len(b))
204 for k to max_iter {
205 x_new = py_bdiv(py_bsub(b, R @ x), d)
206 if py_npmax(abs(py_bsub(x_new, x))) < tol {
207 return x_new, py_int(k + 1)
208 }
209 x = x_new
210 }
211 return x, max_iter
212}
4 solve and check the residual
26TOL = 1e-10
27MAX_ITER = 200
28x, iters = jacobi(A, b, TOL, MAX_ITER)
29print("iterations:", iters)
30print("x =", " ".join(f"{v:.8f}" for v in x))
31print(f"largest residual {np.max(np.abs(A @ x - b)):.1e}")
32TOL = 1e-10
33MAX_ITER = 200
34x, iters = jacobi(A, b, TOL, MAX_ITER)
35print("iterations:", iters)
36print("x =", " ".join(f"{v:.8f}" for v in x))
37print(f"largest residual {max(abs(A @ x - b)):.1e}")
214TOL = 1e-10
215MAX_ITER = 200
216x, iters = jacobi(A, b, TOL, MAX_ITER)
217print("iterations:", fmt(iters + 0.0, "d"))
218print("x =", join(" ", [fmt(py_v3, ".8f") for py_v3 in x]))
219print("largest residual", fmt(py_npmax(abs(py_bsub(A @ x, b))), ".1e"))
5 checks
34assert np.allclose(x, [1.0, 2.0, -1.0, 1.0], atol=1e-9)
40test "solution" {
41 assert_allclose(x, [1.0, 2.0, -1.0, 1.0], 1e-9)
42}
221if not (py_ballclose(x, [1.0, 2.0, -1.0, 1.0], 1e-05, 1e-09)) {
222 stop("AssertionError")
223}

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 jacobi.py

iterations: 29
x = 1.00000000 2.00000000 -1.00000000 1.00000000
largest residual 3.1e-10

$ viv run jacobi.viv

iterations: 29
x = 1.00000000 2.00000000 -1.00000000 1.00000000
largest residual 3.1e-10
PASS solution

$ viv run jacobi.py (py2viv’s translation)

The translation prints exactly what the Python prints. 

iterations: 29
x = 1.00000000 2.00000000 -1.00000000 1.00000000
largest residual 3.1e-10