From 9f0229580fa8418964632438e3fa971a0baa2b5a Mon Sep 17 00:00:00 2001
From: Ross Cutler <46252169+rosscutler@users.noreply.github.com>
Date: Wed, 4 Jun 2025 12:37:34 -0700
Subject: [PATCH] Fix DataFrame columns initialization

---
 src/create_input.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/create_input.py b/src/create_input.py
index 7f156df..16f6039 100644
--- a/src/create_input.py
+++ b/src/create_input.py
@@ -76,10 +76,12 @@ def add_clips_balanced_block(clips, condition_pattern, keys, n_clips_per_session
     :param output_df:
     :return:
     """
-    block_keys = [x.strip() for x in keys.split(',')]
-    if len(block_keys) > 2:
-        raise SystemExit("Error: balanced_block design- only up to 2 keys in 'block_keys' are supported")
-    data = pd.DataFrame(columns=block_keys.copy().insert(0,' url'))
+    block_keys = [x.strip() for x in keys.split(',')]
+    if len(block_keys) > 2:
+        raise SystemExit("Error: balanced_block design- only up to 2 keys in 'block_keys' are supported")
+
+    columns = ['url'] + block_keys.copy()
+    data = pd.DataFrame(columns=columns)
 
     for clip in clips:
         condition = conv_filename_to_condition(clip, condition_pattern)