Suggested changes

Author: eric-czech

setup.cfg

@@ -101,11 +101,9 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-dask.*]
 ignore_missing_imports = True
-<<<<<<< HEAD
 [mypy-fsspec.*]
-=======
+ignore_missing_imports = True
 [mypy-dask_ml.*]
->>>>>>> PCA implementation #95
 ignore_missing_imports = True
 [mypy-numpy.*]
 ignore_missing_imports = True

sgkit/stats/pca.py

@@ -9,6 +9,8 @@
 from typing_extensions import Literal
 from xarray import DataArray, Dataset
 
+from sgkit import variables
+
 from ..typing import ArrayLike, DType, RandomStateType
 from ..utils import conditional_merge_datasets
 from .aggregation import count_call_alleles
@@ -93,8 +95,10 @@ def pca_transform(
     """ Apply PCA estimator to new data """
     AC = _allele_counts(ds, variable, check_missing=check_missing)
     projection = est.transform(da.asarray(AC).T)
-    new_ds = Dataset({"sample_pca_projection": (("samples", "components"), projection)})
-    return conditional_merge_datasets(ds, new_ds, merge)
+    new_ds = Dataset(
+        {variables.sample_pca_projection: (("samples", "components"), projection)}
+    )
+    return conditional_merge_datasets(ds, variables.validate(new_ds), merge)
 
 
 def _get(est: BaseEstimator, attr: str, fn: Any = lambda v: v) -> Optional[ArrayLike]:
@@ -109,25 +113,25 @@ def _get(est: BaseEstimator, attr: str, fn: Any = lambda v: v) -> Optional[Array
 def pca_stats(ds: Dataset, est: BaseEstimator, *, merge: bool = True) -> Dataset:
     """ Extract attributes from PCA estimator """
     new_ds = {
-        "sample_pca_component": (
+        variables.sample_pca_component: (
             ("variants", "components"),
             _get(est, "components_", fn=lambda v: v.T),
         ),
-        "sample_pca_explained_variance": (
+        variables.sample_pca_explained_variance: (
             "components",
             _get(est, "explained_variance_"),
         ),
-        "sample_pca_explained_variance_ratio": (
+        variables.sample_pca_explained_variance_ratio: (
             "components",
             _get(est, "explained_variance_ratio_"),
         ),
     }
     new_ds = Dataset({k: v for k, v in new_ds.items() if v[1] is not None})
     if "sample_pca_component" in new_ds and "sample_pca_explained_variance" in new_ds:
-        new_ds["sample_pca_loading"] = new_ds["sample_pca_component"] * np.sqrt(
-            new_ds["sample_pca_explained_variance"]
-        )
-    return conditional_merge_datasets(ds, new_ds, merge)
+        new_ds[variables.sample_pca_loading] = new_ds[
+            variables.sample_pca_component
+        ] * np.sqrt(new_ds[variables.sample_pca_explained_variance])
+    return conditional_merge_datasets(ds, variables.validate(new_ds), merge)
 
 
 def pca(

sgkit/variables.py

@@ -247,6 +247,41 @@ def _check_field(
     ArrayLikeSpec("sample_pcs", ndim=2, kind="f")
 )
 """Sample PCs (PCxS)."""
+sample_pca_component, sample_pca_component_spec = SgkitVariables.register_variable(
+    ArrayLikeSpec("sample_pca_component", ndim=2, kind="f")
+)
+"""Principal axes defined as eigenvectors for sample covariance matrix.
+In the context of SVD, these are equivalent to the right singular vectors in
+the decomposition of a (N, M) matrix., i.e. ``dask_ml.decomposition.TruncatedSVD.components_``."""
+(
+    sample_pca_explained_variance,
+    sample_pca_explained_variance_spec,
+) = SgkitVariables.register_variable(
+    ArrayLikeSpec("sample_pca_explained_variance", ndim=1, kind="f")
+)
+"""Variance explained by each principal component. These values are equivalent
+to eigenvalues that result from the eigendecomposition of a (N, M) matrix,
+i.e. ``dask_ml.decomposition.TruncatedSVD.explained_variance_``."""
+(
+    sample_pca_explained_variance_ratio,
+    sample_pca_explained_variance_ratio_spec,
+) = SgkitVariables.register_variable(
+    ArrayLikeSpec("sample_pca_explained_variance_ratio", ndim=1, kind="f")
+)
+"""Ratio of variance explained to total variance for each principal component,
+i.e. ``dask_ml.decomposition.TruncatedSVD.explained_variance_ratio_``."""
+sample_pca_loading, sample_pca_loading_spec = SgkitVariables.register_variable(
+    ArrayLikeSpec("sample_pca_loading", ndim=2, kind="f")
+)
+"""PCA loadings defined as principal axes scaled by square root of eigenvalues.
+These values  can also be interpreted  as the correlation between the original variables
+and unit-scaled principal axes."""
+sample_pca_projection, sample_pca_projection_spec = SgkitVariables.register_variable(
+    ArrayLikeSpec("sample_pca_projection", ndim=2, kind="f")
+)
+"""Projection of samples onto principal axes. This array is commonly
+referred to as "scores" or simply "principal components (PCs)" for a set of samples."""
+
 stat_Fst, stat_Fst_spec = SgkitVariables.register_variable(
     ArrayLikeSpec("stat_Fst", ndim=2, kind="f")
 )