新增计算曲线相似度的评价函数

Signed-off-by: allan716 <525223688@qq.com>
ChineseSubFinder · Nov 8, 2021 · 5db8c3e · 5db8c3e
1 parent c5bb0de
commit 5db8c3e
Show file tree

Hide file tree

Showing 2 changed files with 112 additions and 0 deletions.
diff --git a/internal/pkg/sub_timeline_fixer/calculate_curve_correlation.go b/internal/pkg/sub_timeline_fixer/calculate_curve_correlation.go
@@ -0,0 +1,45 @@
+package sub_timeline_fixer
+
+import "math"
+
+// CalculateCurveCorrelation 计算两条曲线的相似度
+// 返回值在区间： [-1,1]
+// 如返回-10，则证明输入参数无效
+// 原始出处不详，《数学之美：判定两个随机信号序列的相似度》
+func CalculateCurveCorrelation(s1, s2 []float64, n int) float64 {
+	var sum_s12 = 0.0
+	var sum_s1 = 0.0
+	var sum_s2 = 0.0
+	var sum_s1s1 = 0.0 //s1^2
+	var sum_s2s2 = 0.0 //s2^2
+	var pxy = 0.0
+	var temp1 = 0.0
+	var temp2 = 0.0
+
+	if s1 == nil || s2 == nil || n <= 0 {
+		return -10
+	}
+
+	for i := 0; i < n; i++ {
+		sum_s12 += s1[i] * s2[i]
+		sum_s1 += s1[i]
+		sum_s2 += s2[i]
+		sum_s1s1 += s1[i] * s1[i]
+		sum_s2s2 += s2[i] * s2[i]
+	}
+
+	temp1 = float64(n)*sum_s1s1 - sum_s1*sum_s1
+	temp2 = float64(n)*sum_s2s2 - sum_s2*sum_s2
+
+	if (temp1 > -delta && temp1 < delta) ||
+		(temp2 > -delta && temp2 < delta) ||
+		(temp1*temp2 <= 0) {
+		return -10
+	}
+
+	pxy = (float64(n)*sum_s12 - sum_s1*sum_s2) / math.Sqrt(temp1*temp2)
+
+	return pxy
+}
+
+const delta = 0.0001
diff --git a/internal/pkg/sub_timeline_fixer/calculate_curve_correlation_test.go b/internal/pkg/sub_timeline_fixer/calculate_curve_correlation_test.go
@@ -0,0 +1,67 @@
+package sub_timeline_fixer
+
+import "testing"
+
+func TestCalculateCurveCorrelation(t *testing.T) {
+	type args struct {
+		s1 []float64
+		s2 []float64
+		n  int
+	}
+	tests := []struct {
+		name string
+		args args
+		want float64
+	}{
+		{name: "00", args: args{
+			s1: []float64{0.309016989, 0.587785244, 0.809016985, 0.95105651, 1, 0.951056526,
+				0.809017016, 0.587785287, 0.30901704, 5.35898e-08, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0},
+			s2: []float64{0.343282816, 0.686491368, 0.874624132, 0.99459642, 1.008448609,
+				1.014252458, 0.884609221, 0.677632906, 0.378334666, 0.077878732,
+				0.050711886, 0.066417083, 0.088759401, 0.005440732, 0.04225661,
+				0.035349939, 0.0631196, 0.007566056, 0.053183895, 0.073143706,
+				0.080285063, 0.030110227, 0.044781145, 0.01875573, 0.08373928,
+				0.04550342, 0.038880858, 0.040611891, 0.046116826, 0.087670453},
+			n: 30,
+		},
+			want: 0.99743484574875},
+		{name: "01", args: args{
+			s1: []float64{0.309016989, 0.587785244, 0.809016985, 0.95105651, 1, 0.951056526,
+				0.809017016, 0.587785287, 0.30901704, 5.35898e-08, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0},
+			s2: []float64{0.309016989, 0.587785244, 0.809016985, 0.95105651, 1, 0.951056526,
+				0.809017016, 0.587785287, 0.30901704, 5.35898e-08, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0},
+			n: 30,
+		},
+			want: 1},
+		{name: "02", args: args{
+			s1: []float64{0.309016989, 0.587785244, 0.809016985, 0.95105651, 1, 0.951056526,
+				0.809017016, 0.587785287, 0.30901704, 5.35898e-08, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0},
+			s2: []float64{-0.309016989, -0.587785244, -0.809016985, -0.95105651, -1, -0.951056526,
+				-0.809017016, -0.587785287, -0.30901704, -5.35898e-08, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0,
+				0, 0, 0, 0, 0, 0},
+			n: 30,
+		},
+			want: -1},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := CalculateCurveCorrelation(tt.args.s1, tt.args.s2, tt.args.n); got != tt.want {
+				t.Errorf("CalculateCurveCorrelation() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}