data_release_builder/filtering.py at main · immerse-eu/data_release_builder · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import os
import pandas as pd


def assessment_window_filtering(assessment_list, source_path):
    '''
    If "Screening" is not needed, this function must be modified since Screening and Baseline share the same value ...
    '''
    print(f"\nFiltering by {assessment_list} assessment window ...")

    window_map = {
        'Screening': 'Screening',
        'Baseline': 'Baseline',
        '2-month post-baseline': 'T1',
        '6-month post-baseline': 'T2',
        '12-month post-baseline': 'T3'
    }
    code_map = {
        'Screening': 0,
        'Baseline': 0,
        '2-month post-baseline': 1,
        '6-month post-baseline': 2,
        '12-month post-baseline': 3
    }

    windows = assessment_list
    if not windows:
        raise ValueError("'assessment_list' does not contain any defined values.")
    if not isinstance(windows, list):
        windows = [windows]

    target_values = [window_map[value] for value in windows if value in window_map]
    target_codes = [code_map[value] for value in windows if value in code_map]
    pattern_values = '|'.join(target_values)

    target_codes_str = [str(code) for code in target_codes]

    filtered_files = {}

    for filename in os.listdir(source_path):
        if not filename.lower().endswith('.csv'):
            continue

        file_path = os.path.join(source_path, filename)
        try:
            df = pd.read_csv(file_path, sep=';')

            if 'VisitCode' in df.columns and target_codes_str:
                df['VisitCode'] = df['VisitCode'].astype(str)
                mask = df['VisitCode'].isin(target_codes_str)

            elif 'visit_name' in df.columns and target_values:
                mask = df['visit_name'].str.contains(pattern_values, case=False, na=False)

            else:
                print(f"{filename}: skipped (no 'visit_name' or 'VisitCode' column found)")
                continue

            filtered_df = df[mask]
            filename = f"{os.path.splitext(filename)[0]}_filter_window.csv"

            filtered_df.to_csv(os.path.join(source_path, filename), index=False, sep=';')
            filtered_files[filename] = filtered_df
            print(f"Saved {filename} ({len(filtered_df)} rows)")

        except Exception as e:
            print(f"Error processing {filename}: {e}")

    return filtered_files


def filtering_excluded_ids(baseline_ids_path, source_path):
    print("\nRemoving excluded ids from baselines...")
    filenames = []
    processed_dataframes = []

    ids_to_exclude = set(pd.read_excel(baseline_ids_path, header=None).iloc[:, 0].astype(str))
    print(f"Excluded {len(ids_to_exclude)}")

    for filename in os.listdir(source_path):
        if filename.endswith('.csv') and filename.endswith('_filter_window.csv'):
            file_path = os.path.join(source_path, filename)
            df = pd.read_csv(file_path, sep=';')

            if 'participant_identifier' not in df.columns:
                print(f"participant_identifier not found in dataframe {filename}")
                continue

            df['participant_identifier'] = df['participant_identifier'].str.strip()
            df = df[~df['participant_identifier'].isin(ids_to_exclude)]

            processed_dataframes.append(df)
            filenames.append(filename)
            filename = f"{filename.replace("_filter_window", "_filter_ids")}"

            out_path_headers = os.path.join(source_path, f"ITEM_{filename}")
            df.to_csv(out_path_headers, index=False, sep=";")
            print(f"Saved {filename} ({len(df)} rows)")

    return processed_dataframes, filenames


def filtering_interesting_ids(baseline_ids_path, source_path):
    print("\nFiltering additional interesting_ids...")
    filenames = []
    processed_dataframes = []

    ids_to_include = set(pd.read_excel(baseline_ids_path, header=None).iloc[:, 0].astype(str))
    print(f"Excluded {len(ids_to_include)}")

    for filename in os.listdir(source_path):
        if filename.endswith('.csv') and filename.startswith('ITEM'):
            file_path = os.path.join(source_path, filename)
            df = pd.read_csv(file_path, sep=';')

            if 'participant_identifier' not in df.columns:
                print(f"participant_identifier not found in dataframe {filename}")
                continue

            df['participant_identifier'] = df['participant_identifier'].str.strip()
            df = df[df['participant_identifier'].isin(ids_to_include)]

            processed_dataframes.append(df)
            filenames.append(filename)
            filename = f"{filename.replace("_filter_ids", "")}"

            out_path_headers = os.path.join(source_path, filename)
            df.to_csv(out_path_headers, index=False, sep=";")
            print(f"Saved {filename} ({len(df)} rows)")

    return processed_dataframes, filenames