雑記帳
僕用勉強ノート 「レイトレーシング」の巻

Haskell でレイトレーシングのチュートリアルを追いかける その8 - 透明の物体

引き続きこのサイトのチュートリアルに則って、レイトレーシングによる画像の生成に挑戦。
進捗状況としては、ひとまず「section 10-3」まで完了。
コードの実行結果
実行結果
(画像はファイルサイズを小さくしたい関係で JPEG 形式にコンバートしているため、各ピクセルの値は実際に得られた値とは微妙に異なる。)
ソースコード
{-# LANGUAGE TypeOperators #-}

module Main where

import Data.Char
import Data.Functor
import Control.Monad
import Control.Lens
import System.Random
import Linear.Vector
import Linear.Metric
import Linear.V3
import Linear.Quaternion

-- https://raytracing.github.io/books/RayTracingInOneWeekend.html
-- section 10-3 Total Internal Reflection with Haskell!!


main :: IO ()
main = do
  let
    -- Image
    aspect_ratio = 16.0 / 9
    image_width = 400
    image_height = round $ fromIntegral image_width / aspect_ratio
    samples_per_pixel = 100
    max_depth = 50
    -- World
    material_ground  = make_shared $ MAT_Lambertian {albedo_Lamb = V3 0.8 0.8 0.0}
    material_center  = make_shared $ MAT_Lambertian {albedo_Lamb = V3 0.7 0.3 0.3}
    material_center2 = make_shared $ MAT_Dielectric {ref_idx = 1.7}
    material_left    = make_shared $ MAT_Metal {albedo_Metal = V3 0.8 0.8 0.8, fuzz = 0.2}
    material_right   = make_shared $ MAT_Metal {albedo_Metal = V3 0.8 0.6 0.2, fuzz = 0.0}
    world = ((((([]
      --`add` RT_Sphere{center = V3 0.0 0.0 (-1.0), radius = 0.5, mat_Sphere = material_center})
      --`add` RT_Sphere{center = V3 (-0.95) 0.0 (-1.0), radius = 0.5, mat_Sphere = material_left})
      --`add` RT_Sphere{center = V3 1.05 0.0 (-1.0), radius = 0.5, mat_Sphere = material_right})
      `add` RT_Torus{
        centerOfTorus = V3 (-0.55) 0.0 (-2.1),
        majorRadius = 0.4,
        minorRadius = 0.1,
        orientationOfTorus = normalize $ V3 (0.5) 1.0 1.5,
        mat_Torus = material_left
        })
      `add` RT_Torus{
        centerOfTorus = V3 1.05 0.0 (-1.0),
        majorRadius = 0.3,
        minorRadius = 0.2,
        orientationOfTorus = normalize $ V3 (-0.5) (-0.2) (-3),
        mat_Torus = material_right
        })
      `add` RT_Sphere{center = V3 (-3.51) 2.4 (-5.9), radius = 2.9, mat_Sphere = material_center})
      `add` RT_Torus{
        centerOfTorus = V3 0.05 0.2 (-1.2),
        majorRadius = 0.35,
        minorRadius = 0.15,
        orientationOfTorus = normalize $ V3 (-0.5) 1.2 1.9,
        mat_Torus = material_center2
        })
      `add` RT_Sphere{center = V3 0 (-10000.5) (-1), radius = 10000, mat_Sphere = material_ground})
    -- Camera
    camera = Camera {
      viewport_height = 2.0,
      viewport_width = aspect_ratio * viewport_height camera,
      focal_length = 1.0,
      origin = zero,
      horizontal = viewport_width camera *^ unit _x,
      vertical = viewport_height camera *^ unit _y,
      lower_left_corner =
        origin camera - horizontal camera ^/2 - vertical camera ^/2
        - focal_length camera *^ unit _z
      }


  -- Render

  img_data <- return $ "P3\n" ++ show image_width ++ " " ++ show image_height ++ "\n255\n"
  putStr $ img_data

  foldr (>>) (return ()) $ (fmap $ ($) $ \((j, i), seed) ->
    let
      f h currentData gen =
        if h < samples_per_pixel then
          let
            (randNum1, newGen1) = random gen :: (Double, StdGen)
            (randNum2, newGen2) = random newGen1 :: (Double, StdGen)
            u = (fromIntegral i + randNum1) / (fromIntegral image_width - 1.0)
            v = (fromIntegral j + randNum2) / (fromIntegral image_height - 1.0)
            r = get_ray camera (u, v)
            pixcel_color = ray_color r world max_depth newGen2
          in
            f (succ h) (currentData + pixcel_color) newGen2
        else
          currentData
    in
      write_color (f 0 zero (mkStdGen seed)) samples_per_pixel) $
      zip ((,) <$> [image_height - 1, image_height - 2 .. 0] <*> [0, 1 .. image_width - 1])
      ((randomRs (0, 536870912) (mkStdGen 21)) :: [Int])


---------------------
-- A Ray Data Type --
---------------------

data Ray = Ray {
  orig :: V3 Double,  -- Origin of this ray (As a position in 3D Euclidean space)
  dir :: V3 Double    -- Direction of this ray (As a direction vector in 3D Euclidean space)
} deriving (Show)


at' :: Ray -> Double -> V3 Double
at' r t = (orig r) + t *^ (dir r)


------------------------
-- A Camera Data Type --
------------------------

data Camera = Camera {
  viewport_height :: Double,
  viewport_width :: Double,
  focal_length :: Double,
  origin :: V3 Double,
  horizontal :: V3 Double,
  vertical :: V3 Double,
  lower_left_corner :: V3 Double
} deriving (Show)

get_ray :: Camera -> (Double, Double) -> Ray
get_ray this (u, v) =
  Ray {
    orig = origin this,
    dir = lower_left_corner this + u *^ horizontal this + v *^ vertical this - origin this
    }


----------------------
-- A Hittable Class --
----------------------

type HittableData = (RT_Sphere + RT_Torus) + RT_Sphere -- Third RT_Sphere is just dummies

class Hittable a where
  toSum :: a -> HittableData
  hit :: a -> Ray -> Double -> Double -> Maybe HitRecord

instance (Hittable a, Hittable b) => Hittable (a + b) where
  toSum = coPair(toSum, toSum)
  hit = coPair(hit, hit)

add :: Hittable a => [HittableData] -> a -> [HittableData]
add list obj = (toSum obj) : list

data HitRecord = HitRecord {
  p :: V3 Double,
  normal :: V3 Double,
  mat :: MaterialData,
  t :: Double,
  front_face :: Bool
} deriving (Show)


set_face_normal :: HitRecord -> Ray -> V3 Double -> HitRecord
set_face_normal this r outward_normal =
  let
    isOutside = (dir r `dot` outward_normal) < 0
  in
    HitRecord {
      p = p this,
      front_face = isOutside,
      normal = if isOutside then outward_normal else -outward_normal,
      t = t this,
      mat = mat this
      }

hitSomething :: [HittableData] -> Ray -> Double -> Double -> Maybe HitRecord
hitSomething list r t_min t_max =
  let
    f (list', r', closest_so_far, currentRecord) =
      case list' of
        x:xs ->
          let
            temp = hit x r' t_min closest_so_far
          in
            case temp of
              Just a ->
                f $ (xs, r', t a, temp)
              Nothing ->
                f $ (xs, r', closest_so_far, currentRecord)
        [] ->
          currentRecord

  in
    f $ (list, r, t_max, Nothing)


----------------------
-- Hittable Objects --
----------------------

-- Sphere
data RT_Sphere = RT_Sphere {
  center :: V3 Double,
  radius :: Double,
  mat_Sphere :: MaterialData
} deriving (Show)

instance Hittable RT_Sphere where
  toSum = Inj1 -: Inj1
  hit obj r t_min t_max =
    let
      p0 = orig r
      c1 = center obj
      r1 = radius obj
      oc = p0 - c1
      a = quadrance (dir r)
      half_b = oc `dot` dir r
      c = quadrance oc - (radius obj) ^ 2
      discriminant = half_b ^ 2 - a*c in

      if discriminant > 0 then
        let
          root = sqrt discriminant
          f k =
            case k of
              x:xs ->
                if t_min < x && x < t_max then
                  return $ set_face_normal HitRecord {
                    p = at' r x,
                    normal = zero,
                    t = x,
                    front_face = False,
                    mat = mat_Sphere obj
                    } r ((at' r x - c1) ^/ r1)
                else
                  f $ xs

              [] ->
                Nothing

        in
          f $ [(-half_b - root) / a, (-half_b + root) / a]
      else
        Nothing

-- Torus
data RT_Torus = RT_Torus {
  centerOfTorus :: V3 Double,
  majorRadius :: Double,
  minorRadius :: Double,
  orientationOfTorus :: V3 Double, -- [BEWARB] this pseudo-vector must be normalized
  mat_Torus :: MaterialData
} deriving (Show)

instance Hittable RT_Torus where
  toSum = Inj2 -: Inj1
  hit obj r t_min t_max =
    let
      p0 = orig r
      c1 = centerOfTorus obj
      r1 = majorRadius obj
      r2 = minorRadius obj
      n = orientationOfTorus obj
      oc = p0 - c1
      a = quadrance (dir r)
      half_b = oc `dot` dir r
      c = quadrance oc - (r1 + r2 + 0.001) ^ 2
      discriminant = half_b ^ 2 - a*c in

      if discriminant > 0 then
        let
          root = sqrt discriminant
          s = newton's_method
            50
            (max 0.0001 ((-half_b - root) / a + 0.0001), max 0.0001 ((-half_b + root) / a))
            (max 0.0001 ((-half_b - root) / a + 0.0001), 0.0001)
            (findIntersection_forTorus obj r)
            (findIntersection_forTorus' obj r)
        in
          s >>= (\k ->
            let
              x = at' r k - c1
            in
              if t_min < k && k < t_max then
                return $ set_face_normal HitRecord {
                  p = c1 + x,
                  normal = zero,
                  t = k,
                  front_face = False,
                  mat = mat_Torus obj
                  } r ((x - (r1 *^ (normalize $ x - (n `dot` x) *^ n))) ^/ r2)

              else
                Nothing
            )
      else
        Nothing


findIntersection_forTorus :: RT_Torus -> Ray -> Double -> Double
findIntersection_forTorus obj r t =
  (quadrance u + r1 ^2 - r2 ^ 2) ^ 2 - 4 * (r1 ^ 2) * (u `dot` (u - (n `dot` u) *^ n))
  -- quadrance u - 2 * r1 * (sqrt $ u `dot` (u - (n `dot` u) *^ n)) + r1 ^2 - r2 ^ 2
  where
    a = dir r
    u = at' r t - centerOfTorus obj
    n = orientationOfTorus obj
    r1 = majorRadius obj
    r2 = minorRadius obj

findIntersection_forTorus' :: RT_Torus -> Ray -> Double -> Double
findIntersection_forTorus' obj r t =
  4 * (a `dot` u) * (quadrance u + r1 ^ 2 - r2 ^ 2)
  - 8 * (r1^2) * (u `dot` (a - (n `dot` a) *^ n))
  -- 2 * (a `dot` u)
  -- - 2 * r1 * (u `dot` (a - (n `dot` a) *^ n)) / (sqrt $ u `dot` (u - (n `dot` u) *^ n))
  where
    a = dir r
    u = at' r t - centerOfTorus obj
    n = orientationOfTorus obj
    r1 = majorRadius obj
    r2 = minorRadius obj


----------------------------------------
-- Computing the color of a given ray --
----------------------------------------

write_color :: V3 Double -> Int -> IO ()
write_color (V3 r g b) spp =
  let
    v' = V3 (sqrt $ r / fromIntegral spp) (sqrt $ g / fromIntegral spp) (sqrt $ b / fromIntegral spp)
    f = show.floor.(256*).(clamp 0 0.999)
  in
    do
      tmp <- return $ f(v' ^._x) ++ " " ++ f(v' ^._y) ++ " " ++ f(v' ^._z) ++ "\n"
      putStr $ tmp

ray_color :: Ray -> [HittableData] -> Int -> StdGen -> V3 Double
ray_color r objects depth gen =
  if depth <= 0 then
    zero
  else
    let
      record = hitSomething objects r 0.0001 infinity
    in
      case record of
        Just record' ->
          let (ret, gen1) = scatter (mat record') r record' gen
          in
            case ret of
              Just (scattered, attenuation) ->
                attenuation * (ray_color scattered objects (pred depth) gen1)

              Nothing ->
                zero

        Nothing ->
          let
            unit_direction = normalize $ (dir r)
            s = 0.5 * (unit_direction ^._y + 1.0)
          in
            lerp s (V3 0.5 0.7 1.0) (V3 1.0 1.0 1.0)


--------------------
-- Random numbers --
--------------------

random_in_unit_sphere :: StdGen -> (V3 Double, StdGen)
random_in_unit_sphere gen0 =
  let
    (rand1,gen1) = randomR (-1, 1) gen0 :: (Double, StdGen)
    (rand2,gen2) = randomR (-1, 1) gen1 :: (Double, StdGen)
    (rand3,gen3) = randomR (-1, 1) gen2 :: (Double, StdGen)
    v = V3 rand1 rand2 rand3
  in
    if quadrance v >= 1 then
      random_in_unit_sphere gen3
    else
      (v, gen3)

random_unit_vector :: StdGen -> (V3 Double, StdGen)
random_unit_vector gen0 =
  let
    (a, gen1) = randomR (0, 2*pi) gen0 :: (Double, StdGen)
    (z, gen2) = randomR (-1, 1) gen1 :: (Double, StdGen)
    r = sqrt $ 1 - z^2
  in
    (V3 (r*cos(a)) (r*sin(a)) z, gen2)


random_in_hemisphere :: V3 Double -> StdGen -> (V3 Double, StdGen)
random_in_hemisphere normal gen0 =
  let
    (in_unit_sphere, gen1) = random_in_unit_sphere gen0
  in
    if in_unit_sphere `dot` normal > 0 then
      (in_unit_sphere, gen1)
    else
      (-in_unit_sphere, gen1)

---------------
-- Utilities --
---------------

infinity :: RealFloat a => a
infinity = encodeFloat (floatRadix 0 - 1) (snd $ floatRange 0)

deg2rad :: Floating a => a -> a
deg2rad degrees = degrees * pi / 180

clamp :: (Ord a, Num a) => a -> a -> a -> a
clamp x y val = (max x).(min y) $ val

isInClosedInterval :: (Ord a, Fractional a) => (a, a) -> a -> Bool
isInClosedInterval (a, b) val = (a <= val && val <= b)

isInOpenInterval :: (Ord a, Fractional a) => (a, a) -> a -> Bool
isInOpenInterval (a, b) val = (a < val && val < b)

newton's_method :: (Ord a, Fractional a) => Int -> (a, a) -> (a, a) -> (a -> a) -> (a -> a) -> Maybe a
newton's_method depth interval (current, prev) f f' =
  if (uncurry (/=)) interval && isInClosedInterval interval current && depth > 0 then
    if abs(current - prev) < 1.0E-10 then
      Just current
    else
      newton's_method (pred depth) interval (current - (f(current) / f'(current)), current) f f'
  else
    Nothing

-- Joke
derivativeOf :: (Floating a) => (a -> a) -> Int -> a -> a
derivativeOf f precision =
  \x -> (f(x + dx) - f(x)) / dx
  where dx = 0.1^precision

reflect :: V3 Double -> V3 Double -> V3 Double
reflect v n = v - (2 * (n `dot` v)) *^ n

refract :: V3 Double -> V3 Double -> Double -> V3 Double
refract uv n eta_over_eta' =
  r_out_perp + r_out_parallel
  where
    cos_theta = (-uv) `dot` n
    r_out_perp = eta_over_eta' *^ (uv + cos_theta *^ n)
    r_out_parallel = (sqrt $ abs (1 - quadrance r_out_perp)) *^ (-n)

--------------------
-- Material Class --
--------------------

type MaterialData = (MAT_Lambertian + MAT_Metal) + MAT_Dielectric

class Material a where
  make_shared :: a -> MaterialData
  scatter :: a -> Ray -> HitRecord -> StdGen -> (Maybe (Ray, V3 Double), StdGen)

instance (Material a, Material b) => Material (a + b) where
  make_shared = coPair(make_shared, make_shared)
  scatter = coPair(scatter, scatter)


-- Lambertian

data MAT_Lambertian = MAT_Lambertian {
  albedo_Lamb :: V3 Double
} deriving (Show)

instance Material MAT_Lambertian where
  make_shared = Inj1 -: Inj1
  scatter this r_in record gen =
    let
      (rand1, gen1) = random_unit_vector gen
      scattered_direction = normal record + rand1
      scattered = Ray{orig = p record, dir = scattered_direction}
      attenuation = albedo_Lamb this
    in
      (Just (scattered, attenuation), gen1)

-- Metal

data MAT_Metal = MAT_Metal {
  albedo_Metal :: V3 Double,
  fuzz :: Double
} deriving (Show)

instance Material MAT_Metal where
  make_shared = Inj2 -: Inj1
  scatter this r_in record gen =
    let
      (rand1, gen1) = random_in_unit_sphere gen
      reflected = reflect (normalize $ dir r_in) (normal record)
      scattered = Ray{orig = p record, dir = reflected + (fuzz this) *^ rand1}
      attenuation = albedo_Metal this
    in
      if (dir scattered `dot` normal record) > 0 then
        (Just (scattered, attenuation), gen1)
      else
        (Nothing, gen1)


-- Dielectric

data MAT_Dielectric = MAT_Dielectric {
  ref_idx :: Double
} deriving (Show)

instance Material MAT_Dielectric where
  make_shared = Inj2
  scatter this r_in record gen =
    let
      attenuation = V3 1 1 1
      eta_over_eta' = if front_face record then 1 / ref_idx this else ref_idx this
      unit_direction = normalize $ dir r_in
      cos_theta = min (-unit_direction `dot` normal record) 1
      sin_theta = sqrt $ 1 - cos_theta ^ 2
    in
      if eta_over_eta' * sin_theta > 1 then
        let
          reflected = reflect unit_direction (normal record)
          scattered = Ray{orig = p record, dir = reflected}
        in
          (Just (scattered, attenuation), gen)
      else
        let
          refracted = refract unit_direction (normal record) eta_over_eta'
          scattered = Ray{orig = p record, dir = refracted}
        in
          (Just (scattered, attenuation), gen)


-----------------------------------
-- Category Theoretic Operations --
-----------------------------------

-- Diagrammatic-order composition
(-:) = flip (.)

-- Sum objects and injections
data (+) a b = Inj1 a | Inj2 b

instance (Show a, Show b) => Show (a + b) where
  show = coPair(show -: (++ ";inj1"), show -: (++ ";inj2"))

-- Dual to pairs
coPair :: (a1 -> b, a2 -> b) -> (a1 + a2 -> b)
coPair (f, g) x = case x of
  Inj1 x -> f x
  Inj2 x -> g x