From 2febb2e88c5cafe990fb40a327ea787b6fdee187 Mon Sep 17 00:00:00 2001
From: angerhang <514140809@qq.com>
Date: Wed, 17 Apr 2024 00:35:24 +0100
Subject: [PATCH] fix visu sleep

---
 pyproject.toml           |  2 +-
 setup.cfg                |  1 +
 src/asleep/visu_sleep.py | 11 ++++++++++-
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 1b332ff..42dc6ad 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,7 +17,7 @@ name = "asleep"  # Required
 #
 # For a discussion on single-sourcing the version, see
 # https://packaging.python.org/guides/single-sourcing-package-version/
-version = "0.4.9"  # Required
+version = "0.4.10"  # Required
 
 # This is a one-line description or tagline of what your project does. This
 # corresponds to the "Summary" metadata field:
diff --git a/setup.cfg b/setup.cfg
index a291c16..cd93cb5 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,3 +1,4 @@
 [options.entry_points]
 console_scripts =
     get_sleep = asleep.get_sleep:main
+    visu_sleep = asleep.visu_sleep:main
diff --git a/src/asleep/visu_sleep.py b/src/asleep/visu_sleep.py
index c2c4ab9..9f3a222 100644
--- a/src/asleep/visu_sleep.py
+++ b/src/asleep/visu_sleep.py
@@ -37,7 +37,7 @@ def main():
     # check all the files are there
     prediction_path = os.path.join(args.filepath, 'predictions.csv')
     sleep_blocks_path = os.path.join(args.filepath, 'sleep_block.csv')
-    raw_x_path = os.path.join('outputs/test', 'data2model.npy')
+    raw_x_path = os.path.join(args.filepath, 'data2model.npy')
     plot_path = os.path.join(args.filepath, 'sleep.png')
 
     if not os.path.exists(prediction_path):
@@ -62,6 +62,8 @@ def main():
     exact_date_df = sleep_block_df.loc[sleep_block_df['is_longest_block']].copy()
     exact_date_df['interval_start'] = pd.to_datetime(exact_date_df['interval_start'])
     exact_date_df['interval_end'] = pd.to_datetime(exact_date_df['interval_end'])
+    exact_date_df['start'] = pd.to_datetime(exact_date_df['start'])
+    exact_date_df['end'] = pd.to_datetime(exact_date_df['end'])
 
     y_df['enmo'] = raw_enmo
     y_df['time'] = pd.to_datetime(y_df['time'])
@@ -81,6 +83,8 @@ def main():
     for _, row in exact_date_df.iterrows():
         day_start = row['interval_start']
         day_end = row['interval_end']
+        long_sleep_start = row['start']
+        long_sleep_end = row['end']
         day_df = y_df[(y_df['time'] >= day_start) & (y_df['time'] <= day_end)].copy()
         day_df['is_nonwear'] = day_df['enmo'] == 0
         day_df.loc[day_df['is_nonwear'], 'sleep_wake'] = 'non_wear'
@@ -95,6 +99,11 @@ def main():
         ax.set_ylim(0, 2)
         ax.set_xlim(day_start, day_end)
 
+        plt.vlines(x=long_sleep_start, ymin=0, ymax=2,
+                   colors='yellow')
+        plt.vlines(x=long_sleep_end, ymin=0, ymax=2,
+                   colors='yellow')
+
         a = np.array([is_sleep, is_wake, is_nonwear]) * MAXRANGE
 
         ax.stackplot(day_df['time'], a,