Source code for suboptimumg.vehicle.aerodynamics.aerodynamics

from collections.abc import Generator
from contextlib import contextmanager

from ...constants import RHO
from ..subsystem_models import AeroModel, DRSModel
from ..vehicle_models import VehicleModel




class DRS:
    def __init__(
        self,
        vehicle_model: VehicleModel,
    ) -> None:
        assert (
            vehicle_model.aero.drs is not None
        ), "Cannot construct DRS for a vehicle whose aero.drs parameters are unset"
        self.vehicle_model: VehicleModel = vehicle_model
        self.params: DRSModel = vehicle_model.aero.drs



    def get_drs_drag(self, v: float) -> float:
        """
        Calculate the drag force based on the current velocity with DRS active.

        Parameters
        ----------
        v : float
            Current velocity (m/s)

        Returns
        -------
        float
            Drag force (N)
        """
        return 0.5 * self.params.drs_cd * self.params.drs_front_area * v**2 * RHO




    def get_drs_lift(self, v: float) -> float:
        """
        Calculate the lift force based on the current velocity with DRS active.

        Parameters
        ----------
        v : float
            Current velocity (m/s)

        Returns
        -------
        float
            Lift force (N)
        """
        return 0.5 * self.params.drs_cl * self.params.drs_front_area * v**2 * RHO




    def get_tot_drs_lift(self, v: float) -> tuple[float, float, float, float]:
        """
        Calculate the total lift force based on the current velocity with DRS active.

        Parameters
        ----------
        v : float
            Current velocity (m/s)

        Returns
        -------
        tuple
            Front and rear lift forces on each tire (N) (fl, fr, rl, rr)
        """
        tot_lift = self.get_drs_lift(v)
        front_lift = tot_lift * self.params.drs_cop
        rear_lift = tot_lift * (1 - self.params.drs_cop)

        return (front_lift / 2, front_lift / 2, rear_lift / 2, rear_lift / 2)




    def drs_activate(self, drs_threshold: float) -> bool:
        """
        Check whether DRS should be activated based on lateral acceleration.

        Parameters
        ----------
        drs_threshold : float
            Current lateral acceleration threshold (m/s^2)

        Returns
        -------
        bool
            True if DRS should be activated, False otherwise
        """
        return self.params.drs_present and drs_threshold <= self.params.drs_accel_thresh






class Aero:
    def __init__(
        self,
        vehicle_model: VehicleModel,
    ) -> None:
        self.vehicle_model: VehicleModel = vehicle_model
        self.params: AeroModel = vehicle_model.aero

        # Construct DRS if it exists in the vehicle model
        self.drs: DRS | None
        if self.params.drs is not None:
            self.drs = DRS(vehicle_model=vehicle_model)
        else:
            self.drs = None

        self.drs_disabled: bool = False



    @contextmanager
    def drs_disabled_override(self) -> Generator[None, None, None]:
        """Force DRS off for the duration of a ``with`` block."""
        previously_disabled = self.drs_disabled
        self.drs_disabled = True
        try:
            yield
        finally:
            self.drs_disabled = previously_disabled




    def drs_is_active(self, drs_threshold: float) -> bool:
        """Check whether DRS is currently active, honoring the disable override.

        Parameters
        ----------
        drs_threshold : float
            Current lateral acceleration threshold (m/s^2)

        Returns
        -------
        bool
            True if DRS is equipped, not overridden off, and within threshold.
        """
        if self.drs is None or self.drs_disabled:
            return False
        return self.drs.drs_activate(drs_threshold)




    def get_drag(self, v: float, drs_threshold: float) -> float:
        """
        Calculates the drag force based on the current velocity.
        Varies the output based on the DRS active state.

        Returns
        -------
        float
            Drag force (N)
        """
        if self.drs_is_active(drs_threshold):
            assert self.drs is not None, "DRS reported active but is not equipped"
            return self.drs.get_drs_drag(v)
        return 0.5 * self.params.cd * self.params.front_area * v**2 * RHO




    def get_lift(self, v: float) -> float:
        """
        Calculates the lift force based on the current velocity without DRS.

        Returns
        -------
        float
            Lift force (N)
        """
        return 0.5 * self.params.cl * self.params.front_area * v**2 * RHO




    def tire_lift_forces(
        self,
        v: float,
        drs_threshold: float,
        pitch: float | None = None,
        roll: float | None = None,
    ) -> tuple[float, float, float, float]:
        """
        Calculates the lift force on each tire based on the current velocity.
        Varies the output based on the DRS active state.

        Returns
        -------
        tuple
            Front and rear lift forces on each tire (N)
        """

        if self.drs_is_active(drs_threshold):
            assert self.drs is not None, "DRS reported active but is not equipped"
            return self.drs.get_tot_drs_lift(v)
        else:
            tot_lift = self.get_lift(v)
            front_lift = tot_lift * self.params.cop
            rear_lift = tot_lift * (1 - self.params.cop)

            return (
                front_lift / 2,
                front_lift / 2,
                rear_lift / 2,
                rear_lift / 2,
            )