Background
A forcefield is a collection of functions and parameters that is used to calculate the potential energy of a complex system. I have text files which contain data about the parameters for a forcefield. The text file is split into many sections, with each section following the same format:
- A section header which is enclosed in square brackets
- On the next line the word
indices:followed by a list of integers. - This is then followed by 1 or more lines of parameters associated with the section
Here is a made-up example file to showcase the format.
############################################
# Comments begin with '#'
############################################
[lj_pairs] # Section 1
indices: 0 2
# ID eps sigma
1 2.344 1.234 5
2 4.423 5.313 5
3 1.573 6.321 5
4 1.921 11.93 5
[bonds]
indices: 0 1
2 4.234e-03 11.2
6 -0.134545 5.7
The goal is to parse such files and store all of the information in a dict.
Code
Main function for review
""" Force-field data reader """
import re
from dataclasses import dataclass, field
from typing import Dict, Iterable, List, TextIO, Tuple, Union, Any
def ff_reader(fname: Union[str, TextIO]) -> Dict[str, "FFSections"]:
""" Reads data from a force-field file """
try:
if _is_string(fname):
fh = open(fname, mode="r")
own = True
else:
fh = iter(fname)
except TypeError:
raise ValueError("fname must be a string or a file handle")
# All the possible section headers
keywords = ("lj_pairs", "bonds") # etc... Long list of possible sections
# Removed for brevity
re_sections = re.compile(r"^\[(%s)\]$" % "|".join(keywords))
ff_data = _strip_comments(fh)
# Empty dict that'll hold all the data.
final_ff_data = {key: FFSections() for key in keywords}
# Get first section header
for line in ff_data:
match = re.match(re_sections, line)
if match:
section = match.group(1)
in_section_for_first_time = True
break
else:
raise FFReaderError("A valid section header must be the first line in file")
else:
raise FFReaderError("No force-field sections exist")
# Read the rest of the file
for line in ff_data:
match = re.match(re_sections, line)
# If we've encounted a section header the next line must be an index list.
if in_section_for_first_time:
if line.split()[0] != "indices:":
raise FFReaderError(f"Missing index list for section: {section}")
idx = _validate_indices(line)
final_ff_data[section].use_idx = idx
in_section_for_first_time = False
in_params_for_first_time = True
continue
if match and in_params_for_first_time:
raise FFReaderError(
f"Section {section} missing parameters"
+ "Sections must contain atleast one type coefficients"
)
if match: # and not in_section_for_first_time and in_params_for_first_time
section = match.group(1)
in_section_for_first_time = True
continue
params = _validate_params(line)
final_ff_data[section].coeffs.update([params])
in_params_for_first_time = False
# Close the file if we opened it
if own:
fh.close()
for section in final_ff_data.values():
# coeff must exist if use_idx does
if section.use_idx is not None:
assert section.coeffs
return final_ff_data
Other stuff for the code to work
def _strip_comments(
instream: TextIO, comments: Union[str, Iterable[str], None] = "#"
) -> Iterable[str]:
""" Strip comments from a text IO stream """
if comments is not None:
if isinstance(comments, str):
comments = [comments]
comments_re = re.compile("|".join(map(re.escape, comments)))
try:
for lines in instream.readlines():
line = re.split(comments_re, lines, 1)[0].strip()
if line != "":
yield line
except AttributeError:
raise TypeError("instream must be a `TextIO` stream") from None
@dataclass(eq=False)
class FFSections:
"""
FFSections(coeffs,use_idx)
Container for forcefield information
"""
coeffs: Dict[int, List[float]] = field(default_factory=dict)
use_idx: List[int] = field(default=None)
class FFReaderError(Exception):
""" Incorrect or badly formatted force-Field data """
def __init__(self, message: str, badline: Optional[str] = None) -> None:
if badline:
message = f"{message}\nError parsing --> ({badline})"
super().__init__(message)
def _validate_indices(line: str) -> List[int]:
"""
Check if given line contains only a whitespace separated
list of integers
"""
# split on indices: followed by whitescape
split = line.split("indices:")[1].split()
# import ipdb; ipdb.set_trace()
if not set(s.isdecimal() for s in split) == {True}:
raise FFReaderError(
"Indices should be integers and seperated by whitespace", line
)
return [int(x) for x in split]
def _validate_params(line: str) -> Tuple[int, List[float]]:
"""
Check if given line is valid param line, which are
an integer followed by one or more floats seperated by whitespace
"""
split = line.split()
id_ = split[0]
coeffs = split[1:]
if not id_.isdecimal():
raise FFReaderError("Invalid params", line)
try:
coeffs = [float(x) for x in coeffs]
except (TypeError, ValueError):
raise FFReaderError("Invalid params", line) from None
return (int(id_), coeffs)
I consider myself a beginner in python and this is my first substantive project. I'd like the review to focus on the ff_reader function, but feel free to comment on the other parts too if there are better ways to do somethings. I feel like the way I've written the ff_reader is kind of ugly and inelegant. I'd be especially interested if there is a better way to read such files, perhaps parsing the whole file instead of line by line.
