def test_create_genotype_dosage_dataset():
variant_contig_names = ["chr1"]
variant_contig = np.array([0, 0], dtype="i1")
variant_position = np.array([1000, 2000], dtype="i4")
variant_alleles = np.array([["A", "C"], ["G", "A"]], dtype="S1")
variant_id = np.array(["rs1", "rs2"], dtype=str)
sample_id = np.array(["sample_1", "sample_2", "sample_3"], dtype=str)
call_dosage = np.array([[0.8, 0.9, 1.0], [1.0, 1.1, 1.2]], dtype="f4")
ds = create_genotype_dosage_dataset(
variant_contig_names=variant_contig_names,
variant_contig=variant_contig,
variant_position=variant_position,
variant_alleles=variant_alleles,
sample_id=sample_id,
call_dosage=call_dosage,
variant_id=variant_id,
)
assert DIM_VARIANT in ds.dims
assert DIM_SAMPLE in ds.dims
assert ds.attrs["contigs"] == variant_contig_names
assert_array_equal(ds["variant_contig"], variant_contig)
assert_array_equal(ds["variant_position"], variant_position)
assert_array_equal(ds["variant_allele"], variant_alleles)
assert_array_equal(ds["variant_id"], variant_id)
assert_array_equal(ds["sample_id"], sample_id)
assert_array_equal(ds["call_dosage"], call_dosage)
assert_array_equal(ds["call_dosage_mask"], np.isnan(call_dosage))
def read_bgen(
path: PathType,
metafile_path: Optional[PathType] = None,
sample_path: Optional[PathType] = None,
chunks: Union[str, int, Tuple[int, int, int]] = "auto",
lock: bool = False,
persist: bool = True,
contig_dtype: DType = "str",
gp_dtype: DType = "float32",
) -> Dataset:
"""Read BGEN dataset.
Loads a single BGEN dataset as dask arrays within a Dataset
from a ``.bgen`` file.
Parameters
----------
path
Path to BGEN file.
metafile_path
Path to companion index file used to determine BGEN byte offsets.
Defaults to ``path`` + ".metafile" if not provided.
This file is necessary for reading BGEN genotype probabilities and it will be
generated the first time the file is read if it does not already exist.
If it needs to be created, it can make the first call to this function
much slower than subsequent calls.
sample_path
Path to ``.sample`` file, by default None. This is used to fetch sample identifiers
and when provided it is preferred over sample identifiers embedded in the ``.bgen`` file.
chunks
Chunk size for genotype probability data (3 dimensions),
by default "auto".
lock
Whether or not to synchronize concurrent reads of
file blocks, by default False. This is passed through to
[dask.array.from_array](https://docs.dask.org/en/latest/array-api.html#dask.array.from_array).
persist
Whether or not to persist variant information in memory, by default True.
This is an important performance consideration as the metadata file for this data will
be read multiple times when False.
contig_dtype
Data type for contig names, by default "str".
This may also be an integer type (e.g. "int"), but will fail if any of the contig names
cannot be converted to integers.
gp_dtype
Data type for genotype probabilities, by default "float32".
Warnings
--------
Only bi-allelic, diploid BGEN files are currently supported.
Returns
-------
A dataset containing the following variables:
- :data:`sgkit.variables.variant_id_spec` (variants)
- :data:`sgkit.variables.variant_contig_spec` (variants)
- :data:`sgkit.variables.variant_position_spec` (variants)
- :data:`sgkit.variables.variant_allele_spec` (variants)
- :data:`sgkit.variables.sample_id_spec` (samples)
- :data:`sgkit.variables.call_dosage_spec` (variants, samples)
- :data:`sgkit.variables.call_dosage_mask_spec` (variants, samples)
- :data:`sgkit.variables.call_genotype_probability_spec` (variants, samples, genotypes)
- :data:`sgkit.variables.call_genotype_probability_mask_spec` (variants, samples, genotypes)
"""
if isinstance(chunks, tuple) and len(chunks) != 3:
raise ValueError(f"`chunks` must be tuple with 3 items, not {chunks}")
if not np.issubdtype(gp_dtype, np.floating):
raise ValueError(
f"`gp_dtype` must be a floating point data type, not {gp_dtype}"
)
if not np.issubdtype(contig_dtype, np.integer) and np.dtype(
contig_dtype
).kind not in {"U", "S"}:
raise ValueError(
f"`contig_dtype` must be of string or int type, not {contig_dtype}"
)
path = Path(path)
sample_path = Path(sample_path) if sample_path else path.with_suffix(".sample")
if sample_path.exists():
sample_id = read_samples(sample_path).sample_id.values.astype("U")
else:
sample_id = _default_sample_ids(path)
bgen_reader = BgenReader(path, metafile_path=metafile_path, dtype=gp_dtype)
df = read_metafile(bgen_reader.metafile_path)
if persist:
df = df.persist()
arrs = dataframe_to_dict(df, METAFILE_DTYPE)
variant_id = arrs["id"]
variant_contig = arrs["chrom"].astype(contig_dtype)
variant_contig, variant_contig_names = encode_contigs(variant_contig)
variant_contig_names = list(variant_contig_names)
variant_position = arrs["pos"]
variant_allele = da.hstack((arrs["a1"][:, np.newaxis], arrs["a2"][:, np.newaxis]))
call_genotype_probability = da.from_array(
bgen_reader,
chunks=chunks,
lock=lock,
fancy=False,
asarray=False,
name=f"{bgen_reader.name}:read_bgen:{path}",
)
call_dosage = _to_dosage(call_genotype_probability)
ds: Dataset = create_genotype_dosage_dataset(
variant_contig_names=variant_contig_names,
variant_contig=variant_contig,
variant_position=variant_position,
variant_allele=variant_allele,
sample_id=sample_id,
call_dosage=call_dosage,
call_genotype_probability=call_genotype_probability,
variant_id=variant_id,
)
return ds
def read_bgen(
path: PathType,
chunks: Union[str, int, Tuple[int, ...]] = "auto",
lock: bool = False,
persist: bool = True,
dtype: Any = "float32",
) -> Dataset:
"""Read BGEN dataset.
Loads a single BGEN dataset as dask arrays within a Dataset
from a bgen file.
Parameters
----------
path : PathType
Path to BGEN file.
chunks : Union[str, int, tuple], optional
Chunk size for genotype probability data (3 dimensions),
by default "auto".
lock : bool, optional
Whether or not to synchronize concurrent reads of
file blocks, by default False. This is passed through to
[dask.array.from_array](https://docs.dask.org/en/latest/array-api.html#dask.array.from_array).
persist : bool, optional
Whether or not to persist variant information in
memory, by default True. This is an important performance
consideration as the metadata file for this data will
be read multiple times when False.
dtype : Any
Genotype probability array data type, by default float32.
Warnings
--------
Only bi-allelic, diploid BGEN files are currently supported.
"""
bgen_reader = BgenReader(path, persist, dtype=dtype)
variant_contig, variant_contig_names = encode_array(bgen_reader.contig.compute())
variant_contig_names = list(variant_contig_names)
variant_contig = variant_contig.astype("int16")
variant_position = np.asarray(bgen_reader.pos, dtype=int)
variant_alleles = np.asarray(bgen_reader.variant_alleles, dtype="S")
variant_id = np.asarray(bgen_reader.variant_id, dtype=str)
sample_id = np.asarray(bgen_reader.sample_id, dtype=str)
call_genotype_probability = da.from_array(
bgen_reader,
chunks=chunks,
lock=lock,
fancy=False,
asarray=False,
name=f"{bgen_reader.name}:read_bgen:{path}",
)
call_dosage = _to_dosage(call_genotype_probability)
ds: Dataset = create_genotype_dosage_dataset(
variant_contig_names=variant_contig_names,
variant_contig=variant_contig,
variant_position=variant_position,
variant_alleles=variant_alleles,
sample_id=sample_id,
call_dosage=call_dosage,
call_genotype_probability=call_genotype_probability,
variant_id=variant_id,
)
return ds