Skip to content

Latest commit

 

History

History
515 lines (369 loc) · 20 KB

user_guide.rst

File metadata and controls

515 lines (369 loc) · 20 KB

User Guide

Important

Use bio2zarr for genetic file format conversion to Zarr, then use :func:`sgkit.load_dataset` for analysis in sgkit. See the bio2zarr documentation for installation and conversion instructions.

BGEN reading and conversion have been removed from sgkit, along with the optional cbgen dependency. See :ref:`bgen_conversion_migration` for conversion guidance and support for existing legacy packed stores. PLINK reading and writing have also been removed; see :ref:`plink_conversion_migration`.

Sgkit reads Zarr datasets directly. For VCF, BGEN and PLINK conversion, install bio2zarr separately. The bgen and plink extras and their direct readers have been removed.

Genetic file format conversion belongs in bio2zarr, rather than sgkit. Convert inputs with bio2zarr before analysis, so the expensive conversion step need only be paid once. BGEN input requires a bio2zarr version with BGEN support (see :ref:`bgen_conversion_migration`).

After converting a file to Zarr format, it can be loaded with sgkit's :func:`sgkit.load_dataset`. This is like Xarray's :func:`xarray.open_zarr` function, with some useful defaults applied for you.

Similarly, the function :func:`sgkit.save_dataset` can be used to save the dataset in Zarr format. This can be used to save a newly-loaded dataset, or a dataset with new variables that are the result of running genetic methods on the data, so it can be loaded again later.

Sgkit requires Zarr Python >= 3 and Xarray >= 2026.9.0, but :func:`sgkit.save_dataset` continues to write Zarr format 2 by default. The Python package version does not determine the on-disk format. Pass zarr_format=3 to opt into format 3; :func:`sgkit.load_dataset` can read either format. The Xarray minimum includes support for the native NumPy string dtype returned by recent Zarr releases, preserving string-array chunking on load.

By default, Zarr format 2 datasets created by sgkit have consolidated metadata, which makes opening the dataset faster.

Warning

sgkit.io.bgen.read_bgen, sgkit.io.bgen.bgen_to_zarr and sgkit.io.bgen.rechunk_bgen have been removed, together with the bgen extra. Sgkit no longer reads BGEN files or depends on cbgen.

Use bio2zarr for BGEN conversion, then load the resulting Zarr dataset with :func:`sgkit.load_dataset`. Install a bio2zarr version that includes BGEN support; availability is tracked in bio2zarr issue 417. An older bio2zarr release without BGEN support cannot replace the removed reader. Sgkit does not provide an alternative BGEN conversion path.

Some older sgkit BGEN converters packed variables to reduce storage size: they saved only the heterozygous and homozygous-alternate probabilities, omitted dosage, and used one missingness flag per variant/sample pair. For these known legacy packed stores only, use :func:`sgkit.io.bgen.unpack_variables` after loading and before methods that expect complete genotype probabilities or dosage. This helper operates on an existing Xarray dataset and needs no BGEN backend:

import sgkit as sg
from sgkit.io.bgen import unpack_variables

packed = sg.load_dataset("legacy-packed.zarr")
ds = unpack_variables(packed)

The expected signature is call_genotype_probability with dimensions (variants, samples, genotypes) and two stored probabilities, plus call_genotype_probability_mask with dimensions (variants, samples). Unpacking restores three diploid, biallelic genotype probabilities, computes alternate-allele dosage, and expands the probability mask. For example, [0.25, 0.5] becomes [0.25, 0.25, 0.5] with dosage 1.25. Missing probabilities remain NaN. The input is unchanged and Dask arrays remain lazy; lost quantization precision cannot be recovered.

Do not unpack ordinary datasets, new bio2zarr output, or an already-unpacked store. A two-state probability array alone is not evidence of legacy packing: other ploidies or allele counts can also produce two states. The helper is not a general BGEN decoder or schema adapter. :func:`sgkit.load_dataset` does not automatically detect or unpack legacy stores. The original sgkit.io.bgen.bgen_reader.unpack_variables import remains supported as a compatibility alias; that module no longer contains a reader.

Use bio2zarr for PLINK-to-Zarr conversion and vcztools for export. The sgkit.io.plink module, including read_plink, plink_to_zarr, write_plink and zarr_to_plink, has been removed, along with the plink extra and the bed-reader dependency.

For a PLINK 1 fileset named input.bed, input.bim and input.fam, convert using the shared path prefix:

plink2zarr convert input output.zarr

Then load the result for analysis:

import sgkit as sg

ds = sg.load_dataset("output.zarr")

This conversion preserves hard calls and variant/sample identifiers; it is not a complete archival copy of the companion files. In particular, bio2zarr does not import FAM family/parent IDs, sex or phenotype, or BIM genetic-map positions. Consult the converter documentation when migrating code that used those fields.

TODO: Show how to read/write Zarr data in cloud storage

Most :ref:`genetic_methods` in sgkit operate on a few variables in an Xarray dataset. Variables have default names, so you can usually just pass in the dataset, but it's also possible to use different variable names.

.. ipython:: python
    :okwarning:

    import sgkit as sg
    ds = sg.simulate_genotype_call_dataset(n_variant=100, n_sample=50, missing_pct=.1)
    ds = ds[['variant_allele', 'call_genotype']]
    ds

    # Use the default variable (call_genotype)
    sg.count_call_alleles(ds).call_allele_count

    # Create a copy of the call_genotype variable, and use that to compute counts
    # (More realistically, this variable would be created from another computation or input.)
    ds["my_call_genotype"] = ds["call_genotype"]
    sg.count_call_alleles(ds, call_genotype="my_call_genotype").call_allele_count

For a full list of variables and their default names, see :ref:`api_variables`.

Methods declare the variables that they use directly. If the variable exists in the dataset, then it will be used for the computation.

If the variable doesn't exist in the dataset, then it will be computed if the variable name is the default one. For example, :func:`sgkit.count_variant_alleles` declares call_allele_count as a variable it needs to perform its computation. If the dataset doesn't contain call_allele_count, then the method will call :func:`sgkit.count_call_alleles` to populate it, before running its own computation.

.. ipython:: python
    :okwarning:

    # The following will create call_allele_count and variant_allele_count
    sg.count_variant_alleles(ds)

If however a non-default variable name is used and it doesn't exist in the dataset, then the intermediate variable is not populated, and an error is raised, since sgkit expects the user to have created the variable in that case.

.. ipython:: python
    :okexcept:

    sg.count_variant_alleles(ds, call_allele_count="my_call_allele_count")

There are also some variables that cannot be automatically defined, such as call_genotype, since it can't be computed from other data.

Generally, method functions in sgkit compute some new variables based on the input dataset, then return a new output dataset that consists of the input dataset plus the new computed variables. The input dataset is unchanged.

This behavior can be controlled using the merge parameter. If set to True (the default), then the function will merge the input dataset and the computed output variables into a single dataset. Output variables will overwrite any input variables with the same name, and a warning will be issued in this case. If False, the function will return only the computed output variables.

Examples:

.. ipython:: python
    :okwarning:

    import sgkit as sg
    ds = sg.simulate_genotype_call_dataset(n_variant=100, n_sample=50, missing_pct=.1)
    ds = ds[['variant_allele', 'call_genotype']]
    ds

    # By default, new variables are merged into a copy of the provided dataset
    ds = sg.count_variant_alleles(ds)
    ds

    # If an existing variable would be re-defined, a warning is thrown
    import warnings
    ds = sg.count_variant_alleles(ds)
    with warnings.catch_warnings(record=True) as w:
        ds = sg.count_variant_alleles(ds)
        print(f"{w[0].category.__name__}: {w[0].message}")

    # New variables can also be returned in their own dataset
    sg.count_variant_alleles(ds, merge=False)

    # This can be useful for merging multiple datasets manually
    ds.merge(sg.count_variant_alleles(ds, merge=False))

Merge can be used to rename output variables too.

.. ipython:: python
    :okwarning:

    import sgkit as sg
    ds = sg.simulate_genotype_call_dataset(n_variant=100, n_sample=50, missing_pct=.1)
    ds = ds[['variant_allele', 'call_genotype']]

    ds.merge(sg.count_variant_alleles(ds, merge=False).rename(variant_allele_count="my_variant_allele_count"))

Note that there is a limitation where intermediate variables (call_allele_count in this case) are not returned if merge=False. See #405.

It's usually easier to pass genetic data between Python libraries as simple NumPy arrays, rather than saving them in files. In sgkit, any data variable can be computed and extracted as a NumPy array using the .values property.

Genetic data is usually stored in a :data:`sgkit.variables.call_genotype_spec` array which has three dimensions (variants, samples and ploidy). This data structure can be difficult to work with in generic statistical libraries and it is often necessary to convert genotypes to a single value per call. The :func:`sgkit.convert_call_to_index` method converts call genotypes into :data:`sgkit.variables.call_genotype_index_spec` which represents each call as a single integer value. For biallelic datasets, this value is simply the count of the alternate allele. Genotype calls with missing alleles will be converted to a -1.

.. ipython:: python
    :okwarning:

    import sgkit as sg
    # Simulate biallelic genotype calls
    ds = sg.simulate_genotype_call_dataset(n_variant=10, n_sample=8, missing_pct=.1, seed=0)
    sg.display_genotypes(ds)

    # Convert genotype calls into a numpy array of alternate allele counts
    sg.convert_call_to_index(ds).call_genotype_index.values

TODO: Show to use a custom naming convention via Xarray renaming features.

TODO: Show how something like sample metadata can be joined to an existing Xarray dataset. Also briefly explain indexing and uniqueness within Xarray/Pandas, since this is critical for understanding joins.

TODO: Finish explaining how Numba works and how users might apply it

Here is an example that demonstrates an alt allele count:

.. ipython:: python

    import numba
    import sgkit as sg
    import numpy as np

    ds = sg.simulate_genotype_call_dataset(5, 3, missing_pct=.2)

    def alt_allele_count(gt):
        out = np.full(gt.shape[:2], -1, dtype=np.int64)
        for i, j in np.ndindex(*out.shape):
            if np.all(gt[i, j] >= 0):
                out[i, j] = np.sum(gt[i, j] > 0)
        return out

    numba.njit(alt_allele_count)(ds.call_genotype.values)

TODO: Describe the upstream tools for PCA (i.e. those in dask-ml/scikit-learn)

For parallelism sgkit uses Dask to distribute work across workers. These workers can either be local processes or remote processes running on a cluster. By default, dask creates a local cluster with one worker thread per CPU core on the machine. This is useful for testing and development, but for larger datasets it is often necessary to run on a cluster. There are several options for this including:

Research institutions often use a job scheduler like Slurm, LFS or PBS to manage compute resources, so here is a worked example of using Dask Jobqueue to run sgkit on an LSF cluster. The process is similar for other schedulers, see the Dask Jobqueue documentation

The first step is to create a dask scheduler that sgkit connects to. It is often desirable to have this running in a separate process, in a python REPL or notebook, so that its scaling can be adjusted on the fly. This needs to be done on a login node or other machine that has job submission privileges.

from dask_jobqueue import LSFCluster
cluster = LSFCluster(
    cores=1,           # Number of cores per-worker
    memory="16GB",     # Amount of memory per-worker
    project="myproject",
    queue="myqueue",   # LSF queue
    walltime="04:00",  # Set this to a reasonable value for your data size
    use_stdin=True,    # This is often required for dask to work with LSF
    lsf_units="MB",    # This may very depending on your clusters config
    # Any bash commands needed to setup the environment
    job_script_prolouge="module load example/3.8.5",
    # This last line is useful to gracefully rotate out workers as they reach
    # the maximum wallclock time for the given queue. This allows a long-running
    # cluster to run on queues with shorter wallclock limits, but that likely
    # have higher priority.
    worker_extra_args=["--lifetime", "350m", "--lifetime-stagger", "5m"],
)

Now that the cluster is created we can view its dashboard at the address from cluster.dashboard_link, at this point it will have zero workers.

To launch workers we can request a fixed amount with:

cluster.scale(10)

At which point the worker jobs will be scheduled and once they are running they will connect to the dask scheduler and be visible in the dashboard.

The other option is to let dask adaptively scale the number of workers based on the amount of task that are queued up. This can prevent idle workers when there are no tasks to run, but there may be a short delay when new tasks are submitted while the workers are being scheduled. To enable this we can use:

cluster.adapt(minimum=0, maximum=100, wait_count=120)

The minimum and maximum are optional, but it is recommended to set them to reasonable values to prevent the cluster launching thousands of jobs. The wait_count is the number of update cycles that dask will wait before scaling down the cluster when there are no tasks to run. By default the update interval is 1s so the above command will wait 120s before scaling down.

Now that the cluster is running we can connect to it from our sgkit code. The scheduler address is available from cluster.scheduler_address. Then in your sgkit code pass this value to the dask Client constructor:

import sgkit as sg
from dask.distributed import Client
client = Client("http://scheduler-address:8786")
# Now run sgkit code as normal
ds = sgkit.load_dataset("mydata.zarr")
ds = sgkit.variant_stats(ds)

TODO: Show CuPy examples

The simplest way to monitor operations when running sgkit on a single host is to use Dask local diagnostics.

As an example, this code shows how to track the progress of a single sgkit function:

.. ipython:: python
    :okwarning:

    import sgkit as sg
    from dask.diagnostics import ProgressBar
    ds = sg.simulate_genotype_call_dataset(n_variant=100, n_sample=50, missing_pct=.1)
    with ProgressBar():
        ac = sg.count_variant_alleles(ds).variant_allele_count.compute()
    ac[:5]

Monitoring resource utilization with ResourceProfiler and profiling task streams with Profiler are other commonly used local diagnostics.

For similar monitoring in a distributed cluster, see Dask distributed diagnostics.

Dask allows you to visualize the task graph of a computation before running it, which can be handy when trying to understand where the bottlenecks are.

In most cases the number of tasks is too large to visualize, so it's useful to restrict the graph just a few chunks, as shown in this example.

.. ipython:: python
    :okwarning:

    import sgkit as sg
    ds = sg.simulate_genotype_call_dataset(n_variant=100, n_sample=50, missing_pct=.1)
    # Rechunk to illustrate multiple tasks
    ds = ds.chunk({"variants": 25, "samples": 25})
    counts = sg.count_call_alleles(ds).call_allele_count.data

    # Restrict to first 3 chunks in variants dimension
    counts = counts[:3*counts.chunksize[0],...]

    counts.visualize(optimize_graph=True)

_static/mydask.png

By passing keyword arguments to visualize we can see the order tasks will run in:

.. ipython:: python

    # Graph where colors indicate task ordering
    counts.visualize(filename="order", optimize_graph=True, color="order", cmap="autumn", node_attr={"penwidth": "4"})

_static/order.png

Task order number is shown in circular boxes, colored from red to yellow.

Internally, sgkit uses the Numba JIT compiler to accelerate some methods. These methods will be compiled the first time that they are used following a new installation. The compiled methods are automatically cached so that recompilation is not required in future sessions. However, this may occasionally cause issues with some setups.

Caching of compiled sgkit methods can be disabled by setting the environment variable SGKIT_DISABLE_NUMBA_CACHE to 1. This variable can also be set from within a python session before loading sgkit.

.. ipython:: python
    :okwarning:

    import os
    os.environ['SGKIT_DISABLE_NUMBA_CACHE']='1'
    import sgkit as sg

With caching disabled, these methods will be compiled the first time that they are called during each session. Refer to the Numba notes on caching for more information.