Index

def with_alpha(self, alpha: float) -> Vec4:
    """Returns the color with a custom transparency level."""
    return Color.hex_to_rgba(self.value, alpha=alpha)

@classmethod
def hex_to_rgba(cls, hex_str: str, alpha: float = 1.0) -> Vec4:
    """Converts '#RRGGBB' or 'RRGGBB' to normalized [0, 1] RGBA."""
    hex_str = hex_str.lstrip("#")
    # Convert hex to integers using bit-shifting
    lv = len(hex_str)
    if lv == 6:
        rgb = tuple(int(hex_str[i : i + 2], 16) for i in (0, 2, 4))
    else:
        raise ValueError(f"Invalid hex color: {hex_str}. Expected 6 characters.")

    return np.array([*(np.array(rgb) / 255.0), alpha])

@classmethod
def rgba255_to_rgba(cls, rgba255: Vec4) -> Vec4:
    """Converts [0-255, 0-255, 0-255, 0-1] to normalized [0, 1] RGBA."""
    res = np.array(rgba255, dtype=float)
    res[:3] /= 255.0
    return res

@classmethod
def rgba_to_hex(cls, rgba: Vec4) -> str:
    """Converts normalized RGBA to '#RRGGBB' (alpha is discarded)."""
    rgba = np.asarray(rgba, dtype=float)
    rgb255 = (rgba[:3] * 255).astype(int)
    return "#{:02x}{:02x}{:02x}".format(*rgb255)

@classmethod
def rgba_to_rgba255(cls, rgba: Vec4) -> Vec4:
    """Converts normalized RGBA to [0-255, 0-255, 0-255, 0-1]."""
    res = np.array(rgba, dtype=float)
    res[:3] *= 255.0
    return res

@classmethod
def from_series(cls, series: pd.Series, kind: InterpOptions = "linear") -> Self:
    """Create an interpolator directly from a Pandas Series."""
    return cls(x=np.asarray(series.index), y=np.asarray(series.values), kind=kind)

def get_sphere_to_sphere_proximity(
    self,
    mj_model: mujoco.MjModel,
    mj_data: mujoco.MjData,
) -> tuple[float, Vec3, Vec3, bool]:
    """
    Calculates the shortest distance between two geometries using their bounding spheres.

    Args:
        mj_model (mujoco.MjModel): The compiled MuJoCo model instance.
        mj_data (mujoco.MjData): The current simulation state.

    Returns:
        tuple[float, Vec3, Vec3, bool]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and if the estimated distance exceeds dist_max.

    """
    # get world orientations and origins
    origin_geom_1 = self.geom_1.rt_xpos(mj_model, mj_data)
    mat_geom_1 = self.geom_1.rt_xmat(mj_model, mj_data)

    origin_geom_2 = self.geom_2.rt_xpos(mj_model, mj_data)
    mat_geom_2 = self.geom_2.rt_xmat(mj_model, mj_data)

    if np.isnan(self.geom_1._rad):
        self.geom_1._rad, self.geom_1._local_centroid = self.geom_1.vertex_max_norm(
            mj_model
        )

    if np.isnan(self.geom_2._rad):
        self.geom_2._rad, self.geom_2._local_centroid = self.geom_2.vertex_max_norm(
            mj_model
        )

    # shift centers to pre-calculated centroids
    pos_geom_1 = origin_geom_1 + (mat_geom_1 @ self.geom_1._local_centroid)
    pos_geom_2 = origin_geom_2 + (mat_geom_2 @ self.geom_2._local_centroid)

    rad_geom_1 = self.geom_1._rad
    rad_geom_2 = self.geom_2._rad

    vec_geom_1_to_geom_2 = pos_geom_2 - pos_geom_1
    d_centers = float(np.linalg.norm(vec_geom_1_to_geom_2))
    dist = d_centers - (rad_geom_1 + rad_geom_2)

    dist = max(0.0, dist)  # clip to zero
    exceeds_dist_max = dist > self.dist_max

    if d_centers > 1e-9:
        unit_vec = vec_geom_1_to_geom_2 / d_centers
        p1 = pos_geom_1 + (unit_vec * rad_geom_1)
        p2 = pos_geom_2 - (unit_vec * rad_geom_2)
    else:
        p1 = pos_geom_1
        p2 = pos_geom_2

    self.update_last(p1, p2, mj_data)
    return dist, p1, p2, exceeds_dist_max

def get_convex_hull_proximity(
    self,
    mj_model: mujoco.MjModel,
    mj_data: mujoco.MjData,
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the shortest distance between two geometries using their convex hull.

    Args:
        mj_model (mujoco.MjModel): The compiled MuJoCo model instance.
        mj_data (mujoco.MjData): The current simulation state.

    Returns:
        tuple[float, Vec3, Vec3, ProximityType]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

    """
    # ========== BROADPHASE ==========
    if self.geom_1._proximity_configured_for != ProximityType.CONVEX_HULL:
        self.geom_1.bake_proximity(mj_model, ProximityType.CONVEX_HULL)

    if self.geom_2._proximity_configured_for != ProximityType.CONVEX_HULL:
        self.geom_2.bake_proximity(mj_model, ProximityType.CONVEX_HULL)

    d_est, p1, p2, skip = self.get_sphere_to_sphere_proximity(
        mj_model=mj_model, mj_data=mj_data
    )

    if skip:
        return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE

    # ========== NARROWPHASE ==========
    # temp buffer for MuJoCo's 6-element output [x1,y1,z1, x2,y2,z2]
    mj_fromto = np.zeros(6)
    min_dist = mujoco.mj_geomDistance(
        m=mj_model,
        d=mj_data,
        geom1=self.geom_1.get_id(mj_model),
        geom2=self.geom_2.get_id(mj_model),
        distmax=self.dist_max,
        fromto=mj_fromto,
    )

    min_dist = max(0.0, min_dist)  # clip from below to zero

    p1 = mj_fromto[:3].copy()
    p2 = mj_fromto[3:6].copy()
    self.update_last(p1, p2, mj_data)
    return min_dist, p1, p2, ProximityType.CONVEX_HULL

def get_vertex_to_face_proximity(
    self,
    mj_model: mujoco.MjModel,
    mj_data: mujoco.MjData,
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the vertex to face distance using a multi-phase Bounding Volume Hierarchy (BVH) query.

    Phases:
        1. Broad Phase: Sphere-Sphere check (object level).
        2. Mid Phase: BVH Traversal (eliminating triangle groups). No exit here.
        3. Narrow Phase: Point-to-Face proximity.

    Args:
        mj_model (mujoco.MjModel): Compiled MuJoCo model.
        mj_data (mujoco.MjData): MuJoCo runtime data.

    Returns:
        tuple[float, Vec3, Vec3, ProximityType]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

    """
    if self.geom_1._proximity_configured_for != ProximityType.VERTEX_TO_FACE:
        self.geom_1.bake_proximity(mj_model, ProximityType.VERTEX_TO_FACE)

    if self.geom_2._proximity_configured_for != ProximityType.VERTEX_TO_FACE:
        self.geom_2.bake_proximity(mj_model, ProximityType.VERTEX_TO_FACE)

    assert self.geom_1._baked_query and self.geom_2._baked_query
    assert (
        self.geom_2._local_verts is not None
        and self.geom_2._local_verts is not None
    )

    # ========== BROADPHASE: Sphere-Sphere check ==========
    # find center to center to center distance and return early if broad phase
    d_est, p1, p2, skip = self.get_sphere_to_sphere_proximity(
        mj_model=mj_model, mj_data=mj_data
    )
    if skip:
        return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE

    # ========== COORDINATE TRANSFORMATION ==========
    pos_geom_1 = self.geom_1.rt_xpos(mj_model, mj_data)
    pos_geom_2 = self.geom_2.rt_xpos(mj_model, mj_data)

    mat_geom_1 = self.geom_1.rt_xmat(mj_model, mj_data)  # already Mat3 (3x3)
    mat_geom_2 = self.geom_2.rt_xmat(mj_model, mj_data)
    rel_pos = pos_geom_2 - pos_geom_1

    # ========== NARROWPHASE A: Geom_1-Surface vs. Geom_2-Vertices ==========
    # trimesh uses a BVH internall here (Mid-phase) to find closest triangles
    # combine transforms from geom_1 to geom_2: V_local_geom_1 = R_geom_1.T @ (R_geom_2 @ V_local_geom_2 + p_geom_2 - p_geom_1)
    geom_2_v_in_geom_1 = (
        self.geom_2._local_verts @ mat_geom_2.T + rel_pos
    ) @ mat_geom_1
    pts_on_geom_1, dist_a, _ = self.geom_1._baked_query.on_surface(
        geom_2_v_in_geom_1
    )
    idx_a = np.argmin(dist_a)
    min_a = dist_a[idx_a]

    # ========== NARROWPHASE B: Geom_1-Vertices vs. Geom_2-Surface  ==========
    # transform geom_1 vertices into geom_2's local frame
    geom_1_v_in_geom_2 = (
        self.geom_1._local_verts @ mat_geom_1.T - rel_pos
    ) @ mat_geom_2
    pts_on_geom_2, dist_b, _ = self.geom_2._baked_query.on_surface(
        geom_1_v_in_geom_2
    )
    idx_b = np.argmin(dist_b)
    min_b = dist_b[idx_b]

    # ========== CLEANUP ==========
    # find global min
    if min_a < min_b:
        min_dist = float(min_a)
        p1 = (pts_on_geom_1[idx_a] @ mat_geom_1.T) + pos_geom_1
        p2 = (geom_2_v_in_geom_1[idx_a] @ mat_geom_1.T) + pos_geom_1

        self.update_last(p1, p2, mj_data)
        return min_dist, p1, p2, ProximityType.VERTEX_TO_FACE
    else:
        min_dist = float(min_b)

        # pt_on_geom_2 was calculated in geom_2's local frame
        p2 = (pts_on_geom_2[idx_b] @ mat_geom_2.T) + pos_geom_2
        p1 = (geom_1_v_in_geom_2[idx_b] @ mat_geom_2.T) + pos_geom_2

        self.update_last(p1, p2, mj_data)
        return min_dist, p1, p2, ProximityType.VERTEX_TO_FACE

def get_face_to_face_proximity(
    self,
    mj_model: mujoco.MjModel,
    mj_data: mujoco.MjData,
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the face to face distance using a multi-phase Bounding Volume Hierarchy (BVH) query.

    This is more accurate than the vertex to face method, but comes at higher computational cost.

    Phases:
        1. Broad Phase: Sphere-Sphere check (object level).
        2. Mid Phase: BVH Traversal (eliminating triangle groups). No exit here.
        3. Narrow Phase: Face-to-Face proximity.

    Args:
        mj_model (mujoco.MjModel): Compiled MuJoCo model.
        mj_data (mujoco.MjData): MuJoCo runtime data.

    Returns:
        tuple[float, Vec3, Vec3, ProximityType]: Unsigned (`>= 0`) minimum distance from geom_1 to geom_2, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which phase the exit occurred in.

    """
    if self.geom_1._proximity_configured_for != ProximityType.FACE_TO_FACE:
        self.geom_1.bake_proximity(mj_model, ProximityType.FACE_TO_FACE)

    if self.geom_2._proximity_configured_for != ProximityType.FACE_TO_FACE:
        self.geom_2.bake_proximity(mj_model, ProximityType.FACE_TO_FACE)

    assert self.geom_1._baked_manager and self.geom_2._baked_manager

    # ========== BROADPHASE: Sphere-Sphere check ==========

    # find center to center to center distance and return early if broad phase
    d_est, p1, p2, skip = self.get_sphere_to_sphere_proximity(
        mj_model=mj_model, mj_data=mj_data
    )
    if skip:
        return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE

    # ========== NARROWPHASE ==========
    # set the other transformation relative to geom_1's local frame
    t_geom_1 = np.eye(4)
    t_geom_1[:3, :3] = self.geom_1.rt_xmat(mj_model, mj_data)
    t_geom_1[:3, 3] = self.geom_1.rt_xpos(mj_model, mj_data)

    t_geom_2 = np.eye(4)
    t_geom_2[:3, :3] = self.geom_2.rt_xmat(mj_model, mj_data)
    t_geom_2[:3, 3] = self.geom_2.rt_xpos(mj_model, mj_data)

    self.geom_1._baked_manager.set_transform(self.geom_1.name, t_geom_1)
    self.geom_2._baked_manager.set_transform(self.geom_2.name, t_geom_2)

    # CollisionManager returns distance and the two closest points
    result = self.geom_1._baked_manager.min_distance_other(
        self.geom_2._baked_manager, return_data=True
    )
    min_dist = float(result[0])  # pyright: ignore[reportIndexIssue]
    data = result[1]  # pyright: ignore[reportIndexIssue]

    assert data
    p1 = data.point(self.geom_1.name)  # pyright: ignore[reportAttributeAccessIssue]
    p2 = data.point(self.geom_2.name)  # pyright: ignore[reportAttributeAccessIssue]

    self.update_last(p1, p2, mj_data)
    return min_dist, p1, p2, ProximityType.FACE_TO_FACE

def get_proximity(
    self, mj_model: mujoco.MjModel, mj_data: mujoco.MjData
) -> tuple[float, Vec3, Vec3, ProximityType]:
    """
    Calculates the shortest distance between two geometries using the specified proximity algorithm.

    This is a general dispatcher method that routes to different proximity calculation algorithms based on the `algorithm` parameter. Each mode offers different tradeoffs between speed and precision:

    **Modes:**
        - `SPHERE_TO_SPHERE`: Fastest. Uses bounding sphere distance only (broadphase).
        - `CONVEX_HULL`: Fast & accurate. Uses MuJoCo's convex hull-based distance (default).
        - `VERTEX_TO_FACE`: Accurate. Multi-phase BVH with vertex-to-surface queries.
        - `FACE_TO_FACE`: Most accurate but slowest. Full mesh-to-mesh distance calculation.

    **Phases (for non-sphere modes):**
        1. Broad Phase: Sphere-Sphere check (object level).
        2. Narrow Phase: Algorithm-specific distance calculation.

    Args:
        mj_model (mujoco.MjModel): Compiled MuJoCo model.
        mj_data (mujoco.MjData): MuJoCo runtime data.

    Returns:
        tuple[float, ProximityType]: If fromto=False, returns the unsigned (`>= 0`) minimum distance and which algorithm produced the result.

        tuple[tuple[float, Vec3, Vec3], ProximityType]: If fromto=True, returns the minimum distance, world location of minimum distance on geom_1, world location of minimum distance on geom_2, and which algorithm produced the result.

    """
    match self.algorithm:
        case ProximityType.SPHERE_TO_SPHERE:
            d_est, p1, p2, _skip = self.get_sphere_to_sphere_proximity(
                mj_model=mj_model, mj_data=mj_data
            )
            return d_est, p1, p2, ProximityType.SPHERE_TO_SPHERE
        case ProximityType.CONVEX_HULL:
            return self.get_convex_hull_proximity(
                mj_model=mj_model, mj_data=mj_data
            )
        case ProximityType.VERTEX_TO_FACE:
            return self.get_vertex_to_face_proximity(
                mj_model=mj_model, mj_data=mj_data
            )
        case ProximityType.FACE_TO_FACE:
            return self.get_face_to_face_proximity(
                mj_model=mj_model, mj_data=mj_data
            )
        case _:
            msg = f"Method for {self.algorithm.name} not implemented."
            logger.error(msg)
            raise NotImplementedError(msg)

def request(
    self,
    signal_manager: SignalManager,
    attrs: list[Literal["dist", "fromto", "prox_type"]] = ["dist", "prox_type"],
):
    """
    Registers specific geom proximity attributes for logging. Please see the `get_proximity` method for how these outputs are calculated.

    Available Requests:
        `dist`: Minimum distance as calculated by the specified algorithm. Tagged with `Proximities/{pair_name}:dist`.
        `fromto`: World coordinates for where the minimum distance is estimated to occur at. Two sets of coordinates will be returned for geom_1 and geom_2. Tagged with `Proximities/{pair_name}/fromto/{(geom_1 | geom_2).name}:(x | y | z)`.
        `prox_type`: What type of proximity calculation the previous values are from. Using `dist_max`, `get_proximity` can return a broadphase estimate (bounding sphere to sphere) if the two geometries are distant (greater than `dist_max`). This telemetry will return what type of proximity calculation was performed for this timestep. It is intended to help debug to understand if a jump in telemetry (specifically sharp declines) are real or comes from the broadphase estimate. The values returned will be integer values associated with their specific ProximityType (see the enumeration for the mapping, in general a larger value will mean a more accurate one). Tagged with `Proximities/{pair_name}:prox_type`.

    """
    pair_name = f"{self.geom_1.name}_to_{self.geom_2.name}"

    def sample(mj_model: mujoco.MjModel, mj_data: mujoco.MjData):
        dist, p1, p2, prox_type = self.get_proximity(
            mj_model=mj_model,
            mj_data=mj_data,
        )

        for attr in attrs:
            match attr:
                case "dist":
                    # "Proximities/{pair_name}:dist"
                    signal_manager.post(
                        value=dist,
                        category=SignalCategory.PROXIMITIES,
                        subgroups=(pair_name,),
                        attr=attr,
                    )
                case "fromto":
                    # "Proximities/{pair_name}/fromto/{(geom_1 | geom_2).name}:(x | y | z)"
                    for i, k in enumerate("xyz"):
                        signal_manager.post(
                            value=float(p1[i]),
                            category=SignalCategory.PROXIMITIES,
                            subgroups=(pair_name, attr, str(self.geom_1.name)),
                            attr=k,
                        )
                    for i, k in enumerate("xyz"):
                        signal_manager.post(
                            value=float(p2[i]),
                            category=SignalCategory.PROXIMITIES,
                            subgroups=(pair_name, attr, str(self.geom_2.name)),
                            attr=k,
                        )
                case "prox_type":
                    # "Proximities/{pair_name}:prox_type"
                    signal_manager.post(
                        value=float(prox_type.value),
                        category=SignalCategory.PROXIMITIES,
                        subgroups=(pair_name,),
                        attr=attr,
                    )
                case _:
                    continue

    signal_manager.register_sampler(sample)

def run(
    self,
    global_overrides: NamedValueDict[NDArray] = NamedValueDict[NDArray](),
    clean_workdir: bool = False,
    cleanup_delay: int = 10,
    execution_mode: ExecutionMode = ExecutionMode.LOCAL,
    trial_nums: list[int] | None = None,
) -> bool:
    """Vectors a job to be either computed locally or to be orchestrated by SLURM."""
    if clean_workdir:
        if self.config.resume:
            msg = "clean_workdir and resume are mutually exclusive with one another. Use one or the other."
            logger.error(msg)
            raise ValueError(msg)

        self.force_remove_dir(countdown_from=cleanup_delay, path=self.workdir)

    self.workdir.mkdir(parents=True, exist_ok=True)
    if not (self.workdir / ".gitignore").exists():
        (self.workdir / ".gitignore").write_text("*", encoding="utf-8")
    self.capture_environment()

    match execution_mode:
        case ExecutionMode.LOCAL:
            return self.run_local(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        case ExecutionMode.SLURM:
            return self.orchestrate_slurm(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        case _:
            msg = f"No run command has been configured for execution mode {execution_mode}"
            logger.error(msg)
            raise NotImplementedError(msg)

def orchestrate_slurm(
    self,
    global_overrides: NamedValueDict[NDArray],
    trial_nums: list[int] | None = None,
) -> bool:
    """Generates an sbatch script and submits the job array to SLURM for a given config."""
    (self.workdir / "logs").mkdir(exist_ok=True)

    if isinstance(self.config, MonteCarloConfig):
        try:
            had_fails = self.orchestrate_slurm_monte_carlo(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        except (BdbQuit, KeyboardInterrupt):
            print("\n")
            logger.error("Aborted SLURM orchestration!")
            had_fails = True
    else:
        msg = f"A SLURM configuration for {self.config.__class__.__name__} has not been implemented."
        logger.error(msg)
        raise NotImplementedError(msg)
    return had_fails

def execute_single_trial(
    self, trial_num: int, seed: int | None, overrides_payload: dict
) -> tuple[
    MojoModel | None, TrialStatus, mujoco.MjModel | None, mujoco.MjData | None
]:
    """Helper to package a Trial and run it."""
    overrides = NamedValueDict[NDArray].model_validate(overrides_payload)

    trial = Trial(
        trial_num=trial_num,
        base_dir=self.workdir,
        xml_name=self.xml_name,
        model_config_name=self.model_config_name,
        padding_style=self.config.padding_style,
    )

    return trial.run(
        generator=self.generator,
        runtime=self.runtime,
        seed=seed,
        overrides=overrides,
        gen_args=self.gen_args,
        gen_kwargs=self.gen_kwargs,
        run_args=self.run_args,
        run_kwargs=self.run_kwargs,
    )

def run_monte_carlo(
    self,
    global_overrides: NamedValueDict[NDArray] = NamedValueDict[NDArray](),
    trial_nums: list[int] | None = None,
) -> bool:
    """Orchestrates a Monte Carlo job."""
    if self.slurm_trial_id is not None:
        tn = self.slurm_trial_id
        logger.info(f"SLURM Worker detected. Executing Trial {tn}")
        _mojo_model, trial_status, _, __ = self.execute_single_trial(
            trial_num=tn,
            seed=self.seed,
            overrides_payload=global_overrides.model_dump(),
        )

        return trial_status.completion == Completion.FAILED

    job_trial_nums = trial_nums if trial_nums else self.config.trial_nums

    # initialize the status tracker
    status_tracker = JobStatus(
        workdir=self.workdir.resolve(),
        job_type=JobType.MONTE_CARLO,
        execution_mode=ExecutionMode.LOCAL,
        n_proc=self.config.n_proc,
        seed=self.seed,
        padding_style=self.config.padding_style,
        generator=MojoRunner.inspect_protocol(self.generator),
        runtime=MojoRunner.inspect_protocol(self.runtime),
        objective=MojoRunner.inspect_protocol(self.objective),
        gen_args_used=bool(self.gen_args),
        gen_kwargs_used=bool(self.gen_kwargs),
        run_args_used=bool(self.run_args),
        run_kwargs_used=bool(self.run_kwargs),
        trial_nums=job_trial_nums,
    )

    # decide which trials to execute
    if self.config.resume:
        self._renumber_trial_folders(
            self.workdir.resolve(), self.config.padding_style
        )
        status_tracker.refresh_from_disk(n_proc=self.config.n_proc)
    status_tracker.dump_to_path(self.workdir / JOB_STATUS_FNAME)

    to_run = status_tracker.pending_trial_nums

    if not to_run:
        logger.info("All trials were already completed. Nothing to do.")
        return bool(status_tracker.failed_trial_nums)

    if self.config.is_parallel:
        logger.info(
            f"Running {len(to_run)} trials with {self.config.n_proc} processors. {status_tracker.n_done}/{self.config.n_trial} ({status_tracker.progress:.2%}) trials completed."
        )
        # needed for logging on Windows
        with multiprocessing.Manager() as m:
            log_queue = m.Queue()

            parent_log_level = logging.getLogger().getEffectiveLevel()
            listener = QueueListener(log_queue, *logging.getLogger().handlers)
            listener.start()

            executor = ProcessPoolExecutor(
                max_workers=self.config.n_proc,
                initializer=worker_init,
                initargs=(log_queue, parent_log_level),
            )
            try:
                future_to_tn = {
                    executor.submit(
                        self.execute_single_trial,
                        tn,
                        self.seed,
                        global_overrides.model_dump(),
                    ): tn
                    for tn in to_run
                }
                for f in as_completed(future_to_tn):
                    tn = future_to_tn[f]
                    try:
                        _mojo_model, trial_status, _, __ = f.result()
                        status_tracker.update_trial(status=trial_status)
                    except (BdbQuit, KeyboardInterrupt):
                        # user is quitting from breakpoint() or CTRL+C
                        raise
                    except Exception as e:
                        logger.exception(f"Trial {tn} failed: {e}")
                        status_tracker.update_trial(
                            status=TrialStatus(
                                trial_num=tn, completion=Completion.FAILED
                            )
                        )
                    status_tracker.generate_report()
            except (BdbQuit, KeyboardInterrupt):
                # allows killing the job with one CTRL+C
                logger.warning("Interrupt recieved. Stopping all trials.")
                executor.shutdown(wait=False, cancel_futures=True)
                raise
            finally:
                listener.stop()
                executor.shutdown(wait=True)
    else:
        for tn in to_run:
            try:
                _mojo_model, trial_status, _, __ = self.execute_single_trial(
                    trial_num=tn,
                    seed=self.seed,
                    overrides_payload=global_overrides.model_dump(),
                )
                status_tracker.update_trial(
                    status=trial_status,
                )
            except (BdbQuit, KeyboardInterrupt):
                # user is quitting from breakpoint() or CTRL+C
                raise
            except Exception as e:
                logger.exception(f"A trial failed with error: {e}")
                status_tracker.update_trial(
                    status=TrialStatus(trial_num=tn, completion=Completion.FAILED)
                )
            status_tracker.generate_report()

    status_tracker.generate_report(alert_generation=True)
    return bool(status_tracker.failed_trial_nums)

@staticmethod
def get_slurm_array_string(ids: list[int]) -> str:
    """Collapses [0, 1, 2, 5, 6] into '0-2,5-6' for SLURM."""
    if not ids:
        return ""

    ranges = []
    # Identify groups of consecutive integers
    for _, group in itertools.groupby(
        enumerate(sorted(ids)), lambda x: x[1] - x[0]
    ):
        group = list(group)
        start = group[0][1]
        end = group[-1][1]
        if start == end:
            ranges.append(str(start))
        else:
            ranges.append(f"{start}-{end}")

    return ",".join(ranges)

@staticmethod
def get_slurm_partitions() -> tuple[list[str], str | None]:
    """Queries sinfo for available partitions and identifies the default."""
    try:
        result = subprocess.run(
            ["sinfo", "-h", "--format=%P"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0:
            raw_partitions = [
                p.strip() for p in result.stdout.splitlines() if p.strip()
            ]

            default_partition = None
            clean_partitions = []

            for p in raw_partitions:
                if p.endswith("*"):
                    name = p.replace("*", "")
                    default_partition = name
                    clean_partitions.append(name)
                else:
                    clean_partitions.append(p)

            return sorted(list(set(clean_partitions))), default_partition
    except Exception:
        pass
    return [], None