update forecast migration script

code/validation/forecast_migration.py

@@ -24,13 +24,20 @@ def remove_cols_2(df):
 
 def remove_rows_2(df):
     rows_to_retain = [f"{_} day ahead inc hosp" for _ in range(131)] + \
-                         [f"{_} wk ahead inc death" for _ in range(1, 21)] + \
-                         [f"{_} wk ahead cum death" for _ in range(1, 21)] + \
+                     [f"{_} wk ahead inc death" for _ in range(1, 21)] + \
+                     [f"{_} wk ahead cum death" for _ in range(1, 21)] + \
                      [f"{_} wk ahead inc case" for _ in range(1, 9)]
     df_rows = df['target'].isin(rows_to_retain)
     return df.loc[df_rows], df_rows.sum() == df.shape[0]
 
 
+def fix_location(df):
+    df_l = df['location'].astype(str).str.zfill(2)
+    change_flag = df['location'] == df_l
+    df['location'] = df_l
+    return df, change_flag.all()
+
+
 # dictionary containing a list of function for migration.
 # Example: for migration to version "2": specify a list of functions that accept
 # a pandas.DataFrame and return a modified pandas.DataFrame and a boolean flag indicating
@@ -43,6 +50,9 @@ def remove_rows_2(df):
     ],
     3: [
         remove_rows_2
+    ],
+    4: [
+        fix_location
     ]
 }