ComputePlayground/shaders/hydrogen_wave/pixelCompute.compute
Dawid Pietrykowski 672758f7ec Initial commit
2022-11-09 23:54:43 +01:00

147 lines
3.7 KiB
Plaintext
Executable File

#version 430 core
#define PI 3.14159265359
layout(local_size_x = 8, local_size_y = 8) in;
layout(rgba32f) uniform image2D img_output;
uniform int width;
uniform int height;
uniform int quantum_n;
uniform int quantum_l;
uniform int quantum_m;
uniform float u_phi = 0;
//vec4 base_color = vec4(1);
vec4 base_color = vec4(209, 227, 255, 255) / 255.0f;
const double a0 = 5.29177210903e-11;
const double h = 6.582119569e-16;
const double mq = 9.10938356e-31;
const double e = 8.8541878128e-12;
const double q = 1.60217662e-19;
int factorial(int n){
int res = 1;
for(int i = 2; i <= n; i++){
res *= i;
}
return res;
}
double laguerre(double x, int n, double a){
if(n == 1)
return 1 + a - x;
else
return 1;
}
double harmonics(double theta, double phi){
double res = 1;
if(quantum_m == 0 && quantum_l == 2){
res *= 0.25;
res *= sqrt(5.0/PI);
res *= (3 * pow(cos(float(theta)), 2) - 1);
}else if(quantum_m == 1 && quantum_l == 2){
res *= 0.5;
res *= sqrt(15.0/PI);
res *= sin(float(theta)) * cos(float(phi)) * cos(float(theta));
}
return res;
}
double wave_function(double r, double theta, double phi){
double p1 = pow(float(2/(quantum_n*a0)),3);
double p21 = factorial(quantum_n-quantum_l-1);
double p22 = (2 * quantum_n * factorial(quantum_n+quantum_l));
double root = sqrt(p1 * (p21/p22));
double exponent = exp(float(-r/(quantum_n*a0)));
double p3 = pow(float((2 * r)/(quantum_n * a0)), quantum_l);
double lag = laguerre((2 * r)/(quantum_n * a0), quantum_n - quantum_l - 1, 2 * quantum_l + 1);
double Y = harmonics(theta, phi);
double wave = exponent * root * p3 * lag * Y;
//return root / 2e13 + (quantum_n+quantum_l+quantum_m+ r) * 0.000001f;
//return wave / 1e14;
//return (wave * wave) / 2.28e27;
return (wave * wave) / 1e27;
//return (sin(float(r/sqrt(2)) * 10) + 1.0) * 0.5 + (quantum_n+quantum_l+quantum_m) * 0.000001f;
}
vec4 heatmap(float x){
vec4 colors[6];
colors[0] = vec4(0);
colors[1] = vec4(58, 20, 78, 255) / 255.0f;
colors[2] = vec4(148, 50, 99, 255) / 255.0f;
colors[3] = vec4(249, 114, 13, 255) / 255.0f;
colors[4] = vec4(255, 198, 76, 255) / 255.0f;
colors[5] = vec4(1);
colors[0].w = 0.0f;
colors[1].w = 0.1f;
colors[2].w = 0.25f;
colors[3].w = 0.5f;
colors[4].w = 0.75f;
colors[5].w = 1.0f;
int index = -1;
for(int i = 1; i < 6; i++){
if(x <= colors[i].w){
index = i;
break;
}
}
vec4 res;
if(index != -1){
float range = colors[index].w - colors[index - 1].w;
res = mix(colors[index - 1], colors[index], (x - colors[index - 1].w)/range);
}else
res = colors[5];
res.w = 1;
return res;
}
void main(){
ivec2 id = ivec2(gl_GlobalInvocationID.xy);
vec2 offset = vec2(0.0f, 0.0f);
float x;
float y;
// map x,y to [-1,1]
x = ((((float(id.x) - 0.5f + offset.x)/float(width))) - 0.5f) * 2.0f;
y = ((((float(id.y) - 0.5f + offset.y)/float(height))) - 0.5f) * 2.0f;
double raw_r = sqrt(pow(x, 2) + pow(y, 2));
double r = raw_r * 35 * a0;
double theta = acos(float(y/raw_r));
float phi = u_phi;
vec4 base_color = vec4(0);
vec4 max_color = vec4(1);
// double n_ang ;
// if(x > 0){
// n_ang = mod(theta + phi, PI);
// }else
// n_ang = mod(theta - phi, PI);
float val = 4.0f * float(wave_function(r, theta, phi));
//val = 1 - exp(-1 * val * 5);
//vec4 pixel = mix(base_color, max_color, val);
vec4 pixel = heatmap(val);
imageStore(img_output, id, pixel);
//imageStore(img_output, id, heatmap(float(raw_r)));
}