Coverage for encodermap/encodermap_tf1/misc.py: 43%

1import os 

2from math import pi 

3 

4import numpy as np 

5import tensorflow.compat.v1 as tf 

6from tensorboard.backend.event_processing.event_accumulator import EventAccumulator 

7 

8 

9def add_layer_summaries(layer, debug=False): 

10 """ 

11 

12 :param layer: 

13 :return: 

14 """ 

15 weights = layer.variables[0] 

16 biases = layer.variables[1] 

17 variable_summaries(layer.name + "_weights", weights, debug) 

18 variable_summaries(layer.name + "_biases", biases, debug) 

19 

20 

21def periodic_distance(a, b, periodicity=2 * pi): 

22 """ 

23 

24 :param a: 

25 :param b: 

26 :param periodicity: 

27 :return: 

28 """ 

29 d = tf.abs(b - a) 

30 return tf.minimum(d, periodicity - d) 

31 

32 

33def periodic_distance_np(a, b, periodicity=2 * pi): 

34 """ 

35 

36 :param a: 

37 :param b: 

38 :param periodicity: 

39 :return: 

40 """ 

41 d = np.abs(b - a) 

42 return np.minimum(d, periodicity - d) 

43 

44 

45def variable_summaries(name, variables, debug=False): 

46 """ 

47 Attach several summaries to a Tensor for TensorBoard visualization. 

48 

49 :param name: 

50 :param variables: 

51 :return: 

52 """ 

53 if not isinstance(variables, list): 

54 variables = [variables] 

55 

56 for i, var in enumerate(variables): 

57 try: 

58 add_index = len(variables) > 1 

59 except TypeError: 

60 add_index = True 

61 if add_index: 

62 name = name + str(i) 

63 with tf.name_scope(name): 

64 mean = tf.reduce_mean(var) 

65 tf.summary.scalar("mean", mean) 

66 with tf.name_scope("stddev"): 

67 stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean))) 

68 tf.summary.scalar("stddev", stddev) 

69 tf.summary.scalar("max", tf.reduce_max(var)) 

70 tf.summary.scalar("min", tf.reduce_min(var)) 

71 tf.summary.histogram("histogram", var) 

72 tf.summary.tensor_summary("values", var) 

73 

74 

75def create_dir(path): 

76 """ 

77 

78 :param path: 

79 :return: 

80 """ 

81 if not os.path.isdir(path): 

82 os.makedirs(path) 

83 return path 

84 

85 

86def distance_cost(r_h, r_l, sig_h, a_h, b_h, sig_l, a_l, b_l, periodicity): 

87 """ 

88 

89 :param r_h: 

90 :param r_l: 

91 :param sig_h: 

92 :param a_h: 

93 :param b_h: 

94 :param sig_l: 

95 :param a_l: 

96 :param b_l: 

97 :param periodicity: 

98 :return: 

99 """ 

100 if periodicity == float("inf"): 

101 dist_h = pairwise_dist(r_h) 

102 else: 

103 dist_h = pairwise_dist_periodic(r_h, periodicity) 

104 dist_l = pairwise_dist(r_l) 

105 

106 sig_h = sigmoid(dist_h, sig_h, a_h, b_h) 

107 sig_l = sigmoid(dist_l, sig_l, a_l, b_l) 

108 

109 cost = tf.reduce_mean(tf.square(sig_h - sig_l)) 

110 return cost 

111 

112 

113def sigmoid(r, sig, a, b): 

114 """ 

115 

116 :param r: 

117 :param sig: 

118 :param a: 

119 :param b: 

120 :return: 

121 """ 

122 return 1 - (1 + (2 ** (a / b) - 1) * (r / sig) ** a) ** (-b / a) 

123 

124 

125def pairwise_dist_periodic(positions, periodicity): 

126 assert len(positions.shape) == 2 

127 with tf.name_scope("pairwise_dist_periodic"): 

128 vecs = periodic_distance( 

129 tf.expand_dims(positions, axis=1), 

130 tf.expand_dims(positions, axis=0), 

131 periodicity, 

132 ) 

133 mask = tf.to_float(tf.equal(vecs, 0.0)) 

134 vecs = vecs + mask * 1e-16 # gradient infinite for 0 

135 dists = tf.norm(vecs, axis=2) 

136 return dists 

137 

138 

139def pairwise_dist(positions, squared=False, flat=False): 

140 # thanks to https://omoindrot.github.io/triplet-loss 

141 

142 with tf.name_scope("pairwise_dist"): 

143 if not tf.is_numeric_tensor(positions): 

144 positions = tf.convert_to_tensor(positions) 

145 if len(positions.get_shape()) == 2: 145 ↛ 150line 145 didn't jump to line 150, because the condition on line 145 was never false

146 positions = tf.expand_dims(positions, 0) 

147 

148 # Get the dot product between all embeddings 

149 # shape (batch_size, batch_size) 

150 dot_product = tf.matmul(positions, tf.transpose(positions, [0, 2, 1])) 

151 

152 # Get squared L2 norm for each embedding. We can just take the diagonal of `dot_product`. 

153 # This also provides more numerical stability (the diagonal of the result will be exactly 0). 

154 # shape (batch_size,) 

155 square_norm = tf.linalg.diag_part(dot_product) 

156 

157 # Compute the pairwise distance matrix as we have: 

158 # ||a - b||^2 = ||a||^2 - 2 <a, b> + ||b||^2 

159 # shape (batch_size, batch_size) 

160 distances = ( 

161 tf.expand_dims(square_norm, 1) 

162 - 2.0 * dot_product 

163 + tf.expand_dims(square_norm, 2) 

164 ) 

165 

166 # Because of computation errors, some distances might be negative so we put everything >= 0.0 

167 distances = tf.maximum(distances, 0.0) 

168 

169 if flat: 169 ↛ 170line 169 didn't jump to line 170, because the condition on line 169 was never true

170 n = int(positions.shape[1]) 

171 mask = np.ones((n, n), dtype=bool) 

172 mask[np.tril_indices(n)] = False 

173 distances = tf.boolean_mask(distances, mask, axis=1) 

174 

175 if not squared: 175 ↛ 186line 175 didn't jump to line 186, because the condition on line 175 was never false

176 # Because the gradient of sqrt is infinite when distances == 0.0 (ex: on the diagonal) 

177 # we need to add a small epsilon where distances == 0.0 

178 mask = tf.to_float(tf.equal(distances, 0.0)) 

179 distances = distances + mask * 1e-16 

180 

181 distances = tf.sqrt(distances) 

182 

183 # Correct the epsilon added: set the distances on the mask to be exactly 0.0 

184 distances = distances * (1.0 - mask) 

185 

186 return distances 

187 

188 

189def search_and_replace( 

190 file_path, search_pattern, replacement, out_path=None, backup=True 

191): 

192 """ 

193 Searches for a pattern in a text file and replaces it with the replacement 

194 

195 :param file_path: (str) 

196 path to the file to search 

197 :param search_pattern: (str) 

198 pattern to search for 

199 :param replacement: (str) 

200 string that replaces the search_pattern in the output file 

201 :param out_path: (str) 

202 path where to write the output file. If no path is given the original file will be replaced. 

203 :param backup: (bool) 

204 if backup is true the original file is renamed to filename.bak before it is overwritten 

205 

206 :return: 

207 

208 """ 

209 with open(file_path, "r") as f: 

210 file_data = f.read() 

211 

212 file_data = file_data.replace(search_pattern, replacement) 

213 

214 if not out_path: 

215 out_path = file_path 

216 if backup: 

217 os.rename(file_path, file_path + ".bak") 

218 

219 with open(out_path, "w") as file: 

220 file.write(file_data) 

221 

222 

223def run_path(path): 

224 """ 

225 Creates a directory at "path/run{i}" where the i is corresponding to the smallest not yet existing path 

226 

227 :param path: (str) 

228 :return: (str) 

229 path of the created folder 

230 

231 """ 

232 i = 0 

233 while True: 

234 current_path = os.path.join(path, "run{}".format(i)) 

235 if not os.path.exists(current_path): 

236 os.makedirs(current_path) 

237 output_path = current_path 

238 break 

239 else: 

240 i += 1 

241 return output_path 

242 

243 

244def random_on_cube_edges(n_points, sigma=0): 

245 x = y = z = 1 

246 r = np.random.uniform(size=n_points) 

247 

248 coordinates = np.zeros((n_points, 3)) 

249 

250 ids = np.zeros(n_points) 

251 

252 a = np.array([[0, 0, 0]] * 3 + [[x, y, 0]] * 3 + [[0, y, z]] * 3 + [[x, 0, z]] * 3) 

253 

254 b = np.array( 

255 [ 

256 [x, 0, 0], 

257 [0, y, 0], 

258 [0, 0, z], 

259 [-x, 0, 0], 

260 [0, -y, 0], 

261 [0, 0, z], 

262 [x, 0, 0], 

263 [0, -y, 0], 

264 [0, 0, -z], 

265 [-x, 0, 0], 

266 [0, y, 0], 

267 [0, 0, -z], 

268 ] 

269 ) 

270 

271 for i in range(12): 

272 mask = (i / 12 < r) & (r < (i + 1) / 12) 

273 coordinates[mask] += ( 

274 a[i] + (np.expand_dims(r[mask], axis=1) - i / 12) * 12 * b[i] 

275 ) 

276 ids[mask] = i 

277 

278 if sigma: 

279 coordinates += np.random.normal(scale=sigma, size=(n_points, 3)) 

280 

281 return coordinates, ids 

282 

283 

284def rotation_matrix(axis_unit_vec, angle): 

285 angle = tf.expand_dims(tf.expand_dims(angle, axis=-1), axis=-1) 

286 i = tf.expand_dims(tf.eye(3), 0) 

287 zeros = tf.zeros(tf.shape(axis_unit_vec)[0]) 

288 cross_prod_matrix = tf.convert_to_tensor( 

289 [ 

290 [zeros, -axis_unit_vec[:, 2], axis_unit_vec[:, 1]], 

291 [axis_unit_vec[:, 2], zeros, -axis_unit_vec[:, 0]], 

292 [-axis_unit_vec[:, 1], axis_unit_vec[:, 0], zeros], 

293 ] 

294 ) 

295 cross_prod_matrix = tf.transpose(cross_prod_matrix, [2, 0, 1]) 

296 r = tf.cos(angle) * i 

297 r += tf.sin(angle) * cross_prod_matrix 

298 axis_unit_vec = tf.expand_dims(axis_unit_vec, 2) 

299 r += (1 - tf.cos(angle)) * tf.matmul( 

300 axis_unit_vec, tf.transpose(axis_unit_vec, [0, 2, 1]) 

301 ) 

302 return r 

303 

304 

305def potential_energy(angles, dihedrals, distances): 

306 energy = 0 

307 energies = tf.where( 

308 distances < 2, 

309 tf.minimum(10 - distances, (1 / distances) ** 12), 

310 tf.zeros(tf.shape(distances)), 

311 ) 

312 energy += tf.reduce_mean(tf.reduce_sum(energies, axis=1)) 

313 return energy 

314 

315 

316def read_from_log(run_path, names): 

317 ea = EventAccumulator(run_path) 

318 ea.Reload() 

319 results = [np.array(ea.Scalars(name)) for name in names] 

320 return results