UMAP (Uniform Manifold Approximation and Projection)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

import numpy as np
import matplotlib.pyplot as plt
import umap
 
def explore_n_neighbors(X, y=None, values=[5, 15, 50, 100]):
    """
    Visualize the effect of n_neighbors on embedding.
    
    Creates a grid of embeddings at different n_neighbors values
    to help identify the appropriate setting for your data.
    
    Parameters:
    -----------
    X : ndarray of shape (n_samples, n_features)
        High-dimensional data
    y : ndarray of shape (n_samples,), optional
        Labels for coloring (ground truth or cluster assignments)
    values : list of int
        n_neighbors values to explore
    """
    n_values = len(values)
    fig, axes = plt.subplots(2, (n_values + 1) // 2, figsize=(5 * (n_values + 1) // 2, 10))
    axes = axes.flatten()
    
    for i, n_neighbors in enumerate(values):
        print(f"Computing UMAP with n_neighbors={n_neighbors}...")
        
        reducer = umap.UMAP(
            n_neighbors=n_neighbors,
            min_dist=0.1,  # Keep fixed to isolate n_neighbors effect
            random_state=42
        )
        embedding = reducer.fit_transform(X)
        
        ax = axes[i]
        scatter = ax.scatter(
            embedding[:, 0], 
            embedding[:, 1], 
            c=y if y is not None else 'steelblue',
            cmap='Spectral' if y is not None else None,
            s=5, 
            alpha=0.6
        )
        ax.set_title(f'n_neighbors = {n_neighbors}')
        ax.set_xticks([])
        ax.set_yticks([])
    
    # Hide unused axes
    for i in range(len(values), len(axes)):
        axes[i].axis('off')
    
    plt.tight_layout()
    plt.suptitle('Effect of n_neighbors on UMAP Embedding', y=1.02, fontsize=14)
    return fig
 
 
def estimate_optimal_n_neighbors(X, k_range=range(5, 101, 5)):
    """
    Heuristic for estimating optimal n_neighbors.
    
    Uses the elbow method on the graph connectivity:
    - Too small k: disconnected graph
    - Too large k: over-smoothed structure
    - Optimal k: connected graph with distinct structure
    
    Metric: number of connected components + edge weight entropy
    """
    from scipy.sparse.csgraph import connected_components
    
    results = []
    
    for k in k_range:
        reducer = umap.UMAP(n_neighbors=k, random_state=42)
        reducer.fit(X)
        
        # Get the fuzzy graph
        graph = reducer.graph_
        
        # Metric 1: Connected components (want 1)
        n_components, _ = connected_components(graph > 0.01)
        
        # Metric 2: Edge weight entropy (measure of structure)
        weights = graph.data
        weights_norm = weights / weights.sum()
        entropy = -np.sum(weights_norm * np.log(weights_norm + 1e-10))
        
        results.append({
            'n_neighbors': k,
            'n_components': n_components,
            'entropy': entropy
        })
    
    return results
 
 
def plot_n_neighbors_diagnostics(results):
    """Plot diagnostic metrics for n_neighbors selection."""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    ks = [r['n_neighbors'] for r in results]
    n_comps = [r['n_components'] for r in results]
    entropies = [r['entropy'] for r in results]
    
    ax1.plot(ks, n_comps, 'b-o')
    ax1.axhline(y=1, color='r', linestyle='--', label='Target: 1 component')
    ax1.set_xlabel('n_neighbors')
    ax1.set_ylabel('Number of Connected Components')
    ax1.set_title('Graph Connectivity')
    ax1.legend()
    
    ax2.plot(ks, entropies, 'g-o')
    ax2.set_xlabel('n_neighbors')
    ax2.set_ylabel('Edge Weight Entropy')
    ax2.set_title('Structure Complexity')
    
    plt.tight_layout()
    return fig

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

import numpy as np
import umap
import matplotlib.pyplot as plt
 
def compare_metrics(X, y=None, metrics=['euclidean', 'cosine', 'manhattan', 'correlation']):
    """
    Compare UMAP embeddings using different distance metrics.
    
    Parameters:
    -----------
    X : ndarray of shape (n_samples, n_features)
        High-dimensional data
    y : ndarray of shape (n_samples,), optional
        Labels for coloring
    metrics : list of str
        Metrics to compare
    """
    n_metrics = len(metrics)
    fig, axes = plt.subplots(1, n_metrics, figsize=(5 * n_metrics, 5))
    
    for i, metric in enumerate(metrics):
        print(f"Computing UMAP with metric={metric}...")
        
        try:
            reducer = umap.UMAP(
                n_neighbors=15,
                min_dist=0.1,
                metric=metric,
                random_state=42
            )
            embedding = reducer.fit_transform(X)
            
            ax = axes[i] if n_metrics > 1 else axes
            ax.scatter(
                embedding[:, 0], 
                embedding[:, 1],
                c=y if y is not None else 'steelblue',
                cmap='Spectral' if y is not None else None,
                s=5, 
                alpha=0.6
            )
            ax.set_title(f'metric = "{metric}"')
            ax.set_xticks([])
            ax.set_yticks([])
            
        except Exception as e:
            ax = axes[i] if n_metrics > 1 else axes
            ax.text(0.5, 0.5, f"Error: {e}", ha='center', va='center')
            ax.set_title(f'metric = "{metric}" (failed)')
    
    plt.tight_layout()
    return fig
 
 
def custom_metric_example():
    """
    Example: Using a custom distance metric.
    
    UMAP supports any function with signature:
    f(x, y) -> float
    
    where x and y are 1D arrays.
    """
    from numba import njit
    
    @njit
    def weighted_euclidean(x, y, weights=None):
        """
        Weighted Euclidean distance.
        
        Allows different features to contribute differently
        to the distance calculation.
        """
        if weights is None:
            weights = np.ones(len(x))
        
        d = 0.0
        for i in range(len(x)):
            d += weights[i] * (x[i] - y[i]) ** 2
        return np.sqrt(d)
    
    # Note: For custom metrics, you may need to specify
    # metric_kwds with any additional parameters
    
    # Example usage (pseudo-code):
    # reducer = umap.UMAP(metric=weighted_euclidean, 
    #                     metric_kwds={'weights': my_weights})
    
    return weighted_euclidean

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

import numpy as np
import umap
from sklearn.neighbors import NearestNeighbors
from sklearn.metrics import silhouette_score
import matplotlib.pyplot as plt
 
def systematic_umap_tuning(X, y=None, n_jobs=-1):
    """
    Systematic UMAP hyperparameter tuning.
    
    Follows a structured approach:
    1. Sweep n_neighbors
    2. For best n_neighbors, sweep min_dist
    3. Evaluate with quantitative metrics and visual inspection
    
    Parameters:
    -----------
    X : ndarray
        High-dimensional data
    y : ndarray, optional
        Labels for silhouette score calculation
    
    Returns:
    --------
    best_params : dict
        Recommended parameters
    results : list
        All evaluation results
    """
    results = []
    
    # ============================================
    # STEP 1: n_neighbors sweep
    # ============================================
    print("Step 1: Sweeping n_neighbors...")
    n_neighbors_values = [5, 10, 15, 20, 30, 50, 75, 100]
    
    for n_neighbors in n_neighbors_values:
        reducer = umap.UMAP(
            n_neighbors=n_neighbors,
            min_dist=0.1,  # Fixed for this sweep
            random_state=42
        )
        embedding = reducer.fit_transform(X)
        
        # Compute evaluation metrics
        metrics = evaluate_embedding(X, embedding, y)
        metrics['n_neighbors'] = n_neighbors
        metrics['min_dist'] = 0.1
        results.append(metrics)
        
        print(f"  n_neighbors={n_neighbors}: "
              f"trustworthiness={metrics['trustworthiness']:.3f}, "
              f"continuity={metrics['continuity']:.3f}")
    
    # Find best n_neighbors based on combined metric
    n_neighbors_results = [r for r in results if r['min_dist'] == 0.1]
    best_nn_result = max(n_neighbors_results, 
                          key=lambda x: x['trustworthiness'] + x['continuity'])
    best_n_neighbors = best_nn_result['n_neighbors']
    print(f"\nBest n_neighbors: {best_n_neighbors}")
    
    # ============================================
    # STEP 2: min_dist sweep with best n_neighbors
    # ============================================
    print("\nStep 2: Sweeping min_dist...")
    min_dist_values = [0.0, 0.05, 0.1, 0.2, 0.3, 0.5, 0.8]
    
    for min_dist in min_dist_values:
        reducer = umap.UMAP(
            n_neighbors=best_n_neighbors,
            min_dist=min_dist,
            random_state=42
        )
        embedding = reducer.fit_transform(X)
        
        metrics = evaluate_embedding(X, embedding, y)
        metrics['n_neighbors'] = best_n_neighbors
        metrics['min_dist'] = min_dist
        results.append(metrics)
        
        print(f"  min_dist={min_dist}: "
              f"trustworthiness={metrics['trustworthiness']:.3f}")
    
    # Find best min_dist
    min_dist_results = [r for r in results 
                        if r['n_neighbors'] == best_n_neighbors]
    best_md_result = max(min_dist_results,
                         key=lambda x: x['trustworthiness'])
    best_min_dist = best_md_result['min_dist']
    print(f"\nBest min_dist: {best_min_dist}")
    
    best_params = {
        'n_neighbors': best_n_neighbors,
        'min_dist': best_min_dist
    }
    
    return best_params, results
 
 
def evaluate_embedding(X, embedding, y=None, k=10):
    """
    Compute quantitative metrics for embedding quality.
    
    Metrics:
    - Trustworthiness: Are embedding neighbors true neighbors?
    - Continuity: Are high-D neighbors still neighbors in embedding?
    - Silhouette (if labels): How well-separated are clusters?
    """
    from sklearn.manifold import trustworthiness as tw
    
    n_samples = X.shape[0]
    metrics = {}
    
    # Trustworthiness
    metrics['trustworthiness'] = tw(X, embedding, n_neighbors=k)
    
    # Continuity (inverse trustworthiness)
    metrics['continuity'] = tw(embedding, X, n_neighbors=k)
    
    # Silhouette score if labels provided
    if y is not None and len(np.unique(y)) > 1:
        metrics['silhouette'] = silhouette_score(embedding, y)
    else:
        metrics['silhouette'] = np.nan
    
    return metrics
 
 
def visualize_tuning_results(results):
    """Visualize the tuning sweep results."""
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    
    # n_neighbors sweep results
    nn_results = [r for r in results if r['min_dist'] == 0.1]
    nn_values = [r['n_neighbors'] for r in nn_results]
    trust = [r['trustworthiness'] for r in nn_results]
    cont = [r['continuity'] for r in nn_results]
    
    axes[0].plot(nn_values, trust, 'b-o', label='Trustworthiness')
    axes[0].plot(nn_values, cont, 'g-o', label='Continuity')
    axes[0].set_xlabel('n_neighbors')
    axes[0].set_ylabel('Score')
    axes[0].set_title('n_neighbors Sweep')
    axes[0].legend()
    
    # min_dist sweep results
    md_results = [r for r in results if r['min_dist'] != 0.1 or r == nn_results[2]]
    md_values = [r['min_dist'] for r in md_results]
    trust = [r['trustworthiness'] for r in md_results]
    
    axes[1].plot(md_values, trust, 'b-o')
    axes[1].set_xlabel('min_dist')
    axes[1].set_ylabel('Trustworthiness')
    axes[1].set_title('min_dist Sweep (best n_neighbors)')
    
    plt.tight_layout()
    return fig