Skip to content
Snippets Groups Projects
Verified Commit 320e7da2 authored by James Graham's avatar James Graham
Browse files

refactor: split up main init method

parent a89d2cee
No related branches found
No related tags found
No related merge requests found
...@@ -7,6 +7,7 @@ ...@@ -7,6 +7,7 @@
*.offsets *.offsets
# Dependencies and runtime # Dependencies and runtime
/.venv/
/env/ /env/
/minenv/ /minenv/
/venv/ /venv/
......
...@@ -7,6 +7,7 @@ import pathlib ...@@ -7,6 +7,7 @@ import pathlib
import sys import sys
import textwrap import textwrap
import time import time
import typing
from mdplus.multiscale import GLIMPS from mdplus.multiscale import GLIMPS
from rich.logging import RichHandler from rich.logging import RichHandler
...@@ -16,6 +17,8 @@ from .mapping import Mapping ...@@ -16,6 +17,8 @@ from .mapping import Mapping
from .bondset import BondSet from .bondset import BondSet
from .forcefield import ForceField from .forcefield import ForceField
PathLike = typing.Union[pathlib.Path, str]
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
...@@ -36,15 +39,7 @@ class PyCGTOOL: ...@@ -36,15 +39,7 @@ class PyCGTOOL:
self.mapping = None self.mapping = None
self.out_frame = self.in_frame self.out_frame = self.in_frame
if self.config.mapping: if self.config.mapping:
self.mapping = Mapping(self.config.mapping, self.mapping, self.out_frame = self.apply_mapping(self.in_frame)
self.config,
itp_filename=self.config.itp)
self.out_frame = self.mapping.apply(self.in_frame)
self.out_frame.save(self.get_output_filepath('gro'),
frame_number=0)
if self.config.backmapper_resname and self.out_frame.n_frames > 1:
self.train_backmapper(self.config.resname)
self.bondset = None self.bondset = None
if self.config.bondset: if self.config.bondset:
...@@ -54,14 +49,24 @@ class PyCGTOOL: ...@@ -54,14 +49,24 @@ class PyCGTOOL:
if self.config.output_xtc: if self.config.output_xtc:
self.out_frame.save(self.get_output_filepath('xtc')) self.out_frame.save(self.get_output_filepath('xtc'))
def get_output_filepath(self, ext: str) -> pathlib.Path: def get_output_filepath(self, ext: PathLike) -> pathlib.Path:
"""Get file path for an output file by extension. """Get file path for an output file by extension."""
:param ext:
"""
out_dir = pathlib.Path(self.config.out_dir) out_dir = pathlib.Path(self.config.out_dir)
return out_dir.joinpath(self.config.output_name + '.' + ext) return out_dir.joinpath(self.config.output_name + '.' + ext)
def apply_mapping(self, in_frame: Frame) -> typing.Tuple[Mapping, Frame]:
"""Map input frame to output using requested mapping file."""
mapping = Mapping(self.config.mapping,
self.config,
itp_filename=self.config.itp)
out_frame = mapping.apply(in_frame)
out_frame.save(self.get_output_filepath('gro'), frame_number=0)
if self.config.backmapper_resname and self.out_frame.n_frames > 1:
self.train_backmapper(self.config.resname)
return mapping, out_frame
def measure_bonds(self) -> None: def measure_bonds(self) -> None:
"""Measure bonds at the end of a run.""" """Measure bonds at the end of a run."""
self.bondset.apply(self.out_frame) self.bondset.apply(self.out_frame)
......
...@@ -193,7 +193,7 @@ class Frame: ...@@ -193,7 +193,7 @@ class Frame:
return self._topology.add_atom(name, element, residue) return self._topology.add_atom(name, element, residue)
def save(self, def save(self,
filename: str, filename: PathLike,
frame_number: typing.Optional[int] = None, frame_number: typing.Optional[int] = None,
**kwargs) -> None: **kwargs) -> None:
"""Write trajctory to file. """Write trajctory to file.
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment