def test_wavefrontify(self, decorator, b): l1, l2, s = 14, 37, 9 minval_open, maxval_open = 10.5, 11.5 minval_extend, maxval_extend = 0.8, 1.2 sim_mat = random_sim_mat(b, l1=l1, l2=l2, emb_dim=3) gap_open = random_gap_penalty(minval_open, maxval_open, b, l1, l2) gap_extend = random_gap_penalty(minval_extend, maxval_extend, b, l1, l2) w = alignment.weights_from_sim_mat(sim_mat, gap_open, gap_extend) wavefrontify_fn = tf_ops.wavefrontify unwavefrontify_fn = tf_ops.unwavefrontify if decorator is not None: wavefrontify_fn = decorator(wavefrontify_fn) unwavefrontify_fn = decorator(unwavefrontify_fn) w_wavefrontified = wavefrontify_fn(w) w_unwavefrontified = unwavefrontify_fn(w_wavefrontified) self.assertEqual(w_wavefrontified.shape, (l1 + l2 - 1, s, l1, b)) self.assertAllEqual(w_unwavefrontified, w) for n in tf.range(b): for a in tf.range(s): for i in tf.range(l1): for j in tf.range(l2): self.assertEqual(w_wavefrontified[i + j, a, i, n], w[n, i, j, a])
def forward(sw_params): # Perturbs the Smith-Waterman LP parameters. sim_mat, gap_open, gap_extend = sw_params[0], sw_params[1], sw_params[2] sw_params = alignment.weights_from_sim_mat(sim_mat, gap_open, gap_extend) shape, dtype = tf.shape(sw_params), sw_params.dtype pert_shape = tf.concat([[num_samples], shape], axis=0) noise, noise_grad = tf.nest.map_structure( lambda t: tf.cast(t, dtype), noise_sampler(pert_shape)) pert_sw_params = tf.expand_dims(sw_params, 0) + sigma * noise pert_sw_params = tf.reshape(pert_sw_params, tf.concat([[-1], shape[1:]], axis=0)) # Computes optimal Smith-Waterman scores and alignments. scores, paths = hard_sw_affine(pert_sw_params) # Average scores and outputs over random perturbations def perturbation_reshape(t): new_shape = tf.concat([[num_samples], [-1], tf.shape(t)[1:]], 0) return tf.reshape(t, new_shape) scores = perturbation_reshape(scores) # [num_samples, batch] paths = perturbation_reshape(paths) # [num_samples, batch, len1, len2, 9] pert_scores = tf.reduce_mean(scores, axis=0) pert_paths = tf.reduce_mean(paths, axis=0) if stop_paths_gradient: pert_paths = tf.stop_gradient(pert_paths) def grad(ds, dp): # Computes grad of scores w.r.t. (packed) sw_params. grad_scores = pert_paths * tf.reshape(ds, [tf.shape(ds)[0], 1, 1, 1]) # Computes grad of scores w.r.t. sim_mat, gap_open and gap_extend. grad_scores = alignment.adjoint_weights_from_sim_mat( grad_scores, gap_open.shape, gap_extend.shape) if stop_paths_gradient: gradients = grad_scores else: # Flattens paths to shape [num_samples, batch, len1 * len2 * 9]. flat_paths = tf.reshape( paths, [num_samples, tf.shape(paths)[1], -1]) flat_noise_grad = tf.reshape( noise_grad, [num_samples, tf.shape(noise_grad)[1], -1]) # Flattens dp to shape [num_samples, len1 * len2 * 9] flat_dp = tf.reshape(dp, [tf.shape(dp)[0], -1]) # Computes grad of paths w.r.t. (packed) sw_params. grad_paths = tf.einsum('nbd,nb->bd', flat_noise_grad, tf.einsum('nbd,bd->nb', flat_paths, flat_dp)) grad_paths /= sigma * num_samples grad_paths = tf.reshape(grad_paths, shape) # [batch, len1, len2, 9] # Computes grad of paths w.r.t. sim_mat, gap_open and gap_extend. grad_paths = alignment.adjoint_weights_from_sim_mat( grad_paths, gap_open.shape, gap_extend.shape) # Adds gradients w.r.t. scores and gradients w.r.t. paths. gradients = tf.nest.map_structure( lambda x, y: x + y, grad_scores, grad_paths) return gradients return (pert_scores, pert_paths), grad
def test_smith_waterman_empty(self, decorator): smith_waterman_fn = tf_ops.hard_sw_affine if decorator is not None: smith_waterman_fn = decorator(smith_waterman_fn) tol = 1e-6 single_sub = - 5*tf.ones((5, 5)) toy_sub = tf.expand_dims(single_sub, 0) toy_gap_open = 0.03 * tf.ones((toy_sub.shape[0],)) toy_gap_extend = 0.02 * tf.ones((toy_sub.shape[0],)) w = alignment.weights_from_sim_mat(toy_sub, toy_gap_open, toy_gap_extend) values, paths = smith_waterman_fn(w, tol=tol) paths_squeeze = alignment.path_label_squeeze(paths) self.assertAllClose(values, [0.0], atol=2 * tol) self.assertAllEqual(paths_squeeze, tf.zeros([1, 5, 5], tf.float32))
def setUp(self): super().setUp() single_sub = 6*tf.eye(5) - 5*tf.ones((5, 5)) second_sub = tf.tensor_scatter_nd_update(single_sub, indices=[[0, 0], [1, 1]], updates=[-5, -5]) third_sub = tf.tensor_scatter_nd_update(single_sub, indices=[[0, 0], [2, 2]], updates=[-5, -5]) fourth_sub = tf.tensor_scatter_nd_update(single_sub, indices=[[0, 0], [4, 4]], updates=[-5, -5]) self._toy_sub = tf.stack([single_sub, second_sub, third_sub, fourth_sub]) self._toy_gap_open = 0.03 * tf.ones((self._toy_sub.shape[0],)) self._toy_gap_extend = 0.02 * tf.ones((self._toy_sub.shape[0],)) self._weights = alignment.weights_from_sim_mat(self._toy_sub, self._toy_gap_open, self._toy_gap_extend)
def test_smith_waterman_termination(self, decorator): smith_waterman_fn = tf_ops.hard_sw_affine if decorator is not None: smith_waterman_fn = decorator(smith_waterman_fn) tol = 1e-6 single_sub = tf.concat([- 5*tf.ones((3, 1)), 6*tf.eye(3) - 5*tf.ones((3, 3))], 1) toy_sub = tf.expand_dims(single_sub, 0) toy_gap_open = 0.03 * tf.ones((toy_sub.shape[0],)) toy_gap_extend = 0.02 * tf.ones((toy_sub.shape[0],)) w = alignment.weights_from_sim_mat(toy_sub, toy_gap_open, toy_gap_extend) values, paths = smith_waterman_fn(w, tol=tol) paths_squeeze = alignment.path_label_squeeze(paths) self.assertAllClose(values, [3.0], atol=2 * tol) paths_squeeze_test = tf.convert_to_tensor([[[0., 1., 0., 0.], [0., 0., 2., 0.], [0., 0., 0., 2.]]], tf.float32) self.assertAllEqual(paths_squeeze, paths_squeeze_test)
def unperturbed_alignment_score(sim_mat, gap_open, gap_extend): """Noiseless alignment score and paths from Smith-Waterman parameters.""" sw_params = alignment.weights_from_sim_mat(sim_mat, gap_open, gap_extend) return hard_sw_affine(sw_params)