-
Notifications
You must be signed in to change notification settings - Fork 72
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
randomized svd draft #3008
Draft
hanbin973
wants to merge
23
commits into
tskit-dev:main
Choose a base branch
from
hanbin973:rsvd
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Draft
randomized svd draft #3008
Changes from 11 commits
Commits
Show all changes
23 commits
Select commit
Hold shift + click to select a range
1f45245
randomized svd draft
hanbin973 e408ab3
modified api remove scipy
hanbin973 a176132
remove scipy
hanbin973 8c662c8
correct docstring and comments
hanbin973 aa13613
space remove
hanbin973 5bf405a
rng to random seed
hanbin973 6e415e5
add windows feature
hanbin973 45ac61e
output shape change when windows=wholegnome
hanbin973 2cdb9dd
make centre work with nodes
hanbin973 1f194aa
remove redundant options from internal functions
hanbin973 fdc5842
start at testing
petrelharp c0a2854
change variable name to n_* to num_*
hanbin973 667d60f
return range sketch matrix Q
hanbin973 36c10e9
fix random sketch option to handle None
hanbin973 0ac88b8
random_sketch needs windows specified
hanbin973 a6fa8ab
input checking for range_sketch to align with windows
hanbin973 791c7da
docstring change to reflect range_sketch
hanbin973 2f4ce2b
linting has a bug; when converting lambda to ordinary function defini…
hanbin973 be2f736
now output is a dataclass
hanbin973 bcbbcf6
docstring change
hanbin973 1a8ff7a
change variable name of PCAResult class
hanbin973 af16340
move internal function of PCA out
hanbin973 e81db15
function rearrangement
hanbin973 File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -8592,6 +8592,195 @@ def genetic_relatedness_vector( | |
) | ||
return out | ||
|
||
def pca( | ||
self, | ||
n_components: int = 10, | ||
windows: list = None, | ||
samples: np.ndarray = None, | ||
individuals: np.ndarray = None, | ||
mode: str = "branch", | ||
centre: bool = True, | ||
iterated_power: int = 3, | ||
n_oversamples: int = 10, | ||
random_seed: int = None, | ||
) -> (np.ndarray, np.ndarray): | ||
""" | ||
Run randomized singular value decomposition (rSVD) to obtain principal | ||
components. | ||
API partially adopted from `scikit-learn`'s | ||
`sklearn.decomposition.PCA.html` | ||
|
||
By default, performs PCA for the samples, so output has one coordinate | ||
for each sample), but alternatively either a list of sample IDs or a | ||
list of individual IDs can be provided (but not both). | ||
|
||
TODO: say exactly what is returned (and relationship to | ||
:meth:`genetic_relatedness <.TreeSequence.genetic_relatedness>`). | ||
|
||
TODO: say what algorithms are used. | ||
|
||
:param int n_components: Number of principal components. | ||
:param list windows: An increasing list of breakpoints between the windows | ||
to compute the statistic in. | ||
:param np.ndarray samples: Samples to perform PCA with. | ||
:param np.ndarray individuals: Individuals to perform PCA with. Cannot specify | ||
both `samples` and `individuals`. | ||
:param str mode: A string giving the "type" of relatedness to be computed | ||
(defaults to "branch"; see | ||
:meth:`genetic_relatedness_vector | ||
<.TreeSequence.genetic_relatedness_vector>`) | ||
:param bool centre: Centre the genetic relatedness matrix. | ||
:param int iterated_power: Number of power iteration of range finder. | ||
:param int n_oversamples: Number of additional test vectors. | ||
:param int random_seed: The random seed. If this is None, a random seed will | ||
be automatically generated. Valid random seeds must be between 1 and | ||
:math:`2^32 − 1`. | ||
:return: A tuple (U, D) of ndarrays, with the principal component loadings in U | ||
and the principal values in D. | ||
""" | ||
|
||
if samples is None and individuals is None: | ||
samples = self.samples() | ||
|
||
if samples is not None and individuals is not None: | ||
raise ValueError("Samples and individuals cannot be used at the same time") | ||
elif samples is not None: | ||
output_type = "node" | ||
dim = len(samples) | ||
else: | ||
assert individuals is not None | ||
output_type = "individual" | ||
dim = len(individuals) | ||
|
||
if n_components > dim: | ||
raise ValueError( | ||
"Number of components must be less than or equal to " | ||
"the number of samples (or individuals, if specified)." | ||
) | ||
|
||
random_state = np.random.default_rng(random_seed) | ||
|
||
def _rand_pow_range_finder( | ||
operator, | ||
operator_dim: int, | ||
rank: int, | ||
depth: int, | ||
num_vectors: int, | ||
rng: np.random.Generator, | ||
) -> np.ndarray: | ||
""" | ||
Algorithm 9 in https://arxiv.org/pdf/2002.01387 | ||
""" | ||
assert num_vectors >= rank > 0 | ||
test_vectors = rng.normal(size=(operator_dim, num_vectors)) | ||
Q = test_vectors | ||
for _ in range(depth): | ||
Q = np.linalg.qr(Q).Q | ||
Q = operator(Q) | ||
Q = np.linalg.qr(Q).Q | ||
return Q[:, :rank] | ||
|
||
def _rand_svd( | ||
operator, | ||
operator_dim: int, | ||
rank: int, | ||
depth: int, | ||
num_vectors: int, | ||
rng: np.random.Generator, | ||
) -> (np.ndarray, np.ndarray, np.ndarray): | ||
""" | ||
Algorithm 8 in https://arxiv.org/pdf/2002.01387 | ||
""" | ||
assert num_vectors >= rank > 0 | ||
Q = _rand_pow_range_finder( | ||
operator, operator_dim, num_vectors, depth, num_vectors, rng | ||
) | ||
C = operator(Q).T | ||
U_hat, D, V = np.linalg.svd(C, full_matrices=False) | ||
U = Q @ U_hat | ||
return U[:, :rank], D[:rank], V[:rank] | ||
|
||
def _genetic_relatedness_vector_individual( | ||
petrelharp marked this conversation as resolved.
Show resolved
Hide resolved
|
||
arr: np.ndarray, | ||
centre: bool = True, | ||
windows=None, | ||
) -> np.ndarray: | ||
ij = np.vstack( | ||
[ | ||
[n, k] | ||
for k, i in enumerate(individuals) | ||
for n in self.individual(i).nodes | ||
] | ||
) | ||
samples, sample_individuals = ( | ||
ij[:, 0], | ||
ij[:, 1], | ||
) # sample node index, individual of those nodes | ||
x = ( | ||
arr - arr.mean(axis=0) if centre else arr | ||
) # centering within index in rows | ||
x = self.genetic_relatedness_vector( | ||
W=x[sample_individuals], | ||
windows=windows, | ||
mode=mode, | ||
centre=False, | ||
nodes=samples, | ||
)[0] | ||
|
||
def bincount_fn(w): | ||
np.bincount(sample_individuals, w) | ||
|
||
x = np.apply_along_axis(bincount_fn, axis=0, arr=x) | ||
x = x - x.mean(axis=0) if centre else x # centering within index in cols | ||
|
||
return x | ||
|
||
def _genetic_relatedness_vector_node( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. same: automatic linting |
||
arr: np.ndarray, | ||
centre: bool = True, | ||
windows=None, | ||
) -> np.ndarray: | ||
x = arr - arr.mean(axis=0) if centre else arr | ||
x = self.genetic_relatedness_vector( | ||
W=x, windows=windows, mode=mode, centre=False, nodes=samples | ||
)[0] | ||
x = x - x.mean(axis=0) if centre else x | ||
|
||
return x | ||
|
||
drop_windows = windows is None | ||
windows = self.parse_windows(windows) | ||
num_windows = len(windows) - 1 | ||
if num_windows < 1: | ||
raise ValueError("Number of windows must be at least 1.") | ||
|
||
U = np.empty((num_windows, dim, n_components)) | ||
D = np.empty((num_windows, n_components)) | ||
for i in range(num_windows): | ||
this_window = windows[i : i + 2] | ||
_f = ( | ||
_genetic_relatedness_vector_node | ||
if output_type == "node" | ||
else _genetic_relatedness_vector_individual | ||
) | ||
|
||
def _G(x): | ||
_f(x, centre=centre, windows=this_window) # NOQA: B023 | ||
|
||
U[i], D[i], _ = _rand_svd( | ||
operator=_G, | ||
operator_dim=dim, | ||
rank=n_components, | ||
depth=iterated_power, | ||
num_vectors=n_components + n_oversamples, | ||
rng=random_state, | ||
) | ||
|
||
if drop_windows: | ||
U, D = U[0], D[0] | ||
|
||
return U, D | ||
|
||
def trait_covariance(self, W, windows=None, mode="site", span_normalise=True): | ||
""" | ||
Computes the mean squared covariances between each of the columns of ``W`` | ||
|
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
perhaps we should not have a default?