【Rust】並列処理でマンデルブロ集合の画像を作るコード

最近Rustを勉強しています。 体系的に学ぶためにO’ReillyのプログラミングRustを買いました。

book

第2章のRustがどんな言語なのか体験するために書くコードが、まさかのマンデルブロ集合で驚きつつも楽しく取り組めました。

この記事では全ては説明しませんが、自分なりに重要なところとかをまとめます。

ToC

何ができるようになったか

Rustでマンデルブロ集合の画像を描画するコードを書きながらRustについて学びました。 さらに、並列処理でレンダリングもしています。

完成すると、以下のコマンドでマンデルブロ集合の画像を出力します。

# ビルドした実行ファイル 画像ファイル名 サイズ 左上複素数 右下複素数
$ mandelbrot.exe  mandelbrot.png 1000x750 -1.20,0.35 -1,0.20

出力 mandelbrot

8ビットのグレイスケールですが模様と淡いグレイが大理石っぽくて綺麗です!

マンデルブロ集合について

マンデルブロ集合(Mandelbrot set)は拡大していくと似たようなモノが繰り返し現れるフラクタルとしても有名で、拡大していく映像はYouTubeでも結構再生されてて人気です。「神の指紋」とも呼ばれている2

↓4Kの解像度のまま1時間かけて拡大していく超大作。無限に見てられる。

マンデルブロ集合は次の漸化式

{zn+1=zn2+cz0=0{\displaystyle {\begin{cases} z_{n+1}=z_{n}^{2}+c\\z_{0}=0 \end{cases}} }

で定義される複素数列 zn{z_n} nN0n∈N∪{0}nn → ∞ の極限で無限大に発散しないという条件を満たす複素数 cc 全体が作る集合がマンデルブロ集合である1

式を展開するとこんな感じ。

zn+1=zn2+czn+1=(Ren+Imn)2+czn+1=(a+bi)2+czn+1=a2b2+2abi+cz_{n+1}=z_{n}^{2}+c\\ z_{n+1}={(Re_n+Im_n)}^{2}+c\\ z_{n+1}={(a+bi)}^{2}+c\\ z_{n+1}={a}^{2}-{b}^{2}+2abi+c

iは1=i\sqrt{-1}=iの複素数です。 zn+1=a2b2+2abi+cz_{n+1}={a}^{2}-{b}^{2}+2abi+cとなってzzが無限大に発散していくが、その速度は複素数によって異なります。 Rustではnum::Complexのクレートを使っているため、複素数の計算もz * zと簡潔に書けます。

画像にするときには横軸を実数、縦軸を虚数と見立て画素の座標における複素数について発散(または一定回数)まで繰り返し計算を求めます。 発散するまでの計算回数を画像の濃淡として表すと、このような美しい模様になります。

参考

コード

コードはクレートのバージョンを新しくしたため、書籍のサンプルとはすこし異なります。 書籍は2018年の第2刷です。

書籍のサンプルから以下の点を変更しています。

  • クレートを最新化した
  • 並列処理のクレートを crossbeamからrayonに変更した

書いたコードのリポジトりはこちら https://github.com/suzu6/trial-code/tree/master/rust/algorithms/mandelbrot

コードは最終的にO'Reillyのリポジトリにあったものに寄せました。

Code examples for the book Programming Rust, from O'Reilly
https://github.com/ProgrammingRust/mandelbrot

クレートのバージョン

Cargo.toml
[dependencies]
# I update to latest version crate.
num = "0.4.0"      # <- 0.1.27
image = "0.23.14"  # <- 0.13.0
rayon = "1.5.0"    # <- crossbeam=0.2.8

O’Reillyのリポジトリでは並列処理がcrossbeamからrayonに変わってました。

main(): 処理の流れ

全て一つのファイルに書いているので少し長めです。 個人的にメイン関数が一番上にあって実行する順に関数があると全体を把握しやすいのでそう並べています。 関数でわけて補足していきます。

全体の処理の流れは以下の通りです。

  1. 引数の受け取り
  2. 引数から画像サイズと複素空間の位置の情報をパースする
  3. マンデルブロ集合の計算
    1. 直列または並列処理
  4. 計算結果を画像として書き出す
main.rs
use std::io::Write;

fn main() {
    let args: Vec<String> = std::env::args().collect();

    if args.len() != 5 {
        writeln!(
            std::io::stderr(),
            "Usage: mandelbrot FILE PIXELS UPPERLEFT LOWERRIGHT"
        )
        .unwrap();
        writeln!(
            std::io::stderr(),
            "Example: {} mandel.png 1000x750 -1.20,0.35 -1,0.20",
            args[0]
        )
        .unwrap();
        std::process::exit(1);
    }

    let bounds = parse_pair(&args[2], 'x').expect("error parsing image dimensions");
    let upper_left = parse_complex(&args[3]).expect("error parsing upper left corner point");
    let lower_right = parse_complex(&args[4]).expect("error parsing lower right corner point");

    
    // Single thread rendering;
    let pixels = single_thread_render(bounds, upper_left, lower_right);
    write_image("no_thread.png", &pixels, bounds).expect("error writing PNG file");
    
    // Multi thread rendering;
    let pixels = multi_threads_render(bounds, upper_left, lower_right);
    write_image(&args[1], &pixels, bounds).expect("error writing PNG file");
}

引数は<ファイル名> <画像サイズ(widthxheight)> <複素数空間の左上(Re,Im)> <複素数空間の右下(Re,Im)>と取っています。

引数の例
# Usage: mandelbrot FILE PIXELS UPPERLEFT LOWERRIGHT
$ cargo run mandel.png 1000x750 -1.20,0.35 -1,0.20

引数の数が合わない場合は処理の前に終了します。

# 引数が多い
> cargo run no_thread.png 500x500 1.0,0.0 -1.0,1.0 xxxx
   Compiling mandelbrot v0.1.0 (C:\git\trial-code\rust\algorithms\mandelbrot)
    Finished dev [unoptimized + debuginfo] target(s) in 2.31s
     Running `target\debug\mandelbrot.exe no_thread.png 500x500 1,0 -1.0,1 xxxx`
Usage: mandelbrot FILE PIXELS UPPERLEFT LOWERRIGHT
Example: target\debug\mandelbrot.exe mandel.png 1000x750 -1.20,0.35 -1,0.20
error: process didn't exit successfully: `target\debug\mandelbrot.exe no_thread.png 500x500 1,0 -1.0,1 xxxx` 
(exit code: 1)

parse_pair(): 引数のパース①

引数で受け取った画像サイズwidthxheightを区切り文字xで分けてタプル()で返します。

面白いのは返すタプルがジェネリック型になっているためparse_pair::<i32>("10,20", ',')で32ビット整数型、parse_pair::<f64>("0.5x", 'x')と書くと64ビット浮動小数点型のタプルが返ってきます。

返り値はOptionです。 型のパースが出来なかった場合はNoneが返るため、その条件分岐をmatchで判定している。

parse_pair()
/// Parse the string `s` as a coordinate pair, like `"400x600"` or `"1.0,0.5"`.
///
/// Specifically, `s` should have the form <left><sep><right>, where <sep> is
/// the character given by the `separator` argument, and <left> and <right> are both
/// strings that can be parsed by `T::from_str`.
///
/// If `s` has the proper form, return `Some<(x, y)>`. If it doesn't parse
/// correctly, return `None`.
use std::str::FromStr;

fn parse_pair<T: FromStr>(s: &str, separator: char) -> Option<(T, T)> {
    match s.find(separator) {
        None => None,
        Some(index) => match (T::from_str(&s[..index]), T::from_str(&s[index + 1..])) {
            (Ok(l), Ok(r)) => Some((l, r)),
            _ => None,
        },
    }
}

#[test]
fn test_parse_pair() {
    assert_eq!(parse_pair::<i32>("", ','), None);
    assert_eq!(parse_pair::<i32>("10,", ','), None);
    assert_eq!(parse_pair::<i32>(",10", ','), None);
    assert_eq!(parse_pair::<i32>("10,20", ','), Some((10, 20)));
    assert_eq!(parse_pair::<i32>("10,20xy", ','), None);
    assert_eq!(parse_pair::<f64>("0.5x", 'x'), None);
    assert_eq!(parse_pair::<f64>("0.5x1.5", 'x'), Some((0.5, 1.5)));
}

テストもこのように書くことができる。

引数のパースに失敗した場合はNoneが返るため、その後の処理で型が合わずpanicが起きて終了します。 型があると何行目でどの引数で失敗したか示されてデバッグがしやすいですね。

# パースに失敗した場合
$ cargo run no_thread.png 500x500 1.0,0.0 -1.0,1.0s    
    Finished dev [unoptimized + debuginfo] target(s) in 0.12s
     Running `target\debug\mandelbrot.exe no_thread.png 500x500 1,0 -1.0,1.0s`
thread 'main' panicked at 'error parsing lower right corner point', src\main.rs:23:47
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
error: process didn't exit successfully: `target\debug\mandelbrot.exe no_thread.png 500x500 1,0 -1.0,1.0s` (exit code: 101)

parse_complex(): 引数のパース②

parse_pair()をラップしてf64,f64に対応した関数。 タプル(f64, f64)からComlex<f64>への変換もしています。

extern crate num;
use num::Complex;

// Parse a pair of floating-point numbers separated by a comma as a complex number.
fn parse_complex(s: &str) -> Option<Complex<f64>> {
    match parse_pair(s, ',') {
        Some((re, im)) => Some(Complex { re, im }),
        None => None,
    }
}

#[test]
fn test_parse_complex() {
    assert_eq!(
        parse_complex("1.25,-0.00625"),
        Some(Complex {
            re: 1.25,
            im: -0.00625
        })
    );
    assert_eq!(parse_complex(", -0.00625"), None);
}

pixel_to_point(): 画素から複素平面の座標への写像

画像と複素数の二つの空間を対応させる必要があります。 この関数では画素から複素平面の座標への写像です。

虚部の計算が引き算になっているのは、上に点が動くと画素は高さが増えるが虚部は減るためです。

pixel_to_point()
/// Given the row and column of a pixel in the output image, return the
/// corresponding point on the complex plane.
///
/// `bounds` is a pair giving the width and height of the image in pixels.
/// `pixel` is a (column, row) pair indicating a particular pixel in that image.
/// The `upper_left` and `lower_right` parameters are points on the complex
/// plane designating the area our image covers.
fn pixel_to_point(
    bounds: (usize, usize),
    pixel: (usize, usize),
    upper_left: Complex<f64>,
    lower_right: Complex<f64>,
) -> Complex<f64> {
    let (width, height) = (
        lower_right.re - upper_left.re,
        upper_left.im - lower_right.im,
    );
    Complex {
        re: upper_left.re + pixel.0 as f64 * width / bounds.0 as f64,
        im: upper_left.im - pixel.1 as f64 * height / bounds.1 as f64,
    }
}

#[test]
fn test_pixel_to_point() {
    assert_eq!(
        pixel_to_point(
            (100, 100),
            (25, 75),
            Complex { re: -1.0, im: 1.0 },
            Complex { re: 1.0, im: -1.0 }
        ),
        Complex { re: -0.5, im: -0.5 }
    );
}

single_thread_render()とrender(): レンダリング

こちらはシングルスレッドでレンダリングをしています。

render()に渡すpixelsは参照渡しであり、画素ごとに独立して計算できるため並列に毛産することも可能です。 後で並列の関数を紹介します。

single_thread_render()
fn single_thread_render(
    bounds: (usize, usize),
    upper_left: Complex<f64>,
    lower_right: Complex<f64>
) -> Vec<u8> {
    let mut pixels = vec![0; bounds.0 * bounds.1];
    render(&mut pixels, bounds, upper_left, lower_right);

    pixels
}

/// Render a rectangle of the Mandelbrot set into a buffer of pixels.
///
/// The `bounds` argument gives the width and height of the buffer `pixels`,
/// which holds one grayscale pixel per byte. The `upper_left` and `lower_right`
/// arguments specify points on the complex plane corresponding to the upper-
/// left and lower-right corners of the pixel buffer.
fn render(
    pixels: &mut [u8],
    bounds: (usize, usize),
    upper_left: Complex<f64>,
    lower_right: Complex<f64>,
) {
    assert!(pixels.len() == bounds.0 * bounds.1);

    for row in 0..bounds.1 {
        for column in 0..bounds.0 {
            let point = pixel_to_point(bounds, (column, row), upper_left, lower_right);
            pixels[row * bounds.0 + column] = match escape_time(point, 255) {
                None => 0,
                Some(count) => 255 - count as u8,
            };
        }
    }
}

escape_time(): 漸化式とリミット

マンデルブロ集合の画像を作るためには、画像上のすべての画素に対する複素平面上の点が発散するまでの計算回数をカウントする必要がある。 グレイスケールは8bitの画素で表現できるため、今回limitには255を受け取る。

#[allow(dead_code)]
fn escape_time(c: Complex<f64>, limit: u32) -> Option<u32> {
    let mut z = Complex { re: 0.0, im: 0.0 };
    for i in 0..limit {
        z = z * z + c;
        if z.norm_sqr() > 4.0 {
            return Some(i);
        }
    }

    None
}

z.norm_sqr()は複素数のノルム(絶対値)を求めている。

write_image(): 画像への書き出し

8ビットの2次元配列を画像として出力します。 iamgeクレートを利用してpngのグレイスケールの画像を書き出します。

write_image()
extern crate image;
use image::png::PngEncoder;
use image::ColorType;
use std::fs::File;

/// Write the buffer `pixels`, whose dimensions are given by
/// `bounds`, to the file named `filename`.
fn write_image(
    filename: &str,
    pixels: &[u8],
    bounds: (usize, usize),
) -> Result<(), image::ImageError> {
    let output = File::create(filename)?;

    let encoder = PngEncoder::new(output);
    let result = encoder.encode(&pixels, bounds.0 as u32, bounds.1 as u32, ColorType::L8);
    match result {
        Ok(()) => {
            return Ok(());
        }
        Err(e) => {
            return Err(e);
        }
    }
}

multi_threads_render(): 並列処理

rayonを使って並列に画素のレンダリングをします。

pixelsは画像の左上から右下までを順に並べた1次元配列です。 これをpixels.chunks_mut(bounds.0)で画像の幅毎に配列を切り分けて、2次元の配列にします。 この幅毎に分けてそれぞれ並列に計算していきます。

std::vec::Vec::chunks_mut

use rayon::prelude::*;

fn multi_threads_render(
    bounds: (usize, usize),
    upper_left: Complex<f64>,
    lower_right: Complex<f64>
) -> Vec<u8> {
    let mut pixels = vec![0; bounds.0 * bounds.1];
    // Scope of slicing up `pixels` into horizontal bands.
    {
        let bands: Vec<(usize, &mut [u8])> = pixels.chunks_mut(bounds.0).enumerate().collect();

        bands.into_par_iter().for_each(|(i, band)| {
            let top = i;
            let band_bounds = (bounds.0, 1);
            let band_upper_left = pixel_to_point(bounds, (0, top), upper_left, lower_right);
            let band_lower_right =
                pixel_to_point(bounds, (bounds.0, top + 1), upper_left, lower_right);
            render(band, band_bounds, band_upper_left, band_lower_right);
        });
    }

    pixels
}

 関連記事

 新着 or 更新記事