Split dataset with some conditions

_Malte: <p>Hey there <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>I have to split a training dataset into training and validation with a ratio of 80/20, but with some conditions. Let me first introduce the dataset. I am using a medical dataset which consists of the following attributes: patient_id, image_name, target and some other features, however they are not important for that moment. Each patient_id is unique, and there exists a mapping from one patient_id to multiple image_names. Each patient has at least two image_names but at most 120 image_names. Now I want to split the training dataset into training and validation such that if a patient belongs to train or / and val there are at least two images of that patient in that dataset. I think I have a solution, but this one is quite slow and has too many conditional statements. Furthermore, this solution does not have a stratified option, yet.</p>
<p>Here is what I implemented so far:</p>
<pre><code class="lang-auto">def split_meta_file_to_train_and_val(path_dataset : str, ratio_train : float = 0.8, ratio_val : float = 0.2, shuffle : bool = True, seed : int = 44):
    if ratio_train+ratio_val != 1:
        raise ValueError(f"The sum of the parameters ratio_train and ratio_val should be 1. However, the sum is actually {ratio_train+ratio_val}.")
    meta = pd.read_csv(os.path.join(path_dataset, 'train.csv'))
    if shuffle == True:
        meta = meta.sample(frac=1.0, random_state = seed).reset_index()
        meta['old_index'] = meta['index']
        meta = meta.drop(columns=['index'])
    length_dataset = len(meta)
    train_length = round(len(meta)*ratio_train)
    val_length = round(len(meta)*ratio_val)
    np.random.seed(42)
    train_idx = []
    val_idx = []
    meta_grouped = meta.groupby(['patient_id']).count().reset_index()
    meta_grouped['in_list'] = 0
    for i, patient in enumerate(meta.values):
        # Check whether the patient has exactly three images and if so whether it fits into the train subset w.r.t. to the split ratio.
        # We need to do so because if two images goes to the train or val subset then we are missing one image in the other subset
        # subset since the model needs two images from one patient. 
        if  (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'].values == 3) and (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'in_list'].values == 0) and (len(train_idx)+3 &lt;= train_length): 
            train_idx.append(patient[0])
            patient_images = meta.loc[(meta['patient_id'] == patient[1]) &amp; (meta['image_name'] != patient[0])] # Get all other images from that same patient
            train_idx.append(patient_images['image_name'].ravel())
            meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'in_list'] = 1 # Set in_list to one since we have images from that particular patient within the train subset
            meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] = meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] - 3 # Subtract the amount of images such that we can track the amount of images which are left from that patient
            continue
        # Same as the one above with the difference that we are checking whether images of that patient already allocated to train_idx
        elif (patient[0] not in train_idx) and (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'].values == 3) and (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'in_list'].values == 1) and (len(train_idx)+3 &lt;= train_length): 
            train_idx.append(patient[0])
            patient_images = meta.loc[(meta['patient_id'] == patient[1]) &amp; (meta['image_name'] != patient[0]) &amp; (meta[~meta['image_name'].isin(train_idx)])]
            images = np.random.choice(patient_images['image_name'].to_numpy().ravel(), size=2)
            train_idx.append(images.ravel())
            meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'in_list'] = 1
            meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] = meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] - 3
            continue
        # Check whether we have two or &gt;= 4 images of that patients left and no image of that patient were already allocated into train_idx. Take two images of that patient and allocate it to train_idx when it fits
        elif (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'in_list'].values == 0) and (len(train_idx)+2 &lt;= train_length) and ((meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'].values == 2) or (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'].values &gt;= 4)):
            train_idx.append(patient[0])
            patient_images = meta.loc[(meta['patient_id'] == patient[1]) &amp; (meta['image_name'] != patient[0]) &amp; (meta[~meta['image_name'].isin(train_idx)])]
            image = np.random.choice(patient_images['image_name'].to_numpy().ravel())            
            train_idx.append(image)
            meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'in_list'] = 1
            meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] = meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] - 2
            continue
        elif (patient[0] not in train_idx) and (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'in_list'].values == 1) and (len(train_idx)+1 &lt;= train_length) and (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'].values &gt;= 3):
            train_idx.append(patient[0])
            meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] = meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] -1 
            continue
        elif (patient[0] not in train_idx) and (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'in_list'].values == 1) and (len(train_idx)+2 &lt;= train_length) and (meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'].values == 2):
            train_idx.append(patient[0])
            patient_images = meta.loc[(meta['patient_id'] == patient[1]) &amp; (meta['image_name'] != patient[0]) &amp; (meta[~meta['image_name'].isin(train_idx)])]
            images = np.random.choice(patient_images['image_name'].to_numpy().ravel(), size=2)
            train_idx.append(images.ravel()) 
            meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] = meta_grouped.loc[meta_grouped['patient_id'] == patient[1], 'image_name'] -2
            continue

    train_image_names = meta[meta['image_name'].isin(train_idx)]
    val_image_names = meta[~meta['image_name'].isin(train_idx)]
    return train_image_names, val_image_names
</code></pre>
<p>Perhaps there exists an easier way to do so. If so, I would appreciate it if someone can give me a hint / tip.</p>
<p>Best regards <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>