finally speedup

Author: codegithubka

results_stress/sweep_metadata.json

@@ -0,0 +1,5 @@
+{
+  "n_sims": 24,
+  "timestamp": "2026-01-21 19:39:59",
+  "grid_size": 200
+}

results_stress/sweep_results.npz

@@ -15,12 +15,12 @@ def run_stress_benchmark():
     cfg = Config()
     # INCREASE WORKLOAD: 150x150 grid and 1000 total steps
     # This ensures the 'race' is long enough to overcome the 'warmup'
-    cfg.default_grid = 150       
+    cfg.default_grid = 200     
     cfg.n_prey_birth = 2        # 2x2 sweep is enough to see scaling
     cfg.n_prey_death = 2
     cfg.n_replicates = 3        
     cfg.warmup_steps = 400
-    cfg.measurement_steps = 600
+    cfg.measurement_steps = 10000
     cfg.n_jobs = 4              
 
     output_dir = Path("results_stress")

scripts/__jnit__.py

@@ -32,57 +32,74 @@ def _pp_async_kernel(
     rows, cols = grid.shape
     n_shifts = len(dr_arr)
 
-    # 1. Identify occupied cells
-    occupied = []
+    # 1. PRE-ALLOCATE COORDINATE BUFFER
+    # Instead of a list, we use a fixed-size array. 
+    # Max possible occupied cells is rows * cols.
+    occupied = np.empty((rows * cols, 2), dtype=np.int32)
+    count = 0
+    
     for r in range(rows):
         for c in range(cols):
             if grid[r, c] != 0:
-                occupied.append((r, c))
+                occupied[count, 0] = r
+                occupied[count, 1] = c
+                count += 1
 
-    # 2. Shuffle (Asynchronous behavior)
-    N = len(occupied)
-    for i in range(N - 1, 0, -1):
+    # 2. IN-PLACE FISHER-YATES SHUFFLE
+    # We only shuffle the part of the buffer we actually filled (up to 'count')
+    for i in range(count - 1, 0, -1):
         j = np.random.randint(0, i + 1)
-        temp = occupied[i]
-        occupied[i] = occupied[j]
-        occupied[j] = temp
+        # Swap row index
+        r_temp = occupied[i, 0]
+        occupied[i, 0] = occupied[j, 0]
+        occupied[j, 0] = r_temp
+        # Swap col index
+        c_temp = occupied[i, 1]
+        occupied[i, 1] = occupied[j, 1]
+        occupied[j, 1] = c_temp
+        
+    # 3. PROCESS ACTIVE CELLS
+    for i in range(count):
+        r = occupied[i, 0]
+        c = occupied[i, 1]
 
-    # 3. Process
-    for pos in occupied:
-        r, c = pos
         state = grid[r, c]
-        if state == 0: continue 
+        if state == 0: 
+            continue 
 
+        # Pick random neighbor
         nbi = np.random.randint(0, n_shifts)
         nr = (r + dr_arr[nbi]) % rows
         nc = (c + dc_arr[nbi]) % cols
 
         if state == 1: # PREY
-            # Death
+            # Death logic
             if np.random.random() < prey_death_arr[r, c]:
                 grid[r, c] = 0
                 prey_death_arr[r, c] = np.nan
-            # Birth into empty cell
+            # Birth logic
             elif grid[nr, nc] == 0:
                 if np.random.random() < p_birth_val:
                     grid[nr, nc] = 1
-                    # Inheritance
                     parent_val = prey_death_arr[r, c]
                     if not evolution_stopped:
                         child_val = parent_val + np.random.normal(0, evolve_sd)
-                        prey_death_arr[nr, nc] = max(evolve_min, min(evolve_max, child_val))
+                        # Manual clip is faster than np.clip in Numba
+                        if child_val < evolve_min: child_val = evolve_min
+                        if child_val > evolve_max: child_val = evolve_max
+                        prey_death_arr[nr, nc] = child_val
                     else:
                         prey_death_arr[nr, nc] = parent_val
 
         elif state == 2: # PREDATOR
-            # Death
+            # Death logic
             if np.random.random() < pred_death_val:
                 grid[r, c] = 0
-            # Birth into prey cell
+            # Birth logic (eating prey)
             elif grid[nr, nc] == 1:
                 if np.random.random() < pred_birth_val:
                     grid[nr, nc] = 2
-                    prey_death_arr[nr, nc] = np.nan # Predator eats prey, clear params
+                    prey_death_arr[nr, nc] = np.nan
 
     return grid
 

scripts/macro_benchmark.py

@@ -33,6 +33,10 @@
 from typing import Dict, List, Tuple, Optional
 import warnings
 
+project_root = str(Path(__file__).parents[1])
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
 try:
     from scripts.numba_optimized import (
         compute_pcf_periodic_fast,
@@ -80,8 +84,8 @@ class Config:
     # Simulation length
     warmup_steps: int = 200
     measurement_steps: int = 300
-    cluster_samples: int = 10
-    cluster_interval: int = 100 
+    cluster_samples: int = 3
+    cluster_interval: int = 1000
 
     # Ecological parameters
     evolve_sd: float = 0.10
@@ -104,7 +108,7 @@ class Config:
     synchronous: bool = False
 
     # Snapshot settings
-    save_snapshots: bool = True
+    save_snapshots: bool = False
     snapshot_times: Tuple[int, ...] = (50, 100, 150, 200, 250)
 
     n_jobs: int = -1  # -1 = all available cores
@@ -209,6 +213,7 @@ def get_evolution_metadata(model)->Dict:
     return metadata
 
 
+
 def collect_comprehensive_metrics(model, step: int) -> Dict:
     """Collect comprehensive metrics using built-in CA methods."""
     metrics = {
@@ -581,6 +586,39 @@ def save_diagnostic_snapshot(
     plt.savefig(output_path / f'diagnostic_{run_id}.png', dpi=120, bbox_inches='tight')
     plt.close()
 
+
+def save_sweep_binary(results: List[Dict], output_path: Path):
+    """Saves high-volume sweep data to compressed .npz to avoid JSON overhead."""
+    # We use a dictionary to map each run to its own key-space
+    data_to_save = {}
+    for i, res in enumerate(results):
+        prefix = f"run_{i}_"
+        for key, val in res.items():
+            # np.savez handles numbers and arrays perfectly; 
+            # objects/dicts are converted to array-wrappers
+            data_to_save[f"{prefix}{key}"] = np.array(val)
+            
+    np.savez_compressed(output_path, **data_to_save)
+    
+def warmup_numba_kernels(grid_size: int):
+    """Compiles kernels in the main thread so workers load from cache instantly."""
+    logging.info(f"Warming up Numba kernels for {grid_size}x{grid_size} grid...")
+    
+    # Create dummy data structures matching real types
+    dummy_grid = np.zeros((grid_size, grid_size), dtype=np.int32)
+    p_death_arr = np.full((grid_size, grid_size), 0.05, dtype=np.float64)
+    dr = np.array([-1, 1, 0, 0], dtype=np.int32)
+    dc = np.array([0, 0, -1, 1], dtype=np.int32)
+    
+    # Trigger Simulation JIT
+    from scripts.numba_optimized import _pp_async_kernel, _compute_distance_histogram
+    _ = _pp_async_kernel(dummy_grid, p_death_arr, 0.2, 0.05, 0.2, 0.1, dr, dc, 0.1, 0.001, 0.1, False)
+
+    # Trigger PCF JIT
+    pos = np.ascontiguousarray(np.argwhere(dummy_grid == 0)[:10], dtype=np.float64)
+    if len(pos) > 0:
+        _ = _compute_distance_histogram(pos, pos, float(grid_size), float(grid_size), 20.0, 20, True)
+ 
 
 ##########################################################################
 # Main Simulation Function
@@ -943,6 +981,8 @@ def run_single_simulation_fss(
 def run_2d_sweep(cfg, output_dir, logger) -> List[Dict]:
     """Run full 2D parameter sweep with optional diagnostic snapshots."""
     from joblib import Parallel, delayed
+    if USE_NUMBA:
+        warmup_numba_kernels(cfg.default_grid)
 
     prey_births = cfg.get_prey_births()
     prey_deaths = cfg.get_prey_deaths()
@@ -987,14 +1027,22 @@ def run_2d_sweep(cfg, output_dir, logger) -> List[Dict]:
         for pb, pd, gs, seed, evo, save_diag in jobs
     )
 
-    output_file = output_dir / "sweep_results.json"
-    with open(output_file, "w") as f:
-        json.dump(results, f)
-    logger.info(f"Saved to {output_file}")
+    # REPLACE THE JSON BLOCK WITH THIS:
+    output_file = output_dir / "sweep_results.npz"
+    save_sweep_binary(results, output_file)
+    
+    # Save a tiny 'metadata' JSON just for quick human reading
+    meta = {
+        "n_sims": len(results),
+        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
+        "grid_size": cfg.default_grid
+    }
+    with open(output_dir / "sweep_metadata.json", "w") as f:
+        json.dump(meta, f, indent=2)
 
+    logger.info(f"SUCCESS: Saved binary sweep results to {output_file}")
     return results
 
-
 def run_sensitivity(
     cfg: Config, output_dir: Path, logger: logging.Logger
 ) -> List[Dict]:

scripts/numba_optimized.py

@@ -0,0 +1,26 @@
+import cProfile, pstats
+from pathlib import Path
+import sys
+
+# Ensure we can find our modules
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from pp_analysis import Config, run_single_simulation
+
+# 1. Setup a single simulation configuration
+cfg = Config()
+cfg.default_grid = 150
+cfg.warmup_steps = 200
+cfg.measurement_steps = 300
+
+# 2. Profile the function
+profiler = cProfile.Profile()
+profiler.enable()
+
+# Run a single simulation (no parallelization)
+run_single_simulation(0.2, 0.05, 150, 42, True, cfg)
+
+profiler.disable()
+
+# 3. Print the top 15 time-consumers
+stats = pstats.Stats(profiler).sort_stats('tottime')
+stats.print_stats(15)