Change Dataset.dims to Dataset.sizes to address Xarray FutureWarning

Author: tomwhite

sgkit/display.py

@@ -125,7 +125,7 @@ def truncate(ds: xr.Dataset, max_sizes: Mapping[Hashable, int]) -> xr.Dataset:
     """
     sel = dict()
     for dim, size in max_sizes.items():
-        if ds.dims[dim] <= size:
+        if ds.sizes[dim] <= size:
             # No truncation required
             pass
         else:
@@ -194,7 +194,7 @@ def display_genotypes(
         ds_calls = set_index_if_unique(ds_calls, "variants", variant_index)
     # convert call genotypes to strings
     calls = ds_calls["call_genotype"].values
-    max_chars = max(2, len(str(ds.dims["alleles"] - 1)))
+    max_chars = max(2, len(str(ds.sizes["alleles"] - 1)))
     if "call_genotype_phased" in ds_calls:
         phased = ds_calls["call_genotype_phased"].values
     else:

sgkit/io/bgen/bgen_reader.py

@@ -486,8 +486,8 @@ def rechunk_bgen(
     if isinstance(output, Path):
         output = str(output)
 
-    chunk_length = min(chunk_length, ds.dims["variants"])
-    chunk_width = min(chunk_width, ds.dims["samples"])
+    chunk_length = min(chunk_length, ds.sizes["variants"])
+    chunk_width = min(chunk_width, ds.sizes["samples"])
 
     if pack:
         ds = pack_variables(ds)

sgkit/io/plink/plink_writer.py

@@ -90,9 +90,9 @@ def write_plink(
         raise ValueError(
             "Either `path` or all 3 of `{bed,bim,fam}_path` must be specified but not both"
         )
-    if "ploidy" in ds.dims and ds.dims["ploidy"] != 2:
+    if "ploidy" in ds.sizes and ds.sizes["ploidy"] != 2:
         raise ValueError("write_plink only works for diploid genotypes")
-    if "alleles" in ds.dims and ds.dims["alleles"] != 2:
+    if "alleles" in ds.sizes and ds.sizes["alleles"] != 2:
         raise ValueError("write_plink only works for biallelic genotypes")
 
     if path:

sgkit/io/vcf/vcf_reader.py

@@ -571,8 +571,8 @@ def vcf_to_zarr_sequential(
 
             if first_variants_chunk:
                 # limit chunk width to actual number of samples seen in first chunk
-                if ds.dims["samples"] > 0:
-                    chunk_width = min(chunk_width, ds.dims["samples"])
+                if ds.sizes["samples"] > 0:
+                    chunk_width = min(chunk_width, ds.sizes["samples"])
 
                 # ensure that booleans are not stored as int8 by xarray https://github.com/pydata/xarray/issues/4386
                 for var in ds.data_vars:

sgkit/io/vcf/vcf_writer.py

@@ -147,7 +147,7 @@ def write_vcf(
 
         print(vcf_header, end="", file=output)
 
-        if input.dims["variants"] == 0:
+        if input.sizes["variants"] == 0:
             return
 
         header_info_fields = _info_fields(vcf_header)
@@ -174,8 +174,8 @@ def dataset_chunk_to_vcf(
 
     ds = ds.load()  # load dataset chunk into memory
 
-    n_variants = ds.dims["variants"]  # number of variants in this chunk
-    n_samples = ds.dims["samples"]  # number of samples in whole dataset
+    n_variants = ds.sizes["variants"]  # number of variants in this chunk
+    n_samples = ds.sizes["samples"]  # number of samples in whole dataset
 
     # fixed fields
 

sgkit/model.py

@@ -169,8 +169,8 @@ def create_genotype_dosage_dataset(
 
 def num_contigs(ds: xr.Dataset) -> ArrayLike:
     """Return the number of contigs in a dataset."""
-    if DIM_CONTIG in ds.dims:
-        return ds.dims[DIM_CONTIG]
+    if DIM_CONTIG in ds.sizes:
+        return ds.sizes[DIM_CONTIG]
     else:
         return len(ds.attrs["contigs"])
 

sgkit/stats/aggregation.py

@@ -75,7 +75,7 @@ def count_call_alleles(
     from .aggregation_numba_fns import count_alleles
 
     variables.validate(ds, {call_genotype: variables.call_genotype_spec})
-    n_alleles = ds.dims["alleles"]
+    n_alleles = ds.sizes["alleles"]
     G = da.asarray(ds[call_genotype])
     shape = (G.chunks[0], G.chunks[1], n_alleles)
     # use numpy array to avoid dask task dependencies between chunks
@@ -170,8 +170,8 @@ def count_variant_alleles(
         from .aggregation_numba_fns import count_alleles
 
         variables.validate(ds, {call_genotype: variables.call_genotype_spec})
-        n_alleles = ds.dims["alleles"]
-        n_variant = ds.dims["variants"]
+        n_alleles = ds.sizes["alleles"]
+        n_variant = ds.sizes["variants"]
         G = da.asarray(ds[call_genotype]).reshape((n_variant, -1))
         shape = (G.chunks[0], n_alleles)
         # use uint64 dummy array to return uin64 counts array
@@ -227,7 +227,7 @@ def count_cohort_alleles(
     >>> ds = sg.simulate_genotype_call_dataset(n_variant=5, n_sample=4)
 
     >>> # Divide samples into two cohorts
-    >>> ds["sample_cohort"] = xr.DataArray(np.repeat([0, 1], ds.dims["samples"] // 2), dims="samples")
+    >>> ds["sample_cohort"] = xr.DataArray(np.repeat([0, 1], ds.sizes["samples"] // 2), dims="samples")
     >>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
     samples    S0   S1   S2   S3
     variants
@@ -364,8 +364,8 @@ def count_variant_genotypes(
     mixed_ploidy = ds[call_genotype].attrs.get("mixed_ploidy", False)
     if mixed_ploidy:
         raise ValueError("Mixed-ploidy dataset")
-    ploidy = ds.dims["ploidy"]
-    n_alleles = ds.dims["alleles"]
+    ploidy = ds.sizes["ploidy"]
+    n_alleles = ds.sizes["alleles"]
     n_genotypes = _comb_with_replacement(n_alleles, ploidy)
     G = da.asarray(ds[call_genotype].data)
     N = np.empty(n_genotypes, np.uint64)
@@ -432,8 +432,8 @@ def genotype_coords(
     """
     from .conversion_numba_fns import _comb_with_replacement, _index_as_genotype
 
-    n_alleles = ds.dims["alleles"]
-    ploidy = ds.dims["ploidy"]
+    n_alleles = ds.sizes["alleles"]
+    ploidy = ds.sizes["ploidy"]
     n_genotypes = _comb_with_replacement(n_alleles, ploidy)
     max_chars = len(str(n_alleles - 1))
     # dummy variable for ploidy dim also specifies output dtype
@@ -553,7 +553,7 @@ def cohort_allele_frequencies(
     >>> ds = sg.simulate_genotype_call_dataset(n_variant=5, n_sample=4)
 
     >>> # Divide samples into two cohorts
-    >>> ds["sample_cohort"] = xr.DataArray(np.repeat([0, 1], ds.dims["samples"] // 2), dims="samples")
+    >>> ds["sample_cohort"] = xr.DataArray(np.repeat([0, 1], ds.sizes["samples"] // 2), dims="samples")
     >>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
     samples    S0   S1   S2   S3
     variants

sgkit/stats/association.py

@@ -210,7 +210,7 @@ def gwas_linear_regression(
 
     if len(covariates) == 0:
         if add_intercept:
-            X = da.ones((ds.dims["samples"], 1), dtype=np.float32)
+            X = da.ones((ds.sizes["samples"], 1), dtype=np.float32)
         else:
             raise ValueError("add_intercept must be True if no covariates specified")
     else:

sgkit/stats/conversion.py

@@ -100,7 +100,7 @@ def convert_call_to_index(
         raise ValueError("Mixed-ploidy dataset")
     G = da.asarray(ds[call_genotype].data)
     shape = G.chunks[0:2]
-    if ds.dims.get("alleles") == 2:  # default to general case
+    if ds.sizes.get("alleles") == 2:  # default to general case
         X = da.map_blocks(
             biallelic_genotype_call_index,
             G,
@@ -169,10 +169,10 @@ def convert_probability_to_call(
     variables.validate(
         ds, {call_genotype_probability: variables.call_genotype_probability_spec}
     )
-    if ds.dims["genotypes"] != 3:
+    if ds.sizes["genotypes"] != 3:
         raise NotImplementedError(
             f"Hard call conversion only supported for diploid, biallelic genotypes; "
-            f"num genotypes in provided probabilities array = {ds.dims['genotypes']}."
+            f"num genotypes in provided probabilities array = {ds.sizes['genotypes']}."
         )
     GP = da.asarray(ds[call_genotype_probability])
     # Remove chunking in genotypes dimension, if present

sgkit/stats/grm.py

@@ -176,7 +176,7 @@ def genomic_relationship(
                [1, 0, 0],
                [1, 1, 2]], dtype=uint8)
     >>> # use sample population frequency as ancestral frequency
-    >>> ds["sample_frequency"] = ds.call_dosage.mean(dim="samples") / ds.dims["ploidy"]
+    >>> ds["sample_frequency"] = ds.call_dosage.mean(dim="samples") / ds.sizes["ploidy"]
     >>> ds = sg.genomic_relationship(ds, ancestral_frequency="sample_frequency")
     >>> ds.stat_genomic_relationship.values # doctest: +NORMALIZE_WHITESPACE
         array([[ 0.93617021, -0.21276596, -0.72340426],
@@ -208,7 +208,7 @@ def genomic_relationship(
            [ 2.,  2.,  0.,  0.]])
     >>> ds["sample_frequency"] = ds.call_dosage.mean(
     ...     dim="samples", skipna=True
-    ... ) / ds.dims["ploidy"]
+    ... ) / ds.sizes["ploidy"]
     >>> ds = sg.genomic_relationship(
     ...     ds, ancestral_frequency="sample_frequency", skipna=True
     ... )
@@ -249,7 +249,7 @@ def genomic_relationship(
            [2.        , 2.        , 0.        , 0.        ]])
     >>> ds["sample_frequency"] = ds.call_dosage.mean(
     ...     dim="samples", skipna=True
-    ... ) / ds.dims["ploidy"]
+    ... ) / ds.sizes["ploidy"]
     >>> ds = sg.genomic_relationship(
     ...     ds,
     ...     call_dosage="call_dosage_imputed",
@@ -293,7 +293,7 @@ def genomic_relationship(
 
     estimator = estimator or EST_VAN_RADEN
     # TODO: raise on mixed ploidy
-    ploidy = ploidy or ds.dims.get("ploidy")
+    ploidy = ploidy or ds.sizes.get("ploidy")
     if ploidy is None:
         raise ValueError("Ploidy must be specified when the ploidy dimension is absent")
     dosage = da.array(ds[call_dosage].data)