def test_toy_smith_waterman(self, decorator, tol):
        smith_waterman_fn = tf_ops.hard_sw_affine
        if decorator is not None:
            smith_waterman_fn = decorator(smith_waterman_fn)

        values, paths = smith_waterman_fn(self._w, tol)
        paths_squeeze = alignment.path_label_squeeze(paths)
        all_matches = tf.where(
            alignment.paths_to_state_indicators(paths, 'match'))

        values_test = tf.constant([5.0, 3.0, 2.94, 3.0], dtype=tf.float32)
        self.assertAllClose(values, values_test, atol=2 * tol)

        paths_squeeze_test = tf.constant(
            [[[1., 0., 0., 0., 0.], [0., 2., 0., 0., 0.], [0., 0., 2., 0., 0.],
              [0., 0., 0., 2., 0.], [0., 0., 0., 0., 2.]],
             [[0., 0., 0., 0., 0.], [0., 0., 0., 0., 0.], [0., 0., 1., 0., 0.],
              [0., 0., 0., 2., 0.], [0., 0., 0., 0., 2.]],
             [[0., 0., 0., 0., 0.], [0., 1., 5., 0., 0.], [0., 0., 8., 0., 0.],
              [0., 0., 0., 4., 0.], [0., 0., 0., 0., 2.]],
             [[0., 0., 0., 0., 0.], [0., 1., 0., 0., 0.], [0., 0., 2., 0., 0.],
              [0., 0., 0., 2., 0.], [0., 0., 0., 0., 0.]]],
            dtype=tf.float32)
        self.assertAllEqual(paths_squeeze, paths_squeeze_test)

        all_matches_test = tf.constant(
            [[0, 0, 0], [0, 1, 1], [0, 2, 2], [0, 3, 3], [0, 4, 4], [1, 2, 2],
             [1, 3, 3], [1, 4, 4], [2, 1, 1], [2, 3, 3], [2, 4, 4], [3, 1, 1],
             [3, 2, 2], [3, 3, 3]],
            dtype=tf.int32)
        self.assertAllEqual(all_matches, all_matches_test)
    def test_alignments_to_paths(self):
        exp_out = tf.convert_to_tensor(
            [[[1., 0., 0., 0., 0., 0., 0.], [0., 2., 5., 0., 0., 0., 0.],
              [0., 0., 0., 3., 0., 0., 0.], [0., 0., 0., 7., 0., 0., 0.],
              [0., 0., 0., 9., 0., 0., 0.], [0., 0., 0., 0., 4., 0., 0.],
              [0., 0., 0., 0., 0., 2., 0.], [0., 0., 0., 0., 0., 0., 2.]]],
            tf.float32)

        paths = alignment.alignments_to_paths(self.alignments, self.len_x,
                                              self.len_y)
        sq_paths = alignment.path_label_squeeze(paths)
        self.assertAllEqual(sq_paths, exp_out)

        # Invariance to padding.
        padded_alignments = tf.concat(
            [self.alignments, tf.zeros([1, 3, 3], tf.int32)], 2)
        paths = alignment.alignments_to_paths(padded_alignments, self.len_x,
                                              self.len_y)
        sq_paths = alignment.path_label_squeeze(paths)
        self.assertAllEqual(sq_paths, exp_out)

        # Pads correctly via length arguments.
        paths = alignment.alignments_to_paths(self.alignments, self.len_x + 3,
                                              self.len_y + 3)
        sq_paths = alignment.path_label_squeeze(paths)
        self.assertAllEqual(sq_paths[:, :self.len_x, :self.len_y], exp_out)
        self.assertAllEqual(sq_paths[:, self.len_x:, :],
                            tf.zeros([1, 3, self.len_y + 3], tf.float32))
        self.assertAllEqual(sq_paths[Ellipsis, self.len_y:],
                            tf.zeros([1, self.len_x + 3, 3], tf.float32))

        # Deals with empty ground-truth alignments.
        paths = alignment.alignments_to_paths(tf.zeros_like(self.alignments),
                                              self.len_x, self.len_y)
        sq_paths = alignment.path_label_squeeze(paths)
        self.assertAllEqual(sq_paths, tf.zeros_like(exp_out))
  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 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)