#!/usr/bin/env python3
# Description
###############################################################################
'''
Starter Workbench spec templates for GUI and CLI definition workflows.
'''
# Imports
###############################################################################
from __future__ import annotations
from pathlib import Path
from typing import Any
from typing import Callable
from typing import Literal
from OCDocker.Workbench.IO import model_to_data
from OCDocker.Workbench.Models import FeaturePolicySelection
from OCDocker.Workbench.Models import OCScoreAblationSpec
from OCDocker.Workbench.Models import OCScoreInputSpec
from OCDocker.Workbench.Models import OCScoreStudySpec
from OCDocker.Workbench.Models import ResourceSpec
from OCDocker.Workbench.Models import SnakemakeWorkflowSpec
from OCDocker.Workbench.Models import VSInputSpec
from OCDocker.Workbench.Models import VSCampaignSpec
from OCDocker.Workbench.Models import WorkbenchSpec
# License
###############################################################################
'''Copyright (c) Federal University of Rio de Janeiro (UFRJ), Artur Duque Rossi, and Pedro Henrique Monteiro Torres.
SPDX-License-Identifier: BSD-3-Clause
See the LICENSE file for full terms.
'''
# Type aliases
###############################################################################
TemplateName = Literal["ocscore_ablation", "ocscore_study", "vs_campaign"]
# Functions
###############################################################################
## Private ##
def _ocscore_input_template() -> OCScoreInputSpec:
'''Build a placeholder OCScore raw-input selection.
Returns
-------
OCScoreInputSpec
Starter OCScore input selection.
'''
return OCScoreInputSpec(raw_input_dir=Path("path/to/raw_prepare"))
def _ocscore_study_template() -> OCScoreStudySpec:
'''Build a starter OCScore study spec.
Returns
-------
OCScoreStudySpec
Starter OCScore study spec.
'''
return OCScoreStudySpec(
name="new-ocscore-study",
protocol="protocol-name",
inputs=_ocscore_input_template(),
output_dir=Path("runs/new-ocscore-study"),
description="Starter OCScore study generated by OCDocker Workbench.",
tags=("workbench", "ocscore"),
)
def _ocscore_ablation_template() -> OCScoreAblationSpec:
'''Build a starter OCScore ablation spec.
Returns
-------
OCScoreAblationSpec
Starter OCScore ablation spec.
'''
return OCScoreAblationSpec(
name="new-ocscore-ablation",
protocol="protocol-name",
inputs=_ocscore_input_template(),
output_dir=Path("runs/new-ocscore-ablation"),
feature_policies=FeaturePolicySelection(run_all=True),
include_full_reference=True,
description="Starter OCScore feature-policy ablation generated by OCDocker Workbench.",
tags=("workbench", "ocscore", "ablation"),
)
def _vs_campaign_template() -> VSCampaignSpec:
'''Build a starter virtual-screening campaign spec.
Returns
-------
VSCampaignSpec
Starter virtual-screening campaign spec.
'''
return VSCampaignSpec(
name="new-vs-campaign",
workspace=Path("runs/new-vs-campaign"),
workflow=SnakemakeWorkflowSpec(
snakefile=Path("Snakefile"),
workdir=Path("."),
targets=("all",),
resources=ResourceSpec(cores=1),
),
inputs=(
VSInputSpec(
sample="sample_001",
receptor=Path("data/receptors/sample_001.pdbqt"),
ligand=Path("data/ligands/sample_001.sdf"),
box=Path("data/boxes/sample_001.txt"),
engines=("vina", "smina", "plants"),
rescoring_engines=("oddt",),
),
),
description="Starter virtual-screening campaign generated by OCDocker Workbench.",
tags=("workbench", "virtual-screening"),
)
_TEMPLATE_BUILDERS: dict[str, Callable[[], WorkbenchSpec]] = {
"ocscore_ablation": _ocscore_ablation_template,
"ocscore_study": _ocscore_study_template,
"vs_campaign": _vs_campaign_template,
}
def _template_builder(name: str) -> Callable[[], WorkbenchSpec]:
'''Return the builder registered for a template name.
Parameters
----------
name : str
Template name.
Returns
-------
Callable[[], WorkbenchSpec]
Template builder.
'''
normalized = str(name).strip()
try:
return _TEMPLATE_BUILDERS[normalized]
except KeyError as exc:
available = ", ".join(available_template_names())
raise ValueError(
f"Unknown Workbench template {normalized!r}. Expected one of: {available}."
) from exc
## Public ##
[docs]
def available_template_names() -> tuple[str, ...]:
'''Return registered starter template names.
Returns
-------
tuple[str, ...]
Registered template names in deterministic order.
'''
return tuple(sorted(_TEMPLATE_BUILDERS))
[docs]
def build_template_spec(name: str) -> WorkbenchSpec:
'''Build a validated starter spec for a Workbench template name.
Parameters
----------
name : str
Registered template name.
Returns
-------
WorkbenchSpec
Validated starter spec.
'''
return _template_builder(name)()
[docs]
def build_template_payload(name: str) -> dict[str, Any]:
'''Build a JSON-compatible starter spec payload.
Parameters
----------
name : str
Registered template name.
Returns
-------
dict[str, Any]
JSON-compatible starter spec payload.
'''
return model_to_data(build_template_spec(name))
__all__ = [
"TemplateName",
"available_template_names",
"build_template_payload",
"build_template_spec",
]