RayDifferentials.cpp

// Ray Tracer with Ray Differentials
// Compile: g++ -O3 -std=c++11 RayDifferentials.cpp -o raydiff
// Run: ./raydiff
// This outputs 3 different scenes to separate files

#include <cmath>
#include <limits>
#include <memory>
#include <random>
#include <iostream>
#include <fstream>
#include <vector>

// Constants
const double infinity = std::numeric_limits<double>::infinity();
const double pi = 3.1415926535897932385;

// Utility Functions
inline double degrees_to_radians(double degrees) {
    return degrees * pi / 180.0;
}

inline double random_double() {
    static std::uniform_real_distribution<double> distribution(0.0, 1.0);
    static std::mt19937 generator;
    return distribution(generator);
}

inline double random_double(double min, double max) {
    return min + (max-min)*random_double();
}

inline double clamp(double x, double min, double max) {
    if (x < min) return min;
    if (x > max) return max;
    return x;
}

// Vec3 class
class vec3 {
public:
    double e[3];

    vec3() : e{0,0,0} {}
    vec3(double e0, double e1, double e2) : e{e0, e1, e2} {}

    double x() const { return e[0]; }
    double y() const { return e[1]; }
    double z() const { return e[2]; }

    vec3 operator-() const { return vec3(-e[0], -e[1], -e[2]); }
    double operator[](int i) const { return e[i]; }
    double& operator[](int i) { return e[i]; }

    vec3& operator+=(const vec3 &v) {
        e[0] += v.e[0];
        e[1] += v.e[1];
        e[2] += v.e[2];
        return *this;
    }

    vec3& operator*=(const double t) {
        e[0] *= t;
        e[1] *= t;
        e[2] *= t;
        return *this;
    }

    vec3& operator/=(const double t) {
        return *this *= 1/t;
    }

    double length() const {
        return sqrt(length_squared());
    }

    double length_squared() const {
        return e[0]*e[0] + e[1]*e[1] + e[2]*e[2];
    }

    bool near_zero() const {
        const auto s = 1e-8;
        return (fabs(e[0]) < s) && (fabs(e[1]) < s) && (fabs(e[2]) < s);
    }

    inline static vec3 random() {
        return vec3(random_double(), random_double(), random_double());
    }

    inline static vec3 random(double min, double max) {
        return vec3(random_double(min,max), random_double(min,max), random_double(min,max));
    }
};

using point3 = vec3;
using color = vec3;

// Vec3 Utility Functions
inline vec3 operator+(const vec3 &u, const vec3 &v) {
    return vec3(u.e[0] + v.e[0], u.e[1] + v.e[1], u.e[2] + v.e[2]);
}

inline vec3 operator-(const vec3 &u, const vec3 &v) {
    return vec3(u.e[0] - v.e[0], u.e[1] - v.e[1], u.e[2] - v.e[2]);
}

inline vec3 operator*(const vec3 &u, const vec3 &v) {
    return vec3(u.e[0] * v.e[0], u.e[1] * v.e[1], u.e[2] * v.e[2]);
}

inline vec3 operator*(double t, const vec3 &v) {
    return vec3(t*v.e[0], t*v.e[1], t*v.e[2]);
}

inline vec3 operator*(const vec3 &v, double t) {
    return t * v;
}

inline vec3 operator/(vec3 v, double t) {
    return (1/t) * v;
}

inline double dot(const vec3 &u, const vec3 &v) {
    return u.e[0] * v.e[0]
         + u.e[1] * v.e[1]
         + u.e[2] * v.e[2];
}

inline vec3 cross(const vec3 &u, const vec3 &v) {
    return vec3(u.e[1] * v.e[2] - u.e[2] * v.e[1],
                u.e[2] * v.e[0] - u.e[0] * v.e[2],
                u.e[0] * v.e[1] - u.e[1] * v.e[0]);
}

inline vec3 unit_vector(vec3 v) {
    return v / v.length();
}

vec3 random_in_unit_sphere() {
    while (true) {
        auto p = vec3::random(-1,1);
        if (p.length_squared() >= 1) continue;
        return p;
    }
}

vec3 random_unit_vector() {
    return unit_vector(random_in_unit_sphere());
}

vec3 random_in_hemisphere(const vec3& normal) {
    vec3 in_unit_sphere = random_in_unit_sphere();
    if (dot(in_unit_sphere, normal) > 0.0)
        return in_unit_sphere;
    else
        return -in_unit_sphere;
}

vec3 reflect(const vec3& v, const vec3& n) {
    return v - 2*dot(v,n)*n;
}

vec3 refract(const vec3& uv, const vec3& n, double etai_over_etat) {
    auto cos_theta = fmin(dot(-uv, n), 1.0);
    vec3 r_out_perp =  etai_over_etat * (uv + cos_theta*n);
    vec3 r_out_parallel = -sqrt(fabs(1.0 - r_out_perp.length_squared())) * n;
    return r_out_perp + r_out_parallel;
}

// RAY DIFFERENTIALS: Extended ray class with differentials
class ray {
public:
    point3 orig;
    vec3 dir;
    
    // Ray differentials for texture filtering
    point3 rx_origin, ry_origin;
    vec3 rx_direction, ry_direction;
    bool has_differentials;

    ray() : has_differentials(false) {}
    ray(const point3& origin, const vec3& direction)
        : orig(origin), dir(direction), has_differentials(false)
    {}

    point3 origin() const  { return orig; }
    vec3 direction() const { return dir; }

    point3 at(double t) const {
        return orig + t*dir;
    }

    // RAY DIFFERENTIALS: Initialize differentials for primary rays
    void init_differentials(const point3& camera_pos, 
                           const vec3& dx, const vec3& dy) {
        has_differentials = true;
        rx_origin = camera_pos;
        ry_origin = camera_pos;
        rx_direction = dir + dx;
        ry_direction = dir + dy;
    }

    // RAY DIFFERENTIALS: Compute differential footprint at hit point
    double get_footprint(double t, const vec3& normal) const {
        if (!has_differentials) return 0.01;
        
        // Compute where differential rays hit the tangent plane
        point3 p = at(t);
        
        // Transfer differentials to hit point
        double dt_dx = -dot(rx_origin - p, normal) / dot(rx_direction, normal);
        double dt_dy = -dot(ry_origin - p, normal) / dot(ry_direction, normal);
        
        point3 px = rx_origin + (t + dt_dx) * rx_direction;
        point3 py = ry_origin + (t + dt_dy) * ry_direction;
        
        // Footprint is the area covered by differentials
        vec3 dpx = px - p;
        vec3 dpy = py - p;
        
        return sqrt(dpx.length_squared() + dpy.length_squared());
    }

    // RAY DIFFERENTIALS: Propagate differentials through reflection
    void propagate_reflection(const vec3& normal, double t) {
        if (!has_differentials) return;
        
        // Update differential origins
        rx_origin = rx_origin + t * rx_direction;
        ry_origin = ry_origin + t * ry_direction;
        
        // Reflect differential directions
        rx_direction = reflect(rx_direction, normal);
        ry_direction = reflect(ry_direction, normal);
    }

    // RAY DIFFERENTIALS: Propagate differentials through refraction
    void propagate_refraction(const vec3& normal, double eta, double t) {
        if (!has_differentials) return;
        
        rx_origin = rx_origin + t * rx_direction;
        ry_origin = ry_origin + t * ry_direction;
        
        rx_direction = refract(unit_vector(rx_direction), normal, eta);
        ry_direction = refract(unit_vector(ry_direction), normal, eta);
    }
};

// Hit Record
class material;

struct hit_record {
    point3 p;
    vec3 normal;
    std::shared_ptr<material> mat_ptr;
    double t;
    bool front_face;
    double footprint; // RAY DIFFERENTIALS: Store footprint for texture filtering

    inline void set_face_normal(const ray& r, const vec3& outward_normal) {
        front_face = dot(r.direction(), outward_normal) < 0;
        normal = front_face ? outward_normal :-outward_normal;
        
        // RAY DIFFERENTIALS: Compute footprint at hit point
        footprint = r.get_footprint(t, normal);
    }
};

// Hittable
class hittable {
public:
    virtual bool hit(const ray& r, double t_min, double t_max, hit_record& rec) const = 0;
};

// Sphere
class sphere : public hittable {
public:
    point3 center;
    double radius;
    std::shared_ptr<material> mat_ptr;

    sphere() {}
    sphere(point3 cen, double r, std::shared_ptr<material> m)
        : center(cen), radius(r), mat_ptr(m) {};

    virtual bool hit(const ray& r, double t_min, double t_max, hit_record& rec) const override;
};

bool sphere::hit(const ray& r, double t_min, double t_max, hit_record& rec) const {
    vec3 oc = r.origin() - center;
    auto a = r.direction().length_squared();
    auto half_b = dot(oc, r.direction());
    auto c = oc.length_squared() - radius*radius;

    auto discriminant = half_b*half_b - a*c;
    if (discriminant < 0) return false;
    auto sqrtd = sqrt(discriminant);

    auto root = (-half_b - sqrtd) / a;
    if (root < t_min || t_max < root) {
        root = (-half_b + sqrtd) / a;
        if (root < t_min || t_max < root)
            return false;
    }

    rec.t = root;
    rec.p = r.at(rec.t);
    vec3 outward_normal = (rec.p - center) / radius;
    rec.set_face_normal(r, outward_normal);
    rec.mat_ptr = mat_ptr;

    return true;
}

// Hittable List
class hittable_list : public hittable {
public:
    std::vector<std::shared_ptr<hittable>> objects;

    hittable_list() {}
    hittable_list(std::shared_ptr<hittable> object) { add(object); }

    void clear() { objects.clear(); }
    void add(std::shared_ptr<hittable> object) { objects.push_back(object); }

    virtual bool hit(const ray& r, double t_min, double t_max, hit_record& rec) const override;
};

bool hittable_list::hit(const ray& r, double t_min, double t_max, hit_record& rec) const {
    hit_record temp_rec;
    bool hit_anything = false;
    auto closest_so_far = t_max;

    for (const auto& object : objects) {
        if (object->hit(r, t_min, closest_so_far, temp_rec)) {
            hit_anything = true;
            closest_so_far = temp_rec.t;
            rec = temp_rec;
        }
    }

    return hit_anything;
}

// Material
class material {
public:
    virtual bool scatter(
        const ray& r_in, const hit_record& rec, color& attenuation, ray& scattered
    ) const = 0;
};

// Lambertian Material
class lambertian : public material {
public:
    color albedo;

    lambertian(const color& a) : albedo(a) {}

    virtual bool scatter(
        const ray& r_in, const hit_record& rec, color& attenuation, ray& scattered
    ) const override {
        auto scatter_direction = rec.normal + random_unit_vector();

        if (scatter_direction.near_zero())
            scatter_direction = rec.normal;

        scattered = ray(rec.p, scatter_direction);
        
        // RAY DIFFERENTIALS: Use footprint for adaptive sampling
        // Larger footprints could trigger texture filtering (simplified here)
        attenuation = albedo * (1.0 / (1.0 + rec.footprint * 0.5));
        
        return true;
    }
};

class metal : public material {
public:
    color albedo;
    double fuzz;

    metal(const color& a, double f) : albedo(a), fuzz(f < 1 ? f : 1) {}

    virtual bool scatter(
        const ray& r_in, const hit_record& rec, color& attenuation, ray& scattered
    ) const override {
        vec3 reflected = reflect(unit_vector(r_in.direction()), rec.normal);
        scattered = ray(rec.p, reflected + fuzz*random_in_unit_sphere());
        attenuation = albedo;
        
        // RAY DIFFERENTIALS: Propagate through reflection
        if (r_in.has_differentials) {
            scattered.has_differentials = true;
            scattered.rx_origin = rec.p;
            scattered.ry_origin = rec.p;
            scattered.rx_direction = reflect(r_in.rx_direction, rec.normal);
            scattered.ry_direction = reflect(r_in.ry_direction, rec.normal);
        }
        
        return (dot(scattered.direction(), rec.normal) > 0);
    }
};

class dielectric : public material {
public:
    double ir;

    dielectric(double index_of_refraction) : ir(index_of_refraction) {}

    virtual bool scatter(
        const ray& r_in, const hit_record& rec, color& attenuation, ray& scattered
    ) const override {
        attenuation = color(1.0, 1.0, 1.0);
        double refraction_ratio = rec.front_face ? (1.0/ir) : ir;

        vec3 unit_direction = unit_vector(r_in.direction());
        double cos_theta = fmin(dot(-unit_direction, rec.normal), 1.0);
        double sin_theta = sqrt(1.0 - cos_theta*cos_theta);

        bool cannot_refract = refraction_ratio * sin_theta > 1.0;
        vec3 direction;

        if (cannot_refract || reflectance(cos_theta, refraction_ratio) > random_double())
            direction = reflect(unit_direction, rec.normal);
        else
            direction = refract(unit_direction, rec.normal, refraction_ratio);

        scattered = ray(rec.p, direction);
        
        // RAY DIFFERENTIALS: Propagate through refraction/reflection
        if (r_in.has_differentials) {
            scattered.has_differentials = true;
            scattered.rx_origin = rec.p;
            scattered.ry_origin = rec.p;
            if (cannot_refract) {
                scattered.rx_direction = reflect(r_in.rx_direction, rec.normal);
                scattered.ry_direction = reflect(r_in.ry_direction, rec.normal);
            } else {
                scattered.rx_direction = refract(r_in.rx_direction, rec.normal, refraction_ratio);
                scattered.ry_direction = refract(r_in.ry_direction, rec.normal, refraction_ratio);
            }
        }
        
        return true;
    }

private:
    static double reflectance(double cosine, double ref_idx) {
        auto r0 = (1-ref_idx) / (1+ref_idx);
        r0 = r0*r0;
        return r0 + (1-r0)*pow((1 - cosine),5);
    }
};

// Camera
class camera {
public:
    point3 origin;
    point3 lower_left_corner;
    vec3 horizontal;
    vec3 vertical;
    double pixel_dx, pixel_dy; // RAY DIFFERENTIALS: Pixel spacing

    camera() {
        auto aspect_ratio = 16.0 / 9.0;
        auto viewport_height = 2.0;
        auto viewport_width = aspect_ratio * viewport_height;
        auto focal_length = 1.0;

        origin = point3(0, 0, 0);
        horizontal = vec3(viewport_width, 0, 0);
        vertical = vec3(0, viewport_height, 0);
        lower_left_corner = origin - horizontal/2 - vertical/2 - vec3(0, 0, focal_length);
        
        // RAY DIFFERENTIALS: Compute pixel spacing
        pixel_dx = viewport_width / 400.0;  // Image width
        pixel_dy = viewport_height / 225.0; // Image height
    }

    ray get_ray(double u, double v) const {
        ray r(origin, lower_left_corner + u*horizontal + v*vertical - origin);
        
        // RAY DIFFERENTIALS: Initialize with pixel spacing
        vec3 dx = pixel_dx * vec3(1, 0, 0);
        vec3 dy = pixel_dy * vec3(0, 1, 0);
        r.init_differentials(origin, dx, dy);
        
        return r;
    }
};

// Color utilities
void write_color(std::ostream& out, color pixel_color, int samples_per_pixel) {
    auto r = pixel_color.x();
    auto g = pixel_color.y();
    auto b = pixel_color.z();

    auto scale = 1.0 / samples_per_pixel;
    r = sqrt(scale * r);
    g = sqrt(scale * g);
    b = sqrt(scale * b);

    out << static_cast<int>(256 * clamp(r, 0.0, 0.999)) << ' '
        << static_cast<int>(256 * clamp(g, 0.0, 0.999)) << ' '
        << static_cast<int>(256 * clamp(b, 0.0, 0.999)) << '\n';
}

// Ray Color
color ray_color(const ray& r, const hittable& world, int depth) {
    hit_record rec;

    if (depth <= 0)
        return color(0,0,0);

    if (world.hit(r, 0.001, infinity, rec)) {
        ray scattered;
        color attenuation;
        if (rec.mat_ptr->scatter(r, rec, attenuation, scattered)) {
            return attenuation * ray_color(scattered, world, depth-1);
        }
        return color(0,0,0);
    }

    vec3 unit_direction = unit_vector(r.direction());
    auto t = 0.5*(unit_direction.y() + 1.0);
    return (1.0-t)*color(1.0, 1.0, 1.0) + t*color(0.5, 0.7, 1.0);
}

// Scene 1: Classic scene
hittable_list create_scene_1() {
    hittable_list world;

    auto material_ground = std::make_shared<lambertian>(color(0.8, 0.8, 0.0));
    auto material_center = std::make_shared<lambertian>(color(0.1, 0.2, 0.5));
    auto material_left   = std::make_shared<dielectric>(1.5);
    auto material_right  = std::make_shared<metal>(color(0.8, 0.6, 0.2), 0.0);

    world.add(std::make_shared<sphere>(point3( 0.0, -100.5, -1.0), 100.0, material_ground));
    world.add(std::make_shared<sphere>(point3( 0.0,    0.0, -1.0),   0.5, material_center));
    world.add(std::make_shared<sphere>(point3(-1.0,    0.0, -1.0),   0.5, material_left));
    world.add(std::make_shared<sphere>(point3(-1.0,    0.0, -1.0),  -0.4, material_left));
    world.add(std::make_shared<sphere>(point3( 1.0,    0.0, -1.0),   0.5, material_right));

    return world;
}

// Scene 2: Colorful metals
hittable_list create_scene_2() {
    hittable_list world;

    auto material_ground = std::make_shared<lambertian>(color(0.5, 0.5, 0.5));
    auto material_red    = std::make_shared<metal>(color(0.9, 0.1, 0.1), 0.0);
    auto material_green  = std::make_shared<metal>(color(0.1, 0.9, 0.1), 0.3);
    auto material_blue   = std::make_shared<metal>(color(0.1, 0.1, 0.9), 0.0);
    auto material_gold   = std::make_shared<metal>(color(1.0, 0.8, 0.2), 0.1);

    world.add(std::make_shared<sphere>(point3( 0.0, -100.5, -1.0), 100.0, material_ground));
    world.add(std::make_shared<sphere>(point3( 0.0,    0.0, -1.0),   0.5, material_gold));
    world.add(std::make_shared<sphere>(point3(-1.2,    0.0, -1.0),   0.5, material_red));
    world.add(std::make_shared<sphere>(point3( 1.2,    0.0, -1.0),   0.5, material_blue));
    world.add(std::make_shared<sphere>(point3( 0.0,    1.0, -1.5),   0.5, material_green));

    return world;
}

// Scene 3: Glass and diffuse
hittable_list create_scene_3() {
    hittable_list world;

    auto material_ground = std::make_shared<lambertian>(color(0.2, 0.6, 0.2));
    auto material_glass  = std::make_shared<dielectric>(1.5);
    auto material_left   = std::make_shared<lambertian>(color(0.8, 0.3, 0.3));
    auto material_right  = std::make_shared<lambertian>(color(0.3, 0.3, 0.8));
    auto material_center = std::make_shared<lambertian>(color(0.9, 0.9, 0.1));

    world.add(std::make_shared<sphere>(point3( 0.0, -100.5, -1.0), 100.0, material_ground));
    world.add(std::make_shared<sphere>(point3( 0.0,    0.0, -1.0),   0.5, material_glass));
    world.add(std::make_shared<sphere>(point3( 0.0,    0.0, -1.0),  -0.45, material_glass));
    world.add(std::make_shared<sphere>(point3(-1.0,    0.0, -1.0),   0.5, material_left));
    world.add(std::make_shared<sphere>(point3( 1.0,    0.0, -1.0),   0.5, material_right));
    world.add(std::make_shared<sphere>(point3( 0.0,    1.0, -2.0),   0.4, material_center));

    return world;
}

void render_scene(const std::string& filename, hittable_list& world, int scene_num) {
    const auto aspect_ratio = 16.0 / 9.0;
    const int image_width = 400;
    const int image_height = static_cast<int>(image_width / aspect_ratio);
    const int samples_per_pixel = 100;
    const int max_depth = 50;

    camera cam;

    std::ofstream outfile(filename);
    outfile << "P3\n" << image_width << ' ' << image_height << "\n255\n";

    for (int j = image_height-1; j >= 0; --j) {
        std::cerr << "\rScene " << scene_num << " - Scanlines remaining: " << j << ' ' << std::flush;
        for (int i = 0; i < image_width; ++i) {
            color pixel_color(0, 0, 0);
            for (int s = 0; s < samples_per_pixel; ++s) {
                auto u = (i + random_double()) / (image_width-1);
                auto v = (j + random_double()) / (image_height-1);
                ray r = cam.get_ray(u, v);
                pixel_color += ray_color(r, world, max_depth);
            }
            write_color(outfile, pixel_color, samples_per_pixel);
        }
    }
    
    outfile.close();
    std::cerr << "\nScene " << scene_num << " complete: " << filename << "\n";
}

int main() {
    // Render Scene 1
    hittable_list world1 = create_scene_1();
    render_scene("raydiff_scene1.ppm", world1, 1);

    // Render Scene 2
    hittable_list world2 = create_scene_2();
    render_scene("raydiff_scene2.ppm", world2, 2);

    // Render Scene 3
    hittable_list world3 = create_scene_3();
    render_scene("raydiff_scene3.ppm", world3, 3);

    std::cerr << "\nAll scenes complete!\n";
}