From 53f8993f2b155228fff5cbee84f10c62eb149a1f Mon Sep 17 00:00:00 2001
From: Deepak Mallubhotla <dmallubhotla+github@gmail.com>
Date: Sun, 29 Dec 2024 21:13:34 -0600
Subject: [PATCH] feat: add configurable max number of dipoles to write

---
 deepdog/direct_monte_carlo/direct_mc.py | 32 ++++++++++++++++---------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/deepdog/direct_monte_carlo/direct_mc.py b/deepdog/direct_monte_carlo/direct_mc.py
index 49acf08..71b5c8b 100644
--- a/deepdog/direct_monte_carlo/direct_mc.py
+++ b/deepdog/direct_monte_carlo/direct_mc.py
@@ -284,31 +284,43 @@ class DirectMonteCarloRun:
 
 					pool_results = sum(result[0] for result in raw_pool_results)
 
+					_logger.debug(f"Pool results: {pool_results}")
+
 					if self.config.write_successes_to_file:
+	
+						_logger.info("Writing dipole results")
+
 						cycle_success_configs = numpy.concatenate(
 							[result[1] for result in raw_pool_results]
 						)
-						if len(cycle_success_configs):
 
+						dipole_count = numpy.array(cycle_success_configs).shape[1]
+
+						max_number_dipoles_to_write = self.config.target_success * 5
+						_logger.debug(
+							f"Limiting to {max_number_dipoles_to_write=}, have {len(cycle_success_configs)}"
+						)
+
+
+						if len(cycle_success_configs):
 							sorted_by_freq = numpy.array(
 								[
 									pdme.subspace_simulation.sort_array_of_dipoles_by_frequency(
 										dipole_config
 									)
-									for dipole_config in cycle_success_configs
+									for dipole_config in cycle_success_configs[:max_number_dipoles_to_write]
 								]
 							)
-							dipole_count = numpy.array(cycle_success_configs).shape[1]
 
-							number_dipoles_to_write = self.config.target_success * 5
-							_logger.info(
-								f"Limiting to {number_dipoles_to_write=}, have {dipole_count}"
-							)
 
 							for n in range(dipole_count):
+								
+								dipole_filename = f"{self.config.tag}_{step_count}_dipole_{n}.csv"
+								_logger.debug(f"Writing {min(len(cycle_success_configs), max_number_dipoles_to_write)} to {dipole_filename}")
+								
 								numpy.savetxt(
-									f"{self.config.tag}_{step_count}_dipole_{n}.csv",
-									sorted_by_freq[::number_dipoles_to_write, n],
+									dipole_filename,
+									sorted_by_freq[:, n],
 									delimiter=",",
 								)
 						else:
@@ -316,8 +328,6 @@ class DirectMonteCarloRun:
 								"Instructed to write results, but none obtained"
 							)
 
-					_logger.debug(f"Pool results: {pool_results}")
-
 					total_success += pool_results
 					total_count += count_per_step
 					_logger.debug(