From 93f136b08062243168054a579259f008a8a339d6 Mon Sep 17 00:00:00 2001
From: junos <junos.lukan@ijs.si>
Date: Thu, 19 Aug 2021 17:32:02 +0200
Subject: [PATCH] Add a method to get communication features.

---
 features/communication.py    | 4 +++-
 machine_learning/pipeline.py | 4 ++++
 test/test_communication.py   | 3 +--
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/features/communication.py b/features/communication.py
index 2fb6bc9..a858830 100644
--- a/features/communication.py
+++ b/features/communication.py
@@ -37,7 +37,7 @@ FEATURES_SMS = (
 #  "no_received_ratio", "no_sent_ratio",
 #  "no_contacts"]
 
-FEATURES_CONTACT = [
+FEATURES_CALLS_SMS_PROP = [
     "proportion_calls_all",
     "proportion_calls_incoming",
     "proportion_calls_outgoing",
@@ -45,6 +45,8 @@ FEATURES_CONTACT = [
     "proportion_calls_missed_sms_received",
 ]
 
+FEATURES_CALLS_SMS_ALL = FEATURES_CALLS + FEATURES_SMS + FEATURES_CALLS_SMS_PROP
+
 
 def get_call_data(usernames: Collection) -> pd.DataFrame:
     """
diff --git a/machine_learning/pipeline.py b/machine_learning/pipeline.py
index 8f91f28..f5c35ae 100644
--- a/machine_learning/pipeline.py
+++ b/machine_learning/pipeline.py
@@ -69,6 +69,10 @@ class SensorFeatures:
             if feature_names == "all":
                 feature_names = proximity.FEATURES_PROXIMITY
             return self.df_proximity_counts[feature_names]
+        elif data_type == "communication":
+            if feature_names == "all":
+                feature_names = communication.FEATURES_CALLS_SMS_ALL
+            return self.df_calls_sms[feature_names]
         else:
             raise KeyError("This data type has not been implemented.")
 
diff --git a/test/test_communication.py b/test/test_communication.py
index 5258a66..2f4f0df 100644
--- a/test/test_communication.py
+++ b/test/test_communication.py
@@ -88,6 +88,5 @@ class CallsFeatures(unittest.TestCase):
         self.features_call_sms = calls_sms_features(self.calls, self.sms)
         self.assertIsInstance(self.features_call_sms, pd.DataFrame)
         self.assertCountEqual(
-            self.features_call_sms.columns.to_list(),
-            FEATURES_CALLS + FEATURES_SMS + FEATURES_CONTACT,
+            self.features_call_sms.columns.to_list(), FEATURES_CALLS_SMS_ALL
         )