Correlation matrix should work on data series with nils

Author: jecisc

src/DataFrame-Math/DataPearsonCorrelationMethod.class.st

@@ -9,12 +9,10 @@ DataPearsonCorrelationMethod class >> between: x and: y [
 	"Calcualte the Pearson correlation coefficient between two data series"
 
 	| xDeviation yDeviation |
+	x size = y size ifFalse: [ SizeMismatch signal: 'Correlation can not be calculated for two series of different size' ].
 
-	x size = y size ifFalse: [
-		SizeMismatch signal: 'Correlation can not be calculated for two series of different size' ].
-
-	xDeviation := x - x average.
-	yDeviation := y - y average.
+	xDeviation := x replaceNilsWithZeros - x average.
+	yDeviation := y replaceNilsWithZeros - y average.
 
 	^ (xDeviation * yDeviation) sum / ((xDeviation ** 2) sum * (yDeviation ** 2) sum) sqrt
 ]

src/DataFrame-Tests/DataFrameStatsTest.class.st

@@ -37,28 +37,36 @@ DataFrameStatsTest >> testCorrelationMatrix [
 
 	| expectedCorrelationMatrix actualCorrelationMatrix |
 	expectedCorrelationMatrix := DataFrame withRows:
-		                             #( #( 1 0.311398 0.538922 0.454601 )
-		                                #( 0.311398 1 -0.321281 -0.308023 )
-		                                #( 0.538922 -0.321281 1 0.982956 )
-		                                #( 0.454601 -0.308023 0.982956
-		                                   1 ) ).
-	expectedCorrelationMatrix columnNames:
-		#( sepalLength sepalWidth petalLength petalWidth ).
-	expectedCorrelationMatrix rowNames:
-		#( sepalLength sepalWidth petalLength petalWidth ).
+		                             #( #( 1 0.311398 0.538922 0.454601 ) #( 0.311398 1 -0.321281 -0.308023 ) #( 0.538922 -0.321281 1 0.982956 )
+		                                #( 0.454601 -0.308023 0.982956 1 ) ).
+	expectedCorrelationMatrix columnNames: #( sepalLength sepalWidth petalLength petalWidth ).
+	expectedCorrelationMatrix rowNames: #( sepalLength sepalWidth petalLength petalWidth ).
 	actualCorrelationMatrix := df correlationMatrix.
-	self
-		assert: actualCorrelationMatrix rowNames
-		equals: expectedCorrelationMatrix rowNames.
-	self
-		assert: actualCorrelationMatrix columnNames
-		equals: expectedCorrelationMatrix columnNames.
+	self assert: actualCorrelationMatrix rowNames equals: expectedCorrelationMatrix rowNames.
+	self assert: actualCorrelationMatrix columnNames equals: expectedCorrelationMatrix columnNames.
 
 	1 to: actualCorrelationMatrix numberOfColumns do: [ :j |
-		1 to: actualCorrelationMatrix numberOfRows do: [ :i |
-			self
-				assert: (actualCorrelationMatrix at: i at: j)
-				closeTo: (expectedCorrelationMatrix at: i at: j) ] ]
+		1 to: actualCorrelationMatrix numberOfRows do: [ :i | self assert: (actualCorrelationMatrix at: i at: j) closeTo: (expectedCorrelationMatrix at: i at: j) ] ]
+]
+
+{ #category : #tests }
+DataFrameStatsTest >> testCorrelationMatrixWithNils [
+
+	| expectedCorrelationMatrix actualCorrelationMatrix |
+	df := DataFrame withRows: #( #( 1 1 ) #( 2 nil ) #( nil 3 ) #( 4 4 ) ).
+	df columnNames: #( dogs cats ).
+
+	expectedCorrelationMatrix := DataFrame withRows: #( #( 1 0.3207134902949093 ) #( 0.3207134902949093 1 ) ).
+	expectedCorrelationMatrix columnNames: #( dogs cats ).
+	expectedCorrelationMatrix rowNames: #( dogs cats ).
+
+	actualCorrelationMatrix := df correlationMatrix.
+
+	self assert: actualCorrelationMatrix rowNames equals: expectedCorrelationMatrix rowNames.
+	self assert: actualCorrelationMatrix columnNames equals: expectedCorrelationMatrix columnNames.
+
+	1 to: actualCorrelationMatrix numberOfColumns do: [ :j |
+		1 to: actualCorrelationMatrix numberOfRows do: [ :i | self assert: (actualCorrelationMatrix at: i at: j) closeTo: (expectedCorrelationMatrix at: i at: j) ] ]
 ]
 
 { #category : #tests }