feat(validation): add visualization output module
This commit is contained in:
74
src/macro_lactone_toolkit/validation/visualization_output.py
Normal file
74
src/macro_lactone_toolkit/validation/visualization_output.py
Normal file
@@ -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
|
||||||
Reference in New Issue
Block a user