fix tvar

Quafadas · Jul 26, 2024 · 70178d9 · 70178d9
1 parent fc798e5
commit 70178d9
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/vecxt/src/tvar.scala b/vecxt/src/tvar.scala
@@ -27,7 +27,7 @@ extension [N <: Int](thisVector: NArray[Double])
     def tVar(alpha: Double): Double =
       val numYears = thisVector.length
       val nte = numYears * (1.0 - alpha);
-      val fte = Math.floor(nte).toInt;
+      val fte = Math.floor(nte+0.001).toInt;
       val sorted = thisVector.sort()
       var i = 0
       var tailSum = 0.0;
@@ -41,7 +41,7 @@ extension [N <: Int](thisVector: NArray[Double])
     def tVarIdx(alpha: Double): NArray[Boolean] =
       val numYears = thisVector.length
       val nte = numYears * (1.0 - alpha);
-      val fte = Math.floor(nte).toInt;
+      val fte = Math.floor(nte+0.001).toInt;
       val sorted = thisVector.toVector.zipWithIndex.sortBy(_._1).map(_._2)
       val idx = NArray.fill[Boolean](numYears)(false)
       var i = 0

diff --git a/vecxt/test/src/arrayExtensions.test.scala b/vecxt/test/src/arrayExtensions.test.scala
@@ -190,5 +190,25 @@ class ArrayExtensionSuite extends munit.FunSuite:
     end for
   }
 
+    test("tvar index 2") {
+    import vecxt.rpt.tVarIdx
+    val v1 = NArray.from(Array(6.0, 2.0, 5.0, 5.0, 10.0, 1.0, 2.0, 3.0, 5.0, 8.0))
+    val tvar = v1.tVarIdx(0.90)
+
+    assertEquals(tvar.countTrue, 1)
+    assert(tvar(5))
+  }
+
+
+    test("tvar index 3") {
+    import vecxt.rpt.tVarIdx
+    val v1 = NArray.from(Array(6.0, 2.0, 5.0, 5.0, 10.0, 1.0, 2.0, 3.0, 5.0, 8.0))
+    val tvar = v1.tVarIdx(0.8)
+
+    assertEquals(tvar.countTrue, 1)
+    assert(tvar(5))
+    assert(tvar(6))
+  }
+
 
 end ArrayExtensionSuite