nbody/gen_data.py

62 lines
1.9 KiB
Python
Executable file

#!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
length: int
depth: int
speed: float
radius: float
count: int
def print_part(data: Any) -> None:
"""
Prints a CSV field
:param data: The data to put in the field
"""
print(str(data), end=",")
def main(args: Args) -> None:
"""
Generates the starting data
:param args: Parameters for the random generation
"""
# Print <count> objects
for _ in range(args.count):
# 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()))