Coverage for src / molecular_simulations / utils / mda_utils.py: 0%

1import MDAnalysis as mda 

2from pathlib import Path 

3from rust_simulation_tools import (kabsch_align, 

4 unwrap_system, 

5 rewrap_system) 

6from typing import Optional 

7 

8def trim_trajectory(u: mda.Universe, 

9 out: Path, 

10 stride: int=1, 

11 align: bool=False, 

12 rewrap: bool=False, 

13 sel: Optional[str]=None, 

14 align_sel: Optional[str]=None,) -> None: 

15 if sel is not None: 

16 selection = u.select_atoms(sel) 

17 else: 

18 selection = u.atoms 

19 

20 positions = np.zeros((u.trajectory.n_frames // stride, selection.n_atoms, 3), dtype=np.float32) 

21 

22 for i, ts in u.trajectory[::stride]: 

23 positions[i, ...] = selection.positions.copy().astype(np.float32) 

24 

25 if align: 

26 if align_sel is None: 

27 align_idx = selection.select_atoms('backbone or nucleicbackbone').ix 

28 else: 

29 align_idx = selection.select_atoms(align_sel).ix 

30 

31 positions = kabsch_align(positions, positions[0], align_idx) 

32 

33 if rewrap: 

34 pass 

35 

36 with mda.Writer(str(out), n_atoms=selection.n_atoms) as W: 

37 for pos in positions: 

38 W.write(pos)