だいぶ前だけど、MojoでPythonが高速化できるという話があって、はーなるほどねぇ、ぐらいだったんだけど、ほとんどそのままな感じでRustにしたらどうなるんだっけ?というのをやってなかったのを急に思い出した。
#!/usr/bin/env python3.9
# -*- coding: utf-8 -*-
import time
import numpy as np
from numba import jit, prange
xmin = -1.75
xmax = 0.75
width = 4096
ymin = -1.25
ymax = 1.25
height = 4096
max_iter = 500
@jit(nopython=True, parallel=True)
def mandelbrot_kernel(c):
z = c
for i in range(max_iter):
z = z * z + c
if abs(z) > 2:
return i
return max_iter
@jit(nopython=True, parallel=True)
def compute_mandelbrot(image, xmin, xmax, ymin, ymax, width, height, max_iter):
dx = (xmax - xmin) / width
dy = (ymax - ymin) / height
for j in prange(height):
for i in range(width):
y = ymin + j * dy
x = xmin + i * dx
image[j, i] = mandelbrot_kernel(complex(x, y))
return image
image = np.zeros((height, width), dtype=np.int32)
t0 = time.time()
image = compute_mandelbrot(image, xmin, xmax, ymin, ymax, width, height, max_iter)
t1 = time.time()
print(f"elapsed time: {(t1-t0) * 1000:.3f} ms")
(mandelbrot) rifujita@macstudio mandelbrot % ./mandelbrot.py
elapsed time: 119282.660 ms
Pythonでnumbaを使っても119秒かかる。
use num_complex::Complex;
use rayon::prelude::*;
use std::time::Instant;
const XMIN: f64 = -1.75;
const XMAX: f64 = 0.75;
const WIDTH: usize = 4096;
const YMIN: f64 = -1.25;
const YMAX: f64 = 1.25;
const HEIGHT: usize = 4096;
const MAX_ITER: usize = 500;
fn mandelbrot_kernel(c: Complex<f64>) -> usize {
let mut z = c;
for i in 0..MAX_ITER {
z = z * z + c;
if z.norm() > 2.0 {
return i;
}
}
MAX_ITER
}
fn compute_mandelbrot(image: &mut Vec<Vec<usize>>) {
let dx = (XMAX - XMIN) / WIDTH as f64;
let dy = (YMAX - YMIN) / HEIGHT as f64;
image.par_iter_mut().enumerate().for_each(|(j, row)| {
for i in 0..WIDTH {
let y = YMIN + j as f64 * dy;
let x = XMIN + i as f64 * dx;
row[i] = mandelbrot_kernel(Complex::new(x, y));
}
});
}
fn main() {
let mut image = vec![vec![0; WIDTH]; HEIGHT];
let start = Instant::now();
compute_mandelbrot(&mut image);
let duration = start.elapsed();
println!("elapsed time: {:.3} ms", duration.as_secs_f64() * 1000.0);
}
rifujita@macstudio mandelbrot % cargo run --release
Finished `release` profile [optimized] target(s) in 0.04s
Running `target/release/mandelbrot`
elapsed time: 8239.660 ms
8.2秒だった。14.48倍ぐらいしか速くならなかった。もっと真面目にやらないとダメか。