ray_marching.py

from math import radians, sin, cos
import matplotlib.pyplot as plt
import numpy as np
from dataclasses import dataclass


def length(vec: np.array):
  return vec.dot(vec) ** 0.5


def normalize(vec: np.array):
  return vec / (vec.dot(vec) ** 0.5)


def norm(vec: np.array, dist):
  return normalize(np.array([
    dist(vec + np.array([1e-4, 0.0, 0.0])) - dist(vec + np.array([-1e-4, 0.0, 0.0])),
    dist(vec + np.array([0.0, 1e-4, 0.0])) - dist(vec + np.array([0.0, -1e-4, 0.0])),
    dist(vec + np.array([0.0, 0.0, 1e-4])) - dist(vec + np.array([0.0, 0.0, -1e-4]))
  ]))


@dataclass
class Camera:
  pos: np.array
  dir: np.array
  up: np.array
  resolution: np.array
  marching_loops: int
  field_of_view: float = 60.0
  light: np.array = np.array([-0.577, 0.577, 0.577])

  @property
  def init_matrix(self):
    if hasattr(self, '.init_matrix'):
      init = getattr(self, '.init_matrix')
      if init.shape[0] == self.resolution[0] or init.shape[1] == self.resolution[1]:
        return init

    init = np.array([
        [[i, j] for j in range(self.resolution[1])]
                for i in range(self.resolution[0])
    ])
    setattr(self, '.init_matrix', init)
    return init

  def ray(self, pix: np.array) -> np.array:
    sin_fov = sin(radians(self.field_of_view))
    cos_fov = cos(radians(self.field_of_view))
    side = np.cross(self.dir, self.up)

    return normalize(
            side * sin_fov * pix[0] +
            self.up * sin_fov * pix[1] +
            self.dir * cos_fov) 

  def march(self, dist_func):
    def march_inner(pixel: np.array):
      ray = self.ray((pixel * 2 - self.resolution) / self.resolution.min())
      length = 0
      position = self.pos
      for i in range(self.marching_loops):
        d = dist_func(position)
        length += d
        position = self.pos + ray * length
  
      if abs(d) < 0.001:
        n = norm(position, dist_func)
        return np.array([255, 255, 255]) * max(min(n.dot(self.light), 1.0), 0.1)

      return np.array([0, 0, 0])
    return march_inner

  def render(self, dist_func):
    fn = self.march(dist_func)
    init = self.init_matrix
    return np.vectorize(
                    np.vectorize(fn, signature='(2)->(3)'),
                    signature='(n,2)->(n,3)'
                )(init)


if __name__ == '__main__':
    def dist(vec: np.array):
        return length(vec) - 2


    cam = Camera(
        pos=np.array([0., 0., 3.]),
        dir=np.array([0., 0., -1.]),
        up=np.array([0., 1., 0.]),
        resolution=np.array([1080//2, 1920//2]),
        marching_loops=32
    )

    plt.imshow(cam.render(dist))
    plt.show()