Python API

tsp-maker can be used as a Python library for programmatic TSP creation.

Installation

pip install git+https://github.com/TeamMacLean/tsp-maker.git

Quick Example

from pathlib import Path
from tsp_maker import AlphaFold3Parser, build_tsp, upload_to_zenodo

# Parse predictions
parser = AlphaFold3Parser(output_dir=Path("/tmp/intermediate"), top_n=5)
parser.parse_batch(Path("/data/af3_predictions"))

# Build TSP
build_tsp(
    input_dir=Path("/tmp/intermediate"),
    output_dir=Path("/tmp/my-dataset"),
    name="my-structures",
    title="My Structure Dataset",
    author="Jane Doe",
)

# Upload
result = upload_to_zenodo(
    tsp_path=Path("/tmp/my-dataset"),
    token="your-zenodo-token",
    sandbox=True,
    publish=True,
)
print(f"DOI: {result['doi']}")

Parsers

AlphaFold2Parser

from tsp_maker import AlphaFold2Parser

parser = AlphaFold2Parser(
    output_dir=Path("/output"),
    top_n=5,                    # Models per protein
    id_pattern=r"[A-Z][0-9]...",  # Protein ID regex
    verbose=True,               # Print progress
)

# Parse single directory
scores = parser.parse_directory(
    input_dir=Path("/data/P12345"),
    structure_id="P12345"
)

# Batch process
results = parser.parse_batch(Path("/data/af2_outputs"))

AlphaFold3Parser

from tsp_maker import AlphaFold3Parser

parser = AlphaFold3Parser(output_dir=Path("/output"), top_n=5)
results = parser.parse_batch(Path("/data/af3_outputs"))

Boltz2Parser

from tsp_maker import Boltz2Parser

parser = Boltz2Parser(output_dir=Path("/output"), top_n=5)
results = parser.parse_batch(Path("/data/boltz2_outputs"))

Custom Parser

Extend BaseParser for new predictors:

from tsp_maker.parsers.base import BaseParser

class MyPredictor(BaseParser):
    PREDICTOR_NAME = "MyPredictor"
    PREDICTOR_SUFFIX = "MYP"

    def is_valid_directory(self, directory):
        return (directory / "my_output.json").exists()

    def parse_directory(self, input_dir, structure_id):
        # Your parsing logic
        scores = [...]
        self.save_scores(scores, structure_id)
        return scores

Builder

build_tsp

from tsp_maker import build_tsp

output_path = build_tsp(
    input_dir=Path("/intermediate"),
    output_dir=Path("/output"),
    name="my-dataset",              # Required
    title="My Dataset",             # Optional
    description="Description",      # Optional
    author="Jane Doe",              # Optional
    affiliation="My Institute",     # Optional
    version="1.0.0",                # Default: 1.0.0
    license="CC-BY-4.0",            # Default: CC-BY-4.0
    batch_size_mb=100,              # Default: 100
    verbose=True,                   # Default: True
)

Uploader

upload_to_zenodo

from tsp_maker import upload_to_zenodo

result = upload_to_zenodo(
    tsp_path=Path("/my-dataset"),
    token="your-api-token",
    sandbox=True,       # True for sandbox.zenodo.org
    publish=False,      # True to publish immediately
    verbose=True,
)

# Result dict
print(result["id"])         # Deposit ID
print(result["doi"])        # DOI (if published)
print(result["record_url"]) # Public URL (if published)
print(result["edit_url"])   # Edit URL (if draft)

Return Values

Parser Results

parse_batch() returns a dict:

{
    "P12345": {
        "source_dir": "/data/P12345",
        "n_models": 5,
        "status": "success"
    },
    "Q67890": {
        "source_dir": "/data/Q67890",
        "n_models": 0,
        "status": "error: No models found"
    }
}

Score Records

parse_directory() returns a list:

[
    {
        "model_id": "P12345_AF3_1",
        "structure_id": "P12345",
        "rank": 1,
        "predictor": "alphafold3",
        "plddt_mean": 85.2,
        "ptm": 0.82,
        "ranking_score": 0.85,
        # ... more fields
    },
    # ...
]

Error Handling

from tsp_maker import AlphaFold3Parser

parser = AlphaFold3Parser(output_dir=Path("/output"))

try:
    results = parser.parse_batch(Path("/invalid/path"))
except FileNotFoundError as e:
    print(f"No valid directories found: {e}")

Combining Predictors

from pathlib import Path
from tsp_maker import AlphaFold2Parser, AlphaFold3Parser, Boltz2Parser, build_tsp

output = Path("/tmp/intermediate")

# All parsers write to same output directory
AlphaFold2Parser(output, top_n=5).parse_batch(Path("/data/af2"))
AlphaFold3Parser(output, top_n=5).parse_batch(Path("/data/af3"))
Boltz2Parser(output, top_n=3).parse_batch(Path("/data/boltz2"))

# Build combined TSP
build_tsp(output, Path("/tmp/my-tsp"), name="combined-predictions")