141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
from typing import List
|
|
|
|
import pandas as pd
|
|
|
|
from config.models import SMS, Call, Participant
|
|
from setup import db_engine, session
|
|
|
|
call_types = {1: "incoming", 2: "outgoing", 3: "missed"}
|
|
sms_types = {1: "received", 2: "sent"}
|
|
|
|
|
|
def get_call_data(usernames: List) -> pd.DataFrame:
|
|
"""
|
|
Read the data from the calls table and return it in a dataframe.
|
|
|
|
Parameters
|
|
----------
|
|
usernames: List
|
|
A list of usernames to put into the WHERE condition.
|
|
|
|
Returns
|
|
-------
|
|
df_calls: pd.DataFrame
|
|
A dataframe of call data.
|
|
"""
|
|
query_calls = (
|
|
session.query(Call, Participant.username)
|
|
.filter(Participant.id == Call.participant_id)
|
|
.filter(Participant.username.in_(usernames))
|
|
)
|
|
with db_engine.connect() as connection:
|
|
df_calls = pd.read_sql(query_calls.statement, connection)
|
|
return df_calls
|
|
|
|
|
|
def get_sms_data(usernames: List) -> pd.DataFrame:
|
|
"""
|
|
Read the data from the sms table and return it in a dataframe.
|
|
|
|
Parameters
|
|
----------
|
|
usernames: List
|
|
A list of usernames to put into the WHERE condition.
|
|
|
|
Returns
|
|
-------
|
|
df_sms: pd.DataFrame
|
|
A dataframe of call data.
|
|
"""
|
|
query_sms = (
|
|
session.query(SMS, Participant.username)
|
|
.filter(Participant.id == SMS.participant_id)
|
|
.filter(Participant.username.in_(usernames))
|
|
)
|
|
with db_engine.connect() as connection:
|
|
df_sms = pd.read_sql(query_sms.statement, connection)
|
|
return df_sms
|
|
|
|
|
|
def enumerate_contacts(comm_df: pd.DataFrame) -> pd.DataFrame:
|
|
"""
|
|
Count contacts (callers, senders) and enumerate them by their frequency.
|
|
|
|
Parameters
|
|
----------
|
|
comm_df: pd.DataFrame
|
|
A dataframe of calls or SMSes.
|
|
|
|
Returns
|
|
-------
|
|
comm_df: pd.DataFrame
|
|
The altered dataframe with the column contact_id, arranged by frequency.
|
|
"""
|
|
contact_counts = (
|
|
comm_df["trace"]
|
|
.value_counts(sort=True, ascending=False)
|
|
.to_frame(name="frequency")
|
|
)
|
|
# A frequency table of different traces (contacts).
|
|
contact_counts["contact_id"] = list(range(len(contact_counts.index)))
|
|
contact_code = contact_counts["contact_id"].to_dict()
|
|
# Create a dictionary translating traces into integers, enumerated by their frequency.
|
|
comm_df["contact_id"] = comm_df["trace"].astype("category")
|
|
# Transform to categorical data instead of a simple character column.
|
|
comm_df["contact_id"] = comm_df["contact_id"].cat.rename_categories(contact_code)
|
|
# Recode the contacts into integers from 0 to n_contacts, so that the first one is contacted the most often.
|
|
return comm_df
|
|
|
|
|
|
def count_comms(comm_df: pd.DataFrame) -> pd.DataFrame:
|
|
"""
|
|
Calculate frequencies (and duration) of messages (or calls), grouped by their types.
|
|
|
|
Parameters
|
|
----------
|
|
comm_df: pd.DataFrame
|
|
A dataframe of calls or SMSes.
|
|
|
|
Returns
|
|
-------
|
|
comm_features: pd.DataFrame
|
|
A list of communication features for every participant.
|
|
These are:
|
|
* the number of messages by type (received, sent),
|
|
* the number of calls by type (incoming, outgoing missed), and
|
|
* the duration of calls by type.
|
|
"""
|
|
if "call_type" in comm_df:
|
|
comm_counts = (
|
|
comm_df.value_counts(subset=["participant_id", "call_type"])
|
|
.unstack()
|
|
.rename(columns=call_types)
|
|
.add_prefix("no_")
|
|
)
|
|
comm_duration = (
|
|
comm_df.groupby(["participant_id", "call_type"])
|
|
.sum()["call_duration"]
|
|
.unstack()
|
|
.rename(columns=call_types)
|
|
.add_prefix("duration_")
|
|
)
|
|
comm_features = comm_counts.join(comm_duration)
|
|
try:
|
|
comm_features.drop(columns="duration_" + call_types[3], inplace=True)
|
|
# The missed calls are always of 0 duration.
|
|
except KeyError:
|
|
pass
|
|
# If there were no missed calls, this exception is raised.
|
|
# But we are dropping the column anyway, so no need to deal with the exception.
|
|
elif "message_type" in comm_df:
|
|
comm_counts = (
|
|
comm_df.value_counts(subset=["participant_id", "message_type"])
|
|
.unstack()
|
|
.rename(columns=sms_types)
|
|
.add_prefix("no_")
|
|
)
|
|
comm_features = comm_counts
|
|
else:
|
|
raise KeyError("The dataframe contains neither call_type or message_type")
|
|
return comm_features
|