Mandelbrot rendering code in R

Can we do anything about that first loop? It's currently written in a very C style, modifying single elements in the matrix Mxy. Ideally, in a vectorised implementation, you'd be able to update the whole of the matrix Mxy in one go.

Since the code is pretty complicated and slow, we might need some tests and some benchmarking to ensure that after any changes we introduce the code i) still works properly, and ii) is more efficient.

I pulled out the first loop into a function:

# Some of the variables have been renamed
compute_mxy_initial <- function(
                                n_iter,
                                n_row,
                                n_col,
                                r_escape) {
  # Initial Mandelbrot array
  m_xy <- matrix(1, nrow = n_row, ncol = n_col)
  m_x <- seq(n_row)
  m_y <- seq(n_col)

  # Re(C) vector
  c_x <- -0.5 + (-1 + m_x / n_row * 2) * 1.6
  c_y <- (-1 + m_y / n_col * 2) * 0.9

  # Compute the matrix
  for (ny in seq(n_col)) {
    for (nx in seq(n_row)) {
      x0 <- 0
      y0 <- 0
      s <- 1
      for (j in seq(n_iter)) {
        # should it be called r2 ?
        r <- x0^2 + y0^2
        if (r < r_escape) {
          x1 <- x0^2 - y0^2 + c_x[nx]
          y1 <- 2 * x0 * y0 + c_y[ny]
        } else {
          # are x1 / y1 just discarded here?
          x1 <- x0
          y1 <- y0
          break
        }
        x0 <- x1
        y0 <- y1
        if (r > 0) {
          s <- s + min(r, 1 / r)
        }
      }
      m_xy[nx, ny] <- s
    }
  }
  m_xy
}

.. and made another function compute_mxy_test with the same body (this one will be updated).

To check that the two functions give the same results you can add a few test_that tests:

set.seed(1)

test_that(
  "mxy values generated by the refactored function match the original values", {
    for (i in 1:30) {
      n_iter <- 1 + rpois(1, 10)
      n_row <- 1 + rpois(1, 10)
      n_col <- 1 + rpois(1, 10)
      r_escape <- rexp(1, 1 / 1000)
      expect_equal(
        compute_mxy_test(n_iter, n_row, n_col, r_escape),
        compute_mxy_initial(n_iter, n_row, n_col, r_escape),
        tolerance = 1e-4
      )
    }
  }
)

Anytime that you modify compute_mxy_test rerun the test. (You could alternatively, and more efficiently, precompute the expected values for the matrix and just compare the values created by compute_mxy_test to these).

You can also use the bench package to compare timings, but do this sparingly because your code is pretty slow:

timings <- bench::press(
  n_iter = c(10, 100, 400),
  n_row = c(10, 100, 1600),
  n_col = c(10, 100, 900),
  {
    bench::mark(
      initial = compute_mxy_initial(n_iter, n_row, n_col, r_escape = 1000),
      vectorised = compute_mxy_test(n_iter, n_row, n_col, r_escape = 1000)
    )
  }
)

Modifications to the first loop:

What I'm trying to do here is get the for (j in seq(n_iter)) outside the positional for-loops. Unfortunately, the only way I could see to do that, was to store lots more information in RAM. You could add a data.frame that stores the coordinates, squared-radius etc for each nx,ny pair such that at each new iteration, those points with small radii get updated.

Here's the basic data structure:

compute_mxy_test <- function(
                             n_iter,
                             n_row,
                             n_col,
                             r_escape) {
  # Initial Mandelbrot array
  m_xy <- matrix(1, nrow = n_row, ncol = n_col)

  # Re(C) vector
  c_x <- -0.5 + (-1 + seq(n_row) / n_row * 2) * 1.6
  c_y <- (-1 + seq(n_col) / n_col * 2) * 0.9

  # construct a data-frame for use in vectorised updates
  coords <- expand.grid(
    nx = seq(n_row),
    ny = seq(n_col)
  ) %>%
    dplyr::mutate(
      # starting values for each nx/ny pair
      x0 = 0,
      y0 = 0,
      s = 1,
      r = 0
    )

  # Compute the matrix

  # iterate over the rows of `coords` and then extract nx/ny from each row
  for (i in seq(nrow(coords))) {
    nx <- coords[i, "nx"]
    ny <- coords[i, "ny"]
    x0 <- coords[i, "x0"]
    y0 <- coords[i, "y0"]
    s <- coords[i, "s"]
    r <- coords[i, "r"]
    for (j in seq(n_iter)) {
      # should it be called r2 ?
      if (r < r_escape) {
        x1 <- x0^2 - y0^2 + c_x[nx]
        y1 <- 2 * x0 * y0 + c_y[ny]
        s <- ifelse(r > 0, s + min(r, 1 / r), s)
        x0 <- x1
        y0 <- y1
        r <- x0^2 + y0^2
      }
    }
    # update the values for x0 etc for this row
    coords[i, c("x0", "y0", "s", "r")] <- c(x0, y0, s, r)
  }
  m_xy[cbind(coords$nx, coords$ny)] <- coords$s
  m_xy
}

THen, still working with one row at a time, we can introduce a function that will update all the coordinate-data for that row:

compute_mxy_test <- function(
                             n_iter,
                             n_row,
                             n_col,
                             r_escape) {
  # Initial Mandelbrot array
  m_xy <- matrix(1, nrow = n_row, ncol = n_col)

  # Re(C) vector
  c_x <- -0.5 + (-1 + seq(n_row) / n_row * 2) * 1.6
  c_y <- (-1 + seq(n_col) / n_col * 2) * 0.9

  # construct a data-frame for use in vectorised updates
  coords <- expand.grid(
    nx = seq(n_row),
    ny = seq(n_col)
  ) %>%
    dplyr::mutate(
      x0 = 0,
      y0 = 0,
      s = 1,
      r = 0,
      # note these additive offsets are used in .update_coords()
      cx = c_x[nx],
      cy = c_y[ny]
    )

  .update_coords <- function(df) {
    # can assume that r >= r_escape
    x1 <- with(df, x0^2 - y0^2 + cx)
    y1 <- with(df, 2 * x0 * y0 + cy)
    df %>%
      dplyr::mutate(
        # pmin does pairwise `min` between matched elements in two equal-sized vectors
        s = ifelse(r > 0, s + pmin(r, 1 / r), s),
        x0 = x1,
        y0 = y1,
        r = x0^2 + y0^2
      )
  }

  # Compute the matrix
  for (i in seq(nrow(coords))) {
    # working with a single row, but will soon generalise
    my_coords <- coords[i, ]
    for (j in seq(n_iter)) {
      # only update a row if it's coords are near zero
      rows <- which(my_coords$r < r_escape)
      my_coords[rows, ] <- .update_coords(my_coords)
    }
    coords[i, ] <- my_coords
  }
  # note the use of MATRIX[cbind(vec1, vec2)], this allows you to update
  # a set of specific entries in a matrix, rather than contiguous blocks of
  # entries
  m_xy[cbind(coords$nx, coords$ny)] <- coords$s
  m_xy
}

Rather than doing for(index) ... for(iteration){} and updating a single row: for a given iteration, we could determine all rows that need to be updated and then update those rows. That is we can do for (iteration){find-the-relevant-rows; update-those-rows}. The final code looks like this:

compute_mxy_test <- function(
                             n_iter,
                             n_row,
                             n_col,
                             r_escape) {
  # Initial Mandelbrot array
  m_xy <- matrix(1, nrow = n_row, ncol = n_col)

  # Re(C) vector
  c_x <- -0.5 + (-1 + seq(n_row) / n_row * 2) * 1.6
  c_y <- (-1 + seq(n_col) / n_col * 2) * 0.9

  # construct a data-frame for use in vectorised updates
  coords <- expand.grid(
    nx = seq(n_row),
    ny = seq(n_col)
  ) %>%
    dplyr::mutate(
      x0 = 0,
      y0 = 0,
      s = 1,
      r = 0,
      cx = c_x[nx],
      cy = c_y[ny]
    )

  .update_coords <- function(df) {
    # can assume
    # - nx, ny, x0, y0, s, r, cx, cy are columns in `df`
    # - that r >= r_escape
    x1 <- with(df, x0^2 - y0^2 + cx)
    y1 <- with(df, 2 * x0 * y0 + cy)
    df %>%
      dplyr::mutate(
        s = ifelse(r > 0, s + pmin(r, 1 / r), s),
        x0 = x1,
        y0 = y1,
        r = x0^2 + y0^2
      )
  }

  # Compute the matrix
  for (j in seq(n_iter)) {
    rows <- which(coords$r < r_escape)
    coords[rows, ] <- .update_coords(coords[rows, ])
  }

  m_xy[cbind(coords$nx, coords$ny)] <- coords$s
  m_xy
}

Now, you can compare the values returned, and you can compare the speed of the original implementation, and the vectorised implementation. IN my hands, for (400, 1600, 900) the vectorised code was ~ 3x as fast as the imperative for-loop code. But it used way more memory. If you run the benchmark code above, you should be able to convince yourself of this. (But you might notice that the vectorised code is considerably slower for less demanding arguments)

Although vectorised code is (usually) faster, because it pushes more work down into a lower level, the data-structures you have to set up to make it work, and the way the code looks sometimes makes it less appealing. As I said in a comment earlier, you might be better moving this code all the way down into the C level, eg by rewriting your original for-loop implementation using Rcpp.