Coverage for dynasor/trajectory/atomic_indices.py: 100%

1import os 

2import re 

3import numpy as np 

4from numpy.typing import NDArray 

5from dynasor.logging_tools import logger 

6 

7 

8def parse_gromacs_index_file(fname: str) -> dict[str, NDArray[int]]: 

9 """ 

10 Parses a gromacs style index (ndx) file. 

11 Returns a dict with key values as 

12 `group-name: [1, 3, 8]`. 

13 

14 Comments (everything following a ``;``) and blank lines are ignored. 

15 

16 Note 

17 ---- 

18 The atomic indices in gromacs-ndx file starts with 1, but the returned dict starts with 0. 

19 """ 

20 

21 if not os.path.isfile(fname): 

22 raise ValueError('Index file not found') 

23 

24 atomic_indices = dict() 

25 name = None 

26 header_re = re.compile(r'^ *\[ *([a-zA-Z0-9_.-]+) *\] *$') 

27 with open(fname, 'r') as fobj: 

28 for line in fobj.readlines(): 

29 line = line.split(';')[0] # strip comments 

30 match = header_re.match(line) 

31 if match is not None: # get name of group 

32 name = match.group(1) 

33 if name in atomic_indices.keys(): 

34 logger.warning(f'Group name {name} appears twice in index file, ' 

35 'the indices of the last occurrence are used.') 

36 atomic_indices[name] = [] 

37 else: # get indices for group 

38 indices = [int(i) for i in line.split()] 

39 if not indices: # blank or comment-only line 

40 continue 

41 if name is None: 

42 raise ValueError('Index file contains indices before any group header.') 

43 atomic_indices[name].extend(indices) 

44 

45 # cast to indices to numpy arrays and shift so indices start with 0 

46 for name, indices in atomic_indices.items(): 

47 atomic_indices[name] = np.array(indices) - 1 

48 

49 return atomic_indices