65 lines
2 KiB
Python
65 lines
2 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
#
|
|
# Copyright (C) 2011 Radim Rehurek <radimrehurek@seznam.cz>
|
|
# Licensed under the GNU LGPL v2.1 - http://www.gnu.org/licenses/lgpl.html
|
|
|
|
"""
|
|
Automated tests for segmentation algorithms in the segmentation module.
|
|
"""
|
|
|
|
|
|
import logging
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
from gensim.topic_coherence import segmentation
|
|
from numpy import array
|
|
|
|
|
|
class TestSegmentation(unittest.TestCase):
|
|
def setUp(self):
|
|
self.topics = [
|
|
array([9, 4, 6]),
|
|
array([9, 10, 7]),
|
|
array([5, 2, 7])
|
|
]
|
|
|
|
def testSOnePre(self):
|
|
"""Test s_one_pre segmentation."""
|
|
actual = segmentation.s_one_pre(self.topics)
|
|
expected = [
|
|
[(4, 9), (6, 9), (6, 4)],
|
|
[(10, 9), (7, 9), (7, 10)],
|
|
[(2, 5), (7, 5), (7, 2)]
|
|
]
|
|
self.assertTrue(np.allclose(actual, expected))
|
|
|
|
def testSOneOne(self):
|
|
"""Test s_one_one segmentation."""
|
|
actual = segmentation.s_one_one(self.topics)
|
|
expected = [
|
|
[(9, 4), (9, 6), (4, 9), (4, 6), (6, 9), (6, 4)],
|
|
[(9, 10), (9, 7), (10, 9), (10, 7), (7, 9), (7, 10)],
|
|
[(5, 2), (5, 7), (2, 5), (2, 7), (7, 5), (7, 2)]
|
|
]
|
|
self.assertTrue(np.allclose(actual, expected))
|
|
|
|
def testSOneSet(self):
|
|
"""Test s_one_set segmentation."""
|
|
actual = segmentation.s_one_set(self.topics)
|
|
expected = [
|
|
[(9, array([9, 4, 6])), (4, array([9, 4, 6])), (6, array([9, 4, 6]))],
|
|
[(9, array([9, 10, 7])), (10, array([9, 10, 7])), (7, array([9, 10, 7]))],
|
|
[(5, array([5, 2, 7])), (2, array([5, 2, 7])), (7, array([5, 2, 7]))]
|
|
]
|
|
for s_i in range(len(actual)):
|
|
for j in range(len(actual[s_i])):
|
|
self.assertEqual(actual[s_i][j][0], expected[s_i][j][0])
|
|
self.assertTrue(np.allclose(actual[s_i][j][1], expected[s_i][j][1]))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
logging.root.setLevel(logging.WARNING)
|
|
unittest.main()
|