--- a
+++ b/notebooks/train_expansion.py
@@ -0,0 +1,51 @@
+# %%
+import json
+import os
+import random
+
+# enable lib loading even if not installed as a pip package or in PYTHONPATH
+# also convenient for relative paths in example config files
+from pathlib import Path
+
+os.chdir(Path(__file__).resolve().parent.parent)
+
+from adpkd_segmentation.data.data_utils import ( # noqa
+    get_labeled,
+    make_dcmdicts,
+)
+from adpkd_segmentation.data.link_data import makelinks  # noqa
+
+# %%
+makelinks()
+# %%
+dcm_paths = sorted(get_labeled())
+dcm2attribs, patient2dcm = make_dcmdicts(tuple(dcm_paths))
+all_patient_IDS = list(patient2dcm.keys())
+
+
+# %%
+with open("./stratification/strat_split_2020_09_24.json", "r") as f:
+    full_split = json.load(f)
+
+# %%
+new_train = [
+    patient_id
+    for patient_id in all_patient_IDS
+    if patient_id not in full_split["val"]
+    and patient_id not in full_split["test"]
+]
+
+# %%
+random.seed(42)
+random.shuffle(new_train)
+
+# %%
+new = {}
+new["val"] = full_split["val"]
+new["test"] = full_split["test"]
+new["train"] = new_train
+# %%
+with open("./stratification/strat_split_2020_09_28_extended.json", "w") as f:
+    json.dump(new, f, indent=4)
+
+# %%