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