ensemble.py

"""Protpardelle and user ensemble generation."""

from pathlib import Path


def _generate_protpardelle_ensemble(
    pdb_path: str,
    num_conformers: int,
    out_dir: Path,
    weights_dir: str,
) -> dict[str, list[str]]:
    """Generate conformers with Protpardelle-1c, return pdb_to_conformers dict."""
    from caliby import generate_ensembles

    pdb_to_conformers = generate_ensembles(
        [pdb_path],
        out_dir=str(out_dir / "protpardelle_ensemble"),
        num_samples_per_pdb=num_conformers,
        model_params_path=weights_dir,
    )
    # generate_ensembles returns only generated conformers — prepend the primary structure.
    pdb_stem = Path(pdb_path).stem
    pdb_to_conformers[pdb_stem] = [pdb_path] + pdb_to_conformers.get(pdb_stem, [])
    return pdb_to_conformers


def _setup_user_ensemble_dir(
    pdb_paths: list[str],
    **_ignored,
) -> dict[str, list[str]]:
    """Build pdb_to_conformers dict from user-uploaded files.

    First file is the primary conformer, rest are additional conformers.
    """
    pdb_key = Path(pdb_paths[0]).stem
    return {pdb_key: list(pdb_paths)}