dendrogram.py 3.23 KB
Newer Older
Anthony Larcher's avatar
Anthony Larcher committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# -*- coding: utf-8 -*-
#
# This file is part of s4d.
#
# s4d is a python package for speaker diarization.
# Home page: http://www-lium.univ-lemans.fr/s4d/
#
# s4d is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of the License,
# or (at your option) any later version.
#
# s4d is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with s4d.  If not, see <http://www.gnu.org/licenses/>.


"""
Anthony Larcher's avatar
Anthony Larcher committed
23
Copyright 2014-2020 Sylvain Meignier
Anthony Larcher's avatar
Anthony Larcher committed
24
"""
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
25

Anthony Larcher's avatar
Anthony Larcher committed
26
import copy
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
27
28
import numpy as np
import scipy.cluster.hierarchy as hac
Anthony Larcher's avatar
Anthony Larcher committed
29
30
31

from matplotlib import pyplot as plot
from matplotlib import ticker
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
32

Sylvain Meignier's avatar
Sylvain Meignier committed
33

Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
34
def plot_dendrogram(merge, thr, size=(25,6), log=False):
Anthony Larcher's avatar
Anthony Larcher committed
35
36
37
38
39
40
41
42
    """

    :param merge:
    :param thr:
    :param size:
    :param log:
    :return:
    """
Sylvain Meignier's avatar
Sylvain Meignier committed
43
44
45
    minv = 0

    def my_formatter(x, pos):
Anthony Larcher's avatar
Anthony Larcher committed
46
47
48
49
50
51
        """

        :param x:
        :param pos:
        :return:
        """
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
52
53
54
        v = x
        if log:
            v = np.exp(v)
Sylvain Meignier's avatar
Sylvain Meignier committed
55
        v -= minv
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
56
57
58
59

        return "{:10.3f}".format(v)

    def link(merge, log):
Anthony Larcher's avatar
Anthony Larcher committed
60
61
62
63
64
65
        """

        :param merge:
        :param log:
        :return:
        """
66
        cluster_lst = list()
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
67
68
69
70
71
        idx = dict()
        qt = dict()
        k=0;
        while merge[k][0] < 0:
            name = merge[k][1]
72
            cluster_lst.append(name)
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
73
74
75
76
77
            idx[name] = k
            qt[name] = 1
            k += 1

        l = k
Sylvain Meignier's avatar
Sylvain Meignier committed
78
        links = np.zeros((l-1, 4))
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
79
80
81
82
83
84
85
86

        while k < len(merge):
            m = merge[k][0]
            name_i = merge[k][1]
            name_j = merge[k][2]
            v_ij = merge[k][3]

            qt[name_i] += 1
Sylvain Meignier's avatar
Sylvain Meignier committed
87
88
89
90
            links[m, 0] = idx[name_i]
            links[m, 1] = idx[name_j]
            links[m, 2] = v_ij
            links[m, 3] = qt[name_i]
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
91
92
93
94
95
            idx[name_i] = m+l
            idx[name_j] = -1

            k+= 1

Sylvain Meignier's avatar
Sylvain Meignier committed
96
97
        min = -np.min(links[:,2])+2
        links[:,2] += min
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
98
99

        if log:
Sylvain Meignier's avatar
Sylvain Meignier committed
100
            links[:,2] = np.log(links[:,2])
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
101

102
        return cluster_lst, links, min
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
103
104
105
106

    merge = copy.deepcopy(merge)

    plot.figure(figsize=size)
107
    cluster_list, link_mat, minv = link(merge, log)
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
108

Sylvain Meignier's avatar
Sylvain Meignier committed
109
    t=thr+minv
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
110
111
112
    if log:
        t = np.log(t)

113
    dendro_data = hac.dendrogram(link_mat, labels=cluster_list, color_threshold=t)
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
114
115
116
117
118
119
    for i, d in zip(dendro_data['icoord'], dendro_data['dcoord']):
        x = 0.5 * sum(i[1:3])
        y = d[1]
        v = y
        if log:
            v = np.exp(v)
Sylvain Meignier's avatar
Sylvain Meignier committed
120
        v -= minv
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
121
122
        plot.plot(x, y, 'o', c="b")
        plot.annotate("{:10.3f}".format(v), (x, y), xytext=(0, -5),
Anthony Larcher's avatar
Anthony Larcher committed
123
124
                      textcoords='offset points',
                      va='top', ha='center')
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
125
126
127

    plot.axhline(y=t, c='r')
    plot.annotate("{:10.3f}".format(thr), (0, t), xytext=(25, 25),
Anthony Larcher's avatar
Anthony Larcher committed
128
129
                  textcoords='offset points',
                  va='top', ha='center')
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
130
131

    ax = plot.gca()
Sylvain Meignier's avatar
Sylvain Meignier committed
132
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(my_formatter))
Sylvain Meignier's avatar
Origin  
Sylvain Meignier committed
133
134
135
136

    plot.plot()

    return link_mat, dendro_data