fix: first optimization

5 files changed, 105 insertions, 14 deletions

diff --git a/assets/shaders/compute_texture_quad.glsl b/assets/shaders/compute_texture_quad.glsl
index 1b83f8a..b7b76d0 100644
--- a/assets/shaders/compute_texture_quad.glsl
+++ b/assets/shaders/compute_texture_quad.glsl
@@ -1,6 +1,7 @@
 #type vertex
 
 #version 330 core
+
 layout (location = 0) in vec2 a_Pos;
 layout (location = 1) in vec2 a_TexCoord;
 
@@ -16,7 +17,7 @@ void main()
 
 #version 330 core
 
-in vec2 v_TexCoord;
+in  vec2 v_TexCoord;
 out vec4 o_FragColor;
 
 uniform sampler2D u_ScreenTexture;
diff --git a/assets/shaders/geodesic.glsl b/assets/shaders/geodesic.glsl
index 88cc2a1..d02eb93 100644
--- a/assets/shaders/geodesic.glsl
+++ b/assets/shaders/geodesic.glsl
@@ -30,13 +30,25 @@ layout(std140, binding = 3) uniform Objects
     float  mass[16]; 
 };
 
+layout(std140, binding = 4) uniform Simulation 
+{
+    int   maxStepsMoving;
+    int   maxStepsStatic;
+    float earlyExitDistance;
+    int   _pad0;
+};
+
 const float  SagA_rs  = 1.269e10;
 const float  D_LAMBDA = 1e7;
 const double ESCAPE_R = 1e30;
 
-const int   MAX_STEPS_MOVING   = 60000;
-const int   MAX_STEPS_STATIC   = 30000;
-const float EARLY_EXIT_DISTANCE = 5e11;
+const int   DEFAULT_MAX_STEPS_MOVING   = 12000;
+const int   DEFAULT_MAX_STEPS_STATIC   = 8000;
+const float DEFAULT_EARLY_EXIT_DISTANCE = 2e12;
+
+const float MIN_STEP_SIZE = 1e6;
+const float MAX_STEP_SIZE = 5e7;
+const float STEP_ADAPTATION_FACTOR = 1.5;
 
 vec4 objectColor = vec4(0.0);
 vec3 hitCenter = vec3(0.0);
@@ -99,8 +111,12 @@ bool InterceptObject(Ray ray)
     {
         vec3  center = objPosRadius[i].xyz;
         float radius = objPosRadius[i].w;
+        
+        float distSq = dot(P - center, P - center);
+        if (distSq > radius * radius * 4.0) 
+            continue;
 
-        if (distance(P, center) <= radius)
+        if (distSq <= radius * radius)
         {
             objectColor = objColor[i];
             hitCenter = center;
@@ -154,6 +170,15 @@ bool CrossesEquatorialPlane(vec3 oldPos, vec3 newPos)
     return crossed && (r >= disk_r1 && r <= disk_r2);
 }
 
+float CalculateAdaptiveStepSize(Ray ray, float baseStepSize) 
+{
+    float r_factor         = clamp(ray.r / (SagA_rs * 10.0), 0.1, 1.0);
+    float curvature        = length(vec3(ray.dr, ray.dtheta * ray.r, ray.dphi * ray.r * sin(ray.theta)));
+    float curvature_factor = clamp(1e12 / (curvature + 1e6), 0.1, 2.0);
+    
+    return clamp(baseStepSize * r_factor * curvature_factor, MIN_STEP_SIZE, MAX_STEP_SIZE);
+}
+
 void main() 
 {
     ivec2 pix  = ivec2(gl_GlobalInvocationID.xy);
@@ -181,17 +206,29 @@ void main()
     bool hitDisk      = false;
     bool hitObject    = false;
 
-    int maxSteps = cam.moving ? MAX_STEPS_MOVING : MAX_STEPS_STATIC;
+    int maxSteps = cam.moving ? maxStepsMoving : maxStepsStatic;
+    
+    if (maxSteps <= 0)
+        maxSteps = cam.moving ? DEFAULT_MAX_STEPS_MOVING : DEFAULT_MAX_STEPS_STATIC;
     
     float cameraDistance = length(cam.camPos);
-    if (cameraDistance > 1e12)
-        maxSteps = maxSteps / 4;
-    else if (cameraDistance > 5e11)
+    if (cameraDistance > 2e12)
+        maxSteps = maxSteps / 2;
+    else if (cameraDistance > 1e12)
+        maxSteps = int(maxSteps * 0.75);
+
+    float initialEscapeVelocity = sqrt(2.0 * SagA_rs / ray.r);
+    if (ray.dr > initialEscapeVelocity * 0.95 && 
+        ray.r  > SagA_rs * 200.0)
         maxSteps = maxSteps / 2;
 
+    float currentStepSize = D_LAMBDA;
+    int objectCheckInterval = 5;
+
     for (int i = 0; i < maxSteps; ++i) 
     {
-        if (ray.r > EARLY_EXIT_DISTANCE) 
+        float exitDistance = earlyExitDistance > 0.0 ? earlyExitDistance : DEFAULT_EARLY_EXIT_DISTANCE;
+        if (ray.r > exitDistance) 
             break;
         if (ray.r > ESCAPE_R) 
             break;
@@ -202,8 +239,10 @@ void main()
             break; 
         }
         
-        RK4Step(ray, D_LAMBDA);
-        lambda += D_LAMBDA;
+        currentStepSize = CalculateAdaptiveStepSize(ray, D_LAMBDA);
+        
+        RK4Step(ray, currentStepSize);
+        lambda += currentStepSize;
 
         vec3 newPos = vec3(ray.x, ray.y, ray.z);
         
@@ -213,13 +252,16 @@ void main()
             break; 
         }
         
-        if (InterceptObject(ray)) 
+        if (i % objectCheckInterval == 0 && InterceptObject(ray)) 
         { 
             hitObject = true; 
             break; 
         }
         
         prevPos = newPos;
+        
+        if (ray.dr > 0.0 && ray.r > SagA_rs * 100.0 && lambda > 2e8)
+            break;
     }
 
     if (hitDisk) 
diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp
index c2a84a8..c9a23ca 100644
--- a/src/engine/engine.cpp
+++ b/src/engine/engine.cpp
@@ -39,6 +39,8 @@ namespace Donut
             + 16 * (sizeof(glm::vec4) + sizeof(glm::vec4))
             + 16 * sizeof(float);
         m_objects_ubo = UniformBuffer::create(obj_ubo_size, 3);
+        
+        m_simulation_ubo = UniformBuffer::create(sizeof(int) * 2 + sizeof(float) * 2, 4);
 
         auto result = quad_vao();
         m_quad_vao = result.first;
@@ -97,6 +99,7 @@ namespace Donut
         upload_camera_ubo(cam);
         upload_disk_ubo();
         upload_objects_ubo(m_objects);
+        upload_simulation_ubo();
         m_texture->bind_as_image(0, false);
 
         uint32_t groups_x = static_cast<uint32_t>(std::ceil(cw / 16.0f));
@@ -173,6 +176,25 @@ namespace Donut
         m_disk_ubo->bind(2);
     }
 
+    void Engine::upload_simulation_ubo()
+    {
+        struct UBOData
+        {
+            int max_steps_moving;
+            int max_steps_static;
+            float early_exit_distance;
+            int padding;
+        } data;
+
+        data.max_steps_moving    = m_max_steps_moving;
+        data.max_steps_static    = m_max_steps_static;
+        data.early_exit_distance = m_early_exit_distance;
+        data.padding           = 0;
+
+        m_simulation_ubo->set_data(&data, sizeof(data));
+        m_simulation_ubo->bind(4);
+    }
+
     void Engine::update_physics(float delta_time)
     {
         for (auto& obj : m_objects)
diff --git a/src/engine/engine.h b/src/engine/engine.h
index 39791c1..0d0fb1b 100644
--- a/src/engine/engine.h
+++ b/src/engine/engine.h
@@ -73,6 +73,7 @@ namespace Donut
         void upload_camera_ubo(const Camera& cam);
         void upload_objects_ubo(const std::vector<ObjectData>& objs);
         void upload_disk_ubo();
+        void upload_simulation_ubo();
         void render_scene();
         void update_physics(float delta_time);
         void update_window_dimensions();
@@ -93,6 +94,13 @@ namespace Donut
         int   get_compute_height()     const { return m_compute_height;                        }
         int   get_compute_width()      const { return (m_width * m_compute_height) / m_height; }
         void  update_compute_dimensions();
+        
+        int   get_max_steps_moving()    const { return m_max_steps_moving; }
+        int   get_max_steps_static()    const { return m_max_steps_static; }
+        float get_early_exit_distance() const { return m_early_exit_distance; }
+        void  set_max_steps_moving(int steps) { m_max_steps_moving = steps; }
+        void  set_max_steps_static(int steps) { m_max_steps_static = steps; }
+        void  set_early_exit_distance(float distance) { m_early_exit_distance = distance; }
     private:
         Ref<Shader> create_compute_program(const char* path);
         std::pair<Ref<VertexArray>, Ref<Texture2D>> quad_vao();
@@ -104,6 +112,7 @@ namespace Donut
         Ref<UniformBuffer> m_camera_ubo;
         Ref<UniformBuffer> m_disk_ubo;
         Ref<UniformBuffer> m_objects_ubo;
+        Ref<UniformBuffer> m_simulation_ubo;
 
         int   m_width;
         int   m_height;
@@ -113,11 +122,15 @@ namespace Donut
         int   m_target_fps     = 60;
         float m_current_fps    = 60.0f;
         float m_last_frame_time = 0.0f;
-        int   m_compute_height = 320;
+        int   m_compute_height = 150;
 
         std::vector<ObjectData> m_objects;
         BlackHole               m_sag_a;
         Camera                  m_camera;
         bool                    m_gravity = false;
+        
+        int   m_max_steps_moving = 60000;
+        int   m_max_steps_static = 30000;
+        float m_early_exit_distance = 5e11f;
     };
 };
diff --git a/src/states/simulation_state.cpp b/src/states/simulation_state.cpp
index e59f1fa..ecb2898 100644
--- a/src/states/simulation_state.cpp
+++ b/src/states/simulation_state.cpp
@@ -124,6 +124,19 @@ namespace Donut
         
         ImGui::Separator();
         
+        ImGui::Text("Simulation Quality:");
+        int max_steps_moving = m_engine.get_max_steps_moving();
+        if (ImGui::SliderInt("Max Steps (Moving)", &max_steps_moving, 1000, 60000))
+            m_engine.set_max_steps_moving(max_steps_moving);
+        int max_steps_static = m_engine.get_max_steps_static();
+        if (ImGui::SliderInt("Max Steps (Static)", &max_steps_static, 1000, 30000))
+            m_engine.set_max_steps_static(max_steps_static);
+        float early_exit_distance = m_engine.get_early_exit_distance();
+        if (ImGui::SliderFloat("Early Exit Distance", &early_exit_distance, 1e11f, 1e13f, "%.2e"))
+            m_engine.set_early_exit_distance(early_exit_distance);
+        
+        ImGui::Separator();
+        
         ImGui::Text("Controls:");
         ImGui::Text("Left Mouse: Orbit camera");
         ImGui::Text("Scroll: Zoom in/out");