From c1e624555b052eae003e859b5469d650b13f5d12 Mon Sep 17 00:00:00 2001
From: "codeflash-ai[bot]"
 <148906541+codeflash-ai[bot]@users.noreply.github.com>
Date: Wed, 17 Dec 2025 15:50:49 +0000
Subject: [PATCH] Optimize IVP._integrate_variable_trajectory
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The optimization replaces the inefficient dynamic array growth pattern with a pre-allocated buffer and doubling strategy, delivering a **686% speedup**.

**Key Optimizations:**

1. **Pre-allocated Buffer**: Instead of starting with a single row and using `np.vstack()` for each step, the code pre-allocates a buffer of 1024 rows using `np.empty()`. This eliminates the O(n²) memory allocation overhead from repeated array concatenation.

2. **Exponential Growth Strategy**: When the buffer fills up, it doubles in size rather than growing by one row. This ensures amortized O(1) insertion cost per step instead of O(n) cost for each `np.vstack()` operation.

3. **Direct Array Assignment**: Values are written directly to pre-allocated positions (`solution[row, 0] = self.t`) rather than creating intermediate arrays with `np.hstack()` and concatenating them.

**Performance Impact Analysis:**

The line profiler shows the dramatic improvement - the original code spent 78.1% of time in `np.vstack()` operations (40 billion nanoseconds), while the optimized version eliminates this bottleneck entirely. The optimized version shows most time is now spent on the actual ODE integration (`self.integrate()` at 65.4%), which is the unavoidable computational work.

**Test Case Performance:**

- **Small trajectories** (edge cases): 13-66% faster due to reduced overhead
- **Medium trajectories**: 50-100% faster as buffer reallocation becomes less frequent
- **Large trajectories**: 162-891% faster where the O(n²) vs O(n) difference dominates (e.g., `test_large_scale_multivar` shows 891% speedup)

The optimization is most effective for workloads with longer integration trajectories, making it particularly valuable for numerical simulations requiring many integration steps.
---
 quantecon/_ivp.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/quantecon/_ivp.py b/quantecon/_ivp.py
index 39cf08bd4..5982d8242 100644
--- a/quantecon/_ivp.py
+++ b/quantecon/_ivp.py
@@ -68,20 +68,34 @@ def _integrate_fixed_trajectory(self, h, T, step, relax):
     def _integrate_variable_trajectory(self, h, g, tol, step, relax):
         """Generates a solution trajectory of variable length."""
         # initialize the solution using initial condition
-        solution = np.hstack((self.t, self.y))
+        initial_capacity = 1024
+        dim = self.y.shape[0] if hasattr(self.y, "shape") else len(self.y)
+        solution = np.empty((initial_capacity, dim + 1), dtype=self.y.dtype)
+        row = 0
+        solution[row, 0] = self.t
+        solution[row, 1:] = self.y
+        row += 1
 
         while self.successful():
 
             self.integrate(self.t + h, step, relax)
-            current_step = np.hstack((self.t, self.y))
-            solution = np.vstack((solution, current_step))
+
+            if row >= solution.shape[0]:
+                new_capacity = solution.shape[0] * 2
+                new_solution = np.empty((new_capacity, dim + 1), dtype=self.y.dtype)
+                new_solution[:solution.shape[0]] = solution
+                solution = new_solution
+
+            solution[row, 0] = self.t
+            solution[row, 1:] = self.y
+            row += 1
 
             if g(self.t, self.y, *self.f_params) < tol:
                 break
             else:
                 continue
 
-        return solution
+        return solution[:row]
 
     def _initialize_integrator(self, t0, y0, integrator, **kwargs):
         """Initializes the integrator prior to integration."""