From 4a03342dadb7aa6d5245cb24b3759c528ed6faa7 Mon Sep 17 00:00:00 2001 From: Ben Lansdell Date: Fri, 2 Feb 2024 14:35:47 -0700 Subject: [PATCH] More type hints (#12) --- ethome/features/generic_features.py | 12 +++++----- ethome/features/mars_features.py | 37 +++++++++++++++-------------- ethome/io.py | 17 ++++++------- 3 files changed, 34 insertions(+), 32 deletions(-) diff --git a/ethome/features/generic_features.py b/ethome/features/generic_features.py index fdcfa97..2be630e 100644 --- a/ethome/features/generic_features.py +++ b/ethome/features/generic_features.py @@ -4,7 +4,7 @@ import numpy as np -def _diff_within_group(df, sort_key, diff_col, **kwargs): +def _diff_within_group(df, sort_key: str, diff_col:str, **kwargs): return df.groupby(sort_key)[diff_col].transform(lambda x: x.diff(**kwargs).bfill()) @@ -45,7 +45,7 @@ def compute_centerofmass_interanimal_distances( def compute_centerofmass_interanimal_speed( - df: pd.DataFrame, raw_col_names: list, n_shifts=5, **kwargs + df: pd.DataFrame, raw_col_names: list, n_shifts:int=5, **kwargs ) -> pd.DataFrame: """Speeds between all animals' centroids""" animal_setup = df.pose.animal_setup @@ -123,7 +123,7 @@ def compute_centerofmass( def compute_centerofmass_velocity( - df: pd.DataFrame, raw_col_names: list, n_shifts=5, bodyparts: list = [], **kwargs + df: pd.DataFrame, raw_col_names: list, n_shifts:int=5, bodyparts: list = [], **kwargs ) -> pd.DataFrame: """Velocity of all animals' centroids""" animal_setup = df.pose.animal_setup @@ -161,7 +161,7 @@ def compute_centerofmass_velocity( def compute_part_velocity( - df: pd.DataFrame, raw_col_names: list, n_shifts=5, bodyparts: list = [], **kwargs + df: pd.DataFrame, raw_col_names: list, n_shifts:int=5, bodyparts: list = [], **kwargs ) -> pd.DataFrame: """Velocity of all animals' bodyparts""" animal_setup = df.pose.animal_setup @@ -198,7 +198,7 @@ def compute_part_velocity( def compute_part_speed( - df: pd.DataFrame, raw_col_names: list, n_shifts=5, bodyparts: list = [], **kwargs + df: pd.DataFrame, raw_col_names: list, n_shifts:int=5, bodyparts: list = [], **kwargs ) -> pd.DataFrame: """Speed of all animals' bodyparts""" animal_setup = df.pose.animal_setup @@ -235,7 +235,7 @@ def compute_part_speed( def compute_speed_features( - df: pd.DataFrame, raw_col_names: list, n_shifts=5, **kwargs + df: pd.DataFrame, raw_col_names: list, n_shifts:int=5, **kwargs ) -> pd.DataFrame: """Speeds between all body parts pairs (within and between animals)""" animal_setup = df.pose.animal_setup diff --git a/ethome/features/mars_features.py b/ethome/features/mars_features.py index c1f8865..99e895d 100644 --- a/ethome/features/mars_features.py +++ b/ethome/features/mars_features.py @@ -1,6 +1,7 @@ import pandas as pd import numpy as np +from typing import Callable from ethome.io import XY_IDS from itertools import product @@ -12,9 +13,9 @@ # The decorator maker, so we can provide arguments -def augment_features(window_size=5, n_shifts=3, mode="shift"): +def augment_features(window_size:int=5, n_shifts:int=3, mode:str="shift"): # The decorator - def decorator(feature_function): + def decorator(feature_function:Callable): # What is called instead of the actual function, assumes the feature making # function returns the names of the columns just made def wrapper(*args, **kwargs): @@ -85,7 +86,7 @@ def wrapper(*args, **kwargs): from pandas.api.types import is_numeric_dtype -def boiler_plate(features_df): +def boiler_plate(features_df:pd.DataFrame): reversemap = None to_drop = ["Unnamed: 0"] @@ -101,7 +102,7 @@ def boiler_plate(features_df): @augment_features() def _compute_centroid( - df, name, animal_setup, body_parts=None, n_shifts=3, mode="shift" + df:pd.DataFrame, name:str, animal_setup:dict, body_parts:list=None, n_shifts:int=3, mode:str="shift" ): bodypart_ids = animal_setup["bodypart_ids"] mouse_ids = animal_setup["mouse_ids"] @@ -121,7 +122,7 @@ def _compute_centroid( @augment_features() def _compute_abs_angle( - df, name, animal_setup, bps, centroid=True, n_shifts=3, mode="shift" + df:pd.DataFrame, name:str, animal_setup:dict, bps:list, centroid:bool=True, n_shifts:int=3, mode:bool="shift" ): mouse_ids = animal_setup["mouse_ids"] df = df.copy() @@ -142,7 +143,7 @@ def _compute_abs_angle( @augment_features() def _compute_rel_angle( - df, name, animal_setup, bps, centroid=False, n_shifts=3, mode="shift" + df:pd.DataFrame, name:str, animal_setup:dict, bps:list, centroid:bool=False, n_shifts:int=3, mode:str="shift" ): mouse_ids = animal_setup["mouse_ids"] df = df.copy() @@ -172,7 +173,7 @@ def _compute_rel_angle( @augment_features() -def _compute_ellipsoid(df, animal_setup, n_shifts=3, mode="shift"): +def _compute_ellipsoid(df:pd.DataFrame, animal_setup:dict, n_shifts:int=3, mode:str="shift"): bodypart_ids = animal_setup["bodypart_ids"] mouse_ids = animal_setup["mouse_ids"] colnames = animal_setup["colnames"] @@ -213,7 +214,7 @@ def _compute_ellipsoid(df, animal_setup, n_shifts=3, mode="shift"): # Recall framerate is 30 fps -def _compute_kinematics(df, names, animal_setup, window_size=5, n_shifts=3): +def _compute_kinematics(df:pd.DataFrame, names:list, animal_setup:dict, window_size:int=5, n_shifts:int=3): bodypart_ids = animal_setup["bodypart_ids"] mouse_ids = animal_setup["mouse_ids"] colnames = animal_setup["colnames"] @@ -235,7 +236,7 @@ def _compute_kinematics(df, names, animal_setup, window_size=5, n_shifts=3): @augment_features() def _compute_relative_body_motions( - df, animal_setup, window_size=3, n_shifts=3, mode="shift" + df:pd.DataFrame, animal_setup:dict, window_size:int=3, n_shifts:int=3, mode:str="shift" ): bodypart_ids = animal_setup["bodypart_ids"] mouse_ids = animal_setup["mouse_ids"] @@ -268,7 +269,7 @@ def _compute_relative_body_motions( @augment_features() -def _compute_relative_body_angles(df, animal_setup, n_shifts=3, mode="shift"): +def _compute_relative_body_angles(df:pd.DataFrame, animal_setup:dict, n_shifts:int=3, mode:str="shift"): bodypart_ids = animal_setup["bodypart_ids"] mouse_ids = animal_setup["mouse_ids"] colnames = animal_setup["colnames"] @@ -313,7 +314,7 @@ def _compute_relative_body_angles(df, animal_setup, n_shifts=3, mode="shift"): @augment_features() -def _compute_iou(df, animal_setup, n_shifts=3, mode="shift"): +def _compute_iou(df:pd.DataFrame, animal_setup:dict, n_shifts:int=3, mode:str="shift"): bodypart_ids = animal_setup["bodypart_ids"] mouse_ids = animal_setup["mouse_ids"] colnames = animal_setup["colnames"] @@ -355,7 +356,7 @@ def _compute_iou(df, animal_setup, n_shifts=3, mode="shift"): # Which can change from video to video, train to test, etc. So perhaps not useful @augment_features() def _compute_cage_distances( - features_df, animal_setup, n_shifts=3, mode="shift" + features_df:pd.DataFrame, animal_setup:dict, n_shifts:int=3, mode:bool="shift" ): # pragma: no cover bodypart_ids = animal_setup["bodypart_ids"] mouse_ids = animal_setup["mouse_ids"] @@ -388,7 +389,7 @@ def _compute_cage_distances( return features_df -def make_features_distances(df, animal_setup): +def make_features_distances(df:pd.DataFrame, animal_setup:dict): bodypart_ids = animal_setup["bodypart_ids"] mouse_ids = animal_setup["mouse_ids"] colnames = animal_setup["colnames"] @@ -430,7 +431,7 @@ def make_features_distances(df, animal_setup): return features_df -def make_features_mars(df, animal_setup, n_shifts=3, mode="shift"): +def make_features_mars(df:pd.DataFrame, animal_setup:dict, n_shifts:int=3, mode:str="shift"): features_df = df.copy() ####################### @@ -541,11 +542,11 @@ def make_features_mars(df, animal_setup, n_shifts=3, mode="shift"): return features_df -def make_features_mars_distr(df, animal_setup): +def make_features_mars_distr(df:pd.DataFrame, animal_setup:dict): return make_features_mars(df, animal_setup, n_shifts=3, mode="distr") -def make_features_mars_reduced(df, animal_setup, n_shifts=2, mode="diff"): +def make_features_mars_reduced(df:pd.DataFrame, animal_setup:dict, n_shifts:int=2, mode:str="diff"): features_df = df.copy() ####################### @@ -641,7 +642,7 @@ def make_features_mars_reduced(df, animal_setup, n_shifts=2, mode="diff"): return features_df -def make_features_velocities(df, animal_setup, n_shifts=5): # pragma: no cover +def make_features_velocities(df:pd.DataFrame, animal_setup:dict, n_shifts:int=5): # pragma: no cover bodypart_ids = animal_setup["bodypart_ids"] mouse_ids = animal_setup["mouse_ids"] colnames = animal_setup["colnames"] @@ -699,7 +700,7 @@ def make_features_velocities(df, animal_setup, n_shifts=5): # pragma: no cover return features_df -def make_features_social(df, animal_setup, n_shifts=3, mode="shift"): +def make_features_social(df:pd.DataFrame, animal_setup:dict, n_shifts:int=3, mode:str="shift"): features_df = df.copy() colnames = animal_setup["colnames"] diff --git a/ethome/io.py b/ethome/io.py index 20f5354..e750001 100644 --- a/ethome/io.py +++ b/ethome/io.py @@ -2,6 +2,7 @@ import pandas as pd import pickle import numpy as np +from typing import Sequence from itertools import product from joblib import dump, load import os @@ -15,14 +16,14 @@ XYLIKELIHOOD_IDS = ["x", "y", "likelihood"] -def uniquifier(seq): +def uniquifier(seq:Sequence): """Return a sequence (e.g. list) with unique elements only, but maintaining original list order""" seen = set() seen_add = seen.add return [x for x in seq if not (x in seen or seen_add(x))] -def _list_replace(ls, renamer): +def _list_replace(ls:list, renamer:dict): """Replace elements in a list according to provided dictionary""" for i, word in enumerate(ls): if word in renamer.keys(): @@ -30,7 +31,7 @@ def _list_replace(ls, renamer): return ls -def save_sklearn_model(model, fn_out): # pragma: no cover +def save_sklearn_model(model, fn_out:str): # pragma: no cover """Save sklearn model to file Args: @@ -40,7 +41,7 @@ def save_sklearn_model(model, fn_out): # pragma: no cover dump(model, fn_out) -def load_sklearn_model(fn_in): # pragma: no cover +def load_sklearn_model(fn_in: str): # pragma: no cover """Load sklearn model from file Args: @@ -152,7 +153,7 @@ def read_sleap_tracks( # POSSIBILITY OF SUCH DAMAGE. -def _load_sleap_data(path, multi_animal=False): +def _load_sleap_data(path:str, multi_animal:bool=False): """loads sleap data h5 format from file path and returns it as pd.DataFrame As sleap body parts (nodes) are not ordered in a particular way, we sort them alphabetically. As sleap tracks do not export a score/likelihood but cut off automatically (nan), we are simulating a likelihood @@ -408,7 +409,7 @@ def _read_DLC_tracks( # SOFTWARE. -def _convert_nwb_to_h5_all(nwbfile): +def _convert_nwb_to_h5_all(nwbfile:str): """ Convert a NWB data file back to DeepLabCut's h5 data format. @@ -546,7 +547,7 @@ def load_data(fn: str): # pragma: no cover # Only used to making a test dataframe for testing and dev purposes -def _make_sample_dataframe(fn_out="sample_dataframe.pkl"): # pragma: no cover +def _make_sample_dataframe(fn_out:str="sample_dataframe.pkl"): # pragma: no cover from ethome import create_dataset, create_metadata cur_dir = os.path.dirname(os.path.abspath(__file__)) @@ -662,7 +663,7 @@ def read_boris_annotation( return ground_truth -def create_behavior_labels(boris_files): +def create_behavior_labels(boris_files:list): """Create behavior labels from BORIS exported csv files. Args: