Python and Viv · Numerics
Runge-Kutta 4 on an oscillator
What to notice
np.array([y[1], -w * w * y[0]])is the array literal[y[1], -w * w * y[0]]: Viv arrays need no library, and+,-,*and/between an array and a number, or two arrays of one shape, work elementwise as NumPy's do.rk4_stepis the same five lines in both.for _ in range(steps)names its counter in Viv (for k to steps).- A variable assigned inside a loop at the top level is a global in both, so the test reads the last
err. np.coson a number is the builtincos.assert err < 1e-6isassert_lt(err, 1e-6), which prints both numbers when it fails.
Python rk4.py · 33 lines
Viv rk4.viv · 40 lines
py2viv translated from rk4.py
What py2viv adds before the program
(nothing)
(nothing)
121 lines: a header and the guard functions its translation calls
1###
2Translated from rk4.py by py2viv 0.2.0 (numeric Python -> Viv).
3Python source sha256: c76476b7b39eaea9721f68958b8a848516499d33df796c0a714a652463a3b6c2
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#| a / b with a float: ZeroDivisionError on a zero divisor.
18def py_div_f(a, b) {
19 if b == 0 {
20 stop("ZeroDivisionError: float division by zero")
21 }
22 return a / b
23}
24
25#| The shape NumPy broadcasts two shapes da and db to (aligned at the last axis, a length 1
26#| stretched); the run stops where NumPy raises ValueError.
27def py_bdims(da: Arr, db: Arr) {
28 n = len(da)
29 if len(db) > n {
30 n = len(db)
31 }
32 e = zeros(n)
33 for i to n {
34 x = 1
35 y = 1
36 if i >= n - len(da) {
37 x = da[i - n + len(da)]
38 }
39 if i >= n - len(db) {
40 y = db[i - n + len(db)]
41 }
42 if x == y or y == 1 {
43 e[i] = x
44 }
45 else {
46 if x == 1 {
47 e[i] = y
48 }
49 else {
50 stop("ValueError: operands could not be broadcast together with shapes")
51 }
52 }
53 }
54 return e
55}
56
57#| The flat indices, in row-major order of the shape e, of the entries of an array of shape d
58#| that broadcasting to e reads (a length 1 of d is read again along that axis of e).
59def py_bidx(d: Arr, e: Arr) {
60 n = len(e)
61 m = len(d)
62 st = zeros(n)
63 s = 1
64 for k to m {
65 j = m - 1 - k
66 if d[j] != 1 {
67 st[j + n - m] = s
68 }
69 s = s * d[j]
70 }
71 idx = zeros(1)
72 for i to n {
73 w = e[i]
74 nxt = zeros(len(idx) * w)
75 for p to len(idx) {
76 for c to w {
77 nxt[p * w + c] = idx[p] + c * st[i]
78 }
79 }
80 idx = nxt
81 }
82 return idx
83}
84
85#| a broadcast to the shape e (a shape py_bdims gave for a and another array): a itself when
86#| it has that shape, else a new array whose entries are copies of a's.
87def py_bto(a: Arr, e: Arr) {
88 d = shape(a)
89 if len(d) == len(e) {
90 if all(d == e) {
91 return a
92 }
93 }
94 return py_shaped(take(a, py_bidx(d, e)), e)
95}
96
97#| a + b of two real arrays, broadcast as NumPy broadcasts them.
98def py_badd(a: Arr, b: Arr) {
99 e = py_bdims(shape(a), shape(b))
100 return py_bto(a, e) + py_bto(b, e)
101}
102
103#| The entries of f, row-major, as an array of shape e (a vector of 1 to 4 lengths).
104def py_shaped(f: Arr, e: Arr) {
105 if len(e) == 1 {
106 return f
107 }
108 if len(e) == 2 {
109 z = zeros(e[0], e[1])
110 }
111 else {
112 if len(e) == 3 {
113 z = zeros(e[0], e[1], e[2])
114 }
115 else {
116 z = zeros(e[0], e[1], e[2], e[3])
117 }
118 }
119 put(z, py_iota(numel(z)), f)
120 return z
121}
1 imports
4import numpy as np
2# nothing to import: arrays and their arithmetic are built in
(nothing)
2 the ODE: a harmonic oscillator as a first-order system
7OMEGA = 2.0
8
9def rhs(t, y):
10 return np.array([y[1], -OMEGA * OMEGA * y[0]])
5OMEGA = 2.0
6
7def rhs(t, y) {
8 return [y[1], -OMEGA * OMEGA * y[0]]
9}
127OMEGA = 2.0
128def rhs(t, y: Arr) {
129 return [y[1], -OMEGA * OMEGA * y[0]]
130}
3 one classical Runge-Kutta step
13def rk4_step(f, t, y, h):
14 k1 = f(t, y)
15 k2 = f(t + h / 2, y + h / 2 * k1)
16 k3 = f(t + h / 2, y + h / 2 * k2)
17 k4 = f(t + h, y + h * k3)
18 return y + h / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
12def rk4_step(f, t, y, h) {
13 k1 = f(t, y)
14 k2 = f(t + h / 2, y + h / 2 * k1)
15 k3 = f(t + h / 2, y + h / 2 * k2)
16 k4 = f(t + h, y + h * k3)
17 return y + h / 6 * (k1 + 2 * k2 + 2 * k3 + k4)
18}
132def rk4_step(f, t, y: Arr, h) {
133 k1 = f(t, y)
134 k2 = f(t + h / 2, py_badd(y, h / 2 * k1))
135 k3 = f(t + h / 2, py_badd(y, h / 2 * k2))
136 k4 = f(t + h, py_badd(y, h * k3))
137 return py_badd(y, h / 6 * py_badd(py_badd(py_badd(k1, 2 * k2), 2 * k3), k4))
138}
4 integrate to t = 1 and compare with cos(2t)
21T_END = 1.0
22for steps in [10, 20, 40]:
23 h = T_END / steps
24 y = np.array([1.0, 0.0])
25 t = 0.0
26 for _ in range(steps):
27 y = rk4_step(rhs, t, y, h)
28 t += h
29 err = abs(y[0] - np.cos(OMEGA * T_END))
30 print(f"{steps:3d} steps: x(1) = {y[0]:.10f}, error {err:.2e}")
21T_END = 1.0
22for steps in [10, 20, 40] {
23 h = T_END / steps
24 y = [1.0, 0.0]
25 t = 0.0
26 for k to steps {
27 y = rk4_step(rhs, t, y, h)
28 t += h
29 }
30 err = abs(y[0] - cos(OMEGA * T_END))
31 print(f"{steps:3d} steps: x(1) = {y[0]:.10f}, error {err:.2e}")
32}
140T_END = 1.0
141py_v1 = [10, 20, 40]
142for py_i2 to 3 {
143 steps = py_v1[py_i2]
144 h = py_div_f(T_END, steps)
145 y = [1.0, 0.0]
146 t = 0.0
147 for _ to steps {
148 y = rk4_step(rhs, t, y, h)
149 t += h
150 }
151 err = abs(y[0] - cos(OMEGA * T_END))
152 print(fmt(steps + 0.0, "3d"), " steps: x(1) = ", fmt(y[0], ".10f"), ", error ", fmt(err, ".2e"), sep="")
153}
5 checks
33assert err < 1e-6
35test "fourth order" {
36 assert_lt(err, 1e-6)
37}
155if not (err < 1e-06) {
156 stop("AssertionError")
157}
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 rk4.py
10 steps: x(1) = -0.4161210938, error 2.57e-0520 steps: x(1) = -0.4161452687, error 1.57e-0640 steps: x(1) = -0.4161467401, error 9.64e-08
$ viv run rk4.viv
10 steps: x(1) = -0.4161210938, error 2.57e-0520 steps: x(1) = -0.4161452687, error 1.57e-0640 steps: x(1) = -0.4161467401, error 9.64e-08PASS fourth order
$ viv run rk4.py (py2viv’s translation)
The translation prints exactly what the Python prints.
10 steps: x(1) = -0.4161210938, error 2.57e-0520 steps: x(1) = -0.4161452687, error 1.57e-0640 steps: x(1) = -0.4161467401, error 9.64e-08