Possible index mismatch in `transform()` with `unique=True`

Hi,

While experimenting with unique=True, I noticed a possible inconsistency in the transform() path for new data.

During fit(),

the search index is built on the deduped dataset,

# Standard case
(self._knn_indices, self._knn_dists, self._knn_search_index, ) = nearest_neighbors(X_indexed,...)

while the final embedding is expanded back to the original dataset length,

self.embedding_ = self.embedding_[inverse]

Later, transform() performs

indices, dists = self._knn_search_index.query(X,...)

followed by

 graph = scipy.sparse.coo_matrix(
       (vals, (rows, cols)), shape=(X.shape[0], self._raw_data.shape[0])
 )
...
 embedding = init_graph_transform(csr_graph, self.embedding_)

So although the CSR graph has shape (len(X_test), len(X_train)), its populated columns appear to correspond to the indices of X_indexed rather than the original training set _raw_data. Am I overlooking a remapping step somewhere, that maps the indices back to original training space (including duplicates) before taking weighted averages of training embeddings inside init_graph_transform() ? Because init_graph_transform() is performing simply

...
normalized = inv_sums @ avg_graph
result[avg_mask] = (normalized @ embedding).astype(np.float32)

Experiment

I have also performed a small experiment using simulated data.

Raw size : (5103, 2) Unique size: (5000, 2) Noise Rate % : 2.02

I then artificially created test samples that are not exactly present in the training data, but differ only by a small jitter from the duplicate points in the training data. The results show that (1) the column index of the closest neighbor in csr_graph matches the index in the deduplicated data, and (2) the points identified as neighbors are quite far away from the query points. I understand that PyNNDescent returns approximate neighbors; however, I would appreciate clarification on whether the observed indexing behavior is expected.

• Example-1

Query : [1099.99976691 1087.02675875]
Expected Closest Neighbor : [1100.     1087.0281] ; deduped position id : 1100 ; original row ids : [1124 1125 1126 1127]
Returned NN indices : [1100 1102 1105] ; Closest NN == deduped pos : True
Graph columns :  [1100 1102 1105]
points correspondong to Graph columns & their weights :
 [(array([1076.    , 1073.6313], dtype=float32), np.float32(0.9999)), (array([1078.    , 1083.6003], dtype=float32), np.float32(0.8137)), (array([1081.    , 1083.2019], dtype=float32), np.float32(0.7713))]

• Example-2

Query : [4799.99607916 4801.7731728 ]
Expected Closest Neighbor : [4800.     4801.7734] ; deduped position id : 4800 ; original row ids : [4901 4902]
Returned NN indices : [4800 4801 4802] ; Closest NN == deduped pos : True
Graph columns :  [4800 4801 4802]
points correspondong to Graph columns & their weights :
 [(array([4700.    , 4699.1504], dtype=float32), np.float32(0.9997)), (array([4700.    , 4699.1504], dtype=float32), np.float32(0.8304)), (array([4701.   , 4694.958], dtype=float32), np.float32(0.7546))]

I've attached a minimal notebook that reproduces my experiment. UMAP_unique.ipynb The deviation increases systematically as the gap between the original row index and the corresponding deduplicated index increases. In contrast, the deviation is approximately equal to zero for unique=False (expected behaviour heuristically; there is training data very close to query point)

NOTE: This does not affect fit_transform(), only transform() path.

Initially I thought the mismatch was limited only to the initialization:

 embedding = init_graph_transform(...)

But a closer inspection reveals that it cascades further. The initialized embedding is followed by

head = graph.row
tail = graph.col
...
embedding = optimize_layout_euclidean(
    embedding,  # head_embedding
    self.embedding_.astype(np.float32, copy=True),  # tail_embedding
    head,
    tail,
    ...
)

where the indices of tail_embedding are not in sync with tail. So, if my understanding is correct, under the unique=True transform() path:

1. We initialize the embedding incorrectly.
2. The initialized points are then pulled/pushed in the wrong direction throughout the subsequent SGD optimization, which is a more serious concern.

A clarification on this would be helpful.