R: Mandelbrot (Rcpp 0.12.6) その2

先日のコードのバグとりのついでに、画像出力まで関数化してみました。

library(Rcpp)
library(inline)

## C++ source
mandelbrot.src <- "
Rcpp::NumericVector x(xx);
Rcpp::NumericVector y(yy);
unsigned short t = Rcpp::as<unsigned short>(threshold);
int nx = x.size(), ny = y.size();
Rcpp::NumericVector u(nx);

if (nx != ny) {
  return Rcpp::wrap(-1);
} else {
  for (int i = 0; i < nx; i++) {
    std::complex<double> z(0.0, 0.0);
    std::complex<double> c(x[i], y[i]);
    unsigned short k = 0;

    while (k < t) {
      z = z * z + c;
      if (abs(z) > 2.0) break;
      k++;
    }
    u[i] = k;
  }
  return Rcpp::wrap(u);
}
"

## Call C++
mandelbrot.c <- rcpp(signature(xx = "numeric",
                               yy = "numeric",
                               threshold = "numeric"),
                     mandelbrot.src, 
                     includes = "#include <complex>")


## R function
mandelbrot <- function(min.x = -2, max.x = 1, min.y = -1.5, max.y = 1.5,
                       px = 256,
                       py = round((max.y - min.y) / (max.x - min.x) * px),
                       threshold = 255,
                       col = terrain.colors(threshold + 1)) {
  if (min.x > max.x) {
    a <- min.x
    min.x <- max.x
    max.x <- a
  }
  if (min.y > max.y) {
    a <- min.y
    min.y <- max.y
    max.y <- a
  }
  if (px < 0) px <- -px
  if (py < 0) py <- -py

  rx <- seq(min.x, max.x, length = px)
  ry <- seq(min.y, max.y, length = py)
  xx <- rep(rx, py)
  yy <- rep(ry, each = px)

  z <- mandelbrot.c(xx, yy, threshold)
  matz <- matrix(z, nrow = px)
  image(seq(min.x, max.x, length = nrow(matz)),
        seq(min.y, max.y, length = ncol(matz)),
        matz,
        col = col,
        xlab = "x", ylab = "y", asp = 1.0)
}

例

mandelbrot(0.2595, 0.2600, 0.0015, 0.0020, 900, 900,
           threshold = 1023,
           col = rep(rainbow(256), (1023 + 1) / 256))

タグ：R