DicePlayer/diceplayer/dice/dice_input.py

from diceplayer.config import PlayerConfig
from diceplayer.logger import logger

from typing_extensions import Self

from abc import ABC
from dataclasses import dataclass, fields
from pathlib import Path
from typing import Any, Sequence, TextIO



DICE_KEYWORD_ORDER = [
    "title",
    "ncores",
    "ljname",
    "outname",
    "nmol",
    "dens",
    "temp",
    "press",
    "seed",
    "init",
    "nstep",
    "vstep",
    "mstop",
    "accum",
    "iprint",
    "isave",
    "irdf",
    "upbuf",
]



@dataclass(slots=True)
class BaseConfig(ABC):
    ncores: int
    ljname: str
    outname: str
    nmol: Sequence[int]
    temp: float
    seed: int
    isave: int

    def write(self, directory: Path, filename: str = "input") -> Path:
        input_path = directory / filename

        if input_path.exists():
            logger.info(
                f"Dice input file {input_path} already exists and will be overwritten"
            )
            input_path.unlink()
        input_path.parent.mkdir(parents=True, exist_ok=True)

        with open(input_path, "w") as io:
            self.write_dice_config(io)

        return input_path

    def write_dice_config(self, io_writer: TextIO) -> None:
        values = {f.name: getattr(self, f.name) for f in fields(self)}

        for key in DICE_KEYWORD_ORDER:
            value = values.pop(key, None)
            if value is None:
                continue
            io_writer.write(f"{key} = {self._serialize_value(value)}\n")

        # write any remaining fields (future extensions)
        for key, value in values.items():
            if value is None:
                continue
            io_writer.write(f"{key} = {self._serialize_value(value)}\n")

        io_writer.write("$end\n")

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        base_fields = cls._extract_base_fields(config)
        return cls(**(base_fields | kwargs))

    @staticmethod
    def _extract_base_fields(config: PlayerConfig) -> dict[str, Any]:
        return dict(
            ncores=int(config.ncores / config.dice.nprocs),
            ljname=config.dice.ljname,
            outname=config.dice.outname,
            nmol=config.dice.nmol,
            temp=config.dice.temp,
            seed=config.dice.seed,
            isave=config.dice.isave,
        )

    @staticmethod
    def _get_nstep(config: PlayerConfig, idx: int) -> int:
        if len(config.dice.nstep) > idx:
            return config.dice.nstep[idx]
        return config.dice.nstep[-1]

    @staticmethod
    def _serialize_value(value: Any) -> str:
        if value is None:
            raise ValueError("DICE configuration cannot serialize None values")

        if isinstance(value, bool):
            return "yes" if value else "no"

        if isinstance(value, (list, tuple)):
            return " ".join(str(v) for v in value)

        return str(value)


# -----------------------------------------------------
# NVT BASE
# -----------------------------------------------------


@dataclass(slots=True)
class NVTConfig(BaseConfig):
    title: str = "Diceplayer Run - NVT"
    dens: float = ...
    nstep: int = ...
    vstep: int = 0

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NVTConfig, cls).from_config(
            config,
            dens=config.dice.dens,
            nstep=cls._get_nstep(config, 0),
        )


# -----------------------------------------------------
# NVT THERMALIZATION
# -----------------------------------------------------


@dataclass(slots=True)
class NVTTerConfig(NVTConfig):
    title: str = "Diceplayer Run - NVT Thermalization"
    upbuf: int = 360
    init: str = "yes"

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NVTTerConfig, cls).from_config(
            config,
            nstep=cls._get_nstep(config, 0),
            upbuf=config.dice.upbuf,
            vstep=0,
            **kwargs,
        )

    def write(self, directory: Path, filename: str = "nvt.ter") -> Path:
        return super(NVTTerConfig, self).write(directory, filename)


# -----------------------------------------------------
# NVT PRODUCTION
# -----------------------------------------------------


@dataclass(slots=True)
class NVTEqConfig(NVTConfig):
    title: str = "Diceplayer Run - NVT Production"
    irdf: int = 0
    init: str = "yesreadxyz"

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NVTEqConfig, cls).from_config(
            config,
            nstep=cls._get_nstep(config, 1),
            irdf=config.dice.irdf,
            vstep=0,
            **kwargs,
        )

    def write(self, directory: Path, filename: str = "nvt.eq") -> Path:
        return super(NVTEqConfig, self).write(directory, filename)


# -----------------------------------------------------
# NPT BASE
# -----------------------------------------------------


@dataclass(slots=True)
class NPTConfig(BaseConfig):
    title: str = "Diceplayer Run - NPT"
    nstep: int = 0
    vstep: int = 5000
    press: float = 1.0

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NPTConfig, cls).from_config(
            config,
            press=config.dice.press,
        )


# -----------------------------------------------------
# NPT THERMALIZATION
# -----------------------------------------------------


@dataclass(slots=True)
class NPTTerConfig(NPTConfig):
    title: str = "Diceplayer Run - NPT Thermalization"
    dens: float | None = None

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NPTTerConfig, cls).from_config(
            config,
            dens=config.dice.dens,
            nstep=cls._get_nstep(config, 1),
            vstep=config.dice.vstep,
            **kwargs,
        )

    def write(self, directory: Path, filename: str = "npt.ter") -> Path:
        return super(NPTTerConfig, self).write(directory, filename)


# -----------------------------------------------------
# NPT PRODUCTION
# -----------------------------------------------------


@dataclass(slots=True)
class NPTEqConfig(NPTConfig):
    title: str = "Diceplayer Run - NPT Production"
    dens: float | None = None

    @classmethod
    def from_config(cls, config: PlayerConfig, **kwargs) -> Self:
        return super(NPTEqConfig, cls).from_config(
            config,
            dens=config.dice.dens,
            nstep=cls._get_nstep(config, 2),
            vstep=config.dice.vstep,
            **kwargs,
        )

    def write(self, directory: Path, filename: str = "npt.eq") -> Path:
        return super(NPTEqConfig, self).write(directory, filename)