Merge pull request #1: Add balloon flight simulation

Author: zuorenchen

rocketpy/rocket/aero_surface/linear_generic_surface.py

@@ -13,7 +13,8 @@ def __init__(
         self,
         reference_area,
         reference_length,
-        coefficients,
+        coefficient_constants=None,
+        coefficients=None,
         center_of_pressure=(0, 0, 0),
         name="Generic Linear Surface",
     ):
@@ -41,6 +42,13 @@ def __init__(
         reference_length : int, float
             Reference length of the aerodynamic surface. Has the unit of meters.
             Commonly defined as the rocket's diameter.
+        coefficient_constants: list, optional
+            List of constants to be populated to the coefficients. This
+            list will overwrite the coefficients dict if not None. The order is
+            as follows: [cL_0, cL_alpha, cL_beta, cL_p, cL_q, cL_r, cQ_0, cQ_alpha,
+            cQ_beta, cQ_p, cQ_q, cQ_r, cD_0, cD_alpha, cD_beta, cD_p, cD_q, cD_r,
+            cm_0, cm_alpha, cm_beta, cm_p, cm_q, cm_r, cn_0, cn_alpha, cn_beta,
+            cn_p, cn_q, cn_r, cl_0, cl_alpha, cl_beta, cl_p, cl_q, cl_r].
         coefficients: dict, optional
             List of coefficients. If a coefficient is omitted, it is set to 0.
             The valid coefficients are:\n
@@ -157,6 +165,36 @@ def __init__(
         name : str
             Name of the aerodynamic surface. Default is 'GenericSurface'.
         """
+        self.coefficient_constants = coefficient_constants
+        # Populate the coefficients from the list of constants if they are not defined
+        if coefficients is None:
+            coefficients = self._get_default_coefficients()
+        if coefficient_constants is not None:
+            # helper to build a 7‑input callable returning a fixed value
+            def _constant_factory(value):
+                def _constant(
+                    alpha, beta, mach, reynolds, pitch_rate, yaw_rate, roll_rate
+                ):
+                    """Return the captured constant regardless of inputs."""
+                    return value
+
+                return _constant
+
+            for i, key in enumerate(self._get_default_coefficients().keys()):
+                const = coefficient_constants[i]
+                coefficients[key] = Function(
+                    _constant_factory(const),
+                    inputs=[
+                        "alpha",
+                        "beta",
+                        "mach",
+                        "reynolds",
+                        "pitch_rate",
+                        "yaw_rate",
+                        "roll_rate",
+                    ],
+                    outputs=[key],
+                )
 
         super().__init__(
             reference_area=reference_area,

rocketpy/rocket/rocket.py

@@ -207,13 +207,16 @@ class Rocket:
         Rocket's inertia tensor 13 component with unloaded motor,in kg*m^2.
     Rocket.dry_I_23 : float
         Rocket's inertia tensor 23 component with unloaded motor,in kg*m^2.
+    Rocket.volume : float
+        Rocket's total volume in m³.
     """
 
     def __init__(  # pylint: disable=too-many-statements
         self,
         radius,
         mass,
         inertia,
+        volume,
         power_off_drag,
         power_on_drag,
         center_of_mass_without_motor,
@@ -240,6 +243,10 @@ def __init__(  # pylint: disable=too-many-statements
             in the direction of e_i x e_j. Alternatively, the inertia tensor can
             be given as (I_11, I_22, I_33), where I_12 = I_13 = I_23 = 0. This
             can also be called as "rocket dry inertia tensor".
+        volume : int, float
+            Rocket's total volume in m³. This can be used for buoyancy calculations.
+            It can also be set to 0 if buoyancy is not being considered.
+            Use None to let the program calculate the volume based on the added components.
         power_off_drag : int, float, callable, string, array
             Rocket's drag coefficient when the motor is off. Can be given as an
             entry to the Function class. See help(Function) for more
@@ -375,6 +382,13 @@ def __init__(  # pylint: disable=too-many-statements
         self.evaluate_reduced_mass()
         self.evaluate_thrust_to_weight()
 
+        if volume is None:
+            # calculate volume properties for buoyancy
+            self.evaluate_volume()
+        else:
+            # Volume and buoyancy properties
+            self.volume = volume
+
         # Evaluate stability (even though no aerodynamic surfaces are present yet)
         self.evaluate_center_of_pressure()
         self.evaluate_stability_margin()
@@ -625,12 +639,14 @@ def evaluate_stability_margin(self):
         self.stability_margin.set_source(
             lambda mach, time: (
                 (
-                    self.center_of_mass.get_value_opt(time)
-                    - self.cp_position.get_value_opt(mach)
+                    (
+                        self.center_of_mass.get_value_opt(time)
+                        - self.cp_position.get_value_opt(mach)
+                    )
+                    / (2 * self.radius)
                 )
-                / (2 * self.radius)
+                * self._csys
             )
-            * self._csys
         )
         return self.stability_margin
 
@@ -647,10 +663,12 @@ def evaluate_static_margin(self):
         # Calculate static margin
         self.static_margin.set_source(
             lambda time: (
-                self.center_of_mass.get_value_opt(time)
-                - self.cp_position.get_value_opt(0)
+                (
+                    self.center_of_mass.get_value_opt(time)
+                    - self.cp_position.get_value_opt(0)
+                )
+                / (2 * self.radius)
             )
-            / (2 * self.radius)
         )
         # Change sign if coordinate system is upside down
         self.static_margin *= self._csys
@@ -662,6 +680,20 @@ def evaluate_static_margin(self):
         )
         return self.static_margin
 
+    def evaluate_volume(self):
+        """Calculates the total volume of the rocket including the motor.
+        The volume must be set manually using set_volume() or calculated from motor.
+
+        Returns
+        -------
+        self.volume : float
+            Total volume of the rocket in m³.
+        """
+        # TODO: calculate volume from radius and surface locations
+        self.volume = 0
+
+        return self.volume
+
     def evaluate_dry_inertias(self):
         """Calculates and returns the rocket's dry inertias relative to
         the rocket's center of dry mass. The inertias are saved and returned

rocketpy/simulation/flight.py

@@ -1714,6 +1714,10 @@ def u_dot(self, t, u, post_processing=False):  # pylint: disable=too-many-locals
         ax, ay, az = K @ Vector(L)
         az -= self.env.gravity.get_value_opt(z)  # Include gravity
 
+        # Include buoyancy force: F_buoyancy = rho * V * g (upward)
+        buoyancy_force = rho * self.rocket.volume * self.env.gravity.get_value_opt(z)
+        az += buoyancy_force / total_mass_at_t  # Add buoyancy acceleration
+
         # Coriolis acceleration
         _, w_earth_y, w_earth_z = self.env.earth_rotation_vector
         ax -= 2 * (vz * w_earth_y - vy * w_earth_z)
@@ -1913,9 +1917,12 @@ def u_dot_generalized(self, t, u, post_processing=False):  # pylint: disable=too
         )
         M3 += self.rocket.cp_eccentricity_x * R2 - self.rocket.cp_eccentricity_y * R1
 
-        weight_in_body_frame = Kt @ Vector(
-            [0, 0, -total_mass * self.env.gravity.get_value_opt(z)]
+        # Calculate weight with buoyancy: F_net = -total_mass * g + rho * V * g
+        gravity_accel = self.env.gravity.get_value_opt(z)
+        net_gravitational_force = (
+            -total_mass * gravity_accel + rho * self.rocket.volume * gravity_accel
         )
+        weight_in_body_frame = Kt @ Vector([0, 0, net_gravitational_force])
 
         T00 = total_mass * r_CM
         T03 = 2 * total_mass_dot * (r_NOZ - r_CM) - 2 * total_mass * r_CM_dot
@@ -2035,9 +2042,13 @@ def u_dot_parachute(self, t, u, post_processing=False):
         Dx = pseudo_drag * freestream_x  # add eta efficiency for wake
         Dy = pseudo_drag * freestream_y
         Dz = pseudo_drag * freestream_z
+
+        # Calculate buoyancy force during parachute phase
+        buoyancy_force = rho * self.rocket.volume * self.env.gravity.get_value_opt(z)
+
         ax = Dx / (mp + ma)
         ay = Dy / (mp + ma)
-        az = (Dz - mp * self.env.gravity.get_value_opt(z)) / (mp + ma)
+        az = (Dz - mp * self.env.gravity.get_value_opt(z) + buoyancy_force) / (mp + ma)
 
         # Add coriolis acceleration
         _, w_earth_y, w_earth_z = self.env.earth_rotation_vector

rocketpy/simulation/monte_carlo.py

@@ -461,6 +461,9 @@ def __run_single_simulation(self):
             initial_solution=self.flight.initial_solution,
             terminate_on_apogee=self.flight.terminate_on_apogee,
             time_overshoot=self.flight.time_overshoot,
+            max_time=self.flight.max_time,
+            max_time_step=self.flight.max_time_step,
+            min_time_step=self.flight.min_time_step,
         )
 
     def __evaluate_flight_inputs(self, sim_idx):

rocketpy/stochastic/stochastic_aero_surfaces.py

@@ -6,6 +6,7 @@
 from rocketpy.rocket.aero_surface import (
     AirBrakes,
     EllipticalFins,
+    LinearGenericSurface,
     NoseCone,
     RailButtons,
     Tail,
@@ -546,3 +547,91 @@ def create_object(self):
         )
         air_brakes.drag_coefficient *= generated_dict["drag_coefficient_curve_factor"]
         return air_brakes
+
+
+class StochasticLinearGenericSurface(StochasticModel):
+    """A Stochastic Linear Generic Surface class that inherits from StochasticModel.
+
+    See Also
+    --------
+    :ref:`stochastic_model` and
+    :class:`LinearGenericSurface <rocketpy.LinearGenericSurface>`
+
+    Attributes
+    ----------
+    object : LinearGenericSurface
+        LinearGenericSurface object to be used for validation.
+    reference_area : tuple, list, int, float
+        Reference area of the aerodynamic surface. Has the unit of meters
+        squared. Commonly defined as the rocket's cross-sectional area.
+    reference_length : tuple, list, int, float
+        Reference length of the aerodynamic surface. Has the unit of meters.
+        Commonly defined as the rocket's diameter.
+    center_of_pressure : tuple, optional
+        Application point of the aerodynamic forces and moments. The
+        center of pressure is defined in the local coordinate system of the
+        aerodynamic surface.
+    coefficient_curve_factor : tuple, list, int, float, optional
+        The drag curve factor of the air brakes. This value scales the
+        drag coefficient curve to introduce stochastic variability.
+    name : list[str]
+        List with the name of the object. This attribute can not be randomized.
+    """
+
+    def __init__(
+        self,
+        linear_generic_surface,
+        reference_area=None,
+        reference_length=None,
+        coefficient_constants=None,
+        center_of_pressure=None,
+    ):
+        """Initializes the Stochastic Linear Generic Surface class.
+
+        See Also
+        --------
+        :ref:`stochastic_model`
+
+        Parameters
+        ----------
+        linear_generic_surface : LinearGenericSurface
+            LinearGenericSurface object to be used for validation.
+        reference_area : int, float
+            Reference area of the aerodynamic surface. Has the unit of meters
+            squared. Commonly defined as the rocket's cross-sectional area.
+        reference_length : int, float
+            Reference length of the aerodynamic surface. Has the unit of meters.
+            Commonly defined as the rocket's diameter.
+        coefficients : dict
+            Dictionary containing the aerodynamic coefficients of the surface.
+        center_of_pressure : tuple, optional
+            Application point of the aerodynamic forces and moments. The center of pressure is defined in the local coordinate system of the
+        coefficient_curve_factor : tuple, list, int, float, optional
+            The drag curve factor of the air brakes. This value scales the
+            drag coefficient curve to introduce stochastic variability.
+        """
+        super().__init__(
+            linear_generic_surface,
+            reference_area=reference_area,
+            reference_length=reference_length,
+            coefficient_constants=coefficient_constants,
+            center_of_pressure=center_of_pressure,
+            name=None,
+        )
+
+    def create_object(self):
+        """Creates and returns a LinearGenericSurface object from the randomly generated input arguments.
+
+        Returns
+        -------
+        linear_generic_surface : LinearGenericSurface
+            LinearGenericSurface object with the randomly generated input arguments.
+        """
+        generated_dict = next(self.dict_generator())
+        linear_generic_surface = LinearGenericSurface(
+            reference_area=generated_dict["reference_area"],
+            reference_length=generated_dict["reference_length"],
+            coefficient_constants=generated_dict["coefficient_constants"],
+            center_of_pressure=generated_dict["center_of_pressure"],
+        )
+        return linear_generic_surface

rocketpy/stochastic/stochastic_flight.py

@@ -32,6 +32,13 @@ class StochasticFlight(StochasticModel):
     time_overshoot : bool
         If False, the simulation will run at the time step defined by the controller
         sampling rate. Be aware that this will make the simulation run much slower.
+    max_time : int, float
+        The maximum time of the flight simulation. If the flight simulation
+        reaches this time, it will terminate. This attribute can not be randomized.
+    min_time_step : int, float
+        The minimum time step of the flight simulation. This attribute can not be randomized.
+    max_time_step : int, float
+        The maximum time step of the flight simulation. This attribute can not be randomized.
     """
 
     def __init__(
@@ -43,6 +50,7 @@ def __init__(
         initial_solution=None,
         terminate_on_apogee=None,
         time_overshoot=None,
+        max_time=None,
     ):
         """Initializes the Stochastic Flight class.
 
@@ -70,6 +78,9 @@ def __init__(
         time_overshoot : bool
             If False, the simulation will run at the time step defined by the controller
             sampling rate. Be aware that this will make the simulation run much slower.
+        max_time : int, float
+            The maximum time of the flight simulation. If the flight simulation
+            reaches this time, it will terminate. This attribute can not be randomized.
         """
         if terminate_on_apogee is not None:
             assert isinstance(terminate_on_apogee, bool), (
@@ -78,6 +89,9 @@ def __init__(
         if time_overshoot is not None:
             if not isinstance(time_overshoot, bool):
                 raise TypeError("`time_overshoot` must be a boolean")
+        if max_time is not None:
+            if not isinstance(max_time, (int, float)):
+                raise TypeError("`max_time` must be a number")
         super().__init__(
             flight,
             rail_length=rail_length,
@@ -87,10 +101,16 @@ def __init__(
 
         self.initial_solution = initial_solution
         self.terminate_on_apogee = terminate_on_apogee
+        if max_time is None:
+            self.max_time = flight.max_time
+        else:
+            self.max_time = max_time
         if time_overshoot is None:
             self.time_overshoot = flight.time_overshoot
         else:
             self.time_overshoot = time_overshoot
+        self.max_time_step = flight.max_time_step
+        self.min_time_step = flight.min_time_step
 
     def _validate_initial_solution(self, initial_solution):
         if initial_solution is not None:
@@ -143,4 +163,5 @@ def create_object(self):
             initial_solution=self.initial_solution,
             terminate_on_apogee=self.terminate_on_apogee,
             time_overshoot=self.time_overshoot,
+            max_time=self.max_time,
         )