1#import bevy_pbr::atmosphere::{
2    types::{Atmosphere, AtmosphereSettings},
3    bindings::{settings, atmosphere},
4    functions::{AtmosphereSample, sample_atmosphere, get_local_r, max_atmosphere_distance, MIDPOINT_RATIO},
5    bruneton_functions::{transmittance_lut_uv_to_r_mu, distance_to_bottom_atmosphere_boundary, distance_to_top_atmosphere_boundary},
6}
7
8
9#import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput
10
11@group(0) @binding(13) var transmittance_lut_out: texture_storage_2d<rgba16float, write>;
12
13@compute 
14@workgroup_size(16, 16, 1)
15fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
16    let uv: vec2<f32> = (vec2<f32>(idx.xy) + 0.5) / vec2<f32>(settings.transmittance_lut_size);
17    // map UV coordinates to view height (r) and zenith cos angle (mu)
18    let r_mu = transmittance_lut_uv_to_r_mu(uv);
19
20    // compute the optical depth from view height r to the top atmosphere boundary
21    let optical_depth = ray_optical_depth(r_mu.x, r_mu.y, settings.transmittance_lut_samples);
22    let transmittance = exp(-optical_depth);
23
24    textureStore(transmittance_lut_out, idx.xy, vec4(transmittance, 1.0));
25}
26
27/// Compute the optical depth of the atmosphere from the ground to the top atmosphere boundary
28/// at a given view height (r) and zenith cos angle (mu)
29fn ray_optical_depth(r: f32, mu: f32, sample_count: u32) -> vec3<f32> {
30    let t_max = max_atmosphere_distance(r, mu);
31    var optical_depth = vec3<f32>(0.0f);
32    var prev_t = 0.0f;
33
34    for (var i = 0u; i < sample_count; i++) {
35        let t_i = t_max * (f32(i) + MIDPOINT_RATIO) / f32(sample_count);
36        let dt = t_i - prev_t;
37        prev_t = t_i;
38
39        let r_i = get_local_r(r, mu, t_i);
40
41        let atmosphere_sample = sample_atmosphere(r_i);
42        let sample_optical_depth = atmosphere_sample.extinction * dt;
43
44        optical_depth += sample_optical_depth;
45    }
46
47    return optical_depth;
48}