/
mol_transform.py
120 lines (102 loc) · 3.38 KB
/
mol_transform.py
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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
import numpy as np
import argparse
def read_xyz(filename):
"""
Utility function to read a xyz file.
It returns a numpy matrix for a posterior transformation.
"""
xyz = open(filename, "r")
n_lns = int(xyz.readline())
p_mat = np.zeros((20, 3))
for l in range(n_lns + 1):
ps = xyz.readline().split(" ")
if l > 0:
ps[-1] = ps[-1][:-1]
ps = list(filter(lambda x: x != "" and x != "\n", ps))
p_mat[l - 1] = np.array([float(ps[1]), float(ps[2]), float(ps[3])])
xyz.close()
return p_mat
def write_xyz(filename, p_mat):
"""
Utility function to write a xyz file.
"""
xyz = open(filename, "w")
xyz.write("{}\n".format(len(p_mat)))
xyz.write("\n")
for line in range(len(p_mat)):
xyz.write("C\t\t {: .5f}\t {: .5f}\t {: .5f}\n".format(
p_mat[line, 0], p_mat[line, 1], p_mat[line, 2]))
xyz.close()
def center(p_mat):
"""
Utility function to center a set of 3D coordinates.
"""
print("(WARNING) Not single transformation selected.")
print("(INFO) Molecule centering to be done.")
min_x, max_x = np.min(p_mat[:, 0]), np.max(p_mat[:, 0])
min_y, max_y = np.min(p_mat[:, 1]), np.max(p_mat[:, 1])
min_z, max_z = np.min(p_mat[:, 2]), np.max(p_mat[:, 2])
mean_x = (max_x + min_x) / 2
mean_y = (max_y + min_y) / 2
mean_z = (max_z + min_z) / 2
t_vec = [-mean_x, -mean_y, -mean_z]
t_mat = translate(p_mat, t_vec)
return t_mat
def translate(p_mat, translation_vec=None):
"""
Utility function to translate a set of 3D coordinates.
"""
T = np.eye(4)
T[:-1, 3] = translation_vec
print("(INFO) Transformation matrix to be used:")
print(T)
t_mat = p_mat.copy()
for row in range(len(p_mat)):
t_mat[row] = (T @ np.hstack((p_mat[row], [1])))[:-1]
return t_mat
def scale(p_mat, scaling_vec=None):
"""
Utility function to rotate a set of 3D coordinates.
"""
T = np.eye(4)
if scaling_vec is None:
return p_mat
else:
T[:-1, :-1] *= scaling_vec
print("(INFO) Transformation matrix to be used:")
print(T)
t_mat = p_mat.copy()
for row in range(len(p_mat)):
t_mat[row] = (T @ np.hstack((p_mat[row], [1])))[:-1]
return t_mat
# Construct the argument parser:
ap = argparse.ArgumentParser()
ap.add_argument("-f", "--file", required=True,
help="Path to input .xyz file.")
ap.add_argument("-o", "--output", required=True,
help="Path to output .xyz file.")
ap.add_argument("-t", "--translate", required=False, nargs='+', type=float,
help="Translation vector.")
ap.add_argument("-s", "--scale", required=False, nargs='+', type=float,
help="Scaling vector.")
args = vars(ap.parse_args())
if __name__ == "__main__":
# Parse args:
p_mat = read_xyz(args["file"])
t_vec = args["translate"]
s_vec = args["scale"]
out_f = args["output"]
# Info about input:
print("(INFO) Input file to be transformed: {}.".format(args["file"]))
# Transformations:
if t_vec:
T_mat = translate(p_mat, t_vec)
elif s_vec:
T_mat = scale(p_mat, s_vec)
else:
T_mat = center(p_mat)
# Write output:
write_xyz(out_f, T_mat)
print("(INFO) Processing done.")
# Info about output:
print("(INFO) The output file was stored in: {}.".format(args["output"]))