From 9b590d894f1e68ca5d7d00cb6d268f7fb6e6730c Mon Sep 17 00:00:00 2001
From: Frederick Muriuki Muriithi
Date: Tue, 9 Nov 2021 09:56:41 +0300
Subject: Fix bug: if three columns, ensure last is "z"

Issue:
https://github.com/genenetwork/gn-gemtext-threads/blob/main/topics/gn1-migration-to-gn2/partial-correlations.gmi

* Fix a bug, caught when the function is called in a recursive form, with the
  "z*" columns reducing for each cycle through the recursive form.

  As it was, the last cycle through the recursive form would end up with a
  DataFrame with the columns "x", "y", and "z0" rather than the columns "x",
  "y", "z".

  This commit handles that edge case to ensure that the column name is changed
  from "z0" to simply "z".
---
 gn3/computations/partial_correlations.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/gn3/computations/partial_correlations.py b/gn3/computations/partial_correlations.py
index bd127a7..9d73197 100644
--- a/gn3/computations/partial_correlations.py
+++ b/gn3/computations/partial_correlations.py
@@ -272,8 +272,11 @@ def build_data_frame(
     x_y_df = pandas.DataFrame({"x": xdata, "y": ydata})
     if isinstance(zdata[0], float):
         return x_y_df.join(pandas.DataFrame({"z": zdata}))
-    return x_y_df.join(pandas.DataFrame(
+    interm_df = x_y_df.join(pandas.DataFrame(
         {"z{}".format(i): val for i, val in enumerate(row)} for row in zdata))
+    if interm_df.shape[1] == 3:
+        return interm_df.rename(columns={"z0": "z"})
+    return interm_df
 
 def partial_correlation_matrix(
         xdata: Tuple[float, ...], ydata: Tuple[float, ...],
-- 
cgit v1.2.3