diff --git a/src/macro_lactone_toolkit/validation/visualization_output.py b/src/macro_lactone_toolkit/validation/visualization_output.py new file mode 100644 index 0000000..d7cc2f6 --- /dev/null +++ b/src/macro_lactone_toolkit/validation/visualization_output.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from pathlib import Path + +from rdkit import Chem + +from macro_lactone_toolkit.visualization import save_fragment_png, save_numbered_molecule_png + + +def get_output_paths(output_dir: Path, source_id: str, ring_size: int, classification: str) -> dict: + """Get organized output paths for a molecule.""" + ring_dir = output_dir / f"ring_size_{ring_size}" + + if classification == "standard_macrolactone": + base_dir = ring_dir / "standard" + numbered_dir = base_dir / "numbered" + sidechains_dir = base_dir / "sidechains" / source_id + elif classification == "non_standard_macrocycle": + base_dir = ring_dir / "non_standard" / "original" + numbered_dir = base_dir + sidechains_dir = None + else: + base_dir = ring_dir / "rejected" / "original" + numbered_dir = base_dir + sidechains_dir = None + + numbered_dir.mkdir(parents=True, exist_ok=True) + if sidechains_dir: + sidechains_dir.mkdir(parents=True, exist_ok=True) + + return { + "numbered_image": numbered_dir / f"{source_id}_numbered.png", + "sidechains_dir": sidechains_dir, + } + + +def save_numbered_molecule( + smiles: str, + output_path: Path, + ring_size: int | None = None, + size: tuple[int, int] = (800, 800), +) -> Path | None: + """Save numbered molecule image.""" + try: + return save_numbered_molecule_png( + smiles, + output_path, + ring_size=ring_size, + size=size, + ) + except Exception as e: + print(f"Failed to save numbered image: {e}") + return None + + +def save_fragment_images( + fragments: list, + output_dir: Path, + source_id: str, + size: tuple[int, int] = (400, 400), +) -> list[str]: + """Save fragment images and return paths.""" + paths = [] + + for i, fragment in enumerate(fragments): + try: + output_path = output_dir / f"{source_id}_frag_{i}_pos{fragment.cleavage_position}.png" + save_fragment_png(fragment.fragment_smiles_plain, output_path, size=size) + paths.append(str(output_path.relative_to(output_dir.parent.parent))) + except Exception as e: + print(f"Failed to save fragment {i}: {e}") + paths.append(None) + + return paths