cbrkit.sim.collections

  1import heapq
  2from collections.abc import Collection, Sequence, Set
  3from dataclasses import asdict, dataclass, field
  4from itertools import product
  5from typing import Any, cast, override
  6
  7import numpy as np
  8
  9from ..helpers import (
 10    dist2sim,
 11    get_metadata,
 12    optional_dependencies,
 13    unbatchify_sim,
 14    unpack_float,
 15)
 16from ..typing import (
 17    AnySimFunc,
 18    Float,
 19    HasMetadata,
 20    JsonDict,
 21    SimFunc,
 22    StructuredValue,
 23)
 24
 25Number = float | int
 26
 27__all__ = [
 28    "dtw",
 29    "smith_waterman",
 30    "jaccard",
 31    "mapping",
 32    "isolated_mapping",
 33    "sequence_mapping",
 34    "sequence_correctness",
 35    "SequenceSim",
 36    "Weight",
 37]
 38
 39with optional_dependencies():
 40    from nltk.metrics import jaccard_distance
 41
 42    @dataclass(slots=True, frozen=True)
 43    class jaccard[V](SimFunc[Collection[V], float]):
 44        """Jaccard similarity function.
 45
 46        Examples:
 47            >>> sim = jaccard()
 48            >>> sim(["a", "b", "c", "d"], ["a", "b", "c"])
 49            0.8
 50        """
 51
 52        @override
 53        def __call__(self, x: Collection[V], y: Collection[V]) -> float:
 54            if not isinstance(x, Set):
 55                x = set(x)
 56            if not isinstance(y, Set):
 57                y = set(y)
 58
 59            return dist2sim(jaccard_distance(x, y))
 60
 61
 62with optional_dependencies():
 63    from minineedle import core, smith
 64
 65    @dataclass(slots=True, frozen=True)
 66    class smith_waterman[V](SimFunc[Sequence[V], float]):
 67        """
 68        Performs the Smith-Waterman alignment with configurable scoring parameters. If no element matches it returns 0.0.
 69
 70        Args:
 71            match_score: Score for matching characters. Defaults to 2.
 72            mismatch_penalty: Penalty for mismatching characters. Defaults to -1.
 73            gap_penalty: Penalty for gaps. Defaults to -1.
 74
 75        Example:
 76            >>> sim = smith_waterman()
 77            >>> sim("abcde", "fghe")
 78            2
 79        """
 80
 81        match_score: int = 2
 82        mismatch_penalty: int = -1
 83        gap_penalty: int = -1
 84
 85        @override
 86        def __call__(self, x: Sequence[V], y: Sequence[V]) -> float:
 87            try:
 88                alignment = smith.SmithWaterman(
 89                    cast(Sequence[Any], x), cast(Sequence[Any], y)
 90                )
 91                alignment.change_matrix(
 92                    core.ScoreMatrix(
 93                        match=self.match_score,
 94                        miss=self.mismatch_penalty,
 95                        gap=self.gap_penalty,
 96                    )
 97                )
 98                alignment.align()
 99
100                return alignment.get_score()
101            except ZeroDivisionError:
102                return 0.0
103
104
105@dataclass(slots=True, frozen=True)
106class SequenceSim[V, S: Float](StructuredValue[float]):
107    """
108    A class representing sequence similarity with optional mapping and similarity scores.
109
110    Attributes:
111        value: The overall similarity score as a float.
112        similarities: Optional local similarity scores as a sequence of floats.
113        mapping: Optional alignment information as a sequence of tuples.
114    """
115
116    similarities: Sequence[S] | None = field(default=None)
117    mapping: Sequence[tuple[V, V]] | None = field(default=None)
118
119
120@dataclass(slots=True, init=False)
121class dtw[V](SimFunc[Collection[V] | np.ndarray, SequenceSim[V, float]]):
122    """
123    Dynamic Time Warping similarity function with optional backtracking for alignment.
124
125    Examples:
126        >>> sim = dtw()
127        >>> sim([1, 2, 3], [1, 2, 3, 4])
128        SequenceSim(value=0.5, similarities=None, mapping=None)
129        >>> sim = dtw(distance_func=lambda a, b: abs(a - b))
130        >>> sim([1, 2, 3], [3, 4, 5])
131        SequenceSim(value=0.14285714285714285, similarities=None, mapping=None)
132        >>> sim = dtw(distance_func=lambda a, b: abs(len(str(a)) - len(str(b))))
133        >>> sim(["a", "bb", "ccc"], ["aa", "bbb", "c"], return_alignment=True)
134        SequenceSim(value=0.25, similarities=[0.5, 1.0, 1.0, 0.3333333333333333], mapping=[('aa', 'a'), ('aa', 'bb'), ('bbb', 'ccc'), ('c', 'ccc')])
135        >>> sim = dtw(distance_func=lambda a, b: abs(a - b))
136        >>> sim([1, 2, 3], [1, 2, 3, 4], return_alignment=True)
137        SequenceSim(value=0.5, similarities=[1.0, 1.0, 1.0, 0.5], mapping=[(1, 1), (2, 2), (3, 3), (4, 3)])
138    """
139
140    distance_func: SimFunc[V, Float] | None
141
142    def __init__(self, distance_func: AnySimFunc[V, Float] | None = None):
143        self.distance_func = unbatchify_sim(distance_func) if distance_func else None
144
145    def __call__(
146        self,
147        x: Collection[V] | np.ndarray,
148        y: Collection[V] | np.ndarray,
149        return_alignment: bool = False,
150    ) -> SequenceSim[V, float]:
151        """
152        Perform DTW and optionally return alignment information.
153
154        Args:
155            x: The first sequence as a collection or numpy array.
156            y: The second sequence as a collection or numpy array.
157            return_alignment: Whether to compute and return the alignment (default: False).
158
159        Returns:
160            A SequenceSim object containing the similarity value, local similarities, and optional alignment.
161        """
162        if not isinstance(x, np.ndarray):
163            x = np.array(x, dtype=object)  # Allow non-numeric types
164        if not isinstance(y, np.ndarray):
165            y = np.array(y, dtype=object)
166
167        # Compute the DTW distance
168        dtw_distance, alignment, local_similarities = self.compute_dtw(
169            x, y, return_alignment
170        )
171
172        # Convert DTW distance to similarity
173        similarity = dist2sim(dtw_distance)
174
175        # Return SequenceSim with updated attributes
176        return SequenceSim(
177            value=float(similarity),
178            similarities=local_similarities,
179            mapping=alignment if return_alignment else None,
180        )
181
182    def compute_dtw(
183        self, x: np.ndarray, y: np.ndarray, return_alignment: bool
184    ) -> tuple[float, list[tuple[V, V]] | None, list[float] | None]:
185        """
186        Compute DTW distance and optionally compute the best alignment and local similarities.
187
188        Args:
189            x: The first sequence as a numpy array.
190            y: The second sequence as a numpy array.
191            return_alignment: Whether to compute the alignment.
192
193        Returns:
194            A tuple of (DTW distance, best alignment or None, local similarities or None).
195        """
196        n, m = len(x), len(y)
197        dtw_matrix = np.full((n + 1, m + 1), np.inf)
198        dtw_matrix[0, 0] = 0
199
200        for i in range(1, n + 1):
201            for j in range(1, m + 1):
202                cost = unpack_float(
203                    self.distance_func(y[j - 1], x[i - 1])
204                    if self.distance_func
205                    else abs(y[j - 1] - x[i - 1])
206                )
207                # Take last min from a square box
208                last_min: float = min(
209                    dtw_matrix[i - 1, j],  # Insertion
210                    dtw_matrix[i, j - 1],  # Deletion
211                    dtw_matrix[i - 1, j - 1],  # Match
212                )
213                dtw_matrix[i, j] = cost + last_min
214
215        # If alignment is not requested, skip backtracking
216        if not return_alignment:
217            return dtw_matrix[n, m], None, None
218
219        # Backtracking to find the best alignment and local similarities
220        mapping, local_similarities = self.backtrack(dtw_matrix, x, y, n, m)
221
222        return dtw_matrix[n, m], mapping, local_similarities
223
224    def backtrack(
225        self, dtw_matrix: np.ndarray, x: np.ndarray, y: np.ndarray, n: int, m: int
226    ) -> tuple[list[tuple[V, V]], list[float]]:
227        """
228        Backtrack through the DTW matrix to find the best alignment and local similarities.
229
230        Args:
231            dtw_matrix: The DTW matrix.
232            x: The first sequence as a numpy array.
233            y: The second sequence as a numpy array.
234            n: The length of the first sequence.
235            m: The length of the second sequence.
236
237        Returns:
238            A tuple of (alignment, local similarities).
239        """
240        i, j = n, m
241        alignment = []
242        local_similarities = []
243
244        while i > 0 and j > 0:
245            alignment.append((y[j - 1], x[i - 1]))  # Align elements as (query, case)
246            cost = unpack_float(
247                self.distance_func(y[j - 1], x[i - 1])
248                if self.distance_func
249                else abs(y[j - 1] - x[i - 1])
250            )
251            local_similarities.append(dist2sim(cost))  # Convert cost to similarity
252            # Move in the direction of the minimum cost
253            if dtw_matrix[i - 1, j] == min(
254                dtw_matrix[i - 1, j], dtw_matrix[i, j - 1], dtw_matrix[i - 1, j - 1]
255            ):
256                i -= 1  # Move up
257            elif dtw_matrix[i, j - 1] == min(
258                dtw_matrix[i - 1, j], dtw_matrix[i, j - 1], dtw_matrix[i - 1, j - 1]
259            ):
260                j -= 1  # Move left
261            else:
262                i -= 1  # Move diagonally
263                j -= 1
264
265        # Handle remaining elements in i or j
266        while i > 0:
267            alignment.append((None, x[i - 1]))  # Unmatched element from x (case)
268            local_similarities.append(0.0)  # No similarity for unmatched
269            i -= 1
270        while j > 0:
271            alignment.append((y[j - 1], None))  # Unmatched element from y (query)
272            local_similarities.append(0.0)  # No similarity for unmatched
273            j -= 1
274
275        return alignment[::-1], local_similarities[
276            ::-1
277        ]  # Reverse to start from the beginning
278
279
280@dataclass(slots=True, frozen=True)
281class isolated_mapping[V](SimFunc[Sequence[V], float]):
282    """
283    Isolated Mapping similarity function that compares each element in 'y' (query)
284    with all elements in 'x' (case)
285    and takes the maximum similarity for each element in 'y', then averages
286    these maximums. Assumes y -> x (query operated onto case).
287
288    Args:
289        element_similarity: A function that takes two elements (query_item, case_item)
290        and returns a similarity score between them.
291
292    Examples:
293        >>> from cbrkit.sim.strings import levenshtein
294        >>> sim = isolated_mapping(levenshtein())
295        >>> sim(["kitten", "sitting"], ["sitting", "fitted"])
296        0.8333333333333334
297    """
298
299    element_similarity: SimFunc[V, float]
300
301    @override
302    def __call__(self, x: Sequence[V], y: Sequence[V]) -> float:
303        # x = case sequence, y = query sequence
304        # Logic: For each element in query y, find its best match in case x. Average these best scores.
305
306        if not y:  # If the query is empty, similarity is undefined or trivially 0 or 1? Let's return 0.
307            return 0.0
308        # Avoid division by zero if x is empty but y is not. max default handles this.
309
310        total_similarity = 0.0
311        for yi in y:  # Iterate through QUERY (y)
312            # Find the best match for the current query element yi within the CASE sequence x
313            # Pass (query_item, case_item) to element_similarity, assuming sim(query, case) convention
314            # Use default=0.0 for max() in case the case sequence x is empty
315            max_similarity = max(
316                (self.element_similarity(yi, xi) for xi in x), default=0.0
317            )
318            total_similarity += max_similarity
319
320        # Normalize by the length of the QUERY sequence (y)
321        average_similarity = total_similarity / len(y)
322
323        return average_similarity
324
325
326@dataclass(slots=True, frozen=True)
327class mapping[V](SimFunc[Sequence[V], float]):
328    """
329    Implements an A* algorithm to find the best matching between query items and case items
330    based on the provided similarity function, maximizing the overall similarity score.
331
332    Args:
333        similarity_function: A function that calculates the similarity between two elements.
334        max_queue_size: Maximum size of the priority queue. Defaults to 1000.
335
336    Returns:
337        A similarity function for sequences.
338
339    Examples:
340        >>> def example_similarity_function(x, y) -> float:
341        ...     return 1.0 if x == y else 0.0
342        >>> sim_func = mapping(example_similarity_function)
343        >>> result = sim_func(["Monday", "Tuesday", "Wednesday"], ["Monday", "Tuesday", "Sunday"])
344        >>> print(f"Normalized Similarity Score: {result}")
345        Normalized Similarity Score: 0.6666666666666666
346    """
347
348    element_similarity: SimFunc[V, float]
349    max_queue_size: int = 1000
350
351    @override
352    def __call__(self, x: Sequence[V], y: Sequence[V]) -> float:
353        # Priority queue to store potential solutions with their scores
354        pq = []
355        initial_solution = (0.0, set(), frozenset(y), frozenset(x))
356        heapq.heappush(pq, initial_solution)
357
358        best_score = 0.0
359
360        while pq:
361            current_score, current_mapping, remaining_query, remaining_case = (
362                heapq.heappop(pq)
363            )
364
365            if not remaining_query:  # All query items are mapped
366                best_score = max(best_score, current_score / len(y))
367                continue  # Continue to process remaining potential mappings
368
369            for query_item in remaining_query:
370                for case_item in remaining_case:
371                    sim_score = self.element_similarity(query_item, case_item)
372                    new_mapping = current_mapping | {(query_item, case_item)}
373                    new_score = current_score + sim_score  # Accumulate score correctly
374                    new_remaining_query = remaining_query - {query_item}
375                    new_remaining_case = remaining_case - {case_item}
376
377                    heapq.heappush(
378                        pq,
379                        (
380                            new_score,
381                            new_mapping,
382                            new_remaining_query,
383                            new_remaining_case,
384                        ),
385                    )
386
387                    if len(pq) > self.max_queue_size:
388                        heapq.heappop(pq)
389
390        return best_score
391
392
393@dataclass
394class Weight:
395    """A weighted interval with bounds for sequence mapping similarity."""
396
397    weight: float
398    lower_bound: float
399    upper_bound: float
400    inclusive_lower: bool
401    inclusive_upper: bool
402    normalized_weight: float | None = None
403
404
405@dataclass(slots=True, frozen=True)
406class sequence_mapping[V, S: Float](
407    SimFunc[Sequence[V], SequenceSim[V, S]], HasMetadata
408):
409    """
410    List Mapping similarity function.
411
412    Parameters:
413        element_similarity: The similarity function to use for comparing elements.
414        exact: Whether to use exact or inexact comparison. Default is False (inexact).
415        weights: Optional list of weights for weighted similarity calculation.
416
417    Examples:
418        >>> sim = sequence_mapping(lambda x, y: 1.0 if x == y else 0.0, True)
419        >>> result = sim(["a", "b", "c"], ["a", "b", "c"])
420        >>> result.value
421        1.0
422        >>> result.similarities
423        [1.0, 1.0, 1.0]
424    """
425
426    element_similarity: SimFunc[V, S]
427    exact: bool = False
428    weights: list[Weight] | None = None
429
430    @property
431    def metadata(self) -> JsonDict:
432        """Return metadata describing the sequence mapping configuration."""
433        return {
434            "element_similarity": get_metadata(self.element_similarity),
435            "exact": self.exact,
436            "weights": [asdict(weight) for weight in self.weights]
437            if self.weights
438            else None,
439        }
440
441    def compute_contains_exact(
442        self, query: Sequence[V], case: Sequence[V]
443    ) -> SequenceSim[V, S]:
444        """Compute element-wise similarity for sequences of equal length."""
445        if len(query) != len(case):
446            return SequenceSim(value=0.0, similarities=None, mapping=None)
447
448        sim_sum = 0.0
449        local_similarities: list[S] = []
450
451        for elem_q, elem_c in zip(query, case, strict=False):
452            sim = self.element_similarity(elem_q, elem_c)
453            sim_sum += unpack_float(sim)
454            local_similarities.append(sim)
455
456        return SequenceSim(
457            value=sim_sum / len(query),
458            similarities=local_similarities,
459            mapping=None,
460        )
461
462    def compute_contains_inexact(
463        self, case_list: Sequence[V], query_list: Sequence[V]
464    ) -> SequenceSim[V, S]:
465        """
466        Slides the *shorter* sequence across the *longer* one and always
467        evaluates   query → case   (i.e. query elements are compared against
468        the current window cut from the case list).
469        """
470        case_is_longer = len(case_list) >= len(query_list)
471        larger, smaller = (
472            (case_list, query_list) if case_is_longer else (query_list, case_list)
473        )
474
475        best_value = -1.0
476        best_sims: Sequence[S] = []
477
478        for start in range(len(larger) - len(smaller) + 1):
479            window = larger[start : start + len(smaller)]
480
481            if case_is_longer:
482                sim_res = self.compute_contains_exact(smaller, window)
483            else:
484                sim_res = self.compute_contains_exact(window, smaller)
485
486            if sim_res.value > best_value:
487                best_value = sim_res.value
488                best_sims = sim_res.similarities or []
489
490        return SequenceSim(value=best_value, similarities=best_sims, mapping=None)
491
492    def __call__(self, x: Sequence[V], y: Sequence[V]) -> SequenceSim[V, S]:
493        # x is the “case”, y is the “query”
494        if self.exact:
495            result = self.compute_contains_exact(y, x)
496        else:
497            result = self.compute_contains_inexact(x, y)
498
499        if self.weights and result.similarities:
500            total_weighted_sim = 0.0
501            total_weight = 0.0
502
503            # Normalize weights
504            for weight in self.weights:
505                weight_range = weight.upper_bound - weight.lower_bound
506                weight.normalized_weight = weight.weight / weight_range
507
508            # Apply weights
509            for sim in result.similarities:
510                sim_val = unpack_float(sim)
511
512                for weight in self.weights:
513                    lower_bound = weight.lower_bound
514                    upper_bound = weight.upper_bound
515                    inclusive_lower = weight.inclusive_lower
516                    inclusive_upper = weight.inclusive_upper
517
518                    if (
519                        (inclusive_lower and lower_bound <= sim_val <= upper_bound)
520                        or (
521                            not inclusive_lower and lower_bound < sim_val <= upper_bound
522                        )
523                    ) and (inclusive_upper or sim_val < upper_bound):
524                        assert weight.normalized_weight is not None
525                        weighted_sim = weight.normalized_weight * sim_val
526                        total_weighted_sim += weighted_sim
527                        total_weight += weight.normalized_weight
528
529            if total_weight > 0:
530                final_similarity = total_weighted_sim / total_weight
531            else:
532                final_similarity = result.value
533
534            return SequenceSim(
535                value=final_similarity,
536                similarities=result.similarities,
537                mapping=result.mapping,
538            )
539
540        return result
541
542
543@dataclass(slots=True, frozen=True)
544class sequence_correctness[V](SimFunc[Sequence[V], float]):
545    """List Correctness similarity function.
546
547    Parameters:
548    worst_case_sim (float): The similarity value to use when all pairs are discordant. Default is 0.0.
549
550    Examples:
551        >>> sim = sequence_correctness(0.5)
552        >>> sim(["Monday", "Tuesday", "Wednesday"], ["Monday", "Wednesday", "Tuesday"])
553        0.3333333333333333
554    """
555
556    worst_case_sim: float = 0.0
557
558    @override
559    def __call__(self, x: Sequence[V], y: Sequence[V]) -> float:
560        if len(x) != len(y):
561            return 0.0
562
563        count_concordant = 0
564        count_discordant = 0
565
566        for i, j in product(range(len(x)), repeat=2):
567            if i >= j:
568                continue
569
570            index_x1 = x.index(x[i])
571            index_x2 = x.index(x[j])
572            index_y1 = y.index(x[i])
573            index_y2 = y.index(x[j])
574
575            if index_y1 == -1 or index_y2 == -1:
576                continue
577            elif (index_x1 < index_x2 and index_y1 < index_y2) or (
578                index_x1 > index_x2 and index_y1 > index_y2
579            ):
580                count_concordant += 1
581            else:
582                count_discordant += 1
583
584        if count_concordant + count_discordant == 0:
585            return 0.0
586
587        correctness = (count_concordant - count_discordant) / (
588            count_concordant + count_discordant
589        )
590
591        if correctness >= 0:
592            return correctness
593        else:
594            return abs(correctness) * self.worst_case_sim