Source code for suboptimumg.sweep.motor_sweep

from ..compsim.models import CompetitionData
from ..vehicle.models import MotorModel
from .models import MotorSweepData, SweepParamConfig
from .motor_sweep_results import MotorSweepResults
from .sweeper_1d import Sweeper1D




class MotorSweeper:
    def __init__(
        self,
        comp_data: CompetitionData,
        motors: list[MotorModel],
        var_1: SweepParamConfig,
    ):
        self.comp_data = comp_data
        self.motors = motors
        self.var_1 = var_1



    def sweep(
        self, verbose: bool = False, num_processes: int | None = None
    ) -> MotorSweepResults:
        """
        Runs a one-parameter sweep across all motors.

        Parameters
        ----------
        verbose : bool, optional
            Whether to show detailed progress (default is False)
        num_processes : int, optional
            Number of processes to use. Defaults to CPU count.

        Returns
        -------
        MotorSweepResults
            Results of the motor sweep operation
        """
        res = {}

        for motor in self.motors:
            # Create car config with new motor
            new_comp = self.comp_data.model_copy(deep=True)
            mass_delta = motor.mass - new_comp.vehicle_model.pwrtn.motor.mass
            new_comp.vehicle_model.pwrtn.motor = motor
            new_comp.vehicle_model.mass += mass_delta

            print(f"Sweeping using {motor.name}")
            sweeper = Sweeper1D(
                comp_data=new_comp,
                var_1=self.var_1,
            )

            # Run sweep and collect results
            sweep_results = sweeper.sweep(verbose=verbose, num_processes=num_processes)
            res[motor.name] = sweep_results.sweep_data

        return MotorSweepResults(sweep_data=MotorSweepData(data_by_motor=res))