Python tricks that make your life easier ¶

In [1]: import pytraj as pt

In [2]: traj = pt.load_sample_data('tz2')[:3]

In [3]: traj
Out[3]: 
pytraj.Trajectory, 3 frames: 
Size: 0.000355 (GB)
<Topology: 5293 atoms, 1704 residues, 1692 mols, PBC with box type = ortho>

Unpacking ¶

In [4]: a, b, c = traj[0], traj[1], traj[2]

In [5]: a, b, c
Out[5]: (<Frame with 5293 atoms>, <Frame with 5293 atoms>, <Frame with 5293 atoms>)

Unpacking dictionary in python function ¶

For new users, it’s confusing when you see f(*args, **kwd). This is Python sytax to unpack list/tuple and dictionary respectively

In [6]: def f(x, y, z=0, a=100):
   ...:     print('x = {}'.format(x))
   ...:     print('y = {}'.format(y))
   ...:     print('z = {}'.format(z))
   ...:     print('a = {}'.format(a))
   ...: 

In [7]: my_list = [4, 6]

In [8]: my_dict = {'z': 8, 'a': 9}

# call your function by unpacking list and dictionary
In [9]: f(*my_list, **my_dict)
x = 4
y = 6
z = 8
a = 9

# which is equal to
In [10]: f(4, 6, z=8, a=9)
x = 4
y = 6
z = 8
a = 9

Iterating over Trajectory index and value pairs (enumerate)¶

In [11]: for i, frame in enumerate(traj):
   ....:     print(i, frame)
   ....: 
0 <Frame with 5293 atoms>
1 <Frame with 5293 atoms>
2 <Frame with 5293 atoms>

Zipping ¶

In [12]: h_indices = pt.select_atoms(traj.top, '@H')

In [13]: n_indices = h_indices - 1

In [14]: list(zip(h_indices, n_indices))[:-5]
Out[14]: [(14, 13), (38, 37), (52, 51), (76, 75), (91, 90), (105, 104)]

Sliding windows ¶

In [15]: for value in pt.tools.n_grams([0, 2, 4, 5, 7, 9], 2):
   ....:     print(value)
   ....: 
(0, 2)
(2, 4)
(4, 5)
(5, 7)
(7, 9)

Flattening lists ¶

In [16]: a = [[0, 54, 6], [[6, 7, 8], [9, 7, 5]]]

In [17]: print(pt.tools.flatten(a))
[0, 54, 6, 6, 7, 8, 9, 7, 5]

List comprehensions ¶

In [18]: [x**2 for x in [3., 5., 7.]]
Out[18]: [9.0, 25.0, 49.0]

Dictionary comprehensions ¶

In [19]: a = ['one', 'two', 'three']

In [20]: b = [1, 2, 3]

In [21]: print(dict(zip(a, b)))
{'three': 3, 'one': 1, 'two': 2}

Get value from Dictionary if key does not exist ¶

In [22]: a = dict(x=3, y=4)

In [23]: a
Out[23]: {'x': 3, 'y': 4}

In [24]: a.get('x')
Out[24]: 3

In [25]: a.get('z', 100)
Out[25]: 100

Set ¶

In [26]: set(res.name for res in traj.top.residues)
Out[26]: {'ASN', 'GLU', 'GLY', 'LYS', 'NHE', 'SER', 'THR', 'TRP', 'WAT'}

Filter data ¶

In [27]: a = [3, 8, 2, 9]

In [28]: list(filter(lambda x: x < 5, a))
Out[28]: [3, 2]

Combinations ¶

In [29]: from itertools import combinations

In [30]: a = [3, 8, 2, 9]

In [31]: list(combinations(a, 3))
Out[31]: [(3, 8, 2), (3, 8, 9), (3, 2, 9), (8, 2, 9)]

Chain several iterators ¶

In [32]: from itertools import chain

In [33]: traj = pt.iterload('tz2.nc', 'tz2.parm7')

In [34]: fi_0 = traj(0, 3)

In [35]: fi_0
Out[35]: 
<FrameIterator with start=0, stop=3, step=1, n_frames=3, 
frame_indices=None, 
mask=None, autoimage=False, rmsfit=None, copy=False> 

In [36]: fi_1 = traj(5, 8)

In [37]: fi_1
Out[37]: 
<FrameIterator with start=5, stop=8, step=1, n_frames=3, 
frame_indices=None, 
mask=None, autoimage=False, rmsfit=None, copy=False> 

In [38]: for frame in chain(fi_0, fi_1): print(frame)
<Frame with 223 atoms>
<Frame with 223 atoms>
<Frame with 223 atoms>
<Frame with 223 atoms>
<Frame with 223 atoms>
<Frame with 223 atoms>