make params customizable

README.md

@@ -18,3 +18,19 @@ Run the main.py script:
 ```shell
 python main.py
 ```
+
+Or alternatively, with your own parameters:
+
+```shell
+python main.py --g 9.8 --C_d 0.02 --v_init 150 --theta 45 --h_init 20 --dt 0.01 --t_end 10
+```
+
+### Arguments
+
+- `--g`: The acceleration due to gravity in m/s^2. Default is 9.81.
+- `--C_d`: The drag coefficient. Default is 0.01.
+- `--v_init`: The initial velocity in m/s. Default is 100.
+- `--theta`: The launch angle in degrees. Default is 180.
+- `--h_init`: The launch height in meters. Default is 10.
+- `--dt`: The time step in seconds. Default is 0.01.
+- `--t_end`: The total time of the simulation in seconds. Default is 10.

main.py

@@ -1,9 +1,40 @@
+import argparse
 from projectile.simulation import ProjectileSimulation
 import numpy as np
 
 
 def main():
-    sim = ProjectileSimulation(v_init=100, theta=np.radians(45))
+    parser = argparse.ArgumentParser(description="Run a projectile simulation.")
+    parser.add_argument(
+        "--g", type=float, default=9.81, help="Acceleration due to gravity (m/s^2)"
+    )
+    parser.add_argument("--C_d", type=float, default=0.01, help="Drag coefficient")
+    parser.add_argument(
+        "--v_init", type=float, default=100, help="Initial velocity (m/s)"
+    )
+    parser.add_argument(
+        "--theta",
+        type=float,
+        default=np.radians(180),
+        help="Launch angle (degrees to radians)",
+    )
+    parser.add_argument("--h_init", type=float, default=10, help="Launch height (m)")
+    parser.add_argument("--dt", type=float, default=0.01, help="Time step (s)")
+    parser.add_argument(
+        "--t_end", type=float, default=10, help="Total time of simulation (s)"
+    )
+
+    args = parser.parse_args()
+
+    sim = ProjectileSimulation(
+        g=args.g,
+        C_d=args.C_d,
+        v_init=args.v_init,
+        theta=args.theta,
+        h_init=args.h_init,
+        dt=args.dt,
+        t_end=args.t_end,
+    )
     sim.run()
     sim.plot()
     sim.print_range()

projectile/simulation.py

@@ -3,6 +3,44 @@ import matplotlib.pyplot as plt
 
 
 class ProjectileSimulation:
+    """
+    A class used to simulate the motion of a projectile with air resistance.
+
+    ...
+
+    Attributes
+    ----------
+    g : float
+        acceleration due to gravity (m/s^2)
+    C_d : float
+        drag coefficient
+    v_init : float
+        initial velocity (m/s)
+    theta : float
+        launch angle (degrees to radians)
+    h_init : float
+        launch height (m)
+    dt : float
+        time step (s)
+    t_end : float
+        total time of simulation (s)
+
+    Methods
+    -------
+    reset():
+        Resets the simulation to its initial state.
+    calculate_acceleration(t, state):
+        Calculates the acceleration at a given time and state.
+    update_state(t, state):
+        Updates the state using the Runge-Kutta method.
+    run():
+        Runs the simulation.
+    plot():
+        Plots the trajectory of the projectile.
+    print_range():
+        Prints the range of the projectile.
+    """
+
     def __init__(
         self,
         g=9.81,
@@ -23,6 +61,7 @@ class ProjectileSimulation:
         self.reset()
 
     def reset(self):
+        """Resets the simulation to its initial state."""
         self.x = [0]
         self.y = [self.h_init]
         self.vx = [self.v_init * np.cos(self.theta)]
@@ -30,6 +69,7 @@ class ProjectileSimulation:
         self.t = [0]
 
     def calculate_acceleration(self, t, state):
+        """Calculates the acceleration at a given time and state."""
         x, y, vx, vy = state
         v = np.sqrt(vx**2 + vy**2)
         F_air_x = -self.C_d * v * vx
@@ -39,6 +79,7 @@ class ProjectileSimulation:
         return [vx, vy, ax, ay]
 
     def update_state(self, t, state):
+        """Updates the state using the Runge-Kutta method."""
         k1 = self.calculate_acceleration(t, state)
         k2 = self.calculate_acceleration(
             t + 0.5 * self.dt, [s + 0.5 * self.dt * k for s, k in zip(state, k1)]
@@ -55,6 +96,7 @@ class ProjectileSimulation:
         ]
 
     def run(self):
+        """Runs the simulation."""
         while self.t[-1] < self.t_end:
             state = [self.x[-1], self.y[-1], self.vx[-1], self.vy[-1]]
             state = self.update_state(self.t[-1], state)
@@ -65,6 +107,7 @@ class ProjectileSimulation:
             self.t.append(self.t[-1] + self.dt)
 
     def plot(self):
+        """Plots the trajectory of the projectile."""
         plt.plot(self.x, self.y)
         plt.xlabel("Horizontal distance (m)")
         plt.ylabel("Vertical distance (m)")
@@ -72,5 +115,6 @@ class ProjectileSimulation:
         plt.show()
 
     def print_range(self):
+        """Prints the range of the projectile."""
         range_projectile = self.x[-1]
         print(f"Range of projectile: {range_projectile:.2f} m")

requirements.txt

@@ -1,2 +1,3 @@
 numpy
 matplotlib
+argparse