Coverage for imcluster/reduction.py: 100.00%

1"""Optional dimensionality reduction for image feature vectors.""" 

2 

3from enum import Enum 

4from typing import Any 

5 

6import numpy as np 

7from numpy.typing import ArrayLike, NDArray 

8from rich.console import Console 

9from sklearn.decomposition import PCA 

10from sklearn.manifold import TSNE 

11from umap import UMAP 

12 

13from .io import ImclusterIO 

14 

15console = Console() 

16 

17 

18class ReductionMethod(str, Enum): 

19 """Supported dimensionality-reduction methods.""" 

20 

21 NONE = "none" 

22 UMAP = "umap" 

23 TSNE = "tsne" 

24 PCA = "pca" 

25 

26 

27def reduce_dimensions( 

28 imcluster_io: ImclusterIO, 

29 feature_vectors: ArrayLike, 

30 method: ReductionMethod | str = ReductionMethod.NONE, 

31 dimensions: int = 50, 

32 force: bool = False, 

33) -> NDArray[Any]: 

34 """Reduce feature dimensions and cache the resulting vectors. 

35 

36 Args: 

37 imcluster_io: Image collection used to persist reduced vectors. 

38 feature_vectors: Feature matrix with one row per image. 

39 method: Reduction algorithm, or ``none`` to retain original vectors. 

40 dimensions: Requested number of output dimensions. 

41 force: Recompute an existing reduced-vector cache. 

42 

43 Returns: 

44 Original or dimension-reduced feature vectors. 

45 

46 Raises: 

47 ValueError: If the method is unsupported or vectors are malformed. 

48 """ 

49 try: 

50 method = ( 

51 method 

52 if isinstance(method, ReductionMethod) 

53 else ReductionMethod(method.lower()) 

54 ) 

55 except (AttributeError, ValueError) as error: 

56 raise ValueError(f"Unsupported reduction method: {method}") from error 

57 

58 vectors = np.asarray(feature_vectors, dtype=float) 

59 if vectors.ndim != 2 or len(vectors) != len(imcluster_io.images): 

60 raise ValueError("feature_vectors must contain one row per image") 

61 if dimensions < 1: 

62 raise ValueError("dimensions must be at least 1") 

63 if method is ReductionMethod.NONE: 

64 return vectors 

65 

66 cache_column = f"reduction_{method.value}_{dimensions}" 

67 if imcluster_io.has_column(cache_column) and not force: 

68 console.print( 

69 f"[green]Using cached {method.value} reduction:[/green] loaded " 

70 f"{len(vectors)} vectors from '{imcluster_io.output}'." 

71 ) 

72 return np.asarray(imcluster_io.get_column(cache_column).to_list(), dtype=float) 

73 

74 n_samples, n_features = vectors.shape 

75 if method is ReductionMethod.PCA: 

76 n_components = min(dimensions, n_samples, n_features) 

77 reducer: Any = PCA(n_components=n_components) 

78 elif method is ReductionMethod.TSNE: 

79 n_components = min(dimensions, 3, n_features) 

80 reducer = TSNE( 

81 n_components=n_components, 

82 perplexity=min(30.0, float(n_samples - 1)), 

83 metric="cosine", 

84 init="random", 

85 learning_rate="auto", 

86 random_state=0, 

87 ) 

88 else: 

89 if n_samples < 3: 

90 raise ValueError("UMAP reduction requires at least three images") 

91 n_components = min(dimensions, n_features, max(1, n_samples - 2)) 

92 reducer = UMAP( 

93 n_components=n_components, 

94 n_neighbors=min(30, n_samples - 1), 

95 min_dist=0.0, 

96 metric="cosine", 

97 init="random", 

98 random_state=0, 

99 ) 

100 

101 console.print( 

102 f"[cyan]Reducing dimensions with {method.value}:[/cyan] " 

103 f"{n_features} dimensions to {n_components}." 

104 ) 

105 with console.status(f"[cyan]Computing {method.value} embedding...[/cyan]"): 

106 reduced = np.asarray(reducer.fit_transform(vectors), dtype=float) 

107 imcluster_io.save_column(cache_column, reduced.tolist()) 

108 console.print( 

109 f"[green]Cached {method.value} reduction:[/green] wrote " 

110 f"{len(reduced)} vectors to '{imcluster_io.output}'." 

111 ) 

112 return reduced