Robotics
code
File size: 1,471 Bytes
aac5fad
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
"""Functions to read and write n2p2 data files."""

__all__ = [
    'read_epre_n2p2',
    'read_fpre_n2p2'
]

import numpy as np

from .utilities import Frame, register_io


@register_io('N2P2_E','read')
def read_epre_n2p2(f_in, column=3):
    """Read the outcome of the energy prediction from file"""

    line = f_in.readline()
    # no more data in the file
    if not line:
        return None
    # Skip comment lines:
    while True:
        if '#' not in line:
            break
        line = f_in.readline()
    energy = float(line.split()[column])
    return Frame(energy=energy)


@register_io('N2P2_F','read')
def read_fpre_n2p2(f_in):
    """Read the outcome of the force prediction from file"""

    line = f_in.readline()
    # no more data in the file
    if not line:
        return None
    # Skip comment lines:
    while True:
        if '#' not in line:
            break
        line = f_in.readline()

    items = line.split()
    config = items[0]

    forces = []
    forces.append(float(items[3]))
    while True:
        last_pos = f_in.tell()
        line = f_in.readline()
        # no more data in the file
        if not line:
            break
        items = line.split()
        # Stop if config changes
        if items[0] != config:
            f_in.seek(last_pos)
            break

        forces.append(float(items[3]))

    forces = np.array(forces)
    forces = forces.reshape((len(forces)//3, 3))
    return Frame(forces=forces)