Skip to content

Commit

Permalink
linting has a bug; when converting lambda to ordinary function defini…
Browse files Browse the repository at this point in the history
…tion, it omits return in the end of the function
  • Loading branch information
hanbin973 committed Oct 14, 2024
1 parent 791c7da commit 2f4ce2b
Showing 1 changed file with 15 additions and 9 deletions.
24 changes: 15 additions & 9 deletions python/tskit/trees.py
Original file line number Diff line number Diff line change
Expand Up @@ -8680,7 +8680,7 @@ def _rand_pow_range_finder(
"""
Algorithm 9 in https://arxiv.org/pdf/2002.01387
"""
assert num_vectors >= rank > 0
assert num_vectors >= rank > 0, "num_vectors should be larger than rank"
if range_sketch is None:
test_vectors = rng.normal(size=(operator_dim, num_vectors))
Q = test_vectors
Expand All @@ -8700,7 +8700,7 @@ def _rand_svd(
num_vectors: int,
rng: np.random.Generator,
range_sketch: np.ndarray = None,
) -> (np.ndarray, np.ndarray, np.ndarray):
) -> (np.ndarray, np.ndarray, np.ndarray, float):
"""
Algorithm 8 in https://arxiv.org/pdf/2002.01387
"""
Expand All @@ -8711,7 +8711,12 @@ def _rand_svd(
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], Q

error_factor = np.power(
1 + 4 * np.sqrt(2 * operator_dim / (rank - 1)),
1 / (2 * depth + 1))
error_bound = D[-1] * (2 + error_factor)
return U[:, :rank], D[:rank], V[:rank], Q, error_bound

def _genetic_relatedness_vector_individual(
arr: np.ndarray,
Expand Down Expand Up @@ -8741,7 +8746,7 @@ def _genetic_relatedness_vector_individual(
)[0]

def bincount_fn(w):
np.bincount(sample_individuals, w)
return 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
Expand Down Expand Up @@ -8769,7 +8774,8 @@ def _genetic_relatedness_vector_node(

U = np.empty((num_windows, dim, num_components))
D = np.empty((num_windows, num_components))
Q = np.empty((num_windows, num_components + num_oversamples))
Q = np.empty((num_windows, dim, num_components + num_oversamples))
E = np.empty(num_windows)
for i in range(num_windows):
this_window = windows[i : i + 2]
_f = (
Expand All @@ -8779,22 +8785,22 @@ def _genetic_relatedness_vector_node(
)

def _G(x):
_f(x, centre=centre, windows=this_window) # NOQA: B023
return _f(x, centre=centre, windows=this_window) # NOQA: B023

U[i], D[i], _, Q[i] = _rand_svd(
U[i], D[i], _, Q[i], E[i] = _rand_svd(
operator=_G,
operator_dim=dim,
rank=num_components,
depth=iterated_power,
num_vectors=num_components + num_oversamples,
rng=random_state,
range_sketch=None if random_sketch is None else range_sketch[i],
range_sketch=None if range_sketch is None else range_sketch[i],
)

if drop_windows:
U, D, Q = U[0], D[0], Q[0]

return U, D, Q
return U, D, Q, E

def trait_covariance(self, W, windows=None, mode="site", span_normalise=True):
"""
Expand Down

0 comments on commit 2f4ce2b

Please sign in to comment.