1#import bevy_pbr::{
2    mesh_view_types::{Lights, DirectionalLight},
3    atmosphere::{
4        types::{Atmosphere, AtmosphereSettings},
5        bindings::{atmosphere, settings},
6        functions::{
7            multiscattering_lut_uv_to_r_mu, sample_transmittance_lut,
8            get_local_r, get_local_up, sample_atmosphere, FRAC_4_PI,
9            max_atmosphere_distance, rayleigh, henyey_greenstein,
10            zenith_azimuth_to_ray_dir,
11        },
12        bruneton_functions::{
13            distance_to_top_atmosphere_boundary, distance_to_bottom_atmosphere_boundary, ray_intersects_ground
14        }
15    }
16}
17
18#import bevy_render::maths::{PI,PI_2}
19
20const PHI_2: vec2<f32> = vec2(1.3247179572447460259609088, 1.7548776662466927600495087);
21
22@group(0) @binding(13) var multiscattering_lut_out: texture_storage_2d<rgba16float, write>;
23
24fn s2_sequence(n: u32) -> vec2<f32> {
25    return fract(0.5 + f32(n) * PHI_2);
26}
27
28// Lambert equal-area projection. 
29fn uv_to_sphere(uv: vec2<f32>) -> vec3<f32> {
30    let phi = PI_2 * uv.y;
31    let sin_lambda = 2 * uv.x - 1;
32    let cos_lambda = sqrt(1 - sin_lambda * sin_lambda);
33
34    return vec3(cos_lambda * cos(phi), cos_lambda * sin(phi), sin_lambda);
35}
36
37// Shared memory arrays for workgroup communication
38var<workgroup> multi_scat_shared_mem: array<vec3<f32>, 64>;
39var<workgroup> l_shared_mem: array<vec3<f32>, 64>;
40
41@compute 
42@workgroup_size(1, 1, 64)
43fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
44    var uv = (vec2<f32>(global_id.xy) + 0.5) / vec2<f32>(settings.multiscattering_lut_size);
45
46    let r_mu = multiscattering_lut_uv_to_r_mu(uv);
47    let light_dir = normalize(vec3(0.0, r_mu.y, -1.0));
48
49    let ray_dir = uv_to_sphere(s2_sequence(global_id.z));
50    let ms_sample = sample_multiscattering_dir(r_mu.x, ray_dir, light_dir);
51    
52    // Calculate the contribution for this sample
53    let sphere_solid_angle = 4.0 * PI;
54    let sample_weight = sphere_solid_angle / 64.0;
55    multi_scat_shared_mem[global_id.z] = ms_sample.f_ms * sample_weight;
56    l_shared_mem[global_id.z] = ms_sample.l_2 * sample_weight;
57
58    workgroupBarrier();
59
60    // Parallel reduction bitshift to the right to divide by 2 each step
61    for (var step = 32u; step > 0u; step >>= 1u) {
62        if global_id.z < step {
63            multi_scat_shared_mem[global_id.z] += multi_scat_shared_mem[global_id.z + step];
64            l_shared_mem[global_id.z] += l_shared_mem[global_id.z + step];
65        }
66        workgroupBarrier();
67    }
68
69    if global_id.z > 0u {
70        return;
71    }
72
73    // Apply isotropic phase function
74    let f_ms = multi_scat_shared_mem[0] * FRAC_4_PI;
75    let l_2 = l_shared_mem[0] * FRAC_4_PI;
76    
77    // Equation 10 from the paper: Geometric series for infinite scattering
78    let psi_ms = l_2 / (1.0 - f_ms);
79    textureStore(multiscattering_lut_out, global_id.xy, vec4<f32>(psi_ms, 1.0));
80}
81
82struct MultiscatteringSample {
83    l_2: vec3<f32>,
84    f_ms: vec3<f32>,
85};
86
87fn sample_multiscattering_dir(r: f32, ray_dir: vec3<f32>, light_dir: vec3<f32>) -> MultiscatteringSample {
88    // get the cosine of the zenith angle of the view direction with respect to the light direction
89    let mu_view = ray_dir.y;
90    let t_max = max_atmosphere_distance(r, mu_view);
91
92    let dt = t_max / f32(settings.multiscattering_lut_samples);
93    var optical_depth = vec3<f32>(0.0);
94
95    var l_2 = vec3(0.0);
96    var f_ms = vec3(0.0);
97    var throughput = vec3(1.0);
98    for (var i: u32 = 0u; i < settings.multiscattering_lut_samples; i++) {
99        let t_i = dt * (f32(i) + 0.5);
100        let local_r = get_local_r(r, mu_view, t_i);
101        let local_up = get_local_up(r, t_i, ray_dir);
102
103        let local_atmosphere = sample_atmosphere(local_r);
104        let sample_optical_depth = local_atmosphere.extinction * dt;
105        let sample_transmittance = exp(-sample_optical_depth);
106        optical_depth += sample_optical_depth;
107
108        let mu_light = dot(light_dir, local_up);
109        let scattering_no_phase = local_atmosphere.rayleigh_scattering + local_atmosphere.mie_scattering;
110
111        let ms = scattering_no_phase;
112        let ms_int = (ms - ms * sample_transmittance) / local_atmosphere.extinction;
113        f_ms += throughput * ms_int;
114
115        let transmittance_to_light = sample_transmittance_lut(local_r, mu_light);
116        let shadow_factor = transmittance_to_light * f32(!ray_intersects_ground(local_r, mu_light));
117
118        let s = scattering_no_phase * shadow_factor * FRAC_4_PI;
119        let s_int = (s - s * sample_transmittance) / local_atmosphere.extinction;
120        l_2 += throughput * s_int;
121
122        throughput *= sample_transmittance;
123        if all(throughput < vec3(0.001)) {
124            break;
125        }
126    }
127
128    //include reflected luminance from planet ground 
129    if ray_intersects_ground(r, mu_view) {
130        let transmittance_to_ground = exp(-optical_depth);
131        let local_up = get_local_up(r, t_max, ray_dir);
132        let mu_light = dot(light_dir, local_up);
133        let transmittance_to_light = sample_transmittance_lut(0.0, mu_light);
134        let ground_luminance = transmittance_to_light * transmittance_to_ground * max(mu_light, 0.0) * atmosphere.ground_albedo;
135        l_2 += ground_luminance;
136    }
137
138    return MultiscatteringSample(l_2, f_ms);
139}