X-Git-Url: https://git.arvados.org/lightning.git/blobdiff_plain/b513c8156ae1a20207479cab714a91733187a4fd..29666c930fc3f6f81edc7eb306c1cb2765e77c58:/glm.go

diff --git a/glm.go b/glm.go
index cc06f39e2c..2f2327479b 100644
--- a/glm.go
+++ b/glm.go
@@ -18,8 +18,12 @@ var glmConfig = &glm.Config{
 	ConcurrentIRLS: 1000,
 }
 
-func pvalueGLM(sampleInfo []sampleInfo, onehotPair [][]bool) float64 {
-	nPCA := len(sampleInfo[0].pcaComponents)
+// Logistic regression.
+//
+// onehot is the observed outcome, in same order as sampleInfo, but
+// shorter because it only has entries for samples with
+// isTraining==true.
+func pvalueGLM(sampleInfo []sampleInfo, onehot []bool, nPCA int) float64 {
 	pcaNames := make([]string, 0, nPCA)
 	data := make([][]statmodel.Dtype, 0, nPCA)
 	for pca := 0; pca < nPCA; pca++ {
@@ -35,9 +39,10 @@ func pvalueGLM(sampleInfo []sampleInfo, onehotPair [][]bool) float64 {
 
 	variant := make([]statmodel.Dtype, 0, len(sampleInfo))
 	outcome := make([]statmodel.Dtype, 0, len(sampleInfo))
-	for row, si := range sampleInfo {
+	row := 0
+	for _, si := range sampleInfo {
 		if si.isTraining {
-			if onehotPair[0][row] {
+			if onehot[row] {
 				variant = append(variant, 1)
 			} else {
 				variant = append(variant, 0)
@@ -47,6 +52,7 @@ func pvalueGLM(sampleInfo []sampleInfo, onehotPair [][]bool) float64 {
 			} else {
 				outcome = append(outcome, 0)
 			}
+			row++
 		}
 	}
 	data = append(data, variant, outcome)