snputils.snp.io

 8def read_snp(filename: Union[str, pathlib.Path], **kwargs) -> SNPObject:
 9    """
10    Automatically detect the file format and read it into a SNPObject.
11
12    Args:
13        filename: Filename of the file to read.
14        **kwargs: Additional arguments passed to the reader method.
15
16    Raises:
17        ValueError: If the filename does not have an extension or the extension is not supported.
18    """
19    from snputils.snp.io.read.auto import SNPReader
20
21    return SNPReader(filename).read(**kwargs)

24def read_bed(filename: Union[str, pathlib.Path], **kwargs) -> SNPObject:
25    """
26    Read a BED fileset into a SNPObject.
27
28    Args:
29        filename: Filename of the BED fileset to read.
30        **kwargs: Additional arguments passed to the reader method. See :class:`snputils.snp.io.read.bed.BEDReader` for possible parameters.
31    """
32    from snputils.snp.io.read.bed import BEDReader
33
34    return BEDReader(filename).read(**kwargs)

37def read_pgen(filename: Union[str, pathlib.Path], **kwargs) -> SNPObject:
38    """
39    Read a PGEN fileset into a SNPObject.
40
41    Args:
42        filename: Filename of the PGEN fileset to read.
43        **kwargs: Additional arguments passed to the reader method. See :class:`snputils.snp.io.read.pgen.PGENReader` for possible parameters.
44    """
45    from snputils.snp.io.read.pgen import PGENReader
46
47    return PGENReader(filename).read(**kwargs)

50def read_vcf(filename: Union[str, pathlib.Path], 
51             backend: str = 'polars',
52             **kwargs) -> SNPObject:
53    """
54    Read a VCF fileset into a SNPObject.
55
56    Args:
57        filename: Filename of the VCF fileset to read.
58        backend: Backend to use for reading the VCF file. Options are 'polars' or 'scikit-allel'.
59        **kwargs: Additional arguments passed to the reader method. See :class:`snputils.snp.io.read.vcf.VCFReader` for possible parameters.
60    """
61    from snputils.snp.io.read.vcf import VCFReader, VCFReaderPolars
62    if backend == 'polars':
63        print(f"Reading {filename} with polars backend")
64        return VCFReaderPolars(filename).read(**kwargs)
65    else:
66        print(f"Reading {filename} with scikit-allel backend")
67        return VCFReader(filename).read(**kwargs)

 8class SNPReader:
 9    def __new__(cls,
10                filename: Union[str, pathlib.Path],
11                vcf_backend: str = 'polars') -> SNPReader:
12        """
13        Automatically detect the SNP file format from the file extension, and return its corresponding reader.
14
15        Args:
16            filename: Filename of the file to read.
17            vcf_backend: Backend to use for reading the VCF file. Options are 'polars' or 'scikit-allel'. Default is 'polars'.
18
19        Raises:
20            ValueError: If the filename does not have an extension or the extension is not supported.
21        """
22        filename = pathlib.Path(filename)
23        suffixes = filename.suffixes
24        if not suffixes:
25            raise ValueError("The filename should have an extension when using SNPReader.")
26
27        extension = suffixes[-2] if suffixes[-1].lower() in (".zst", ".gz") else suffixes[-1]
28        extension = extension.lower()
29
30        if extension == ".vcf":
31            if vcf_backend == 'polars':
32                from snputils.snp.io.read.vcf import VCFReaderPolars
33
34                return VCFReaderPolars(filename)
35            elif vcf_backend == 'scikit-allel':
36                from snputils.snp.io.read.vcf import VCFReader
37
38                return VCFReader(filename)
39            else:
40                raise ValueError(f"VCF backend not supported: {vcf_backend}")
41        elif extension in (".bed", ".bim", ".fam"):
42            from snputils.snp.io.read.bed import BEDReader
43
44            return BEDReader(filename)
45        elif extension in (".pgen", ".pvar", ".psam", ".pvar.zst"):
46            from snputils.snp.io.read.pgen import PGENReader
47
48            return PGENReader(filename)
49        else:
50            raise ValueError(f"File format not supported: {filename}")

 16@SNPBaseReader.register
 17class BEDReader(SNPBaseReader):
 18    def read(
 19        self,
 20        fields: Optional[List[str]] = None,
 21        exclude_fields: Optional[List[str]] = None,
 22        sample_ids: Optional[np.ndarray] = None,
 23        sample_idxs: Optional[np.ndarray] = None,
 24        variant_ids: Optional[np.ndarray] = None,
 25        variant_idxs: Optional[np.ndarray] = None,
 26        sum_strands: bool = False,
 27        separator: Optional[str] = None,
 28    ) -> SNPObject:
 29        """
 30        Read a bed fileset (bed, bim, fam) into a SNPObject.
 31
 32        Args:
 33            fields (str, None, or list of str, optional): Fields to extract data for that should be included in the returned SNPObject.
 34                Available fields are 'GT', 'IID', 'REF', 'ALT', '#CHROM', 'ID', 'POS'.
 35                To extract all fields, set fields to None. Defaults to None.
 36            exclude_fields (str, None, or list of str, optional): Fields to exclude from the returned SNPObject.
 37                Available fields are 'GT', 'IID', 'REF', 'ALT', '#CHROM', 'ID', 'POS'.
 38                To exclude no fields, set exclude_fields to None. Defaults to None.
 39            sample_ids: List of sample IDs to read. If None and sample_idxs is None, all samples are read.
 40            sample_idxs: List of sample indices to read. If None and sample_ids is None, all samples are read.
 41            variant_ids: List of variant IDs to read. If None and variant_idxs is None, all variants are read.
 42            variant_idxs: List of variant indices to read. If None and variant_ids is None, all variants are read.
 43            sum_strands: True if the maternal and paternal strands are to be summed together,
 44                False if the strands are to be stored separately. Note that due to the pgenlib backend, when sum_strands is False,
 45                8 times as much RAM is required. Nonetheless, the calldata_gt will only be double the size.
 46                WARNING: bed files do not store phase information. If you need it, use vcf or pgen.
 47            separator: Separator used in the pvar file. If None, the separator is automatically detected.
 48                If the automatic detection fails, please specify the separator manually.
 49
 50        Returns:
 51            snpobj: SNPObject containing the data from the pgen fileset.
 52                If sum_strands is False, calldata_gt is stored as a numpy array of shape
 53                (num_variants, num_samples, 2) and dtype int8 containing 0, 1.
 54                If sum_strands is True, calldata_gt is stored as a numpy array of shape
 55                (num_variants, num_samples) and dtype int8 containing 0, 1, 2.
 56
 57        Raises:
 58            AssertionError: If both sample_idxs and sample_ids are specified.
 59            AssertionError: If both variant_idxs and variant_ids are specified.
 60        """
 61        assert (
 62            sample_idxs is None or sample_ids is None
 63        ), "Only one of sample_idxs and sample_ids can be specified"
 64        assert (
 65            variant_idxs is None or variant_ids is None
 66        ), "Only one of variant_idxs and variant_ids can be specified"
 67
 68        if isinstance(fields, str):
 69            fields = [fields]
 70        if isinstance(exclude_fields, str):
 71            exclude_fields = [exclude_fields]
 72
 73        fields = fields or ["GT", "IID", "REF", "ALT", "#CHROM", "ID", "POS"]
 74        exclude_fields = exclude_fields or []
 75        fields = [field for field in fields if field not in exclude_fields]
 76        only_read_bed = fields == ["GT"] and variant_idxs is None and sample_idxs is None
 77
 78        filename_noext = str(self.filename)
 79        if filename_noext[-4:].lower() in (".bed", ".bim", ".fam"):
 80            filename_noext = filename_noext[:-4]
 81
 82        if only_read_bed:
 83            with open(filename_noext + '.fam', 'r') as f:
 84                file_num_samples = sum(1 for _ in f)  # Get sample count from fam file
 85            file_num_variants = None  # Not needed
 86        else:
 87            log.info(f"Reading {filename_noext}.bim")
 88
 89            if separator is None:
 90                with open(filename_noext + ".bim", "r") as file:
 91                    separator = csv.Sniffer().sniff(file.readline()).delimiter
 92
 93            bim = pl.read_csv(
 94                filename_noext + ".bim",
 95                separator=separator,
 96                has_header=False,
 97                new_columns=["#CHROM", "ID", "CM", "POS", "ALT", "REF"],
 98                schema_overrides={
 99                    "#CHROM": pl.String,
100                    "ID": pl.String,
101                    "CM": pl.Float64,
102                    "POS": pl.Int64,
103                    "ALT": pl.String,
104                    "REF": pl.String
105                },
106                null_values=["NA"]
107            ).with_row_index()
108            file_num_variants = bim.height
109
110            if variant_ids is not None:
111                variant_idxs = bim.filter(pl.col("ID").is_in(variant_ids)).select("index").to_series().to_numpy()
112
113            if variant_idxs is None:
114                num_variants = file_num_variants
115                variant_idxs = np.arange(num_variants, dtype=np.uint32)
116            else:
117                num_variants = np.size(variant_idxs)
118                variant_idxs = np.array(variant_idxs, dtype=np.uint32)
119                bim = bim.filter(pl.col("index").is_in(variant_idxs))
120
121            log.info(f"Reading {filename_noext}.fam")
122
123            fam = pl.read_csv(
124                filename_noext + ".fam",
125                separator=separator,
126                has_header=False,
127                new_columns=["Family ID", "IID", "Father ID",
128                             "Mother ID", "Sex code", "Phenotype value"],
129                schema_overrides={
130                    "Family ID": pl.String,
131                    "IID": pl.String,
132                    "Father ID": pl.String,
133                    "Mother ID": pl.String,
134                    "Sex code": pl.String,
135                },
136                null_values=["NA"]
137            ).with_row_index()
138            file_num_samples = fam.height
139
140            if sample_ids is not None:
141                sample_idxs = fam.filter(pl.col("IID").is_in(sample_ids)).select("index").to_series().to_numpy()
142
143            if sample_idxs is None:
144                num_samples = file_num_samples
145            else:
146                num_samples = np.size(sample_idxs)
147                sample_idxs = np.array(sample_idxs, dtype=np.uint32)
148                fam = fam.filter(pl.col("index").is_in(sample_idxs))
149
150        if "GT" in fields:
151            log.info(f"Reading {filename_noext}.bed")
152            pgen_reader = pg.PgenReader(
153                str.encode(filename_noext + ".bed"),
154                raw_sample_ct=file_num_samples,
155                variant_ct=file_num_variants,
156                sample_subset=sample_idxs,
157            )
158
159            if only_read_bed:
160                num_samples = pgen_reader.get_raw_sample_ct()
161                num_variants = pgen_reader.get_variant_ct()
162                variant_idxs = np.arange(num_variants, dtype=np.uint32)
163
164            # required arrays: variant_idxs + sample_idxs + genotypes
165            if not sum_strands:
166                required_ram = (num_samples + num_variants + num_variants * 2 * num_samples) * 4
167            else:
168                required_ram = (num_samples + num_variants) * 4 + num_variants * num_samples
169            log.info(f">{required_ram / 1024**3:.2f} GiB of RAM are required to process {num_samples} samples with {num_variants} variants each")
170
171            if not sum_strands:
172                genotypes = np.empty((num_variants, 2 * num_samples), dtype=np.int32)  # cannot use int8 because of pgenlib
173                pgen_reader.read_alleles_list(variant_idxs, genotypes)
174                genotypes = genotypes.astype(np.int8).reshape((num_variants, num_samples, 2))
175            else:
176                genotypes = np.empty((num_variants, num_samples), dtype=np.int8)
177                pgen_reader.read_list(variant_idxs, genotypes)
178            pgen_reader.close()
179        else:
180            genotypes = None
181
182        log.info("Constructing SNPObject")
183
184        snpobj = SNPObject(
185            calldata_gt=genotypes if "GT" in fields else None,
186            samples=fam.get_column("IID").to_numpy() if "IID" in fields and "IID" in fam.columns else None,
187            **{f'variants_{k.lower()}': bim.get_column(v).to_numpy() if v in fields and v in bim.columns else None
188               for k, v in {'ref': 'REF', 'alt': 'ALT', 'chrom': '#CHROM', 'id': 'ID', 'pos': 'POS'}.items()}
189        )
190
191        log.info("Finished constructing SNPObject")
192        return snpobj

 17@SNPBaseReader.register
 18class PGENReader(SNPBaseReader):
 19    def read(
 20        self,
 21        fields: Optional[List[str]] = None,
 22        exclude_fields: Optional[List[str]] = None,
 23        sample_ids: Optional[np.ndarray] = None,
 24        sample_idxs: Optional[np.ndarray] = None,
 25        variant_ids: Optional[np.ndarray] = None,
 26        variant_idxs: Optional[np.ndarray] = None,
 27        sum_strands: bool = False,
 28        separator: str = None,
 29    ) -> SNPObject:
 30        """
 31        Read a pgen fileset (pgen, psam, pvar) into a SNPObject.
 32
 33        Args:
 34            fields (str, None, or list of str, optional): Fields to extract data for that should be included in the returned SNPObject.
 35                Available fields are 'GT', 'IID', 'REF', 'ALT', '#CHROM', 'ID', 'POS', 'FILTER', 'QUAL'.
 36                To extract all fields, set fields to None. Defaults to None.
 37            exclude_fields (str, None, or list of str, optional): Fields to exclude from the returned SNPObject.
 38                Available fields are 'GT', 'IID', 'REF', 'ALT', '#CHROM', 'ID', 'POS', 'FILTER', 'QUAL'.
 39                To exclude no fields, set exclude_fields to None. Defaults to None.
 40            sample_ids: List of sample IDs to read. If None and sample_idxs is None, all samples are read.
 41            sample_idxs: List of sample indices to read. If None and sample_ids is None, all samples are read.
 42            variant_ids: List of variant IDs to read. If None and variant_idxs is None, all variants are read.
 43            variant_idxs: List of variant indices to read. If None and variant_ids is None, all variants are read.
 44            sum_strands: True if the maternal and paternal strands are to be summed together,
 45                False if the strands are to be stored separately. Note that due to the pgenlib backend, when sum_strands is False,
 46                8 times as much RAM is required. Nonetheless, the calldata_gt will only be double the size.
 47            separator: Separator used in the pvar file. If None, the separator is automatically detected.
 48                If the automatic detection fails, please specify the separator manually.
 49
 50        Returns:
 51            snpobj: SNPObject containing the data from the pgen fileset.
 52                If sum_strands is False, calldata_gt is stored as a numpy array of shape
 53                (num_variants, num_samples, 2) and dtype int8 containing 0, 1.
 54                If sum_strands is True, calldata_gt is stored as a numpy array of shape
 55                (num_variants, num_samples) and dtype int8 containing 0, 1, 2.
 56
 57        Raises:
 58            AssertionError: If both sample_idxs and sample_ids are specified.
 59            AssertionError: If both variant_idxs and variant_ids are specified.
 60        """
 61        assert (
 62            sample_idxs is None or sample_ids is None
 63        ), "Only one of sample_idxs and sample_ids can be specified"
 64        assert (
 65            variant_idxs is None or variant_ids is None
 66        ), "Only one of variant_idxs and variant_ids can be specified"
 67
 68        if isinstance(fields, str):
 69            fields = [fields]
 70        if isinstance(exclude_fields, str):
 71            exclude_fields = [exclude_fields]
 72
 73        fields = fields or ["GT", "IID", "REF", "ALT", "#CHROM", "ID", "POS", "FILTER", "QUAL"]
 74        exclude_fields = exclude_fields or []
 75        fields = [field for field in fields if field not in exclude_fields]
 76        only_read_pgen = fields == ["GT"] and variant_idxs is None and sample_idxs is None
 77
 78        filename_noext = str(self.filename)
 79        for ext in [".pgen", ".pvar", ".pvar.zst", ".psam"]:
 80            if filename_noext.endswith(ext):
 81                filename_noext = filename_noext[:-len(ext)]
 82                break
 83
 84        if only_read_pgen:
 85            file_num_samples = None  # Not needed for pgen
 86            file_num_variants = None  # Not needed
 87        else:
 88            pvar_extensions = [".pvar", ".pvar.zst"]
 89            pvar_filename = None
 90            for ext in pvar_extensions:
 91                possible_pvar = filename_noext + ext
 92                if os.path.exists(possible_pvar):
 93                    pvar_filename = possible_pvar
 94                    break
 95            if pvar_filename is None:
 96                raise FileNotFoundError(f"No .pvar or .pvar.zst file found for {filename_noext}")
 97
 98            log.info(f"Reading {pvar_filename}")
 99
100            def open_textfile(filename):
101                if filename.endswith('.zst'):
102                    import zstandard as zstd
103                    return zstd.open(filename, 'rt')
104                else:
105                    return open(filename, 'rt')
106
107            pvar_has_header = True
108            pvar_header_line_num = 0
109            with open_textfile(pvar_filename) as file:
110                for line_num, line in enumerate(file):
111                    if line.startswith("##"):  # Metadata
112                        continue
113                    else:
114                        if separator is None:
115                            separator = csv.Sniffer().sniff(file.readline()).delimiter
116                        if line.startswith("#CHROM"):  # Header
117                            pvar_header_line_num = line_num
118                            header = line.strip().split()
119                            break
120                        elif not line.startswith("#"):  # If no header, look at line 1
121                            pvar_has_header = False
122                            cols_in_pvar = len(line.strip().split(separator))
123                            if cols_in_pvar == 5:
124                                header = ["#CHROM", "ID", "POS", "ALT", "REF"]
125                            elif cols_in_pvar == 6:
126                                header = ["#CHROM", "ID", "CM", "POS", "ALT", "REF"]
127                            else:
128                                raise ValueError(
129                                    f"{pvar_filename} is not a valid pvar file."
130                                )
131                            break
132
133            pvar_reading_args = {
134                'separator': separator,
135                'skip_rows': pvar_header_line_num,
136                'has_header': pvar_has_header,
137                'new_columns': None if pvar_has_header else header,
138                'schema_overrides': {
139                    "#CHROM": pl.String,
140                    "POS": pl.UInt32,
141                    "ID": pl.String,
142                    "REF": pl.String,
143                    "ALT": pl.String,
144                },
145                'null_values': ["NA"],
146            }
147            if pvar_filename.endswith('.zst'):
148                pvar = pl.read_csv(pvar_filename, **pvar_reading_args).lazy()
149            else:
150                pvar = pl.scan_csv(pvar_filename, **pvar_reading_args)
151
152            # We need to track row positions of variants but doing this with lazy operations would
153            # defeat predicate pushdown optimization and force loading the entire DataFrame.
154            # Solution: We only materialize the ID column with indices, keeping memory usage minimal
155            # while maintaining the ability to map between variant IDs and their file positions.
156            variant_id_idxs = pvar.select("ID").with_row_index().collect()
157
158            file_num_variants = variant_id_idxs.height
159
160            if variant_ids is not None:
161                num_variants = np.size(variant_ids)
162                pvar = pvar.filter(pl.col("ID").is_in(variant_ids)).collect()
163                variant_idxs = variant_id_idxs.filter(pl.col("ID").is_in(variant_ids)).select("index").to_series().to_numpy()
164            elif variant_idxs is not None:
165                num_variants = np.size(variant_idxs)
166                variant_idxs = np.array(variant_idxs, dtype=np.uint32)
167                variant_ids = variant_id_idxs.filter(pl.col("index").is_in(variant_idxs)).select("ID")
168                pvar = pvar.filter(pl.col("ID").is_in(variant_ids)).collect()
169            else:
170                num_variants = file_num_variants
171                variant_idxs = np.arange(num_variants, dtype=np.uint32)
172                pvar = pvar.collect()
173
174            log.info(f"Reading {filename_noext}.psam")
175
176            with open(filename_noext + ".psam") as file:
177                first_line = file.readline().strip()
178                psam_has_header = first_line.startswith(("#FID", "FID", "#IID", "IID"))
179
180            psam = pl.read_csv(
181                filename_noext + ".psam",
182                separator=separator,
183                has_header=psam_has_header,
184                new_columns=None if psam_has_header else ["FID", "IID", "PAT", "MAT", "SEX", "PHENO1"],
185                null_values=["NA"],
186            ).with_row_index()
187            if "#IID" in psam.columns:
188                psam = psam.rename({"#IID": "IID"})
189            if "#FID" in psam.columns:
190                psam = psam.rename({"#FID": "FID"})
191
192            file_num_samples = psam.height
193
194            if sample_ids is not None:
195                num_samples = np.size(sample_ids)
196                psam = psam.filter(pl.col("IID").is_in(sample_ids))
197                sample_idxs = psam.select("index").to_series().to_numpy()
198            elif sample_idxs is not None:
199                num_samples = np.size(sample_idxs)
200                sample_idxs = np.array(sample_idxs, dtype=np.uint32)
201                psam = psam.filter(pl.col("index").is_in(sample_idxs))
202            else:
203                num_samples = file_num_samples
204
205        if "GT" in fields:
206            log.info(f"Reading {filename_noext}.pgen")
207            pgen_reader = pg.PgenReader(
208                str.encode(filename_noext + ".pgen"),
209                raw_sample_ct=file_num_samples,
210                variant_ct=file_num_variants,
211                sample_subset=sample_idxs,
212            )
213
214            if only_read_pgen:
215                num_samples = pgen_reader.get_raw_sample_ct()
216                num_variants = pgen_reader.get_variant_ct()
217                variant_idxs = np.arange(num_variants, dtype=np.uint32)
218
219            # required arrays: variant_idxs + sample_idxs + genotypes
220            if not sum_strands:
221                required_ram = (num_samples + num_variants + num_variants * 2 * num_samples) * 4
222            else:
223                required_ram = (num_samples + num_variants) * 4 + num_variants * num_samples
224            log.info(f">{required_ram / 1024**3:.2f} GiB of RAM are required to process {num_samples} samples with {num_variants} variants each")
225
226            if not sum_strands:
227                genotypes = np.empty((num_variants, 2 * num_samples), dtype=np.int32)  # cannot use int8 because of pgenlib
228                pgen_reader.read_alleles_list(variant_idxs, genotypes)
229                genotypes = genotypes.astype(np.int8).reshape((num_variants, num_samples, 2))
230            else:
231                genotypes = np.empty((num_variants, num_samples), dtype=np.int8)
232                pgen_reader.read_list(variant_idxs, genotypes)
233            pgen_reader.close()
234        else:
235            genotypes = None
236
237        log.info("Constructing SNPObject")
238
239        snpobj = SNPObject(
240            calldata_gt=genotypes if "GT" in fields else None,
241            samples=psam.get_column("IID").to_numpy() if "IID" in fields and "IID" in psam.columns else None,
242            **{f'variants_{k.lower()}': pvar.get_column(v).to_numpy() if v in fields and v in pvar.columns else None
243               for k, v in {'ref': 'REF', 'alt': 'ALT', 'chrom': '#CHROM', 'id': 'ID', 'pos': 'POS', 'filter_pass': 'FILTER', 'qual': 'QUAL'}.items()}
244        )
245
246        log.info("Finished constructing SNPObject")
247        return snpobj