[fix]: First optimization

5 files changed, 106 insertions, 14 deletions

diff --git a/Assets/Shaders/ComputeTextureQuad.glsl b/Assets/Shaders/ComputeTextureQuad.glsl
index 1b83f8a..b7b76d0 100644
--- a/Assets/Shaders/ComputeTextureQuad.glsl
+++ b/Assets/Shaders/ComputeTextureQuad.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 7d3cc03..9f871a7 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 ac0e891..799a516 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_ObjectsUBO = UniformBuffer::Create(objUBOSize, 3);
+        
+        m_SimulationUBO = UniformBuffer::Create(sizeof(int) * 2 + sizeof(float) * 2, 4);
 
         auto result = QuadVAO();
         m_QuadVAO = result.first;
@@ -97,6 +99,7 @@ namespace Donut
         UploadCameraUBO(cam);
         UploadDiskUBO();
         UploadObjectsUBO(m_Objects);
+        UploadSimulationUBO();
         m_Texture->BindAsImage(0, false);
 
         uint32_t groupsX = static_cast<uint32_t>(std::ceil(cw / 16.0f));
@@ -173,6 +176,25 @@ namespace Donut
         m_DiskUBO->Bind(2);
     }
 
+    void Engine::UploadSimulationUBO()
+    {
+        struct UBOData
+        {
+            int maxStepsMoving;
+            int maxStepsStatic;
+            float earlyExitDistance;
+            int padding;
+        } data;
+
+        data.maxStepsMoving    = m_MaxStepsMoving;
+        data.maxStepsStatic    = m_MaxStepsStatic;
+        data.earlyExitDistance = m_EarlyExitDistance;
+        data.padding           = 0;
+
+        m_SimulationUBO->SetData(&data, sizeof(data));
+        m_SimulationUBO->Bind(4);
+    }
+
     void Engine::UpdatePhysics(float deltaTime)
     {
         for (auto& obj : m_Objects)
diff --git a/src/Engine/Engine.h b/src/Engine/Engine.h
index a004717..e1b4ad8 100644
--- a/src/Engine/Engine.h
+++ b/src/Engine/Engine.h
@@ -73,6 +73,7 @@ namespace Donut
         void UploadCameraUBO(const Camera& cam);
         void UploadObjectsUBO(const std::vector<ObjectData>& objs);
         void UploadDiskUBO();
+        void UploadSimulationUBO();
         void RenderScene();
         void UpdatePhysics(float deltaTime);
         void UpdateWindowDimensions();
@@ -93,6 +94,14 @@ namespace Donut
         int   GetComputeHeight()     const { return m_ComputeHeight;                        }
         int   GetComputeWidth()      const { return (m_Width * m_ComputeHeight) / m_Height; }
         void  UpdateComputeDimensions();
+        
+        // Simulation parameters
+        int   GetMaxStepsMoving()    const { return m_MaxStepsMoving; }
+        int   GetMaxStepsStatic()    const { return m_MaxStepsStatic; }
+        float GetEarlyExitDistance() const { return m_EarlyExitDistance; }
+        void  SetMaxStepsMoving(int steps) { m_MaxStepsMoving = steps; }
+        void  SetMaxStepsStatic(int steps) { m_MaxStepsStatic = steps; }
+        void  SetEarlyExitDistance(float distance) { m_EarlyExitDistance = distance; }
     private:
         Ref<Shader> CreateComputeProgram(const char* path);
         std::pair<Ref<VertexArray>, Ref<Texture2D>> QuadVAO();
@@ -104,6 +113,7 @@ namespace Donut
         Ref<UniformBuffer> m_CameraUBO;
         Ref<UniformBuffer> m_DiskUBO;
         Ref<UniformBuffer> m_ObjectsUBO;
+        Ref<UniformBuffer> m_SimulationUBO;
 
         int   m_Width;
         int   m_Height;
@@ -113,11 +123,15 @@ namespace Donut
         int   m_TargetFPS     = 60;
         float m_CurrentFPS    = 60.0f;
         float m_LastFrameTime = 0.0f;
-        int   m_ComputeHeight = 320;
+        int   m_ComputeHeight = 150;
 
         std::vector<ObjectData> m_Objects;
         BlackHole               m_SagA;
         Camera                  m_Camera;
         bool                    m_Gravity = false;
+        
+        int   m_MaxStepsMoving = 60000;
+        int   m_MaxStepsStatic = 30000;
+        float m_EarlyExitDistance = 5e11f;
     };
 };
diff --git a/src/States/SimulationState.cpp b/src/States/SimulationState.cpp
index 98640e0..7a45bd1 100644
--- a/src/States/SimulationState.cpp
+++ b/src/States/SimulationState.cpp
@@ -124,6 +124,19 @@ namespace Donut
         
         ImGui::Separator();
         
+        ImGui::Text("Simulation Quality:");
+        int maxStepsMoving = m_Engine.GetMaxStepsMoving();
+        if (ImGui::SliderInt("Max Steps (Moving)", &maxStepsMoving, 1000, 60000))
+            m_Engine.SetMaxStepsMoving(maxStepsMoving);
+        int maxStepsStatic = m_Engine.GetMaxStepsStatic();
+        if (ImGui::SliderInt("Max Steps (Static)", &maxStepsStatic, 1000, 30000))
+            m_Engine.SetMaxStepsStatic(maxStepsStatic);
+        float earlyExitDistance = m_Engine.GetEarlyExitDistance();
+        if (ImGui::SliderFloat("Early Exit Distance", &earlyExitDistance, 1e11f, 1e13f, "%.2e"))
+            m_Engine.SetEarlyExitDistance(earlyExitDistance);
+        
+        ImGui::Separator();
+        
         ImGui::Text("Controls:");
         ImGui::Text("Left Mouse: Orbit camera");
         ImGui::Text("Scroll: Zoom in/out");