Remove (c)bgen dependency for unpack_variables?

[jpfeuffer]

Hi!
Thanks for creating sgkit. Very useful package.
We are using it for performing calculations on a dataset that was converted from bgen.

For reading, our code uses something like this:

    ds_region = ds.isel(variants=region_indices, samples=sample_indices)
    ds_region = bgen.bgen_reader.unpack_variables(ds_region)

However due to the import of the bgen module, this pulls in the whole (btw. conda-incompatible) (c)bgen stack, although this function seems to be working perfectly fine without that dependency.

a) Are we using unpack_variables correctly or are we missing something?
b) If so, would you be fine with refactoring such that this function (and potentially others that do not actually depend on cbgen) can be used without that additional dependency?

I might be able to do a PR unless you have some plans for that module already.

[jeromekelleher]

Hi @jpfeuffer 👋 Great to hear you're using sgkit!

I'm not in a position to comment about how you're using the bgen interface (@tomwhite knows this much better than me, and I think he's away for a bit), but I just wanted to point out that we generally want to move away from "views" on the data for plink and bgen in sgkit, towards one-time conversion to VCF Zarr using bio2zarr. That way we can simplify sgkit to just focus on the Zarr backend, and move the data conversion logic to bio2zarr.

Would conversion to Zarr work for you, or do you need to access the BGEN files directly?

I've created issue sgkit-dev/bio2zarr#417 over at bio2zarr to track adding bgen support there, if you'd like to comment.

[jpfeuffer]

Hey! Thank you for the quick answer.
We are using zarr but want to restore dosages from all the separate attributes stored in the zarr that was converted from bgen.

The function which you are providing to do this is:

def unpack_variables(ds: xr.Dataset, dtype: str= "float32") -> xr.Dataset:
    # Restore homozygous reference GP
    gp = ds["call_genotype_probability"].astype(dtype)
    if gp.sizes["genotypes"] != 2:
        raise ValueError(
            "Expecting variable 'call_genotype_probability' to have genotypes "
            f"dimension of size 2 (received sizes = {dict(gp.sizes)})"
        )
    ds = ds.drop_vars("call_genotype_probability")
    ds["call_genotype_probability"] = xr.concat(
        [1 - gp.sum(dim="genotypes", skipna=False), gp], dim="genotypes"
    )

    # Restore dosage
    ds["call_dosage"] = gp[..., 0] + 2 * gp[..., 1]
    ds["call_dosage_mask"] = ds["call_genotype_probability_mask"]
    ds["call_genotype_probability_mask"] = ds[
        "call_genotype_probability_mask"
    ].broadcast_like(ds["call_genotype_probability"])
    return ds

Unfortunately it is guarded by (for this function) unnecessary imports of cbgen since it is coming from your bgen module.
Currently we are therefore just copying this function into our code.

[jpfeuffer]

I guess what I am proposing is, maybe this could move into a utility module to work with "bgen-style" zarrs? Or avoid cbgen import for functions in the bgen module where cbgen is actually not necessary.

[jeromekelleher]

We are using zarr but want to restore dosages from all the separate attributes stored in the zarr that was converted from bgen.

Ah, interesting! How do you get the Zarr then? Is the root of the issue that the Zarr doesn't contain call_dosage? Would it be a solution if bio2zarr did the conversion and included call_dosage?

Coming back to the original point though, I'm sure we can move this function to somewhere else in the library if its useful.

[jpfeuffer]

Yes, right. An alternative solution might indeed be to store the call_dosages as additional or alternative data variable during conversion already. I am not too familiar with the specifications though and what speaks against/for having one or the other. Having both incurs quite some storage overhead of course.

I think we are still using bgen.bgen_to_zarr but we wanted to look into bio2zarr anyway.