fix: better handling of files in stages and remove crufty code

kalpaa/stages/stage01.py

@@ -201,7 +201,9 @@ class Stage01Runner:
 	def run(self):
 		seed_index = 0
 		if self.config.generation_config.override_measurement_filesets is not None:
-			for override_name in self.config.generation_config.override_measurement_filesets.keys():
+			for (
+				override_name
+			) in self.config.generation_config.override_measurement_filesets.keys():
 				# don't need to do anything with the files, just create the out dir
 				out = self.config.get_out_dir_path()
 				directory = out / f"{override_name}"
@@ -273,50 +275,3 @@ def parse_args():
 # 	tantri.cli._generate_dipoles(config_json, dipoles_json, (index, replica, 1))
 
 # 	tantri.cli._write_apsd(dipoles_json, DOTS, X_ELECTRIC_FIELD, DELTA_T, NUM_ITERATIONS, NUM_BIN_TS, (index, replica, 2), output_csv, binned_csv, BIN_WIDTH_LOG, True)
-
-
-def main():
-	args = parse_args()
-
-	tantri_configs = [
-		kalpaa.TantriConfig(31415, 100, 5, 100000),
-		kalpaa.TantriConfig(314, 100, 0.00005, 100000),
-	]
-	generation_config = kalpaa.GenerationConfig(
-		tantri_configs=tantri_configs,
-		counts=[1],
-		num_replicas=3,
-		orientations=[tantri.dipoles.types.Orientation.Z],
-	)
-
-	config = kalpaa.Config(generation_config=generation_config)
-
-	kalpaa.common.set_up_logging(config, args.log_file)
-
-	_logger.info("Generating our data, for the following iterations")
-
-	_logger.info(config)
-	# _logger.info(f"{COUNTS=}")
-	# _logger.info(f"{ORIENTATIONS=}")
-	# _logger.info(f"{NUM_REPLICAS=}")
-
-	# _logger.info("Our parameters used: ")
-
-	# _logger.info(f"\t{INDEX_STARTER=}")
-
-	# _logger.info(f"\t{NUM_SEEDS=}")
-	# # these are obviously not independent but it's just easier than thinking about floats to define them both here
-	# _logger.info(f"\t{DELTA_T=}")
-	# _logger.info(f"\t{SAMPLE_RATE=}")
-	# _logger.info(f"\t{NUM_ITERATIONS=}")
-
-	# # for binnng
-	# _logger.info(f"\t{NUM_BIN_TS=}")
-	# _logger.info(f"\t{BIN_WIDTH_LOG=}")
-
-	runner = Stage01Runner(config)
-	runner.run()
-
-
-if __name__ == "__main__":
-	main()

kalpaa/stages/stage02.py

@@ -115,7 +115,12 @@ class Stage02Runner:
 						_logger.info(f"Working on {trial_name=}")
 						_logger.debug(f"Have {seed_index=}")
 						self.single_run_in_subdir(
-							job_index, cost, dot.label, trial_name, seed_index, override_name=override_key
+							job_index,
+							cost,
+							dot.label,
+							trial_name,
+							seed_index,
+							override_name=override_key,
 						)
 
 	def single_run_in_subdir(
@@ -188,6 +193,7 @@ class Stage02Runner:
 		seed = seed_index
 
 		# TODO find way to store this as a global config file
+		# TODO refactor to account for missing entries, (ex. if occupancy of 150 should use next highest value)
 		occupancies_dict = {
 			1: (500, 1000),
 			2: (250, 2000),
@@ -198,7 +204,9 @@ class Stage02Runner:
 			17: (50, 10000),
 			31: (50, 10000),
 			56: (25, 20000),
-			100: (5, 100000),
+			100: (2, 250000),
+			161: (1, 500000),
+			200: (1, 500000),
 		}
 
 		mccount, mccountcycles = occupancies_dict[avg_filled]
@@ -257,16 +265,17 @@ class Stage02Runner:
 				dipoles_dir.mkdir(exist_ok=True, parents=False)
 				self.run_in_subdir(dipoles_dir, override_key=override_name)
 
-		"""Going to iterate over every folder in out_dir, and execute the subdir stuff inside dirs like <kalpa>/out/z-10-2/dipoles"""
+		else:
-		out_dir_path = self.config.get_out_dir_path()
+			"""Going to iterate over every folder in out_dir, and execute the subdir stuff inside dirs like <kalpa>/out/z-10-2/dipoles"""
-		subdirs = [child for child in out_dir_path.iterdir() if child.is_dir]
+			out_dir_path = self.config.get_out_dir_path()
-		# _logger.info(f"Going to execute within each of the directories in {subdirs=}")
+			subdirs = [child for child in out_dir_path.iterdir() if child.is_dir]
-		for subdir in subdirs:
+			# _logger.info(f"Going to execute within each of the directories in {subdirs=}")
-			# skip try finally for now just blow up if problem
+			for subdir in subdirs:
-			_logger.debug(f"Running for {subdir=}")
+				# skip try finally for now just blow up if problem
-			dipoles_dir = subdir / "dipoles"
+				_logger.debug(f"Running for {subdir=}")
-			dipoles_dir.mkdir(exist_ok=True, parents=False)
+				dipoles_dir = subdir / "dipoles"
-			self.run_in_subdir(subdir / "dipoles")
+				dipoles_dir.mkdir(exist_ok=True, parents=False)
+				self.run_in_subdir(subdir / "dipoles")
 
 
 def parse_args():
@@ -283,28 +292,3 @@ def parse_args():
 	)
 	args = parser.parse_args()
 	return args
-
-
-def main():
-	args = parse_args()
-
-	tantri_configs = [
-		kalpaa.TantriConfig(31415, 100, 5, 100000),
-		kalpaa.TantriConfig(314, 100, 0.00005, 100000),
-	]
-	generation_config = kalpaa.GenerationConfig(tantri_configs=tantri_configs)
-
-	config = kalpaa.Config(generation_config=generation_config)
-
-	kalpaa.common.set_up_logging(config, args.log_file)
-
-	_logger.info("Generating our data, for the following iterations")
-
-	_logger.info(config)
-
-	runner = Stage02Runner(config)
-	runner.run()
-
-
-if __name__ == "__main__":
-	main()

kalpaa/stages/stage03.py

@@ -161,33 +161,3 @@ def parse_args():
 	)
 	args = parser.parse_args()
 	return args
-
-
-def main():
-	args = parse_args()
-
-	tantri_configs = [
-		kalpaa.TantriConfig(31415, 100, 5, 100000),
-		kalpaa.TantriConfig(314, 100, 0.00005, 100000),
-	]
-	generation_config = kalpaa.GenerationConfig(
-		tantri_configs=tantri_configs,
-		counts=[1],
-		num_replicas=3,
-		orientations=[tantri.dipoles.types.Orientation.Z],
-	)
-
-	config = kalpaa.Config(generation_config=generation_config)
-
-	kalpaa.common.set_up_logging(config, args.log_file)
-
-	_logger.info("Generating our data, for the following iterations")
-
-	_logger.info(config)
-
-	runner = Stage03Runner(config)
-	runner.run()
-
-
-if __name__ == "__main__":
-	main()

kalpaa/stages/stage04.py

@@ -156,13 +156,19 @@ class Stage04Runner:
 			writer.writeheader()
 
 			if self.config.generation_config.override_dipole_configs is not None:
-				override_names = self.config.generation_config.override_dipole_configs.keys()
+				override_names = (
-			elif self.config.generation_config.override_measurement_filesets is not None:
+					self.config.generation_config.override_dipole_configs.keys()
-				override_names = self.config.generation_config.override_measurement_filesets.keys()
+				)
+			elif (
+				self.config.generation_config.override_measurement_filesets is not None
+			):
+				override_names = (
+					self.config.generation_config.override_measurement_filesets.keys()
+				)
 			else:
 				override_names = None
 
-			if (override_names is not None):
+			if override_names is not None:
 				_logger.debug(
 					f"We had overridden dipole config, using override {override_names}"
 				)
@@ -237,33 +243,3 @@ def parse_args():
 	)
 	args = parser.parse_args()
 	return args
-
-
-def main():
-	args = parse_args()
-
-	tantri_configs = [
-		kalpaa.TantriConfig(31415, 100, 5, 100000),
-		kalpaa.TantriConfig(314, 100, 0.00005, 100000),
-	]
-	generation_config = kalpaa.GenerationConfig(
-		tantri_configs=tantri_configs,
-		counts=[1],
-		num_replicas=3,
-		orientations=[tantri.dipoles.types.Orientation.Z],
-	)
-
-	config = kalpaa.Config(generation_config=generation_config)
-
-	kalpaa.common.set_up_logging(config, args.log_file)
-
-	_logger.info("Generating our data, for the following iterations")
-
-	_logger.info(config)
-
-	runner = Stage04Runner(config)
-	runner.run()
-
-
-if __name__ == "__main__":
-	main()