9248755285

Committed 17 May 2023 05:43AM UTC coverage: 81.097% (+0.1%) from 80.955%

Build # 9248755285

Build Type

push

github

Committed by

web-flow

Commit Message

use pre-commit.ci (#597)

* use pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* disable codespell, enable ruff autofix

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* disable some pre-commits

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* reset

* fix comment

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Run Details

934 of 1263 branches covered (73.95%)

Branch coverage included in aggregate %.

2 of 2 new or added lines in 2 files covered. (100.0%)

2 existing lines in 2 files now uncovered.

4124 of 4974 relevant lines covered (82.91%)

0.83 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

94.68

/mdsuite/file_io/chemfiles_read.py

"""Test MDSuites chemfiles read module."""
import pathlib
import typing

import chemfiles
import numpy as np
import tqdm

import mdsuite.database.mdsuite_properties
import mdsuite.file_io.file_read
import mdsuite.utils.meta_functions
from mdsuite.database.mdsuite_properties import mdsuite_properties
from mdsuite.database.simulation_database import TrajectoryMetadata


class ChemfilesRead(mdsuite.file_io.file_read.FileProcessor):
    """
    Read trajectory files via chemfiles.
    See https://chemfiles.org/chemfiles/0.10.0/formats.html for supported data formats.
    """

    def __init__(
        self,
        traj_file_path: typing.Union[str, pathlib.Path],
        topol_file_path: typing.Union[str, pathlib.Path] = None,
    ):
        """

        Parameters
        ----------
        traj_file_path
            Path to the trajectory file you want to read.
        topol_file_path : optional
            If the trajectory file does not contain all information about the topology of
            the system (i.e. which data in the trajectory file belongs to which particle),
             you can provide the topology here.
        """
        self.traj_file_path = pathlib.Path(traj_file_path).resolve()

        if topol_file_path is not None:
            topol_file_path = pathlib.Path(topol_file_path).resolve()
        self.topol_file_path = topol_file_path

        # until now, chemfiles only supports these 2 properties.
        # If more are added, link here the mdsuite property to the name of the property
        # attribute in chemfiles.Frame
        self.properties_to_chemfile_attr_dict = {
            # mdsuite_properties.unwrapped_positions: "positions",
            mdsuite_properties.positions: "positions",
            mdsuite_properties.velocities: "velocities",
        }

        # not all properties are guaranteed to be in the given file.
        # A subset of self.properties_to_chemfile_attr_dict will be extracted during
        # self._get_metadata
        self.properties_in_file = None
        self.species_name_to_line_idxs_dict = None
        self.batch_size = None

    def __str__(self):
        return str(self.traj_file_path)

    def _get_metadata(self) -> TrajectoryMetadata:
        """Get the necessary metadata out of chemfiles.

        Trajectory and the first chemfiles.Frame
        """
        with chemfiles.Trajectory(str(self.traj_file_path)) as traj:
            if self.topol_file_path is not None:
                traj.set_topology(str(self.topol_file_path))

            n_configs = traj.nsteps
            frame = traj.read()

        # get the box lengths from the first frame
        box_l = frame.cell.lengths

        # extract which lines in chemfiles.Frame.<property> belong to which species
        # by going through the atoms list
        self.species_name_to_line_idxs_dict = {}
        for line_idx, atm in enumerate(frame.atoms):
            name = atm.name
            if name not in self.species_name_to_line_idxs_dict.keys():
                self.species_name_to_line_idxs_dict[name] = []
            self.species_name_to_line_idxs_dict[name].append(line_idx)

        # see which properties are in the file. there is no way of finding out except
        # trying to access them
        self.properties_in_file = {}
        for (
            mds_prop,
            chemfile_attr_name,
        ) in self.properties_to_chemfile_attr_dict.items():
            try:
                frame.__getattribute__(chemfile_attr_name)
            except chemfiles.ChemfilesError:
                pass
            else:
                self.properties_in_file[mds_prop] = chemfile_attr_name

        species_list = []
        for key, val in self.species_name_to_line_idxs_dict.items():
            species_list.append(
                mdsuite.database.simulation_database.SpeciesInfo(
                    name=key,
                    n_particles=len(val),
                    properties=list(self.properties_in_file.keys()),
                )
            )

        return TrajectoryMetadata(
            n_configurations=n_configs,
            species_list=species_list,
            box_l=box_l,
        )

    def get_configurations_generator(
        self,
    ) -> typing.Iterator[mdsuite.file_io.file_read.TrajectoryChunkData]:
        """Implement parent abstract method."""
        batch_size = mdsuite.utils.meta_functions.optimize_batch_size(
            filepath=self.traj_file_path,
            number_of_configurations=self.metadata.n_configurations,
        )
        n_batches, n_configs_remainder = divmod(
            int(self.metadata.n_configurations), int(batch_size)
        )

        with chemfiles.Trajectory(str(self.traj_file_path)) as traj:
            if self.topol_file_path is not None:
                traj.set_topology(str(self.topol_file_path))
            for _ in tqdm.tqdm(range(n_batches), ncols=70):
                yield self._read_process_n_configurations(traj, batch_size)
            if n_configs_remainder > 0:
                yield self._read_process_n_configurations(traj, n_configs_remainder)

    def _read_process_n_configurations(
        self, traj: chemfiles.Trajectory, n_configs: int
    ) -> mdsuite.database.simulation_database.TrajectoryChunkData:
        """
        Read n configurations and package them into a trajectory chunk of the
        right format.

        Parameters
        ----------
        traj : chemfiles.Trajectory
            An open chemfiles Trajectory
        n_configs : int
            Number of configurations to read in.
        """
        species_list = self.metadata.species_list
        chunk = mdsuite.database.simulation_database.TrajectoryChunkData(
            species_list, n_configs
        )

        for i in range(n_configs):
            frame = traj.read()
            # slice by species
            for sp_info in species_list:
                for mds_prop, chemfile_attrname in self.properties_in_file.items():
                    data = frame.__getattribute__(chemfile_attrname)
                    idxs = self.species_name_to_line_idxs_dict[sp_info.name]
                    write_data = data[idxs, :]
                    # add 'time' axis. we only have one configuration to write
                    write_data = write_data[np.newaxis, :, :]
                    chunk.add_data(write_data, i, sp_info.name, mds_prop.name)
        return chunk

1	"""Test MDSuites chemfiles read module."""
2	import pathlib	1✔
3	import typing	1✔
4
5	import chemfiles	1✔
6	import numpy as np	1✔
7	import tqdm	1✔
8
9	import mdsuite.database.mdsuite_properties	1✔
10	import mdsuite.file_io.file_read	1✔
11	import mdsuite.utils.meta_functions	1✔
12	from mdsuite.database.mdsuite_properties import mdsuite_properties	1✔
13	from mdsuite.database.simulation_database import TrajectoryMetadata	1✔
14
15
16	class ChemfilesRead(mdsuite.file_io.file_read.FileProcessor):	1✔
17	"""
18	Read trajectory files via chemfiles.
19	See https://chemfiles.org/chemfiles/0.10.0/formats.html for supported data formats.
20	"""
21
22	def __init__(	1✔
23	self,
24	traj_file_path: typing.Union[str, pathlib.Path],
25	topol_file_path: typing.Union[str, pathlib.Path] = None,
26	):
27	"""
28
29	Parameters
30	----------
31	traj_file_path
32	Path to the trajectory file you want to read.
33	topol_file_path : optional
34	If the trajectory file does not contain all information about the topology of
35	the system (i.e. which data in the trajectory file belongs to which particle),
36	you can provide the topology here.
37	"""
38	self.traj_file_path = pathlib.Path(traj_file_path).resolve()	1✔
39
40	if topol_file_path is not None:	1!
41	topol_file_path = pathlib.Path(topol_file_path).resolve()	1✔
42	self.topol_file_path = topol_file_path	1✔
43
44	# until now, chemfiles only supports these 2 properties.
45	# If more are added, link here the mdsuite property to the name of the property
46	# attribute in chemfiles.Frame
47	self.properties_to_chemfile_attr_dict = {	1✔
48	# mdsuite_properties.unwrapped_positions: "positions",
49	mdsuite_properties.positions: "positions",
50	mdsuite_properties.velocities: "velocities",
51	}
52
53	# not all properties are guaranteed to be in the given file.
54	# A subset of self.properties_to_chemfile_attr_dict will be extracted during
55	# self._get_metadata
56	self.properties_in_file = None	1✔
57	self.species_name_to_line_idxs_dict = None	1✔
58	self.batch_size = None	1✔
59
60	def __str__(self):	1✔
61	return str(self.traj_file_path)	1✔
62
63	def _get_metadata(self) -> TrajectoryMetadata:	1✔
64	"""Get the necessary metadata out of chemfiles.
65
66	Trajectory and the first chemfiles.Frame
67	"""
68	with chemfiles.Trajectory(str(self.traj_file_path)) as traj:	1✔
69	if self.topol_file_path is not None:	1!
70	traj.set_topology(str(self.topol_file_path))	1✔
71
72	n_configs = traj.nsteps	1✔
73	frame = traj.read()	1✔
74
75	# get the box lengths from the first frame
76	box_l = frame.cell.lengths	1✔
77
78	# extract which lines in chemfiles.Frame.<property> belong to which species
79	# by going through the atoms list
80	self.species_name_to_line_idxs_dict = {}	1✔
81	for line_idx, atm in enumerate(frame.atoms):	1✔
82	name = atm.name	1✔
83	if name not in self.species_name_to_line_idxs_dict.keys():	1✔
84	self.species_name_to_line_idxs_dict[name] = []	1✔
85	self.species_name_to_line_idxs_dict[name].append(line_idx)	1✔
86
87	# see which properties are in the file. there is no way of finding out except
88	# trying to access them
89	self.properties_in_file = {}	1✔
90	for (	1✔
91	mds_prop,
92	chemfile_attr_name,
93	) in self.properties_to_chemfile_attr_dict.items():
94	try:	1✔
95	frame.__getattribute__(chemfile_attr_name)	1✔
96	except chemfiles.ChemfilesError:	1✔
97	pass	1✔
98	else:
99	self.properties_in_file[mds_prop] = chemfile_attr_name	1✔
100
101	species_list = []	1✔
102	for key, val in self.species_name_to_line_idxs_dict.items():	1✔
103	species_list.append(	1✔
104	mdsuite.database.simulation_database.SpeciesInfo(
105	name=key,
106	n_particles=len(val),
107	properties=list(self.properties_in_file.keys()),
108	)
109	)
110
111	return TrajectoryMetadata(	1✔
112	n_configurations=n_configs,
113	species_list=species_list,
114	box_l=box_l,
115	)
116
117	def get_configurations_generator(	1✔
118	self,
119	) -> typing.Iterator[mdsuite.file_io.file_read.TrajectoryChunkData]:
120	"""Implement parent abstract method."""
121	batch_size = mdsuite.utils.meta_functions.optimize_batch_size(	1✔
122	filepath=self.traj_file_path,
123	number_of_configurations=self.metadata.n_configurations,
124	)
125	n_batches, n_configs_remainder = divmod(	1✔
126	int(self.metadata.n_configurations), int(batch_size)
127	)
128
129	with chemfiles.Trajectory(str(self.traj_file_path)) as traj:	1✔
130	if self.topol_file_path is not None:	1!
131	traj.set_topology(str(self.topol_file_path))	1✔
132	for _ in tqdm.tqdm(range(n_batches), ncols=70):	1✔
133	yield self._read_process_n_configurations(traj, batch_size)	1✔
134	if n_configs_remainder > 0:	1!
UNCOV 135	yield self._read_process_n_configurations(traj, n_configs_remainder)	×
136
137	def _read_process_n_configurations(	1✔
138	self, traj: chemfiles.Trajectory, n_configs: int
139	) -> mdsuite.database.simulation_database.TrajectoryChunkData:
140	"""
141	Read n configurations and package them into a trajectory chunk of the
142	right format.
143
144	Parameters
145	----------
146	traj : chemfiles.Trajectory
147	An open chemfiles Trajectory
148	n_configs : int
149	Number of configurations to read in.
150	"""
151	species_list = self.metadata.species_list	1✔
152	chunk = mdsuite.database.simulation_database.TrajectoryChunkData(	1✔
153	species_list, n_configs
154	)
155
156	for i in range(n_configs):	1✔
157	frame = traj.read()	1✔
158	# slice by species
159	for sp_info in species_list:	1✔
160	for mds_prop, chemfile_attrname in self.properties_in_file.items():	1✔
161	data = frame.__getattribute__(chemfile_attrname)	1✔
162	idxs = self.species_name_to_line_idxs_dict[sp_info.name]	1✔
163	write_data = data[idxs, :]	1✔
164	# add 'time' axis. we only have one configuration to write
165	write_data = write_data[np.newaxis, :, :]	1✔
166	chunk.add_data(write_data, i, sp_info.name, mds_prop.name)	1✔
167	return chunk	1✔

zincware / MDSuite / 9248755285

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous