Files
DicePlayer/diceplayer/dice/dice_wrapper.py
Vitor Hideyoshi 2802f10013
Some checks failed
build and upload / test (3.10) (push) Failing after 1m45s
build and upload / pypi-upload (push) Has been skipped
refactor: restructure dice environment handling and update Python version requirement
2026-03-29 17:57:51 -03:00

104 lines
3.1 KiB
Python

import diceplayer.dice.dice_input as dice_input
from diceplayer.config import DiceConfig
from pydantic import TypeAdapter
from pydantic.dataclasses import dataclass
import subprocess
from itertools import islice
from pathlib import Path
from typing import Final, List, Self
@dataclass(slots=True, frozen=True)
class DiceEnvironmentItem:
atom: str
x: float
y: float
z: float
DiceEnvironmentItemAdapter = TypeAdapter(DiceEnvironmentItem)
@dataclass(slots=True)
class DiceEnvironment:
number_of_sites: int
thickness: List[float]
items: List[DiceEnvironmentItem]
@classmethod
def new(cls, thickness: List[float]) -> Self:
return cls(number_of_sites=0, thickness=thickness, items=[])
def add_site(self, site: DiceEnvironmentItem):
self.items.append(site)
self.number_of_sites += 1
DICE_FLAG_LINE: Final[int] = -2
DICE_END_FLAG: Final[str] = "End of simulation"
class DiceWrapper:
def __init__(self, dice_config: DiceConfig, working_directory: Path):
self.dice_config = dice_config
self.working_directory = working_directory
def run(self, dice_config: dice_input.DiceInputConfig) -> None:
input_path = dice_input.write_config(dice_config, self.working_directory)
output_path = input_path.parent / (input_path.name + ".out")
with open(output_path, "w") as outfile, open(input_path, "r") as infile:
exit_status = subprocess.call(
self.dice_config.progname,
stdin=infile,
stdout=outfile,
cwd=self.working_directory,
)
if exit_status != 0:
raise RuntimeError(f"Dice simulation failed with exit status {exit_status}")
with open(output_path, "r") as outfile:
line = outfile.readlines()[DICE_FLAG_LINE]
if line.strip() == DICE_END_FLAG:
return
raise RuntimeError(f"Dice simulation failed with exit status {exit_status}")
def parse_results(self) -> list[DiceEnvironment]:
results = []
positions_file = self.working_directory / "phb.xyz"
if not positions_file.exists():
raise RuntimeError(f"Positions file not found at {self.working_directory}")
with open(positions_file, "r") as f:
while True:
line = f.readline()
if not line.startswith(" "):
break
environment = DiceEnvironment(
number_of_sites=int(line.strip()),
thickness=[float(n) for n in f.readline().split()[-3:]],
items=[],
)
# Skip the comment line
environment.items.extend(
[
DiceEnvironmentItemAdapter.validate_python(
{"atom": site[0], "x": site[1], "y": site[2], "z": site[3]}
)
for atom in islice(f, environment.number_of_sites)
if (site := atom.strip().split()) and len(site) == 4
]
)
results.append(environment)
return results