spark: add PCAtransformer binding and test #1

opencb · Jan 22, 2019 · a6c3be2 · a6c3be2
1 parent 8044aba
commit a6c3be2
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 0 deletions.
diff --git a/oskar-spark/src/main/python/pyoskar/analysis.py b/oskar-spark/src/main/python/pyoskar/analysis.py
@@ -509,6 +509,29 @@ def setMissingAsReference(self, value):
         return self._set(missingAsReference=value)
 
 
+class PCATransformer(AbstractTransformer):
+    studyId = Param(Params._dummy(), "studyId", "", typeConverter=TypeConverters.toString)
+    k = Param(Params._dummy(), "k", "", typeConverter=TypeConverters.toInt)
+
+    @keyword_only
+    def __init__(self, studyId=None, k=None):
+        super(PCATransformer, self).__init__()
+        self._java_obj = self._new_java_obj("org.opencb.oskar.spark.variant.analysis.PCATransformer", self.uid)
+        self.setParams(**self._input_kwargs)
+
+    def getStudyId(self):
+        return self.getOrDefault(self.studyId)
+
+    def setStudyId(self, value):
+        return self._set(studyId=value)
+
+    def getK(self):
+        return self.getOrDefault(self.k)
+
+    def setK(self, value):
+        return self._set(k=value)
+
+
 class TdtTransformer(AbstractTransformer):
     studyId = Param(Params._dummy(), "studyId", "", typeConverter=TypeConverters.toString)
     phenotype = Param(Params._dummy(), "phenotype", "", typeConverter=TypeConverters.toString)

diff --git a/oskar-spark/src/main/python/pyoskar/core.py b/oskar-spark/src/main/python/pyoskar/core.py
@@ -285,6 +285,13 @@ def modeOfInheritance(self, df, family, modeOfInheritance, phenotype, studyId=No
         return ModeOfInheritanceTransformer(family=family, modeOfInheritance=modeOfInheritance, phenotype=phenotype, studyId=studyId,
                                             incompletePenetrance=incompletePenetrance, missingAsReference=missingAsReference).transform(df)
 
+    def PCA(self, df, studyId, k):
+        """
+
+        :type df: DataFrame
+        """
+        return PCATransformer(studyId=studyId, k=k).transform(df)
+
     def tdt(self, df, studyId, phenotype):
         """
 

diff --git a/oskar-spark/src/test/python/pyoskar_tests/test_transformers.py b/oskar-spark/src/test/python/pyoskar_tests/test_transformers.py
@@ -77,6 +77,9 @@ def test_mode_of_inheritance(self):
 
         self.oskar.modeOfInheritance(dataframe, family=family, modeOfInheritance="biallelic", phenotype=phenotype, studyId="hgvauser@platinum:illumina_platinum").show(LIMIT)
 
+    def test_pca(self):
+        self.oskar.PCA(self.df, "hgvauser@platinum:illumina_platinum", 2).show(LIMIT, truncate=False)
+
     def test_tdt(self):
         self.oskar.tdt(self.df, "hgvauser@platinum:illumina_platinum", "KK").show(LIMIT)