rm groupby(axis=1) using 2 transposes and some logic changes

ryan-williams · ryan-williams · commit 337665561c32 · 2024-02-02T16:45:59.000-05:00
diff --git a/proj.py b/proj.py
@@ -16,25 +16,26 @@
 def project(df):
     """Compute current year-end projection.
 
-    Assumes original DataFrame shape, suitable for calling like:
+    Assumes transposed DataFrame shape, suitable for calling like:
     ```
-    df.groupby(lambda x:x, axis=1, level=0).apply(project))
+    df.groupby(lambda x:x, level=0).apply(project))
     ```
+    Transposes all operations from `project` (e.g. `df[k]` → `df.loc[k]`)
     """
-    [top_level] = df.columns.get_level_values(0).drop_duplicates()
-    d = df.loc[:, top_level]
-    prv_ytd = d.prv_ytd
-    cur_ytd = d.cur_ytd
+    [top_level] = df.index.get_level_values(0).drop_duplicates()
+    d = df.loc[top_level]
+    prv_ytd = d.loc['prv_ytd']
+    cur_ytd = d.loc['cur_ytd']
     zero_mask = prv_ytd == 0
-    prv_roy = d.prv_end - prv_ytd
-    d.loc[ zero_mask, 'roy'] = (prv_roy + cur_ytd) * cur_roy_frac
-    d.loc[~zero_mask, 'roy'] = prv_roy * (1 + cur_ytd_frac * (cur_ytd / prv_ytd - 1))
-    d['roy'] = round(d.roy).astype(int)
-    d['projected'] = cur_ytd + d.roy
-    return d[['roy', 'projected']]
+    prv_roy = d.loc['prv_end'] - prv_ytd
+    d.loc['roy',  zero_mask] = (prv_roy + cur_ytd) * cur_roy_frac
+    d.loc['roy', ~zero_mask] = prv_roy * (1 + cur_ytd_frac * (cur_ytd / prv_ytd - 1))
+    d.loc['roy'] = round(d.loc['roy'])
+    d.loc['projected'] = cur_ytd + d.loc['roy']
+    return d.loc[['roy', 'projected']].astype(int)  # casting d.loc['roy'] above doesn't work, for some reason
 
 
 # Compute projections
-proj = df.groupby(lambda x: x, axis=1, level=0).apply(project)
+proj = df.T.groupby(lambda x: x, level=0).apply(project).T
 print(proj)
 proj.to_csv('proj.csv')
