CNAM-410: Add start delay for LimitedExposurePeriodAdder strategy

Author: Youcef Sebiat

src/main/scala/fr/polytechnique/cmap/cnam/etl/transformers/exposures/ExposuresTransformer.scala

@@ -56,6 +56,7 @@ class ExposuresTransformer(config: ExposuresTransformerConfig)
       when(col(FollowUpEnd) < col(ExposureEnd), col(FollowUpEnd)).otherwise(col(ExposureEnd))
     ).drop(ExposureEnd) // This makes sure that all the exposures end at the followup end date
       .withColumnRenamed("Correct_Exposure_End", ExposureEnd)
+      .where(col(ExposureStart) <= col(ExposureEnd))
       .aggregateWeight(
         cumulativeExposureWindow,
         cumulativeStartThreshold,

src/main/scala/fr/polytechnique/cmap/cnam/etl/transformers/exposures/LimitedExposurePeriodAdder.scala

@@ -22,13 +22,13 @@ private class LimitedExposurePeriodAdder(data: DataFrame) extends ExposurePeriod
     *   A. The first DrugPurchase defines a new Exposure.
     *   B. If there is a DrugPurchase within the defined window of the first DrugPurchase, then expand the current
     * Exposure with the DrugPurchase.
-    *   C. Else, close and set the end of the Exposure as the reach of the current and create a new Exposure with the
-    * DrugPurchase as the new Exposure.
+    *   C. Else, close and set the end of the Exposure as the reach of the latest Drug Purchase and create a new
+    *   Exposure with the next DrugPurchase as the new Exposure.
     * This strategy is suited for short term effects.
     * !!! WARNING: THIS ONLY RETURNS EXPOSURES.
     *
     * @param minPurchases : Not used.
-    * @param startDelay : Not used.
+    * @param  startDelay : period to be added to the start of each DrugPurchase.
     * @param purchasesWindow : Not used.
     * @param endThresholdGc : the period that defines the reach for Grand Condtionnement.
     * @param endThresholdNgc : the period that defines the reach for Non Grand Condtionnement.
@@ -37,7 +37,7 @@ private class LimitedExposurePeriodAdder(data: DataFrame) extends ExposurePeriod
     */
   def withStartEnd(
     minPurchases: Int = 2,
-    startDelay: Period = 3.months,
+    startDelay: Period = 5.days,
     purchasesWindow: Period = 4.months,
     endThresholdGc: Option[Period] = Some(120.days),
     endThresholdNgc: Option[Period] = Some(40.days),
@@ -46,11 +46,19 @@ private class LimitedExposurePeriodAdder(data: DataFrame) extends ExposurePeriod
 
     val outputColumns = (data.columns.toList ++ List(ExposureStart, ExposureEnd)).map(col)
 
-    val firstLastPurchase = getFirstAndLastPurchase(data, endThresholdGc.get, endThresholdNgc.get, endDelay.get)
+    val delayedDrugPurchases = delayStart(data, startDelay)
+
+    val firstLastPurchase = getFirstAndLastPurchase(delayedDrugPurchases, endThresholdGc.get, endThresholdNgc.get, endDelay.get)
 
     toExposure(firstLastPurchase).select(outputColumns: _*)
   }
 
+  def delayStart(data: DataFrame, startDelay: Period): DataFrame = {
+    data.withColumn("NewStart", col("start").addPeriod(startDelay))
+      .drop(col("start"))
+      .withColumnRenamed("NewStart", "start")
+  }
+
   def toExposure(firstLastPurchase: DataFrame): DataFrame = {
     val condition = (col("Status") === "first"
       && lead(col("Status"), 1).over(orderedWindow) === "last")

src/test/scala/fr/polytechnique/cmap/cnam/etl/transformers/exposures/LimitedExposurePeriodAdderSuite.scala

@@ -12,6 +12,42 @@ class LimitedExposurePeriodAdderSuite extends SharedContext {
   // instance created from a mock DataFrame, to allow testing the InnerImplicits implicit class
   private val mockInstance = new LimitedExposurePeriodAdder(mock(classOf[DataFrame]))
 
+
+  "delayStart" should "delay the start of events by the given period" in {
+    val sqlCtx = sqlContext
+    import sqlCtx.implicits._
+
+    // Given
+    val input = Seq(
+      ("A", "P", makeTS(2008, 1, 1), 1.0),
+      ("A", "P", makeTS(2008, 2, 1), 1.0),
+      ("A", "P", makeTS(2008, 5, 1), 1.0),
+      ("A", "P", makeTS(2008, 6, 1), 1.0),
+      ("A", "P", makeTS(2008, 7, 1), 1.0),
+      ("A", "P", makeTS(2009, 1, 1), 1.0),
+      ("A", "P", makeTS(2009, 7, 1), 1.0),
+      ("A", "P", makeTS(2009, 8, 1), 1.0),
+      ("A", "S", makeTS(2008, 2, 1), 1.0)
+    ).toDF(PatientID, Value, Start, Weight)
+
+
+    val expected = Seq(
+      ("A", "P", makeTS(2008, 1, 11), 1.0),
+      ("A", "P", makeTS(2008, 2, 11), 1.0),
+      ("A", "P", makeTS(2008, 5, 11), 1.0),
+      ("A", "P", makeTS(2008, 6, 11), 1.0),
+      ("A", "P", makeTS(2008, 7, 11), 1.0),
+      ("A", "P", makeTS(2009, 1, 11), 1.0),
+      ("A", "P", makeTS(2009, 7, 11), 1.0),
+      ("A", "P", makeTS(2009, 8, 11), 1.0),
+      ("A", "S", makeTS(2008, 2, 11), 1.0)
+    ).toDF(PatientID, Value, Start, Weight).select(PatientID, Value, Weight, Start)
+
+    val result = mockInstance.delayStart(input, 10.days)
+
+    assertDFs(expected, result)
+  }
+
   "getFirstAndLastPurchase" should "return the first and the last purchase of each potential exposure" in {
 
     val sqlCtx = sqlContext
@@ -93,15 +129,15 @@ class LimitedExposurePeriodAdderSuite extends SharedContext {
     ).toDF(PatientID, Value, Start, Weight)
 
     val expected = Seq(
-      ("A", "P", makeTS(2008, 1, 1), 1.0, makeTS(2008, 1, 1), makeTS(2008, 3, 11)),
-      ("A", "P", makeTS(2008, 5, 1), 1.0, makeTS(2008, 5, 1), makeTS(2008, 8, 11)),
-      ("A", "P", makeTS(2009, 1, 1), 1.0, makeTS(2009, 1, 1), makeTS(2009, 2, 11)),
-      ("A", "P", makeTS(2009, 7, 1), 1.0, makeTS(2009, 7, 1), makeTS(2009, 9, 11)),
-      ("A", "S", makeTS(2008, 2, 1), 1.0, makeTS(2008, 2, 1), makeTS(2008, 3, 11))
+      ("A", "P", makeTS(2008, 1, 6), 1.0, makeTS(2008, 1, 6), makeTS(2008, 3, 16)),
+      ("A", "P", makeTS(2008, 5, 6), 1.0, makeTS(2008, 5, 6), makeTS(2008, 8, 16)),
+      ("A", "P", makeTS(2009, 1, 6), 1.0, makeTS(2009, 1, 6), makeTS(2009, 2, 16)),
+      ("A", "P", makeTS(2009, 7, 6), 1.0, makeTS(2009, 7, 6), makeTS(2009, 9, 16)),
+      ("A", "S", makeTS(2008, 2, 6), 1.0, makeTS(2008, 2, 6), makeTS(2008, 3, 16))
     ).toDF(PatientID, Value, Start, Weight, ExposureStart, ExposureEnd)
 
     val result = new LimitedExposurePeriodAdder(input)
-      .withStartEnd(0, 0.months, 0.months, Some(1.months), Some(3.months), Some(10.days))
+      .withStartEnd(0, 5.days, 0.months, Some(1.months), Some(3.months), Some(10.days))
 
     assertDFs(expected, result)
   }