Clean up and comment

gen_data.py

@@ -1,38 +1,62 @@
-#!./venv/bin/python
-import argparse
-from random import uniform, randint
+#!venv/bin/python
+"""Generates random CSV data to be read by the simulator"""
+
+from argparse import ArgumentParser
+from random import uniform
+from typing import Any, cast
 
 
 class Args:
+	"""
+	The types of the arguments retrieved from the user
+	"""
 	width: int
-	height: int
+	length: int
 	depth: int
-	velocity: float
-	mass: float
+	speed: float
+	radius: float
 	count: int
 
 
-if __name__ == "__main__":
-	parser = argparse.ArgumentParser(
-		prog="n-body data generator",
-		description="Generates data for the n-body simulator.",
-		add_help=False
-	)
+def print_part(data: Any) -> None:
+	"""
+	Prints a CSV field
+	:param data: The data to put in the field
+	"""
+	print(str(data), end=",")
 
-	parser.add_argument("-w", "--width", type=int, default=1900)
-	parser.add_argument("-h", "--height", type=int, default=1000)
-	parser.add_argument("-d", "--depth", type=int, default=0)
-	parser.add_argument("-v", "--velocity", type=float, default=1.)
-	parser.add_argument("-m", "--mass", type=float, default=1.)
-	parser.add_argument("-c", "--count", type=int, default=500)
-
-	args: Args = parser.parse_args()
 
+def main(args: Args) -> None:
+	"""
+	Generates the starting data
+	:param args: Parameters for the random generation
+	"""
+	# Print <count> objects
 	for _ in range(args.count):
-		print(f"{randint(-args.width // 2, args.width // 2)},"
-			  f"{randint(-args.height // 2, args.height // 2)},"
-			  f"{f'{randint(-args.depth // 2, args.depth // 2)},' if args.depth else ''}"
-			  f"{uniform(-args.velocity, args.velocity)},"
-			  f"{uniform(-args.velocity, args.velocity)},"
-			  f"{f'{uniform(-args.velocity, args.velocity)},' if args.depth else ''}"
-			  f"{uniform(1e-2, args.mass)}")
+		# Object location
+		print_part(uniform(-args.width / 2, args.width / 2))
+		print_part(uniform(-args.length / 2, args.length / 2))
+		if args.depth:
+			print_part(uniform(-args.depth / 2, args.depth / 2))
+
+		# Object velocity
+		print_part(uniform(-args.speed, args.speed))
+		print_part(uniform(-args.speed, args.speed))
+		if args.depth:
+			print_part(uniform(-args.speed, args.speed))
+
+		# Finish line with a positive radius
+		print(uniform(1e-2, args.radius))
+
+
+if __name__ == "__main__":
+	parser = ArgumentParser(description="Generates data for the n-body simulator", epilog="You should redirect the output to a file")
+
+	parser.add_argument("-w", "--width", type=float, default=1900., help="The width of the spawning area")
+	parser.add_argument("-l", "--length", type=float, default=1000., help="The length of the spawning area")
+	parser.add_argument("-d", "--depth", type=float, default=0., help="The depth of the spawning area, where 0 implies only 2 dimensions")
+	parser.add_argument("-s", "--speed", type=float, default=1., help="The maximum initial starting speed of an object in any dimension")
+	parser.add_argument("-r", "--radius", type=float, default=1., help="The maximum radius of an object")
+	parser.add_argument("-c", "--count", type=int, default=500, help="How many objects to create")
+
+	main(cast(Args, parser.parse_args()))