Skip to content

Improve sampling #302

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.

Sign up for GitHub

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

Open
florian-huber wants to merge 9 commits into
base: main
Choose a base branch
from
Open

Improve sampling #302

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
d98cd9c
fix test
florian-huber Sep 4, 2026
5bd11fb
add caching ans more memory-efficient data types/storage
florian-huber Sep 4, 2026
47bde82
fix wrong loss computation
florian-huber Sep 4, 2026
d9282ee
add caching
florian-huber Sep 4, 2026
22cd61a
larger refactor
florian-huber Sep 4, 2026
5e994fb
seed refactoring (to make it consistent with other new code)
florian-huber Sep 4, 2026
50504d6
update/refactor tests
florian-huber Sep 4, 2026
e6a080e
Merge branch 'main' into improve_sampling
florian-huber Sep 4, 2026
bd8f207
linting
florian-huber Sep 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions ms2deepscore/MS2DeepScoreEvaluated.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,6 @@ def _matrix_components(
return result

if MATCHMS_V1_API:

def matrix(
self,
spectra_1: List[Spectrum],
Expand All @@ -219,7 +218,6 @@ def matrix(
return as_matchms_scores(score_arrays)

else:

def matrix(
self,
references: List[Spectrum],
Expand Down
114 changes: 95 additions & 19 deletions ms2deepscore/fingerprint_similarity_computations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from typing import Tuple
import numpy as np
from numba import jit, prange
from numba import jit, njit, prange
from chemap.metrics import (
tanimoto_similarity_dense,
tanimoto_similarity_sparse,
Expand Down Expand Up @@ -84,6 +84,15 @@ def compute_fingerprint_similarity_row(
)[0]


@njit
def _derive_seed(random_seed, row_index, bin_index):
return (
random_seed
+ 1_000_003 * row_index
+ 97 * bin_index
) % 4_294_967_295


# Add row based similarity computations
# -------------------------------------

Expand Down Expand Up @@ -141,24 +150,50 @@ def _fill_pairs_for_row_same_set(
max_pairs_per_bin,
selection_bins,
include_diagonal,
random_seed,
):
"""Reservoir-sample at most ``max_pairs_per_bin`` candidates per bin.

This avoids allocating/shuffling a potentially O(N) index array for every
score bin and row. Each qualifying pair has equal probability of ending up
in the bounded reservoir.
"""
num_bins = len(selection_bins)

for bin_number in range(num_bins):
selection_bin = selection_bins[bin_number]
indices = np.nonzero(
(tanimoto_scores > selection_bin[0]) & (tanimoto_scores <= selection_bin[1])
)[0]
# Derive an independent deterministic seed for each row/bin pair.
# This avoids dependence on prange/thread execution order.
if random_seed >= 0:
np.random.seed(
_derive_seed(
random_seed,
idx_fingerprint_i,
bin_number,
)
)

if not include_diagonal and idx_fingerprint_i in indices:
indices = indices[indices != idx_fingerprint_i]
seen = 0
filled = 0
for candidate_index in range(len(tanimoto_scores)):
if not include_diagonal and candidate_index == idx_fingerprint_i:
continue

np.random.shuffle(indices)
indices = indices[:max_pairs_per_bin]
num_indices = len(indices)
score = tanimoto_scores[candidate_index]
if not (score > selection_bin[0] and score <= selection_bin[1]):
continue

selected_pairs_per_bin[bin_number, idx_fingerprint_i, :num_indices] = indices
selected_scores_per_bin[bin_number, idx_fingerprint_i, :num_indices] = tanimoto_scores[indices]
seen += 1
if filled < max_pairs_per_bin:
slot = filled
filled += 1
else:
slot = np.random.randint(0, seen)
if slot >= max_pairs_per_bin:
continue

selected_pairs_per_bin[bin_number, idx_fingerprint_i, slot] = candidate_index
selected_scores_per_bin[bin_number, idx_fingerprint_i, slot] = score


@jit(nopython=True, parallel=True)
Expand All @@ -167,6 +202,7 @@ def _compute_tanimoto_similarity_per_bin_dense(
max_pairs_per_bin,
selection_bins=np.array([(x / 10, x / 10 + 0.1) for x in range(10)], dtype=np.float32),
include_diagonal=True,
random_seed=-1,
) -> Tuple[np.ndarray, np.ndarray]:
size = fingerprints.shape[0]
num_bins = len(selection_bins)
Expand All @@ -186,6 +222,7 @@ def _compute_tanimoto_similarity_per_bin_dense(
max_pairs_per_bin,
selection_bins,
include_diagonal,
random_seed,
)

return selected_pairs_per_bin, selected_scores_per_bin
Expand All @@ -197,6 +234,7 @@ def _compute_tanimoto_similarity_per_bin_sparse_binary(
max_pairs_per_bin,
selection_bins=np.array([(x / 10, x / 10 + 0.1) for x in range(10)], dtype=np.float32),
include_diagonal=True,
random_seed=-1,
) -> Tuple[np.ndarray, np.ndarray]:
size = len(fingerprints)
num_bins = len(selection_bins)
Expand All @@ -216,6 +254,7 @@ def _compute_tanimoto_similarity_per_bin_sparse_binary(
max_pairs_per_bin,
selection_bins,
include_diagonal,
random_seed,
)

return selected_pairs_per_bin, selected_scores_per_bin
Expand All @@ -228,6 +267,7 @@ def _compute_tanimoto_similarity_per_bin_sparse_count(
max_pairs_per_bin,
selection_bins=np.array([(x / 10, x / 10 + 0.1) for x in range(10)], dtype=np.float32),
include_diagonal=True,
random_seed=-1,
) -> Tuple[np.ndarray, np.ndarray]:
size = len(fingerprints_bins)
num_bins = len(selection_bins)
Expand All @@ -250,6 +290,7 @@ def _compute_tanimoto_similarity_per_bin_sparse_count(
max_pairs_per_bin,
selection_bins,
include_diagonal,
random_seed,
)

return selected_pairs_per_bin, selected_scores_per_bin
Expand All @@ -267,8 +308,10 @@ def compute_tanimoto_similarity_per_bin(
fingerprint_type: str,
selection_bins=np.array([(x / 10, x / 10 + 0.1) for x in range(10)], dtype=np.float32),
include_diagonal=True,
random_seed=None,
) -> Tuple[np.ndarray, np.ndarray]:
"""Dispatch to the appropriate pairwise-per-bin Tanimoto implementation."""
jit_random_seed = -1 if random_seed is None else int(random_seed)
if fingerprint_type not in SUPPORTED_FINGERPRINT_TYPES:
raise ValueError(f"Unsupported fingerprint type: {fingerprint_type}")

Expand All @@ -278,6 +321,7 @@ def compute_tanimoto_similarity_per_bin(
max_pairs_per_bin=max_pairs_per_bin,
selection_bins=selection_bins,
include_diagonal=include_diagonal,
random_seed=jit_random_seed,
)

if is_unfolded_binary_fingerprint_type(fingerprint_type):
Expand All @@ -286,6 +330,7 @@ def compute_tanimoto_similarity_per_bin(
max_pairs_per_bin=max_pairs_per_bin,
selection_bins=selection_bins,
include_diagonal=include_diagonal,
random_seed=jit_random_seed,
)

if is_unfolded_count_fingerprint_type(fingerprint_type):
Expand All @@ -296,6 +341,7 @@ def compute_tanimoto_similarity_per_bin(
max_pairs_per_bin=max_pairs_per_bin,
selection_bins=selection_bins,
include_diagonal=include_diagonal,
random_seed=jit_random_seed,
)

raise ValueError(f"Unsupported fingerprint type: {fingerprint_type}")
Expand All @@ -310,21 +356,37 @@ def _fill_pairs_for_row_between_sets(
target_offset,
max_pairs_per_bin,
selection_bins,
random_seed,
):
"""Reservoir-sample bounded candidates for one cross-set row."""
num_bins = len(selection_bins)

for bin_number in range(num_bins):
selection_bin = selection_bins[bin_number]
indices = np.nonzero(
(tanimoto_scores > selection_bin[0]) & (tanimoto_scores <= selection_bin[1])
)[0]
if random_seed >= 0:
np.random.seed(
(random_seed + 1_000_003 * row_index + 97 * bin_number)
% 4_294_967_295
)

seen = 0
filled = 0
for candidate_index in range(len(tanimoto_scores)):
score = tanimoto_scores[candidate_index]
if not (score > selection_bin[0] and score <= selection_bin[1]):
continue

np.random.shuffle(indices)
indices = indices[:max_pairs_per_bin]
num_indices = len(indices)
seen += 1
if filled < max_pairs_per_bin:
slot = filled
filled += 1
else:
slot = np.random.randint(0, seen)
if slot >= max_pairs_per_bin:
continue

selected_pairs_per_bin[bin_number, row_index, :num_indices] = indices + target_offset
selected_scores_per_bin[bin_number, row_index, :num_indices] = tanimoto_scores[indices]
selected_pairs_per_bin[bin_number, row_index, slot] = candidate_index + target_offset
selected_scores_per_bin[bin_number, row_index, slot] = score


@jit(nopython=True, parallel=True)
Expand All @@ -333,6 +395,7 @@ def _compute_tanimoto_similarity_per_bin_between_sets_dense(
fingerprints_2,
max_pairs_per_bin,
selection_bins=np.array([(x / 10, x / 10 + 0.1) for x in range(10)], dtype=np.float32),
random_seed=-1,
) -> Tuple[np.ndarray, np.ndarray]:
size_1 = fingerprints_1.shape[0]
size_2 = fingerprints_2.shape[0]
Expand All @@ -353,6 +416,7 @@ def _compute_tanimoto_similarity_per_bin_between_sets_dense(
size_1,
max_pairs_per_bin,
selection_bins,
random_seed,
)

for idx_fingerprint_j in prange(size_2):
Expand All @@ -368,6 +432,7 @@ def _compute_tanimoto_similarity_per_bin_between_sets_dense(
0,
max_pairs_per_bin,
selection_bins,
random_seed,
)

return selected_pairs_per_bin, selected_scores_per_bin
Expand All @@ -379,6 +444,7 @@ def _compute_tanimoto_similarity_per_bin_between_sets_sparse_binary(
fingerprints_2,
max_pairs_per_bin,
selection_bins=np.array([(x / 10, x / 10 + 0.1) for x in range(10)], dtype=np.float32),
random_seed=-1,
) -> Tuple[np.ndarray, np.ndarray]:
size_1 = len(fingerprints_1)
size_2 = len(fingerprints_2)
Expand All @@ -399,6 +465,7 @@ def _compute_tanimoto_similarity_per_bin_between_sets_sparse_binary(
size_1,
max_pairs_per_bin,
selection_bins,
random_seed,
)

for idx_fingerprint_j in prange(size_2):
Expand All @@ -414,6 +481,7 @@ def _compute_tanimoto_similarity_per_bin_between_sets_sparse_binary(
0,
max_pairs_per_bin,
selection_bins,
random_seed,
)

return selected_pairs_per_bin, selected_scores_per_bin
Expand All @@ -427,6 +495,7 @@ def _compute_tanimoto_similarity_per_bin_between_sets_sparse_count(
fingerprints_2_counts,
max_pairs_per_bin,
selection_bins=np.array([(x / 10, x / 10 + 0.1) for x in range(10)], dtype=np.float32),
random_seed=-1,
) -> Tuple[np.ndarray, np.ndarray]:
size_1 = len(fingerprints_1_bins)
size_2 = len(fingerprints_2_bins)
Expand All @@ -453,6 +522,7 @@ def _compute_tanimoto_similarity_per_bin_between_sets_sparse_count(
size_1,
max_pairs_per_bin,
selection_bins,
random_seed,
)

for idx_fingerprint_j in prange(size_2):
Expand All @@ -474,6 +544,7 @@ def _compute_tanimoto_similarity_per_bin_between_sets_sparse_count(
0,
max_pairs_per_bin,
selection_bins,
random_seed,
)

return selected_pairs_per_bin, selected_scores_per_bin
Expand All @@ -485,8 +556,10 @@ def compute_tanimoto_similarity_per_bin_between_sets(
max_pairs_per_bin,
fingerprint_type: str,
selection_bins=np.array([(x / 10, x / 10 + 0.1) for x in range(10)], dtype=np.float32),
random_seed=None,
) -> Tuple[np.ndarray, np.ndarray]:
"""Compute cross-set Tanimoto per bin for all supported fingerprint types."""
jit_random_seed = -1 if random_seed is None else int(random_seed)
if fingerprint_type not in SUPPORTED_FINGERPRINT_TYPES:
raise ValueError(f"Unsupported fingerprint type: {fingerprint_type}")

Expand All @@ -496,6 +569,7 @@ def compute_tanimoto_similarity_per_bin_between_sets(
fingerprints_2,
max_pairs_per_bin=max_pairs_per_bin,
selection_bins=selection_bins,
random_seed=jit_random_seed,
)

if is_unfolded_binary_fingerprint_type(fingerprint_type):
Expand All @@ -504,6 +578,7 @@ def compute_tanimoto_similarity_per_bin_between_sets(
fingerprints_2,
max_pairs_per_bin=max_pairs_per_bin,
selection_bins=selection_bins,
random_seed=jit_random_seed,
)

if is_unfolded_count_fingerprint_type(fingerprint_type):
Expand All @@ -516,6 +591,7 @@ def compute_tanimoto_similarity_per_bin_between_sets(
fingerprints_2_counts,
max_pairs_per_bin=max_pairs_per_bin,
selection_bins=selection_bins,
random_seed=jit_random_seed,
)

raise ValueError(f"Unsupported fingerprint type: {fingerprint_type}")
15 changes: 13 additions & 2 deletions ms2deepscore/models/EmbeddingEvaluatorModel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,15 +152,26 @@ def train_evaluator(
with no_grad():
self.eval()
val_losses = []

for sample in val_generator:
tanimoto_scores, ms2ds_scores, embeddings = sample

outputs = self(embeddings.reshape(-1, 1, embeddings.shape[-1]).to(device))

mse_per_embedding = ((tanimoto_scores - ms2ds_scores) ** 2).mean(dim=1)
mse_per_embedding = mse_per_embedding.reshape(-1, 1).clone().detach()

loss = criterion(outputs.to(device), mse_per_embedding.to(device, dtype=float32))
val_losses.append(loss_value)
loss = criterion(
outputs,
mse_per_embedding.to(
device,
dtype=float32,
),
)

val_loss_value = loss.detach().item()
val_losses.append(val_loss_value)

print(f">>> Val_loss: {np.mean(val_losses):.6f}")

self.train()
Expand Down
Loading