diff options
| -rw-r--r-- | gn3/computations/partial_correlations.py | 16 | ||||
| -rw-r--r-- | tests/unit/computations/test_partial_correlations.py | 22 |
2 files changed, 38 insertions, 0 deletions
diff --git a/gn3/computations/partial_correlations.py b/gn3/computations/partial_correlations.py index 07dc16d..ffdf0c5 100644 --- a/gn3/computations/partial_correlations.py +++ b/gn3/computations/partial_correlations.py @@ -10,6 +10,8 @@ from typing import Any, Tuple, Sequence from scipy.stats import pearsonr, spearmanr from gn3.settings import TEXTDIR +import pandas + from gn3.data_helpers import parse_csv_line def control_samples(controls: Sequence[dict], sampleslist: Sequence[str]): @@ -258,6 +260,20 @@ def compute_partial_correlations_fast(# pylint: disable=[R0913, R0914] else fetched_correlations[corr[0]][0:2]) for idx, corr in enumerate(all_correlations)) +def build_data_frame( + xdata: Tuple[float, ...], ydata: Tuple[float, ...], + zdata: Union[ + Tuple[float, ...], + Tuple[Tuple[float, ...], ...]]) -> pandas.DataFrame: + """ + Build a pandas DataFrame object from xdata, ydata and zdata + """ + 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( + {"z{}".format(i): val for i, val in enumerate(row)} for row in zdata)) + def partial_correlation_matrix( xdata: Tuple[float, ...], ydata: Tuple[float, ...], zdata: Tuple[float, ...], method: str = "pearsons", diff --git a/tests/unit/computations/test_partial_correlations.py b/tests/unit/computations/test_partial_correlations.py index c5c35d1..b22bc62 100644 --- a/tests/unit/computations/test_partial_correlations.py +++ b/tests/unit/computations/test_partial_correlations.py @@ -2,9 +2,12 @@ import csv from unittest import TestCase + +import pandas from gn3.computations.partial_correlations import ( fix_samples, control_samples, + build_data_frame, dictify_by_samples, tissue_correlation, find_identical_traits, @@ -297,6 +300,25 @@ class TestPartialCorrelations(TestCase): ("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l")), (0, 4, 8, 10)) + def test_build_data_frame(self): + """ + Check that the function builds the correct data frame. + """ + for xdata, ydata, zdata, expected in ( + ((0.1, 1.1, 2.1), (2.1, 3.1, 4.1), (5.1, 6.1 ,7.1), + pandas.DataFrame({ + "x": (0.1, 1.1, 2.1), "y": (2.1, 3.1, 4.1), + "z": (5.1, 6.1 ,7.1)})), + ((0.1, 1.1, 2.1), (2.1, 3.1, 4.1), + ((5.1, 6.1 ,7.1), (5.2, 6.2, 7.2), (5.3, 6.3, 7.3)), + pandas.DataFrame({ + "x": (0.1, 1.1, 2.1), "y": (2.1, 3.1, 4.1), + "z0": (5.1, 5.2 ,5.3), "z1": (6.1, 6.2 ,6.3), + "z2": (7.1, 7.2 ,7.3)}))): + with self.subTest(xdata=xdata, ydata=ydata, zdata=zdata): + self.assertTrue( + build_data_frame(xdata, ydata, zdata).equals(expected)) + def test_partial_correlation_matrix(self): """ Test that `partial_correlation_matrix` computes the appropriate |