From 0d342d8e78bc34a23e1070424067f575b33bb6a3 Mon Sep 17 00:00:00 2001
From: Ignacio Spiousas <ignaciospiousas@gmail.com>
Date: Tue, 30 Jul 2024 05:53:47 -0700
Subject: [PATCH] Added mean temperature and light report at predictions.csv
 (#53)

Thanks @spiousas for making this PR for adding light output from asleep.
---
 src/asleep/get_sleep.py | 60 +++++++++++++++++++++++++++++++++++------
 1 file changed, 52 insertions(+), 8 deletions(-)

diff --git a/src/asleep/get_sleep.py b/src/asleep/get_sleep.py
index 646115e..bd1ee39 100644
--- a/src/asleep/get_sleep.py
+++ b/src/asleep/get_sleep.py
@@ -4,6 +4,7 @@
 import pandas as pd
 import json
 import os
+import sys
 import joblib
 import urllib
 import shutil
@@ -149,6 +150,26 @@ def transform_data2model_input(
     return data2model, times, non_wear
 
 
+def mean_temp_and_light(data):
+    # It stops processing if data does not include temperature and light columns
+    if not {'temperature', 'light'}.issubset(data.columns):
+        sys.exit('There is no temperature and light columns in the raw data.')
+
+    # Calculates mean temperature and light for each 30s window
+    print("Calculating mean temperature and light for each 30s interval.")
+
+    data['time'] = pd.to_datetime(data['time'])
+    # Grouping by time
+    grouped = data.groupby(pd.Grouper(key='time', freq='30S', origin=data['time'].min()))
+    # Calculating the mean temperature and light for each group
+    mean_values = grouped[['temperature', 'light']].mean().iloc[:-1]
+
+    temp = mean_values["temperature"].to_numpy()
+    light = mean_values["light"].to_numpy()
+
+    return temp, light
+
+
 def get_sleep_windows(data2model, times, non_wear, args):
     # data2model: N x 3 x 900
     # non_wear_flag: N x 1
@@ -240,6 +261,10 @@ def main():
         action="store_true",
         help="Remove intermediate files to save space but it "
              "will take longer to run the next time.")
+    parser.add_argument(
+        "--report_light_and_temp",
+        action="store_true",
+        help="If true, it adds mean temp. and light columns to the predictions.csv file.")
     parser.add_argument(
         "--pytorch_device",
         "-d",
@@ -305,6 +330,12 @@ def main():
     print("times shape: {}".format(times.shape))
     print("Non_wear flag shape: {}".format(non_wear.shape))
 
+    # 1.2 Get the mean temperature and light (optional)
+    if args.report_light_and_temp:
+        temp, light = mean_temp_and_light(data)
+        print("temperature shape: {}".format(temp.shape))
+        print("light shape: {}".format(light.shape))
+
     # times and non-wear flag need to be stored for visualization
     if args.remove_intermediate_files:
         os.remove(data2model_path)
@@ -346,14 +377,27 @@ def main():
     sleep_stage_predictions = np.vectorize(
         SLEEPNET_THRE_CLASS_LABELS.get)(sleepnet_output)
 
-    predictions_df = pd.DataFrame(
-        {
-            'time': times,
-            'sleep_wake': sleep_wake_predictions,
-            'sleep_stage': sleep_stage_predictions,
-            'raw_label': sleepnet_output,
-        }
-    )
+    if args.report_light_and_temp:
+        predictions_df = pd.DataFrame(
+            {
+                'time': times,
+                'sleep_wake': sleep_wake_predictions,
+                'sleep_stage': sleep_stage_predictions,
+                'raw_label': sleepnet_output,
+                'temperature': temp,
+                'light': light,
+            }
+        )
+    else:
+        predictions_df = pd.DataFrame(
+            {
+                'time': times,
+                'sleep_wake': sleep_wake_predictions,
+                'sleep_stage': sleep_stage_predictions,
+                'raw_label': sleepnet_output,
+            }
+        )
+
     final_prediction_path = os.path.join(args.outdir, 'predictions.csv')
     print("predictions_df shape: {}".format(predictions_df.shape))
     print(predictions_df.head())