1
2struct Nishita {
3    ray_origin: vec3<f32>,
4    sun_position: vec3<f32>,
5    sun_intensity: f32,
6    planet_radius: f32,
7    atmosphere_radius: f32,
8    rayleigh_coefficient: vec3<f32>,
9    rayleigh_scale_height: f32,
10    mie_coefficient: f32,
11    mie_scale_height: f32,
12    mie_direction: f32,
13}
14
15const PI: f32 = 3.141592653589793;
16const ISTEPS: u32 = 16u;
17const JSTEPS: u32 = 8u;
18
19fn rsi(rd: vec3<f32>, r0: vec3<f32>, sr: f32) -> vec2<f32> {
20    // ray-sphere intersection that assumes
21    // the sphere is centered at the origin.
22    // No intersection when result.x > result.y
23    let a = dot(rd, rd);
24    let b = 2.0 * dot(rd, r0);
25    let c = dot(r0, r0) - (sr * sr);
26    let d = (b * b) - (4.0 * a * c);
27
28    if d < 0.0 {
29        return vec2<f32>(1e5, -1e5);
30    } else {
31        return vec2<f32>(
32            (-b - sqrt(d)) / (2.0 * a),
33            (-b + sqrt(d)) / (2.0 * a)
34        );
35    }
36}
37
38fn render_nishita(r_full: vec3<f32>, r0: vec3<f32>, p_sun_full: vec3<f32>, i_sun: f32, r_planet: f32, r_atmos: f32, k_rlh: vec3<f32>, k_mie: f32, sh_rlh: f32, sh_mie: f32, g: f32) -> vec3<f32> {
39    // Normalize the ray direction and sun position.
40    let r = normalize(r_full);
41    let p_sun = normalize(p_sun_full);
42
43    // Calculate the step size of the primary ray.
44    var p = rsi(r, r0, r_atmos);
45    if p.x > p.y { return vec3<f32>(0f); }
46    p.y = min(p.y, rsi(r, r0, r_planet).x);
47    let i_step_size = (p.y - p.x) / f32(ISTEPS);
48
49    // Initialize the primary ray depth.
50    var i_depth = 0.0;
51
52    // Initialize accumulators for Rayleigh and Mie scattering.
53    var total_rlh = vec3<f32>(0f);
54    var total_mie = vec3<f32>(0f);
55
56    // Initialize optical depth accumulators for the primary ray.
57    var i_od_rlh = 0f;
58    var i_od_mie = 0f;
59
60    // Calculate the Rayleigh and Mie phases.
61    let mu = dot(r, p_sun);
62    let mumu = mu * mu;
63    let gg = g * g;
64    let p_rlh = 3.0 / (16.0 * PI) * (1.0 + mumu);
65    let p_mie = 3.0 / (8.0 * PI) * ((1.0 - gg) * (mumu + 1.0)) / (pow(1.0 + gg - 2.0 * mu * g, 1.5) * (2.0 + gg));
66
67    // Sample the primary ray.
68    for (var i = 0u; i < ISTEPS; i++) {
69        // Calculate the primary ray sample position.
70        let i_pos = r0 + r * (i_depth + i_step_size * 0.5);
71
72        // Calculate the height of the sample.
73        let i_height = length(i_pos) - r_planet;
74
75        // Calculate the optical depth of the Rayleigh and Mie scattering for this step.
76        let od_step_rlh = exp(-i_height / sh_rlh) * i_step_size;
77        let od_step_mie = exp(-i_height / sh_mie) * i_step_size;
78
79        // Accumulate optical depth.
80        i_od_rlh += od_step_rlh;
81        i_od_mie += od_step_mie;
82
83        // Calculate the step size of the secondary ray.
84        let j_step_size = rsi(p_sun, i_pos, r_atmos).y / f32(JSTEPS);
85
86        // Initialize the secondary ray depth.
87        var j_depth = 0f;
88
89        // Initialize optical depth accumulators for the secondary ray.
90        var j_od_rlh = 0f;
91        var j_od_mie = 0f;
92
93        // Sample the secondary ray.
94        for (var j = 0u; j < JSTEPS; j++) {
95
96            // Calculate the secondary ray sample position.
97            let j_pos = i_pos + p_sun * (j_depth + j_step_size * 0.5);
98
99            // Calculate the height of the sample.
100            let j_height = length(j_pos) - r_planet;
101
102            // Accumulate the optical depth.
103            j_od_rlh += exp(-j_height / sh_rlh) * j_step_size;
104            j_od_mie += exp(-j_height / sh_mie) * j_step_size;
105
106            // Increment the secondary ray depth.
107            j_depth += j_step_size;
108        }
109
110        // Calculate attenuation.
111        let attn = exp(-(k_mie * (i_od_mie + j_od_mie) + k_rlh * (i_od_rlh + j_od_rlh)));
112
113        // Accumulate scattering.
114        total_rlh += od_step_rlh * attn;
115        total_mie += od_step_mie * attn;
116
117        // Increment the primary ray depth.
118        i_depth += i_step_size;
119    }
120
121    // Calculate and return the final color.
122    return i_sun * (p_rlh * k_rlh * total_rlh + p_mie * k_mie * total_mie);
123}
124
125@group(0) @binding(0)
126var<uniform> nishita: Nishita;
127
128@group(1) @binding(0)
129var image: texture_storage_2d_array<rgba16float, write>;
130
131@compute @workgroup_size(8, 8, 1)
132fn main(@builtin(global_invocation_id) invocation_id: vec3<u32>, @builtin(num_workgroups) num_workgroups: vec3<u32>) {
133    let size = textureDimensions(image).x;
134    let scale = f32(size) / 2f;
135
136    let dir = vec2<f32>((f32(invocation_id.x) / scale) - 1f, (f32(invocation_id.y) / scale) - 1f);
137
138    var ray: vec3<f32>;
139
140    switch invocation_id.z {
141        case 0u {
142            ray = vec3<f32>(1f, -dir.y, -dir.x); // +X
143        }
144        case 1u {
145            ray = vec3<f32>(-1f, -dir.y, dir.x);// -X
146        }
147        case 2u {
148            ray = vec3<f32>(dir.x, 1f, dir.y); // +Y
149        }
150        case 3u {
151            ray = vec3<f32>(dir.x, -1f, -dir.y);// -Y
152        }
153        case 4u {
154            ray = vec3<f32>(dir.x, -dir.y, 1f); // +Z
155        }
156        default: {
157            ray = vec3<f32>(-dir.x, -dir.y, -1f);// -Z
158        }
159    }
160
161    let render = render_nishita(
162        ray,
163        nishita.ray_origin,
164        nishita.sun_position,
165        nishita.sun_intensity,
166        nishita.planet_radius,
167        nishita.atmosphere_radius,
168        nishita.rayleigh_coefficient,
169        nishita.mie_coefficient,
170        nishita.rayleigh_scale_height,
171        nishita.mie_scale_height,
172        nishita.mie_direction,
173    );
174
175    textureStore(
176        image,
177        vec2<i32>(invocation_id.xy),
178        i32(invocation_id.z),
179        vec4<f32>(render, 1.0)
180    );
181}