Implement 6-DOF motor clustering and dynamic inertia

rocketpy/motors/ClusterMotor

@@ -0,0 +1,340 @@
+import numpy as np
+from rocketpy.mathutils.function import Function
+from rocketpy.mathutils.vector_matrix import Vector
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from rocketpy.tools import (
+    parallel_axis_theorem_I11,
+    parallel_axis_theorem_I22,
+    parallel_axis_theorem_I33,
+    parallel_axis_theorem_I12,
+    parallel_axis_theorem_I13,
+    parallel_axis_theorem_I23,
+)
+
+
+class ClusterMotor:
+    """
+    Manages a cluster of motors.
+    This class behaves like a single motor by aggregating the properties
+    of several individual motors and implementing the full interface
+    expected by the Rocket class.
+    """
+
+    def __init__(self, motors, positions, orientations=None):
+        if not motors:
+            raise ValueError("The list of motors cannot be empty.")
+
+        self.motors = motors
+        self.positions = [np.array(pos) for pos in positions]
+
+        if orientations is None:
+            self.orientations = [np.array([0, 0, 1.0]) for _ in motors]
+        else:
+            self.orientations = [
+                np.array(ori) / np.linalg.norm(ori) for ori in orientations
+            ]
+
+        if not len(self.motors) == len(self.positions) == len(self.orientations):
+            raise ValueError(
+                "The 'motors', 'positions', and 'orientations' lists must have the same length."
+            )
+
+        self._csys = 1
+        self.coordinate_system_orientation = "tail_to_nose"
+
+        # 1. Basic properties
+        self.propellant_initial_mass = sum(
+            m.propellant_initial_mass for m in self.motors
+        )
+        self.dry_mass = sum(m.dry_mass for m in self.motors)
+
+        # 2. Burn time
+        self.burn_start_time = min(motor.burn_start_time for motor in self.motors)
+        self.burn_out_time = max(motor.burn_out_time for motor in self.motors)
+        self.burn_time = (self.burn_start_time, self.burn_out_time)
+
+        # 3. Thrust and Impulse
+        self.thrust = Function(
+            self._calculate_total_thrust_scalar, inputs="t", outputs="Thrust (N)"
+        )
+        if self.burn_time[1] > self.burn_time[0]:
+            self.thrust.set_discrete(self.burn_start_time, self.burn_out_time, 100)
+
+        self.total_impulse = sum(m.total_impulse for m in self.motors)
+
+        # 4. Mass flow rate (Independent calculation, based on impulse)
+        if self.propellant_initial_mass > 0:
+            average_exhaust_velocity = self.total_impulse / self.propellant_initial_mass
+            self.total_mass_flow_rate = self.thrust / -average_exhaust_velocity
+        else:
+            self.total_mass_flow_rate = Function(0)
+
+        self.mass_flow_rate = self.total_mass_flow_rate  # Alias
+
+        # 5. Propellant mass (based on flow rate)
+        self.propellant_mass = Function(
+            self._calculate_propellant_mass, inputs="t", outputs="Propellant Mass (kg)"
+        )
+
+        # 6. Total mass (based on propellant mass)
+        self.total_mass = Function(
+            self._calculate_total_mass, inputs="t", outputs="Total Mass (kg)"
+        )
+
+        self.nozzle_radius = np.sqrt(sum(m.nozzle_radius**2 for m in self.motors))
+        self.throat_radius = np.sqrt(sum(m.throat_radius**2 for m in self.motors))
+        self.nozzle_position = np.mean([pos[2] for pos in self.positions])
+
+        self.center_of_mass = Function(
+            self._calculate_center_of_mass, inputs="t", outputs="Cluster CoM Vector (m)"
+        )
+        self.center_of_propellant_mass = Function(
+            self._calculate_center_of_propellant_mass,
+            inputs="t",
+            outputs="Cluster CoPM Vector (m)",
+        )
+
+        self.center_of_dry_mass_position = self._calculate_center_of_mass(
+            self.burn_out_time
+        )
+
+        (
+            self.dry_I_11,
+            self.dry_I_22,
+            self.dry_I_33,
+            self.dry_I_12,
+            self.dry_I_13,
+            self.dry_I_23,
+        ) = self._calculate_total_dry_inertia()
+
+        def propellant_inertia_func(t):
+            return self._calculate_total_propellant_inertia(t)
+
+        self.propellant_I_11 = Function(
+            lambda t: propellant_inertia_func(t)[0], inputs="t"
+        )
+        self.propellant_I_22 = Function(
+            lambda t: propellant_inertia_func(t)[1], inputs="t"
+        )
+        self.propellant_I_33 = Function(
+            lambda t: propellant_inertia_func(t)[2], inputs="t"
+        )
+        self.propellant_I_12 = Function(
+            lambda t: propellant_inertia_func(t)[3], inputs="t"
+        )
+        self.propellant_I_13 = Function(
+            lambda t: propellant_inertia_func(t)[4], inputs="t"
+        )
+        self.propellant_I_23 = Function(
+            lambda t: propellant_inertia_func(t)[5], inputs="t"
+        )
+
+    def _calculate_total_mass(self, t):
+        return self.dry_mass + self._calculate_propellant_mass(t)
+
+    def _calculate_propellant_mass(self, t):
+        """Calculates the total propellant mass at time t by integrating the total mass flow rate."""
+        # Note: mass flow rate is negative (mass leaving), so the integral is negative.
+        return self.propellant_initial_mass + self.total_mass_flow_rate.integral(
+            self.burn_start_time, t
+        )
+
+    def _calculate_total_mass_flow_rate(self, t):
+        return -self.propellant_mass.differentiate(t, dx=1e-6)
+
+    def get_total_thrust_vector(self, t):
+        total_thrust = np.zeros(3)
+        for motor, orientation in zip(self.motors, self.orientations):
+            scalar_thrust = motor.thrust.get_value_opt(t)
+            total_thrust += scalar_thrust * orientation
+        return Vector(total_thrust)
+
+    def _calculate_total_thrust_scalar(self, t):
+        return abs(self.get_total_thrust_vector(t))
+
+    def get_total_moment(self, t, ref_point):
+        total_moment = np.zeros(3, dtype=np.float64)
+        ref_point_arr = np.array(ref_point, dtype=np.float64)
+
+        for motor, pos, orientation in zip(
+            self.motors, self.positions, self.orientations
+        ):
+            force_magnitude = motor.thrust.get_value_opt(t)
+            force = force_magnitude * orientation
+            arm = pos - ref_point_arr
+            total_moment += np.cross(arm, force)
+
+            if hasattr(motor, "thrust_eccentricity_y") and hasattr(
+                motor, "thrust_eccentricity_x"
+            ):
+                total_moment[0] += motor.thrust_eccentricity_y * force_magnitude
+                total_moment[1] -= motor.thrust_eccentricity_x * force_magnitude
+
+        return Vector(total_moment)
+
+    def pressure_thrust(self, pressure):
+        return sum(motor.pressure_thrust(pressure) for motor in self.motors)
+
+    def _calculate_center_of_mass(self, t):
+        total_mass_val = self._calculate_total_mass(t)
+        if total_mass_val == 0:
+            return Vector(
+                np.mean(self.positions, axis=0) if self.positions else np.zeros(3)
+            )
+
+        weighted_sum = np.zeros(3, dtype=np.float64)
+        for motor, pos in zip(self.motors, self.positions):
+            motor_com_local = motor.center_of_mass.get_value_opt(t)
+            motor_com_global = pos + np.array([0, 0, motor_com_local * motor._csys])
+            weighted_sum += motor.total_mass.get_value_opt(t) * motor_com_global
+
+        return Vector(weighted_sum / total_mass_val)
+
+    def _calculate_center_of_propellant_mass(self, t):
+        total_prop_mass = self._calculate_propellant_mass(t)
+        if total_prop_mass == 0:
+            return self.center_of_dry_mass_position
+
+        weighted_sum = np.zeros(3, dtype=np.float64)
+        for motor, pos in zip(self.motors, self.positions):
+            prop_com_local = motor.center_of_propellant_mass.get_value_opt(t)
+            prop_com_global = pos + np.array([0, 0, prop_com_local * motor._csys])
+            weighted_sum += motor.propellant_mass.get_value_opt(t) * prop_com_global
+
+        return Vector(weighted_sum / total_prop_mass)
+
+    def _calculate_total_dry_inertia(self):
+        I_11, I_22, I_33 = 0, 0, 0
+        I_12, I_13, I_23 = 0, 0, 0
+
+        ref_point = self.center_of_dry_mass_position
+
+        for motor, pos in zip(self.motors, self.positions):
+            m = motor.dry_mass
+            motor_cdm_global = pos + np.array(
+                [0, 0, motor.center_of_dry_mass_position * motor._csys]
+            )
+            r_vec = Vector(motor_cdm_global) - ref_point
+
+            I_11 += parallel_axis_theorem_I11(motor.dry_I_11, m, r_vec)
+            I_22 += parallel_axis_theorem_I22(motor.dry_I_22, m, r_vec)
+            I_33 += parallel_axis_theorem_I33(motor.dry_I_33, m, r_vec)
+
+            I_12 += parallel_axis_theorem_I12(motor.dry_I_12, m, r_vec)
+            I_13 += parallel_axis_theorem_I13(motor.dry_I_13, m, r_vec)
+            I_23 += parallel_axis_theorem_I23(motor.dry_I_23, m, r_vec)
+
+        return I_11, I_22, I_33, I_12, I_13, I_23
+
+    def _calculate_total_propellant_inertia(self, t):
+        I_11, I_22, I_33 = 0, 0, 0
+        I_12, I_13, I_23 = 0, 0, 0
+
+        ref_point = self._calculate_center_of_propellant_mass(t)
+
+        for motor, pos in zip(self.motors, self.positions):
+            m = motor.propellant_mass.get_value_opt(t)
+            if m == 0:
+                continue
+
+            prop_com_local = motor.center_of_propellant_mass.get_value_opt(t)
+            prop_com_global = pos + np.array([0, 0, prop_com_local * motor._csys])
+            r_vec = Vector(prop_com_global) - ref_point
+
+            I_11 += parallel_axis_theorem_I11(
+                motor.propellant_I_11.get_value_opt(t), m, r_vec
+            )
+            I_22 += parallel_axis_theorem_I22(
+                motor.propellant_I_22.get_value_opt(t), m, r_vec
+            )
+            I_33 += parallel_axis_theorem_I33(
+                motor.propellant_I_33.get_value_opt(t), m, r_vec
+            )
+
+            I_12 += parallel_axis_theorem_I12(
+                motor.propellant_I_12.get_value_opt(t), m, r_vec
+            )
+            I_13 += parallel_axis_theorem_I13(
+                motor.propellant_I_13.get_value_opt(t), m, r_vec
+            )
+            I_23 += parallel_axis_theorem_I23(
+                motor.propellant_I_23.get_value_opt(t), m, r_vec
+            )
+
+        return I_11, I_22, I_33, I_12, I_13, I_23
+
+    def draw_rear_view(self, rocket_radius, tail_radius=None, filename=None):
+        """
+        Plots a 2D rear view of the motor cluster, showing the main
+        rocket body and optionally the tail cone diameter.
+
+        Args:
+            rocket_radius (float): The main radius of the rocket body.
+            tail_radius (float, optional): The rocket's radius at the tail (boattail).
+                                           If provided, a second circle will be drawn.
+            filename (str, optional): If provided, saves the plot to a file.
+                                      Otherwise, shows the plot.
+        """
+        _, ax = plt.subplots(figsize=(8.5, 8.5))
+
+        rocket_body = patches.Circle(
+            (0, 0),
+            radius=rocket_radius,
+            facecolor="lightgrey",
+            edgecolor="black",
+            linewidth=2,
+            label=f"Main Body ({rocket_radius * 200:.0f} cm)",
+        )
+        ax.add_patch(rocket_body)
+
+        if tail_radius is not None:
+            tail_body = patches.Circle(
+                (0, 0),
+                radius=tail_radius,
+                facecolor="silver",
+                edgecolor="black",
+                linestyle="--",
+                linewidth=1.5,
+                label=f"Shrunk ({tail_radius * 200:.0f} cm)",
+            )
+            ax.add_patch(tail_body)
+
+        for i, (motor, pos) in enumerate(zip(self.motors, self.positions)):
+            motor_body = patches.Circle(
+                (pos[0], pos[1]),
+                radius=motor.grain_outer_radius,
+                facecolor="dimgrey",
+                edgecolor="black",
+                linewidth=1.5,
+            )
+            ax.add_patch(motor_body)
+            ax.text(
+                pos[0],
+                pos[1],
+                f"M{i + 1}",
+                ha="center",
+                va="center",
+                color="white",
+                fontsize=12,
+                fontweight="bold",
+            )
+
+        ax.set_aspect("equal", adjustable="box")
+        plot_limit = rocket_radius * 1.2
+        ax.set_xlim(-plot_limit, plot_limit)
+        ax.set_ylim(-plot_limit, plot_limit)
+        ax.set_title("Detailed Rear View of the Cluster", fontsize=16)
+        ax.set_xlabel("Axis X (m)")
+        ax.set_ylabel("Axis Y (m)")
+        ax.grid(True, linestyle="--", alpha=0.6)
+        ax.axhline(0, color="black", linewidth=0.5)
+        ax.axvline(0, color="black", linewidth=0.5)
+
+        plt.legend()
+
+        if filename:
+            plt.savefig(filename)
+            print(f"Rear view plot saved to: {filename}")
+        else:
+            plt.show()