Skip to content

Commit

Permalink
fmt auxiliary (#4802)
Browse files Browse the repository at this point in the history
  • Loading branch information
@RMeli
RMeli authored Nov 27, 2024
1 parent abc9806 commit 441e2c6
Show file tree
Hide file tree
Showing 11 changed files with 630 additions and 409 deletions.
62 changes: 39 additions & 23 deletions package/MDAnalysis/auxiliary/EDR.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -231,11 +231,15 @@ class EDRStep(base.AuxStep):
:class:`MDAnalysis.auxiliary.base.AuxStep`
"""

def __init__(self, time_selector: str = "Time",
data_selector: Optional[str] = None, **kwargs):
super(EDRStep, self).__init__(time_selector=time_selector,
data_selector=data_selector,
**kwargs)
def __init__(
self,
time_selector: str = "Time",
data_selector: Optional[str] = None,
**kwargs,
):
super(EDRStep, self).__init__(
time_selector=time_selector, data_selector=data_selector, **kwargs
)

def _select_time(self, key: str) -> np.float64:
"""'Time' is one of the entries in the dict returned by pyedr.
Expand All @@ -249,12 +253,14 @@ def _select_data(self, key: Union[str, None]) -> np.float64:
try:
return self._data[key]
except KeyError:
raise KeyError(f"'{key}' is not a key in the data_dict dictionary."
" Check the EDRReader.terms attribute")
raise KeyError(
f"'{key}' is not a key in the data_dict dictionary."
" Check the EDRReader.terms attribute"
)


class EDRReader(base.AuxReader):
""" Auxiliary reader to read data from an .edr file.
"""Auxiliary reader to read data from an .edr file.
`EDR files`_
are created by GROMACS during a simulation. They are binary files which
Expand Down Expand Up @@ -310,8 +316,9 @@ class EDRReader(base.AuxReader):

def __init__(self, filename: str, convert_units: bool = True, **kwargs):
if not HAS_PYEDR:
raise ImportError("EDRReader: To read EDR files please install "
"pyedr.")
raise ImportError(
"EDRReader: To read EDR files please install " "pyedr."
)
self._auxdata = Path(filename).resolve()
self.data_dict = pyedr.edr_to_dict(filename)
self.unit_dict = pyedr.get_unit_dictionary(filename)
Expand Down Expand Up @@ -340,8 +347,10 @@ def _convert_units(self):
self.data_dict[term] = units.convert(data, unit, target_unit)
self.unit_dict[term] = units.MDANALYSIS_BASE_UNITS[unit_type]
if unknown_units:
warnings.warn("Could not find unit type for the following "
f"units: {unknown_units}")
warnings.warn(
"Could not find unit type for the following "
f"units: {unknown_units}"
)

def _memory_usage(self):
size = 0
Expand All @@ -365,8 +374,10 @@ def _read_next_step(self) -> EDRStep:
auxstep = self.auxstep
new_step = self.step + 1
if new_step < self.n_steps:
auxstep._data = {term: self.data_dict[term][self.step + 1]
for term in self.terms}
auxstep._data = {
term: self.data_dict[term][self.step + 1]
for term in self.terms
}
auxstep.step = new_step
return auxstep
else:
Expand All @@ -375,7 +386,7 @@ def _read_next_step(self) -> EDRStep:
raise StopIteration

def _go_to_step(self, i: int) -> EDRStep:
""" Move to and read i-th auxiliary step.
"""Move to and read i-th auxiliary step.
Parameters
----------
Expand All @@ -392,14 +403,16 @@ def _go_to_step(self, i: int) -> EDRStep:
If step index not in valid range.
"""
if i >= self.n_steps or i < 0:
raise ValueError("Step index {0} is not valid for auxiliary "
"(num. steps {1})".format(i, self.n_steps))
raise ValueError(
"Step index {0} is not valid for auxiliary "
"(num. steps {1})".format(i, self.n_steps)
)
self.auxstep.step = i - 1
self.next()
return self.auxstep

def read_all_times(self) -> np.ndarray:
""" Get list of time at each step.
"""Get list of time at each step.
Returns
-------
Expand All @@ -408,9 +421,10 @@ def read_all_times(self) -> np.ndarray:
"""
return self.data_dict[self.time_selector]

def get_data(self, data_selector: Union[str, List[str], None] = None
) -> Dict[str, np.ndarray]:
""" Returns the auxiliary data contained in the :class:`EDRReader`.
def get_data(
self, data_selector: Union[str, List[str], None] = None
) -> Dict[str, np.ndarray]:
"""Returns the auxiliary data contained in the :class:`EDRReader`.
Returns either all data or data specified as `data_selector` in form
of a str or a list of any of :attr:`EDRReader.terms`. `Time` is
always returned to allow easy plotting.
Expand Down Expand Up @@ -438,8 +452,10 @@ def _get_data_term(term, datadict):
try:
return datadict[term]
except KeyError:
raise KeyError(f"data selector {term} is invalid. Check the "
"EDRReader's `terms` attribute.")
raise KeyError(
f"data selector {term} is invalid. Check the "
"EDRReader's `terms` attribute."
)

data_dict = {"Time": self.data_dict["Time"]}

Expand Down
69 changes: 38 additions & 31 deletions package/MDAnalysis/auxiliary/XVG.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@


def uncomment(lines):
""" Remove comments from lines in an .xvg file
"""Remove comments from lines in an .xvg file
Parameters
----------
Expand All @@ -92,10 +92,10 @@ def uncomment(lines):
if not stripped_line:
continue
# '@' must be at the beginning of a line to be a grace instruction
if stripped_line[0] == '@':
if stripped_line[0] == "@":
continue
# '#' can be anywhere in the line, everything after is a comment
comment_position = stripped_line.find('#')
comment_position = stripped_line.find("#")
if comment_position > 0 and stripped_line[:comment_position]:
yield stripped_line[:comment_position]
elif comment_position < 0 and stripped_line:
Expand All @@ -104,7 +104,7 @@ def uncomment(lines):


class XVGStep(base.AuxStep):
""" AuxStep class for .xvg file format.
"""AuxStep class for .xvg file format.
Extends the base AuxStep class to allow selection of time and
data-of-interest fields (by column index) from the full set of data read
Expand All @@ -127,9 +127,9 @@ class XVGStep(base.AuxStep):
"""

def __init__(self, time_selector=0, data_selector=None, **kwargs):
super(XVGStep, self).__init__(time_selector=time_selector,
data_selector=data_selector,
**kwargs)
super(XVGStep, self).__init__(
time_selector=time_selector, data_selector=data_selector, **kwargs
)

def _select_time(self, key):
if key is None:
Expand All @@ -138,7 +138,7 @@ def _select_time(self, key):
if isinstance(key, numbers.Integral):
return self._select_data(key)
else:
raise ValueError('Time selector must be single index')
raise ValueError("Time selector must be single index")

def _select_data(self, key):
if key is None:
Expand All @@ -148,15 +148,17 @@ def _select_data(self, key):
try:
return self._data[key]
except IndexError:
errmsg = (f'{key} not a valid index for data with '
f'{len(self._data)} columns')
errmsg = (
f"{key} not a valid index for data with "
f"{len(self._data)} columns"
)
raise ValueError(errmsg) from None
else:
return np.array([self._select_data(i) for i in key])


class XVGReader(base.AuxReader):
""" Auxiliary reader to read data from an .xvg file.
"""Auxiliary reader to read data from an .xvg file.
Default reader for .xvg files. All data from the file will be read and
stored on initialisation.
Expand Down Expand Up @@ -188,24 +190,25 @@ def __init__(self, filename, **kwargs):
auxdata_values = []
# remove comments before storing
for i, line in enumerate(uncomment(lines)):
if line.lstrip()[0] == '&':
if line.lstrip()[0] == "&":
# multiple data sets not supported; stop at end of first set
break
auxdata_values.append([float(val) for val in line.split()])
# check the number of columns is consistent
if len(auxdata_values[i]) != len(auxdata_values[0]):
raise ValueError('Step {0} has {1} columns instead of '
'{2}'.format(i, auxdata_values[i],
auxdata_values[0]))
raise ValueError(
"Step {0} has {1} columns instead of "
"{2}".format(i, auxdata_values[i], auxdata_values[0])
)
self._auxdata_values = np.array(auxdata_values)
self._n_steps = len(self._auxdata_values)
super(XVGReader, self).__init__(**kwargs)

def _memory_usage(self):
return(self._auxdata_values.nbytes)
return self._auxdata_values.nbytes

def _read_next_step(self):
""" Read next auxiliary step and update ``auxstep``.
"""Read next auxiliary step and update ``auxstep``.
Returns
-------
Expand All @@ -228,7 +231,7 @@ def _read_next_step(self):
raise StopIteration

def _go_to_step(self, i):
""" Move to and read i-th auxiliary step.
"""Move to and read i-th auxiliary step.
Parameters
----------
Expand All @@ -245,14 +248,16 @@ def _go_to_step(self, i):
If step index not in valid range.
"""
if i >= self.n_steps or i < 0:
raise ValueError("Step index {0} is not valid for auxiliary "
"(num. steps {1})".format(i, self.n_steps))
self.auxstep.step = i-1
raise ValueError(
"Step index {0} is not valid for auxiliary "
"(num. steps {1})".format(i, self.n_steps)
)
self.auxstep.step = i - 1
self.next()
return self.auxstep

def read_all_times(self):
""" Get NumPy array of time at each step.
"""Get NumPy array of time at each step.
Returns
-------
Expand All @@ -263,7 +268,7 @@ def read_all_times(self):


class XVGFileReader(base.AuxFileReader):
""" Auxiliary reader to read (one step at a time) from an .xvg file.
"""Auxiliary reader to read (one step at a time) from an .xvg file.
An alternative XVG reader which reads each step from the .xvg file as
needed (rather than reading and storing all from the start), for a lower
Expand All @@ -286,14 +291,14 @@ class XVGFileReader(base.AuxFileReader):
The default reader for .xvg files is :class:`XVGReader`.
"""

format = 'XVG-F'
format = "XVG-F"
_Auxstep = XVGStep

def __init__(self, filename, **kwargs):
super(XVGFileReader, self).__init__(filename, **kwargs)

def _read_next_step(self):
""" Read next recorded step in xvg file and update ``auxstep``.
"""Read next recorded step in xvg file and update ``auxstep``.
Returns
-------
Expand All @@ -307,7 +312,7 @@ def _read_next_step(self):
"""
line = next(self.auxfile)
while True:
if not line or (line.strip() and line.strip()[0] == '&'):
if not line or (line.strip() and line.strip()[0] == "&"):
# at end of file or end of first set of data (multiple sets
# currently not supported)
self.rewind()
Expand All @@ -325,15 +330,17 @@ def _read_next_step(self):
# haven't set n_cols yet; set now
auxstep._n_cols = len(auxstep._data)
if len(auxstep._data) != auxstep._n_cols:
raise ValueError(f'Step {self.step} has '
f'{len(auxstep._data)} columns instead '
f'of {auxstep._n_cols}')
raise ValueError(
f"Step {self.step} has "
f"{len(auxstep._data)} columns instead "
f"of {auxstep._n_cols}"
)
return auxstep
# line is comment only - move to next
line = next(self.auxfile)

def _count_n_steps(self):
""" Iterate through all steps to count total number.
"""Iterate through all steps to count total number.
Returns
-------
Expand All @@ -358,7 +365,7 @@ def _count_n_steps(self):
return count

def read_all_times(self):
""" Iterate through all steps to build times list.
"""Iterate through all steps to build times list.
Returns
-------
Expand Down
Loading

0 comments on commit 441e2c6

Please sign in to comment.