need to fix seeding rng issue for multiprocessing. refactored numba class->we get double speed overall for the hpc script. Ran unit tests for optimizations and numba methods. Need to add a pp_analysis.py unit test module as well after fixing the seeding issue

Author: codegithubka

docs/kimon_prompts.md

@@ -46,9 +46,24 @@
 
 ---
 
-### Script Optimization\
-
+### Script Optimization
 
 1. I am considering using numba for optimization and faster runs in the HPC. Outline an implementation plan, practical considerations, and feasibility within a logical timeframe.
 
-2. Walk me through modifying tg
+2. Walk me through modifying the update_async method for simulation speedup
+
+3. Provide me with a script using cProfile to profile the analysis script and target optimization goals specifically.
+
+4. Write a small benchmark script for a large simulation to test how much the performance using numba has imrpoved (or no)
+
+5. Identify the major bottlenecks in the script and suggest possible (realisic) fixes. if you lie about performance improvement expecatations, you will be replaced.
+
+6. I am currently storing anslysis data in a .json file. Help me refactor the 2D sweep function to store binary data instead to minimize communication overhead.
+
+7. The PCF is O(N^2). Show me how to use cell lists instead and spatial hashing. This is the biggest bottlneck so we must prioritize this. Walk me through reftorings in numba_optimized.py
+
+8. We should also pre-allocate the buffers outside the kernel. Right now we are allocating memoery on every call.
+
+9. Help me write a benchmark script to test the updated version against a non-numba optimized version. Create a summary txt with benchmark results.
+
+10. Create unit tests for pp_analysis and numba_optimized files. Test all methods and make sure the scripts are HPC ready. Provide me with an outline of tests.

docs/prompts.md

@@ -690,5 +690,3 @@ Lastly, add functionality to stop evolution after a certain time-step. This shou
 6. Create a time series analysis plot of the evolution of prey and predator density vs. time. Make sure enough time steps all visible to see how the system eventually stabilizes.
 
 7. Create a bifuracation diagram to confirm the monotonic relationship for a varying prey death rate vs. equilibrium density. 
-
-###

models/CA.py

@@ -8,7 +8,7 @@
 
 import numpy as np
 import logging
-from scripts.numba_optimized import _pp_async_kernel
+from scripts.numba_optimized import PPKernel
 from numba import njit
 
 # Module logger
@@ -891,6 +891,8 @@ def __init__(
 		if seed is not None:
 			# This sets the seed for all @njit functions globally
 			set_numba_seed(seed)
+   
+		self._kernel = PPKernel(rows, cols, neighborhood)
 
 
 	# Remove PP-specific evolve wrapper; use CA.evolve with optional species
@@ -1175,30 +1177,26 @@ def _process_reproduction(sources, birth_param_key, birth_prob, target_state_req
 		_process_reproduction(pred_sources, "predator_birth", self.params["predator_birth"], 1, 2)
 
 	def update_async(self) -> None:
-		dr_arr, dc_arr, _ = self._neighbor_shifts()
-		
 		# Get the evolved prey death map
 		# Fallback to a full array of the global param if it doesn't exist yet
 		p_death_arr = self.cell_params.get("prey_death")
 		if p_death_arr is None:
-			p_death_arr = np.full(self.grid.shape, self.params["prey_death"])
+			p_death_arr = np.full(self.grid.shape, self.params["prey_death"], dtype=np.float64)
 
 		meta = self._evolve_info.get("prey_death", {"sd": 0.05, "min": 0.001, "max": 0.1})
 
-		# Call the JIT kernel
-		self.grid = _pp_async_kernel(
+		# Call the optimized kernel (uses pre-allocated buffers)
+		self._kernel.update(
 			self.grid,
 			p_death_arr,
 			float(self.params["prey_birth"]),
 			float(self.params["prey_death"]),
 			float(self.params["predator_birth"]),
 			float(self.params["predator_death"]),
-			dr_arr.astype(np.int32),
-			dc_arr.astype(np.int32),
 			float(meta["sd"]),
 			float(meta["min"]),
 			float(meta["max"]),
-			self._evolution_stopped
+			self._evolution_stopped,
 		)
 
 	def update(self) -> None:

scripts/__jnit__.py

@@ -0,0 +1 @@
+from .numba_optimized import *