Repository: apriha/lineage

Issue:
Return discrepant SNPs
Return discrepant SNPs when adding SNPs to an ``Individual``.

Discussion:

Patch:
diff --git a/README.rst b/README.rst
index ff76836..d49a33c 100644
--- a/README.rst
+++ b/README.rst
@@ -97,24 +97,23 @@ Merge Raw Data Files
 The dataset for ``User662`` consists of three raw data files from two different DNA testing
 companies. Let's load the remaining two files.
 
-As the data gets added, it's compared to the existing data and discrepancies are saved to CSV
-files. (The discrepancy thresholds can be tuned via parameters.)
+As the data gets added, it's compared to the existing data, and SNP position and genotype
+discrepancies are identified. (The discrepancy thresholds can be tuned via parameters.)
 
+>>> user662.snp_count
+708092
 >>> user662.load_snps(['resources/662.23andme.304.txt.gz', 'resources/662.23andme.340.txt.gz'],
 ...                   discrepant_genotypes_threshold=160)
 Loading resources/662.23andme.304.txt.gz
 3 SNP positions being added differ; keeping original positions
-Saving output/User662_discrepant_positions_1.csv
 8 genotypes were discrepant; marking those as null
-Saving output/User662_discrepant_genotypes_1.csv
 Loading resources/662.23andme.340.txt.gz
 27 SNP positions being added differ; keeping original positions
-Saving output/User662_discrepant_positions_2.csv
 156 genotypes were discrepant; marking those as null
-Saving output/User662_discrepant_genotypes_2.csv
-
-All `output files <https://lineage.readthedocs.io/en/latest/output_files.html>`_ are saved to the output
-directory.
+>>> len(user662.discrepant_positions)
+30
+>>> user662.snp_count
+1006960
 
 Save SNPs
 `````````
@@ -125,6 +124,9 @@ SNPs to a CSV file:
 >>> saved_snps = user662.save_snps()
 Saving output/User662.csv
 
+All `output files <https://lineage.readthedocs.io/en/latest/output_files.html>`_ are saved to the output
+directory.
+
 Compare Individuals
 ```````````````````
 Let's create another ``Individual`` for the ``User663`` dataset:
diff --git a/lineage/individual.py b/lineage/individual.py
index 1a2c07b..3165a9e 100644
--- a/lineage/individual.py
+++ b/lineage/individual.py
@@ -23,6 +23,7 @@ import os
 import re
 
 import numpy as np
+import pandas as pd
 
 import lineage
 from lineage.snps import (SNPs, get_assembly_name, get_chromosomes, get_chromosomes_summary,
@@ -56,6 +57,8 @@ class Individual(object):
         self._source = []
         self._discrepant_positions_file_count = 0
         self._discrepant_genotypes_file_count = 0
+        self._discrepant_positions = pd.DataFrame()
+        self._discrepant_genotypes = pd.DataFrame()
 
         if raw_data is not None:
             self.load_snps(raw_data)
@@ -159,8 +162,28 @@ class Individual(object):
         """
         return determine_sex(self._snps)
 
+    @property
+    def discrepant_positions(self):
+        """ SNPs with discrepant positions discovered while loading SNPs.
+
+        Returns
+        -------
+        pandas.DataFrame
+        """
+        return self._discrepant_positions
+
+    @property
+    def discrepant_genotypes(self):
+        """ SNPs with discrepant genotypes discovered while loading SNPs.
+
+        Returns
+        -------
+        pandas.DataFrame
+        """
+        return self._discrepant_genotypes
+
     def load_snps(self, raw_data, discrepant_snp_positions_threshold=100,
-                  discrepant_genotypes_threshold=10000):
+                  discrepant_genotypes_threshold=10000, save_output=False):
         """ Load raw genotype data.
 
         Parameters
@@ -173,19 +196,31 @@ class Individual(object):
         discrepant_genotypes_threshold : int
             threshold for discrepant genotype data between existing data and data to be loaded,
             a large value could indicated mismatched individuals
+        save_output : bool
+            specifies whether to save discrepant SNP output to CSV files in the output directory
         """
         if type(raw_data) is list:
             for file in raw_data:
-                print('Loading ' + os.path.relpath(file))
-                self._add_snps(SNPs(file), discrepant_snp_positions_threshold,
-                               discrepant_genotypes_threshold)
+                self._load_snps_helper(file, discrepant_snp_positions_threshold,
+                                       discrepant_genotypes_threshold, save_output)
         elif type(raw_data) is str:
-            print('Loading ' + os.path.relpath(raw_data))
-            self._add_snps(SNPs(raw_data), discrepant_snp_positions_threshold,
-                           discrepant_genotypes_threshold)
+            self._load_snps_helper(raw_data, discrepant_snp_positions_threshold,
+                                   discrepant_genotypes_threshold, save_output)
         else:
             raise TypeError('invalid filetype')
 
+    def _load_snps_helper(self, file, discrepant_snp_positions_threshold,
+                          discrepant_genotypes_threshold, save_output):
+        print('Loading ' + os.path.relpath(file))
+        discrepant_positions, discrepant_genotypes = \
+            self._add_snps(SNPs(file), discrepant_snp_positions_threshold,
+                           discrepant_genotypes_threshold, save_output)
+
+        self._discrepant_positions = self._discrepant_positions.append(discrepant_positions,
+                                                                       sort=True)
+        self._discrepant_genotypes = self._discrepant_genotypes.append(discrepant_genotypes,
+                                                                       sort=True)
+
     def save_snps(self):
         """ Save SNPs to file.
 
@@ -309,7 +344,8 @@ class Individual(object):
         self._snps = snps
         self._assembly = assembly
 
-    def _add_snps(self, snps, discrepant_snp_positions_threshold, discrepant_genotypes_threshold):
+    def _add_snps(self, snps, discrepant_snp_positions_threshold,
+                  discrepant_genotypes_threshold, save_output):
         """ Add SNPs to this Individual.
 
         Parameters
@@ -320,9 +356,19 @@ class Individual(object):
             see above
         discrepant_genotypes_threshold : int
             see above
+        save_output
+            see above
+
+        Returns
+        -------
+        discrepant_positions : pandas.DataFrame
+        discrepant_genotypes : pandas.DataFrame
         """
+        discrepant_positions = pd.DataFrame()
+        discrepant_genotypes = pd.DataFrame()
+
         if snps.snps is None:
-            return
+            return discrepant_positions, discrepant_genotypes
 
         assembly = snps.assembly
         source = snps.source
@@ -352,13 +398,14 @@ class Individual(object):
                 print(str(len(discrepant_positions)) + ' SNP positions being added differ; '
                       'keeping original positions')
 
-                self._discrepant_positions_file_count += 1
-                lineage.save_df_as_csv(discrepant_positions, self._output_dir,
-                                       self.get_var_name() + '_discrepant_positions_' + str(
-                                           self._discrepant_positions_file_count) + '.csv')
+                if save_output:
+                    self._discrepant_positions_file_count += 1
+                    lineage.save_df_as_csv(discrepant_positions, self._output_dir,
+                                           self.get_var_name() + '_discrepant_positions_' + str(
+                                               self._discrepant_positions_file_count) + '.csv')
             elif len(discrepant_positions) >= discrepant_snp_positions_threshold:
                 print('too many SNPs differ in position; ensure same genome build is being used')
-                return
+                return discrepant_positions, discrepant_genotypes
 
             # remove null genotypes
             common_snps = common_snps.loc[~common_snps['genotype'].isnull() &
@@ -381,14 +428,15 @@ class Individual(object):
                 print(str(len(discrepant_genotypes)) + ' genotypes were discrepant; '
                       'marking those as null')
 
-                self._discrepant_genotypes_file_count += 1
-                lineage.save_df_as_csv(discrepant_genotypes, self._output_dir,
-                                       self.get_var_name() + '_discrepant_genotypes_' + str(
-                                           self._discrepant_genotypes_file_count) + '.csv')
+                if save_output:
+                    self._discrepant_genotypes_file_count += 1
+                    lineage.save_df_as_csv(discrepant_genotypes, self._output_dir,
+                                           self.get_var_name() + '_discrepant_genotypes_' + str(
+                                               self._discrepant_genotypes_file_count) + '.csv')
             elif len(discrepant_genotypes) >= discrepant_genotypes_threshold:
                 print('too many SNPs differ in their genotype; ensure file is for same '
                       'individual')
-                return
+                return discrepant_positions, discrepant_genotypes
 
             # add new SNPs
             self._source.append(source)
@@ -400,6 +448,8 @@ class Individual(object):
 
         self._snps = sort_snps(self._snps)
 
+        return discrepant_positions, discrepant_genotypes
+
     @staticmethod
     def _double_single_alleles(df, chrom):
         """ Double any single alleles in the specified chromosome.


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 3a313e9..3f4f8a4 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -16,10 +16,27 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 """
 
+import os
+import shutil
+
 import pytest
 
 from lineage import Lineage
 
+
+def del_output_dir_helper():
+    if os.path.exists('output'):
+        shutil.rmtree('output')
+
+
 @pytest.fixture(scope='module')
 def l():
     return Lineage()
+
+
+@pytest.fixture
+def del_output_dir():
+    """ Delete output directory if it exists during setup / teardown. """
+    del_output_dir_helper()
+    yield
+    del_output_dir_helper()
diff --git a/tests/input/discrepant_snps.csv b/tests/input/discrepant_snps.csv
new file mode 100644
index 0000000..3ad9609
--- /dev/null
+++ b/tests/input/discrepant_snps.csv
@@ -0,0 +1,53 @@
+﻿rsid,chromosome,position_file1,position_file2,genotype_file1,genotype_file2,discrepant_position,discrepant_genotype,expected_position,expected_genotype
+rs1,1,1,1,--,--,False,False,1,--
+rs2,1,2,2,--,AA,False,False,2,AA
+rs3,1,3,3,AA,--,False,False,3,AA
+rs4,1,4,4,AA,AA,False,False,4,AA
+rs5,1,5,5,--,--,False,False,5,--
+rs6,1,6,6,--,AB,False,False,6,AB
+rs7,1,7,7,AB,--,False,False,7,AB
+rs8,1,8,8,AB,AB,False,False,8,AB
+rs9,1,9,9,AB,BA,False,False,9,AB
+rs10,1,10,10,BA,AB,False,False,10,BA
+rs11,1,11,11,BA,BA,False,False,11,BA
+rs12,1,12,12,AB,AC,False,True,12,--
+rs13,1,13,13,AB,CA,False,True,13,--
+rs14,1,14,14,BA,AC,False,True,14,--
+rs15,1,15,15,BA,CA,False,True,15,--
+rs16,1,16,16,AB,CD,False,True,16,--
+rs17,1,17,17,AB,DC,False,True,17,--
+rs18,1,18,18,BA,CD,False,True,18,--
+rs19,1,19,19,BA,DC,False,True,19,--
+rs20,MT,20,20,--,--,False,False,20,--
+rs21,MT,21,21,--,A,False,False,21,A
+rs22,MT,22,22,A,--,False,False,22,A
+rs23,MT,23,23,A,A,False,False,23,A
+rs24,MT,24,24,A,B,False,True,24,--
+rs25,MT,25,25,B,A,False,True,25,--
+rs26,MT,26,26,B,B,False,False,26,B
+rs27,1,27,1,--,--,True,False,27,--
+rs28,1,28,2,--,AA,True,False,28,AA
+rs29,1,29,3,AA,--,True,False,29,AA
+rs30,1,30,4,AA,AA,True,False,30,AA
+rs31,1,31,5,--,--,True,False,31,--
+rs32,1,32,6,--,AB,True,False,32,AB
+rs33,1,33,7,AB,--,True,False,33,AB
+rs34,1,34,8,AB,AB,True,False,34,AB
+rs35,1,35,9,AB,BA,True,False,35,AB
+rs36,1,36,10,BA,AB,True,False,36,BA
+rs37,1,37,11,BA,BA,True,False,37,BA
+rs38,1,38,12,AB,AC,True,True,38,--
+rs39,1,39,13,AB,CA,True,True,39,--
+rs40,1,40,14,BA,AC,True,True,40,--
+rs41,1,41,15,BA,CA,True,True,41,--
+rs42,1,42,16,AB,CD,True,True,42,--
+rs43,1,43,17,AB,DC,True,True,43,--
+rs44,1,44,18,BA,CD,True,True,44,--
+rs45,1,45,19,BA,DC,True,True,45,--
+rs46,MT,46,20,--,--,True,False,46,--
+rs47,MT,47,21,--,A,True,False,47,A
+rs48,MT,48,22,A,--,True,False,48,A
+rs49,MT,49,23,A,A,True,False,49,A
+rs50,MT,50,24,A,B,True,True,50,--
+rs51,MT,51,25,B,A,True,True,51,--
+rs52,MT,52,26,B,B,True,False,52,B
diff --git a/tests/test_individual.py b/tests/test_individual.py
index a53c009..c849047 100644
--- a/tests/test_individual.py
+++ b/tests/test_individual.py
@@ -25,6 +25,7 @@ import numpy as np
 import pandas as pd
 import pytest
 
+from lineage.snps import sort_snps
 from tests.test_lineage import simulate_snps
 
 
@@ -48,6 +49,13 @@ def snps_NCBI36():
                          genotype=['AA', np.nan, 'ID', 'AG'])
 
 
+@pytest.fixture(scope='module')
+def snps_NCBI36_discrepant_snps():
+    return create_snp_df(rsid=['rs3094315', 'rs2500347', 'rsIndelTest', 'rs11928389'],
+                         chrom=['1', '1', '1', '3'], pos=[742429, 143649677, 143649678, 50908372],
+                         genotype=['AA', np.nan, 'ID', np.nan])
+
+
 @pytest.fixture(scope='module')
 def snps_GRCh37():
     return create_snp_df(rsid=['rs3094315', 'rs2500347', 'rsIndelTest', 'rs11928389'],
@@ -224,25 +232,85 @@ def test_load_snps_invalid_file(l, snps_GRCh37):
     pd.testing.assert_frame_equal(ind.snps, snps_GRCh37)
 
 
-def test_load_snps_assembly_mismatch(l):
-    ind = l.create_individual('')
+def test_load_snps_assembly_mismatch(del_output_dir, l, snps_NCBI36_discrepant_snps):
+    ind = l.create_individual('ind')
     ind.load_snps(['tests/input/NCBI36.csv', 'tests/input/GRCh37.csv'])
-    assert True  # TODO: assert based on discrepant SNPs; see #14
-
-
-def test_load_snps_assembly_mismatch_exceed_discrepant_positions_threshold(l):
-    ind = l.create_individual('')
+    assert not os.path.exists('output/ind_discrepant_positions_1.csv')
+    assert not os.path.exists('output/ind_discrepant_genotypes_1.csv')
+    assert len(ind.discrepant_positions) == 4
+    assert len(ind.discrepant_genotypes) == 1
+    pd.testing.assert_frame_equal(ind.snps, snps_NCBI36_discrepant_snps)
+
+
+def test_load_snps_assembly_mismatch_save_output(del_output_dir, l, snps_NCBI36_discrepant_snps):
+    ind = l.create_individual('ind')
+    ind.load_snps(['tests/input/NCBI36.csv', 'tests/input/GRCh37.csv'], save_output=True)
+    assert os.path.exists('output/ind_discrepant_positions_1.csv')
+    assert os.path.exists('output/ind_discrepant_genotypes_1.csv')
+    assert len(ind.discrepant_positions) == 4
+    assert len(ind.discrepant_genotypes) == 1
+    pd.testing.assert_frame_equal(ind.snps, snps_NCBI36_discrepant_snps)
+
+
+def test_load_snps_assembly_mismatch_exceed_discrepant_positions_threshold(del_output_dir, l,
+                                                                           snps_NCBI36):
+    ind = l.create_individual('ind')
     ind.load_snps(['tests/input/NCBI36.csv', 'tests/input/GRCh37.csv'],
                   discrepant_snp_positions_threshold=0)
-    assert True  # TODO: assert based on discrepant SNPs; see #14
+    assert not os.path.exists('output/ind_discrepant_positions_1.csv')
+    assert not os.path.exists('output/ind_discrepant_genotypes_1.csv')
+    assert len(ind.discrepant_positions) == 4
+    assert len(ind.discrepant_genotypes) == 0
+    pd.testing.assert_frame_equal(ind.snps, snps_NCBI36)
 
 
-def test_load_snps_assembly_mismatch_exceed_discrepant_genotypes_threshold(l):
-    ind = l.create_individual('')
+def test_load_snps_assembly_mismatch_exceed_discrepant_genotypes_threshold(del_output_dir, l,
+                                                                           snps_NCBI36):
+    ind = l.create_individual('ind')
     ind.load_snps(['tests/input/NCBI36.csv', 'tests/input/GRCh37.csv'],
                   discrepant_genotypes_threshold=0)
-    assert True  # TODO: assert based on discrepant SNPs; see #14
+    assert not os.path.exists('output/ind_discrepant_positions_1.csv')
+    assert not os.path.exists('output/ind_discrepant_genotypes_1.csv')
+    assert len(ind.discrepant_positions) == 4
+    assert len(ind.discrepant_genotypes) == 1
+    pd.testing.assert_frame_equal(ind.snps, snps_NCBI36)
+
+
+def test_merging_files_discrepant_snps(l):
+    df = pd.read_csv('tests/input/discrepant_snps.csv', skiprows=1, na_values='--',
+                     names=['rsid', 'chrom', 'pos_file1', 'pos_file2',
+                            'genotype_file1', 'genotype_file2', 'discrepant_position',
+                            'discrepant_genotype', 'expected_position', 'expected_genotype'],
+                     index_col=0,
+                     dtype={'chrom': object, 'pos_file1': np.int64, 'pos_file2': np.int64,
+                            'discrepant_position': bool, 'discrepant_genotype': bool})
+
+    df1 = df[['chrom', 'pos_file1', 'genotype_file1']]
+    df2 = df[['chrom', 'pos_file2', 'genotype_file2']]
+
+    df1.to_csv('tests/input/discrepant_snps1.csv', na_rep='--',
+               header=['chromosome', 'position', 'genotype'])
+
+    df2.to_csv('tests/input/discrepant_snps2.csv', na_rep='--',
+               header=['chromosome', 'position', 'genotype'])
+
+    ind = l.create_individual('', ['tests/input/discrepant_snps1.csv',
+                                   'tests/input/discrepant_snps2.csv'])
+
+    expected = df[['chrom', 'discrepant_position', 'discrepant_genotype', 'expected_position',
+              'expected_genotype']]
+    expected = expected.rename(columns={'expected_position': 'pos',
+                                        'expected_genotype': 'genotype'})
+    expected = sort_snps(expected)
+
+    pd.testing.assert_index_equal(ind.discrepant_positions.index,
+                                  expected.loc[expected['discrepant_position'] == True].index)
+
+    pd.testing.assert_index_equal(ind.discrepant_genotypes.index,
+                                  expected.loc[expected['discrepant_genotype'] == True].index)
 
+    pd.testing.assert_series_equal(ind.snps['pos'], expected['pos'])
+    pd.testing.assert_series_equal(ind.snps['genotype'], expected['genotype'])
 
 def test_save_snps(l, snps_GRCh37):
     ind = l.create_individual('test save snps', 'tests/input/GRCh37.csv')