Optimized ray tracer implementation, originally written by @jeaye https://gist.github.com/jeaye/77e1d8874c8e76e7335ccf71ef53785c

ray_tracer.clj

(ns ray-tracer
  (:import java.util.List
           java.awt.Color
           java.awt.image.BufferedImage
           javax.imageio.ImageIO
           java.io.File))

(alter-var-root  #'*compiler-options* (constantly {:direct-linking true}) )
(set! *unchecked-math* :warn-on-boxed)

(let [seeded (java.util.Random. 1)]
  (defn rand ^double []
    (.nextDouble ^java.util.Random seeded)))

(defn rand-real ^double [^double min ^double max]
  (+ min (* (- max min) (rand))))

(defn degrees->radians ^double [^double deg]
  (/ (* deg Math/PI) 180.0))

(defrecord Vec3 [^double r, ^double g, ^double b])

(def ^Vec3 zero-vec3 (Vec3. 0.0 0.0 0.0))
(def ^Vec3 unit-vec3 (Vec3. 1.0 1.0 1.0))

(defn vec3-scale [^Vec3 l, ^double n]
  (Vec3. (* (.r l) n) (* (.g l) n) (* (.b l) n)))

(defn vec3-add [^Vec3 l, ^Vec3 r]
  (Vec3. (+ (.r l) (.r r))
         (+ (.g l) (.g r))
         (+ (.b l) (.b r))))

(defn vec3-sub [^Vec3 l, ^Vec3 r]
  (Vec3. (- (.r l) (.r r))
         (- (.g l) (.g r))
         (- (.b l) (.b r))))

(defn vec3-mul [^Vec3 l ^Vec3 r]
  (Vec3. (* (.r l) (.r r))
         (* (.g l) (.g r))
         (* (.b l) (.b r))))

(defn vec3-div ^Vec3 [^Vec3 l ^double n]
  (Vec3. (/ (.r l) n)
         (/ (.g l) n)
         (/ (.b l) n)))

(defn vec3-length-squared ^double [^Vec3 v]
  (+ (* (.r v) (.r v))
     (* (.g v) (.g v))
     (* (.b v) (.b v))))

(defn vec3-length ^double [v]
  (Math/sqrt (vec3-length-squared v)))

(defn vec3-dot ^double [^Vec3 l ^Vec3 r]
  (+ (* (.r l) (.r r))
     (* (.g l) (.g r))
     (* (.b l) (.b r))))

(defn vec3-cross [^Vec3 l ^Vec3 r]
  (Vec3. (- (* (.g l) (.b r))
            (* (.b l) (.g r)))
         (- (* (.b l) (.r r))
            (* (.r l) (.b r)))
         (- (* (.r l) (.g r))
            (* (.g l) (.r r)))))

(defn vec3-normalize ^Vec3 [v]
  (vec3-div v (vec3-length v)))

(defn vec3-rand []
  (Vec3. (rand) (rand) (rand)))

(defn vec3-rand+clamp [^double min, ^double max]
  (Vec3. (rand-real min max) (rand-real min max) (rand-real min max)))

(defn vec3-rand-in-sphere []
  (let [v (vec3-rand+clamp -1.0 1.0)]
    (if (>= (vec3-length-squared v) 1.0)
      (recur)
      v)))

(defn vec3-rand-in-unit-disk []
  (let [p (Vec3. (rand-real -1.0 1.0) (rand-real -1.0 1.0) 0.0)]
    (if (>= (vec3-length-squared p) 1.0)
      (recur)
      p)))

(defn vec3-near-zero? [^Vec3 v]
  (let [epsilon 0.0000024]
    (< (+ (Math/abs (.r v)) (Math/abs (.g v)) (Math/abs (.b v))) epsilon)))

(defn vec3-reflect [v n]
  (vec3-sub v (vec3-scale n (* 2.0 (vec3-dot v n)))))

(defn vec3-refract [uv n ^double etai-over-etat]
  (let [cos-theta (min (vec3-dot (vec3-sub zero-vec3 uv) n)
                       1.0)
        r-out-perp (vec3-scale (vec3-add uv (vec3-scale n cos-theta))
                               etai-over-etat)
        r-out-parallel (vec3-scale n (- (Math/sqrt (Math/abs (- 1.0 (vec3-length-squared r-out-perp))))))]
    (vec3-add r-out-perp r-out-parallel)))

(defrecord Ray [^Vec3 origin ^Vec3 direction])

(defn ray-at [^Ray r ^double t]
  (vec3-add (.origin r) (vec3-scale (.direction r) t)))

(defn reflectance ^double [^double cosine ^double ref-idx]
  (let [r (/ (- 1.0 ref-idx)
             (+ 1.0 ref-idx))
        r2 (* r r)]
    (* (+ r2 (- 1.0 r2))
       (Math/pow (- 1.0 cosine) 5.0))))

(defrecord HitInfo [point normal ^double t material front-face?])
(defrecord Hittable [^Vec3 center ^double radius material])

(defn hit-sphere ^HitInfo [^Hittable hittable ^double t-min ^double t-max ^Ray ray]
  (let [center (.center hittable)
        radius (.radius hittable)
        oc (vec3-sub (.origin ray) center)
        a (vec3-length-squared (.direction ray))
        half-b (vec3-dot oc (.direction ray))
        c (- (vec3-length-squared oc) (* radius radius))
        discriminant (- (* half-b half-b) (* a c))]
    (when (> discriminant 0.0)
      (let [sqrt-d (Math/sqrt discriminant)
            root (let [root (/ (- (- half-b) sqrt-d) a)]
                   (if (or (< root t-min) (< t-max root))
                     (/ (+ (- half-b) sqrt-d) a)
                     root))]
        (when (and (> root t-min) (< root t-max))
          (let [point (ray-at ray root)
                outward-normal (vec3-div (vec3-sub point center) radius)
                front-face? (< (vec3-dot (.direction ray) outward-normal) 0.0)]
            (->HitInfo point
                       (if front-face?
                         outward-normal
                         (vec3-sub zero-vec3 outward-normal))
                       root
                       (.material hittable)
                       front-face?)))))))

(defn hit-all [^double t-min, ^double t-max, ray, ^List hittables]
  (let [it (.iterator hittables)]
    (loop [^HitInfo hit-info nil]
      (if (.hasNext it)
        (if-some [^HitInfo new-hit-info (hit-sphere (.next it)
                                                    t-min
                                                    (if (nil? hit-info) t-max (.t hit-info))
                                                    ray)]
          (recur new-hit-info)
          (recur hit-info))
        hit-info))))

(defn scatter-lambertian [ray, ^HitInfo hit-info]
  (let [normal (.normal hit-info)
        dir (vec3-add normal
                      (vec3-normalize (vec3-rand-in-sphere)))
        scatter-direction (if (vec3-near-zero? dir)
                            normal
                            dir)
        scattered (Ray. (.point hit-info) scatter-direction)
        attenuation (:albedo (.material hit-info))]
    {:ray scattered
     :attenuation attenuation}))

(defn scatter-metal [^Ray ray, ^HitInfo hit-info]
  (let [normal (.normal hit-info)
        material (.material hit-info)
        reflected (vec3-reflect (vec3-normalize (.direction ray))
                                normal)
        scattered (Ray. (.point hit-info)
                        (vec3-add reflected
                                  (vec3-scale (vec3-normalize (vec3-rand-in-sphere))
                                              (:fuzz material))))
        attenuation (:albedo material)]
    (when (> (vec3-dot (.direction scattered) normal) 0.0)
      {:ray scattered
       :attenuation attenuation})))

(defn scatter-dialetric [^Ray ray, ^HitInfo hit-info]
  (let [material (.material hit-info)
        attenuation unit-vec3
        index-of-refraction (double (:index-of-refraction material))
        refraction-ratio (if (.front-face? hit-info)
                           (/ 1.0 index-of-refraction)
                           index-of-refraction)
        unit-direction (vec3-normalize (.direction ray))

        normal (.normal hit-info)
        cos-theta (min (vec3-dot (vec3-sub zero-vec3 unit-direction)
                                 normal)
                       1.0)
        sin-theta (Math/sqrt (- 1.0 (* cos-theta cos-theta)))
        direction (if (or (< 1.0 (* refraction-ratio sin-theta)) ;; cannot refract
                          (< (rand) (reflectance cos-theta refraction-ratio)))
                    (vec3-reflect unit-direction normal)
                    (vec3-refract unit-direction normal refraction-ratio))]
    {:ray (Ray. (.point hit-info) direction)
     :attenuation attenuation}))

(defn ray-cast [^Ray ray, ^long max-ray-bounces, hittables]
  (if (< max-ray-bounces 0)
    zero-vec3
    (let [normalize-direction (vec3-normalize (.direction ray))
          t (* 0.5 (+ (.g normalize-direction) 1.0))]
      (if-some [^HitInfo hit-info (hit-all 0.001 99999999.0 ray hittables)]
        (let [material (.material hit-info)
              scatter-fn (:scatter material)]
          (if-some [scattered (scatter-fn ray hit-info)]
            (vec3-mul (ray-cast (:ray scattered)
                                (dec max-ray-bounces)
                                hittables)
                      (:attenuation scattered))
            zero-vec3))
        (vec3-add (vec3-scale unit-vec3 (- 1.0 t))
                  (vec3-scale (Vec3. 0.5 0.7 1.0) t))))))

(defn rand-scene []
  (into [(->Hittable  (Vec3. 0 -1000 0)
                      1000.0
                      {:albedo (Vec3. 0.5 0.5 0.5)
                       :scatter scatter-lambertian})
         (->Hittable  (Vec3. -4 1 0)
                      1.0
                      {:albedo (Vec3. 0.4 0.2 0.1)
                       :scatter scatter-lambertian})
         (->Hittable  (Vec3. 0 1 0)
                      1.0
                      {:index-of-refraction 1.5
                       :scatter scatter-dialetric})
         (->Hittable  (Vec3. 4 1 0)
                      1.0
                      {:albedo (Vec3. 0.7 0.6 0.5)
                       :fuzz 0
                       :scatter scatter-metal})]
        (keep (fn [^long i]
                (let [x (double (- (rem i 21) 10))
                      z (double (- (quot i 21) 6))
                      choose-mat (rand)
                      center (Vec3. (+ x (* 0.9 (rand)))
                                    0.2
                                    (+ z (* 0.9 (rand))))]
                  (when (> (vec3-length (vec3-sub center (Vec3. 4.0 0.2 0.0)))
                           0.9)
                    (cond (< choose-mat 0.8)
                          (->Hittable center
                                      0.2
                                      {:albedo (vec3-mul (vec3-rand) (vec3-rand))
                                       :scatter scatter-lambertian})

                          (< choose-mat 0.95)
                          (->Hittable center
                                      0.2
                                      {:albedo (vec3-rand+clamp 0.5 1.0)
                                       :fuzz (rand-real 0.0 0.5)
                                       :scatter scatter-metal})

                          :else
                          (->Hittable center
                                      0.2
                                      {:index-of-refraction 1.5
                                       :scatter scatter-dialetric}))))))
        (range 0 200)))

(defn clamp-to-byte ^long [^double n ^double min ^double max]
  (unchecked-long
   (* 256.0 (if (< n min)
              min
              (if (< max n)
                max
                n)))))

(defn put-pixel [^BufferedImage img x y ^Vec3 v samples-per-pixel]
  (let [scale (/ 1.0 samples-per-pixel)
        r (Math/sqrt (* scale (.r v)))
        g (Math/sqrt (* scale (.g v)))
        b (Math/sqrt (* scale (.b v)))
        clr (bit-or 0xFF000000
                    (bit-shift-left (clamp-to-byte r 0.0 0.999) 16)
                    (bit-shift-left (clamp-to-byte g 0.0 0.999) 8)
                    (clamp-to-byte b 0.0 0.999))]
    (.setRGB img x y clr)))

(defn -main []
  (let [aspect-ratio (/ 3.0 2.0)
        image-width 400
        image-height (long (/ image-width aspect-ratio))
        samples-per-pixel 10
        max-ray-bounces 10

        look-from (Vec3. 13.0 2.0 3.0)
        look-at zero-vec3
        aperture 0.1
        lens-radius (/ aperture 2)
        focus-distance 10
        camera-up (Vec3. 0.0 1.0 0.0)
        field-of-view 20
        field-of-view-theta (degrees->radians field-of-view)
        viewport-height (* 2 (Math/tan (/ field-of-view-theta 2.0)))
        viewport-width (* aspect-ratio viewport-height)
        camera-w (vec3-normalize (vec3-sub look-from look-at))
        camera-u (vec3-normalize (vec3-cross camera-up camera-w))
        camera-v (vec3-cross camera-w camera-u)

        origin look-from
        horizontal (vec3-scale camera-u (* viewport-width focus-distance))
        vertical (vec3-scale camera-v (* viewport-height focus-distance))
        lower-left-corner (vec3-sub (vec3-sub (vec3-sub origin (vec3-div horizontal 2))
                                              (vec3-div vertical 2))
                                    (vec3-scale camera-w focus-distance))

        hittables (rand-scene)
        sample-counter (range 0 samples-per-pixel)

        img (BufferedImage. image-width image-height BufferedImage/TYPE_INT_ARGB)]

    (dotimes [y image-height]
      (dotimes [x image-width]
        (loop [i 0, acc zero-vec3]
          (if (< i samples-per-pixel)
            (let [u (/ (+ (unchecked-double x) (rand)) (- image-width 1))
                  v (/ (+ (unchecked-double y) (rand)) (- image-height 1))
                  rd (vec3-scale (vec3-rand-in-unit-disk) lens-radius)
                  offset zero-vec3
                  ray (Ray. (vec3-add origin offset)
                            (vec3-sub (vec3-add (vec3-add lower-left-corner
                                                          (vec3-scale horizontal u))
                                                (vec3-scale vertical v))
                                      (vec3-sub origin offset)))]
              (recur (inc i)
                     (vec3-add acc (ray-cast ray max-ray-bounces hittables))))

            (put-pixel img x (- image-height y 1) acc samples-per-pixel)))))

    (ImageIO/write img "png" (File. "test.png"))
    img))

#_(time (-main))