Tree rendering

tree.rb

require 'chunky_png'

case ARGV[0]
when 'red'
  FLOOR_COLOR = [0xff, 0x80, 0x80]
  LEAF_COLOR = 0xff0044
  TRUNK_COLOR = 0xff8888
  RAY_COLOR = [2, 1, 1]
  LIGHT_THETA = 0.1
when 'green'
  FLOOR_COLOR = [0x80, 0x80, 0x80]
  LEAF_COLOR = 0x44ff44
  TRUNK_COLOR = 0x888888
  RAY_COLOR = [1, 1, 1]
  LIGHT_THETA = 1.2
when 'blue'
  FLOOR_COLOR = [0x80, 0x80, 0xa0]
  LEAF_COLOR = 0x44ffaa
  TRUNK_COLOR = 0x8888aa
  RAY_COLOR = [1, 1, 2]
  LIGHT_THETA = 2.6
else
  puts "Usage: ruby #{__FILE__} red|green|blue"
  exit
end

class Canvas
  def initialize(size)
    @size = size
    @color = size.times.map { [0] * @size }
    @depth = size.times.map { [10] * @size }
    @lightnorm = size.times.map { [0] * @size }
    @cz = Math.cos(0.2)
    @sz = Math.sin(0.2)
    @cl = Math.cos(LIGHT_THETA)
    @sl = Math.sin(LIGHT_THETA)
    @shadow_size = 64
    @shadow = @shadow_size.times.map do
      @shadow_size.times.map { [0] * @shadow_size }
    end
  end

  def add_shadow(x, y, z, r, value = 1)
    y, z = y * @cz - z * @sz, z * @cz + y * @sz
    x, z = x * @cl + z * @sl, z * @cl - x * @sl
    r = 2 * (r * @shadow_size / 2)
    r = [r, 2].max
    x = @shadow_size * (1 + x) / 2
    y = @shadow_size * (1 + y) / 2
    z = @shadow_size * (1 + z) / 2
    xrange = [(x - r).ceil, 0].max..[(x + r).floor, @shadow_size - 1].min
    yrange = [(y - r).ceil, 0].max..[(y + r).floor, @shadow_size - 1].min
    zrange = [(z - r).ceil, 0].max..[(z + r).floor, @shadow_size - 1].min
    xrange.each do |ix|
      yrange.each do |iy|
        zrange.each do |iz|
          v = 1 - ((ix - x) ** 2 + (iy - y) ** 2 + (iz - z) ** 2) / r / r
          if v > 0
            @shadow[ix][iy][iz] += v ** 2 * value
          end
        end
      end
    end
  end

  def get_shadow(x, y, z)
    y, z = y * @cz - z * @sz, z * @cz + y * @sz
    x, z = x * @cl + z * @sl, z * @cl - x * @sl
    interpolate3(@shadow, x, y, z)
  end

  def interpolate3(sh, x, y, z)
    x = (@shadow_size * (1 + x) / 2).clamp(0, @shadow_size - 1)
    y = (@shadow_size * (1 + y) / 2).clamp(0, @shadow_size - 1)
    z = (@shadow_size * (1 + z) / 2).clamp(0, @shadow_size - 1)
    ix = [x.floor, @shadow_size - 2].min
    iy = [y.floor, @shadow_size - 2].min
    iz = [z.floor, @shadow_size - 2].min
    x -= ix
    y -= iy
    z -= iz
    (1 - x) * (1 - y) * (1 - z) * sh[ix][iy][iz] +
    x * (1 - y) * (1 - z) * sh[ix + 1][iy][iz] +
    (1 - x) * y * (1 - z) * sh[ix][iy + 1][iz] +
    x * y * (1 - z) * sh[ix + 1][iy + 1][iz] +
    (1 - x) * (1 - y) * z * sh[ix][iy][iz + 1] +
    x * (1 - y) * z * sh[ix + 1][iy][iz + 1] +
    (1 - x) * y * z * sh[ix][iy + 1][iz + 1] +
    x * y * z * sh[ix + 1][iy + 1][iz + 1]
  end

  def convert_shadow
    @shadow_size.times do |y|
      @shadow_size.times do |z|
        r = (2.0 * y / (@shadow_size - 1) - 1) ** 2 + (2.0 * z / (@shadow_size - 1) - 1) ** 2
        light = r > 1 ? 0 : [(1 - r) ** 2, 1].min
        @shadow_size.times.reverse_each do |x|
          light *= 0.999 ** @shadow[x][y][z]
          @shadow[x][y][z] = light
        end
      end
    end
    @ray = @shadow_size.times.map do
      @shadow_size.times.map { [0] * @shadow_size }
    end
    @shadow_size.times do |x|
      @shadow_size.times do |z|
        sum = 0
        @shadow_size.times.each do |y|
          sum += @shadow[x][y][z] * 0.5 ** (y.fdiv @shadow_size)
          @ray[x][y][z] = sum
        end
      end
    end
  end

  def draw(cx, cy, cz, cr, color, opacity = 1)
    add_shadow(cx, cy, cz, cr, opacity)
    cy, cz = cy * @cz - cz * @sz, cy * @sz + cz * @cz
    x = @size * (1 + cx) / 2.0
    y = @size * (1 - cz) / 2.0
    r = cr * @size / 2.0
    depth = cy
    yrange = [(y - r).ceil, 0].max..[(y + r).floor, @size - 1].min
    xrange = [(x - r).ceil, 0].max..[(x + r).floor, @size - 1].min
    xrange.each do |px|
      yrange.each do |py|
        d2 = r * r - (px - x) ** 2 - (py - y) ** 2
        if d2 > 0 && (d = depth - Math.sqrt(d2) / @size) < @depth[px][py]
          @depth[px][py] = d
          @color[px][py] = color
          @lightnorm[px][py] = (px - x) / r * @cl + (y - py) / r * @sl
        end
      end
    end
  end

  def get_screen_shadow(sx, sy, depth)
    x = 2.0 * sx / @size - 1
    z = 1 - 2.0 * sy / @size
    y = depth
    y, z = y * @cz + z * @sz, -y * @sz + z * @cz
    get_shadow(x, y, z)
  end

  def get_screen_ray(sx, sy, depth)
    x = 2.0 * sx / @size - 1
    z = 1 - 2.0 * sy / @size
    y = depth
    x, z = x * @cl + z * @sl, z * @cl - x * @sl
    interpolate3(@ray, x, y, z)
  end

  def save(file)
    image = ChunkyPNG::Image.new(@size, @size)
    @size.times do |x|
      @size.times do |y|
        d = @depth[x][y]
        c = @color[x][y]
        col = [(c >> 16) & 0xff, c >> 8 & 0xff, c & 0xff]
        floor_z = -0.5
        floor_depth = ((1 - 2.0 * y / @size) * @cz - floor_z) / @sz
        if floor_depth < d
          d = floor_depth
          col = FLOOR_COLOR
          @lightnorm[x][y] = 0
        end
        ray = get_screen_ray(x, y, d)
        shadow = 0.1 + 0.9 * get_screen_shadow(x, y, d)
        shade = 0.7 + 0.3 * @lightnorm[x][y]
        r, g, b = col.zip(RAY_COLOR).map { (_1 * shadow * shade + _2 * ray).round.clamp(0, 0xff) }
        image[x, y] = (r << 24) | (g << 16) | (b << 8) | 0xff
      end
    end
    image.save(file)
  end
end

canvas = Canvas.new(1024)

def sphere_rand
  while true
    x = 2 * rand - 1
    y = 2 * rand - 1
    z = 2 * rand - 1
    r = x * x + y * y + z * z
    return [x / r, y / r, z / r] if r < 1
  end
end
draw_tree = ->(x, y, z, w, dx, dy, dz, t) {
  r = w * 0.03
  if w < 0.05
    10.times do
      sx, sy, sz = sphere_rand
      l = 0.02
      canvas.draw(x + l * (dx + sx), y + l * (dy + sy), z + l * (dz + sz), l / 4, LEAF_COLOR)
    end
    return
  end
  if t < 0
    sx, sy, sz = sphere_rand
    [-1, 1].each do |dir|
      dx2 = dx + sx / 2 * dir
      dy2 = dy + sy / 2 * dir
      dz2 = dz + sz / 2 * dir
      dl = Math.sqrt(dx2 * dx2 + dy2 * dy2 + dz2 * dz2)
      w2 = w * (0.7 + 0.3 * rand)
      draw_tree.call(x, y, z, w2, dx2 / dl, dy2 / dl, dz2 / dl, (w < 0.2 ? 0.2 : 1) * rand)
    end
  else
    5.times do |i|
      canvas.draw(x + r * dx * i / 4, y + r * dy * i / 4, z + r * dz * i / 4, r, TRUNK_COLOR, 1 + 40 * w)
    end
    x += dx * r
    y += dy * r
    z += dz * r
    sx, sy, sz = sphere_rand
    dx += sx / 10
    dy += sy / 10
    dz += sz / 10 + w / 20.0
    dl = Math.sqrt(dx * dx + dy * dy + dz * dz)
    draw_tree.call(x, y, z, w * 0.99, dx / dl, dy / dl, dz / dl, t - 0.1 * w)
  end
}
srand 0
draw_tree.call(0, 0, -0.5, 1, 0, 0, 1, 1)
canvas.convert_shadow
canvas.save 'tree.png'