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