Add substitution, insertion, and deletion cost functions

README.md

@@ -81,6 +81,57 @@ Even if the alignment of the two sequences is identical to `difflib`,
 sequences.  The opcodes returned by this library represent individual character 
 operations, and thus should never span two or more characters.
 
+Custom costs
+------------
+By default every edit costs one and a match costs nothing.  Three optional cost
+functions let you change that:
+
+```python
+sm = edit_distance.SequenceMatcher(
+    a=ref, b=hyp,
+    substitution_cost=lambda x, y: 0 if x == y else 1,
+    insertion_cost=lambda y: 1,
+    deletion_cost=lambda x: 1,
+)
+```
+
+They are also accepted by `edit_distance()` and `edit_distance_backpointer()`,
+as keyword arguments.  Costs may be fractional, in which case the distance is a
+float.
+
+`substitution_cost` is consulted for *every* aligned pair, including pairs that
+compare equal, so a custom function has to handle both cases -- returning a flat
+penalty would charge for matches too:
+
+```python
+# charges 0.5 even when x == y
+substitution_cost=lambda x, y: 0.5
+```
+
+Costs and equality are separate concerns.  `test` decides whether two elements
+are equal, which determines the opcode label and the match count; the cost
+functions only decide the price.  So a free substitution that `test` rejects is
+still a `replace` and still isn't a match, and a substitution that costs
+something but that `test` accepts is still an `equal` and still is.
+
+Two consequences worth knowing:
+
+- A substitution costing more than an insertion plus a deletion is never
+  chosen, because the search gets the same result more cheaply by deleting and
+  then inserting.  Substitution costs are effectively capped at `ins + del`.
+- `distance(a, a) == 0` only holds if a match costs nothing, and
+  `distance(a, b) == distance(b, a)` only holds if insertions and deletions
+  cost the same.  Keeping the result a true metric is up to you.
+
+Setting the substitution cost to the gap-pair total maximizes the number of
+matches rather than minimizing edits, which is what the removed
+`highest_match_action` used to do:
+
+```python
+sm = edit_distance.SequenceMatcher(a=ref, b=hyp,
+     substitution_cost=lambda x, y: 0 if x == y else 2)
+```
+
 Notes
 -----
 This doesn't implement the 'junk' matching features in difflib.

edit_distance/edit_distance.py

@@ -20,55 +20,106 @@
 
 import operator
 import sys
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
+from typing import Any
 
 INSERT: str = "insert"
 DELETE: str = "delete"
 EQUAL: str = "equal"
 REPLACE: str = "replace"
 
+TestFunction = Callable[[Any, Any], Any]
+SubstitutionCostFunction = Callable[[Any, Any], float]
+GapCostFunction = Callable[[Any], float]
+
 
-# Cost is basically: was there a match or not.
-# The other numbers are cumulative costs and matches.
-#
 # At each cell the search takes whichever of the three actions is cheapest,
 # breaking ties in favor of substitution, then insertion, then deletion.
+#
+# Two properties of custom costs are worth knowing:
+#
+# * A substitution costing more than an insertion plus a deletion is never
+#   chosen -- the search routes around it by deleting and then inserting -- so
+#   substitution costs are effectively capped at ``ins + del``.
+# * ``distance(a, a) == 0`` only holds if a match costs nothing, and symmetry
+#   only holds if insertions and deletions cost the same.  Keeping the result a
+#   true metric is the caller's responsibility.
 
 
-def edit_distance(seq1: Sequence, seq2: Sequence, test=operator.eq):
+def edit_distance(
+    seq1: Sequence,
+    seq2: Sequence,
+    *,
+    test: TestFunction = operator.eq,
+    substitution_cost: SubstitutionCostFunction | None = None,
+    insertion_cost: GapCostFunction | None = None,
+    deletion_cost: GapCostFunction | None = None,
+) -> tuple[float, int]:
     """
     Computes the edit distance between the two given sequences.  This uses the
     relatively fast method that only constructs two columns of the 2d array
     for edits.  This function actually uses four columns because we track the
-    number of matches too.
+    number of matches too.  Returns a ``(distance, matches)`` tuple.
+
+    ``test`` decides whether two elements are equal.  It governs the opcode
+    label (``equal`` versus ``replace``) and the match count, and defaults to
+    :py:func:`operator.eq`.
+
+    The three cost functions decide only what each operation *costs*:
+
+    * ``substitution_cost(x, y)`` -- the cost of aligning ``x`` with ``y``.
+      This is consulted for **every** aligned pair, including pairs that
+      ``test`` considers equal, so a custom function must handle both cases;
+      the usual form is ``lambda x, y: 0 if x == y else penalty(x, y)``.
+      Defaults to ``0`` for a match and ``1`` otherwise.
+    * ``insertion_cost(y)`` -- the cost of inserting ``y``.  Defaults to ``1``.
+    * ``deletion_cost(x)`` -- the cost of deleting ``x``.  Defaults to ``1``.
+
+    Costs and equality are independent: a zero-cost pair that ``test`` rejects
+    is still a ``replace`` and does not count as a match, and a nonzero-cost
+    pair that ``test`` accepts is still an ``equal`` and does.
     """
     m = len(seq1)
     n = len(seq2)
-    # Special, easy cases:
-    if test is operator.eq and seq1 == seq2:
+    # Identical sequences are only free when a match costs nothing, which a
+    # custom substitution cost is under no obligation to honor.
+    if substitution_cost is None and test is operator.eq and seq1 == seq2:
         return 0, n
-    if m == 0:
-        return n, 0
-    if n == 0:
-        return m, 0
-    v0 = [0] * (n + 1)  # The two 'error' columns
-    v1 = [0] * (n + 1)
+
+    # Insertion costs are needed once per row, so pay for them only once.
+    ins_costs = [1] * n if insertion_cost is None else [insertion_cost(y) for y in seq2]
+
+    v0: list[float] = [0] * (n + 1)  # The two 'error' columns
+    v1: list[float] = [0] * (n + 1)
     m0 = [0] * (n + 1)  # The two 'match' columns
     m1 = [0] * (n + 1)
-    for i in range(1, n + 1):
-        v0[i] = i
+
+    # The first row: the running cost of inserting all of seq2 up to here.
+    for j in range(1, n + 1):
+        v0[j] = v0[j - 1] + ins_costs[j - 1]
+    if m == 0:
+        return v0[n], 0
+
     for i in range(1, m + 1):
-        v1[0] = i
+        a_elem = seq1[i - 1]
+        elem_del_cost = 1 if deletion_cost is None else deletion_cost(a_elem)
+        # The first column: the running cost of deleting all of seq1 up to here.
+        v1[0] = v0[0] + elem_del_cost
         for j in range(1, n + 1):
-            cost = 0 if test(seq1[i - 1], seq2[j - 1]) else 1
+            b_elem = seq2[j - 1]
+            equal = test(a_elem, b_elem)
+            if substitution_cost is None:
+                cost: float = 0 if equal else 1
+            else:
+                cost = substitution_cost(a_elem, b_elem)
             # The costs
-            ins_cost = v1[j - 1] + 1
-            del_cost = v0[j] + 1
+            ins_cost = v1[j - 1] + ins_costs[j - 1]
+            del_cost = v0[j] + elem_del_cost
             sub_cost = v0[j - 1] + cost
             # Match counts
             ins_match = m1[j - 1]
             del_match = m0[j]
-            sub_match = m0[j - 1] + int(not cost)
+            sub_match = m0[j - 1] + (1 if equal else 0)
 
             # Ties break in favor of substitution, then insertion, then deletion.
             if sub_cost <= ins_cost and sub_cost <= del_cost:
@@ -87,51 +138,73 @@ def edit_distance(seq1: Sequence, seq2: Sequence, test=operator.eq):
     return v1[n], m1[n]
 
 
-def edit_distance_backpointer(seq1, seq2, test=operator.eq):
+def edit_distance_backpointer(
+    seq1: Sequence,
+    seq2: Sequence,
+    *,
+    test: TestFunction = operator.eq,
+    substitution_cost: SubstitutionCostFunction | None = None,
+    insertion_cost: GapCostFunction | None = None,
+    deletion_cost: GapCostFunction | None = None,
+) -> tuple[float, int, list]:
     """
     Similar to :py:func:`~edit_distance.edit_distance` except that this
     function keeps backpointers during the search.  This allows us to return
     the opcodes (i.e. the specific edits that were used to change from one
     string to another).  This function contructs the full 2d array for the
-    backpointers only.
+    backpointers only.  Returns a ``(distance, matches, opcodes)`` tuple.
+
+    ``test`` and the three cost functions mean the same thing they do for
+    :py:func:`~edit_distance.edit_distance`.
     """
     m: int = len(seq1)
     n: int = len(seq2)
     # backpointer array:
-    bp = [[None for _ in range(n + 1)] for _ in range(m + 1)]
+    bp: list[list[str | None]] = [[None for _ in range(n + 1)] for _ in range(m + 1)]
+
+    # Insertion costs are needed once per row, so pay for them only once.
+    ins_costs = [1] * n if insertion_cost is None else [insertion_cost(y) for y in seq2]
 
     # Two columns of the distance and match arrays
-    d0 = [0] * (n + 1)  # The two 'distance' columns
-    d1 = [0] * (n + 1)
+    d0: list[float] = [0] * (n + 1)  # The two 'distance' columns
+    d1: list[float] = [0] * (n + 1)
     m0 = [0] * (n + 1)  # The two 'match' columns
     m1 = [0] * (n + 1)
 
-    # Fill in the first column
-    for i in range(1, n + 1):
-        d0[i] = i
-        bp[0][i] = INSERT
+    # Fill in the first row: the running cost of inserting all of seq2.
+    for j in range(1, n + 1):
+        d0[j] = d0[j - 1] + ins_costs[j - 1]
+        bp[0][j] = INSERT
 
     for i in range(1, m + 1):
-        d1[0] = i
+        a_elem = seq1[i - 1]
+        elem_del_cost = 1 if deletion_cost is None else deletion_cost(a_elem)
+        # The first column: the running cost of deleting all of seq1 up to here.
+        d1[0] = d0[0] + elem_del_cost
         bp[i][0] = DELETE
 
         for j in range(1, n + 1):
-            cost = 0 if test(seq1[i - 1], seq2[j - 1]) else 1
+            b_elem = seq2[j - 1]
+            equal = test(a_elem, b_elem)
+            if substitution_cost is None:
+                cost: float = 0 if equal else 1
+            else:
+                cost = substitution_cost(a_elem, b_elem)
             # The costs of each action...
-            ins_cost = d1[j - 1] + 1  # insertion
-            del_cost = d0[j] + 1  # deletion
+            ins_cost = d1[j - 1] + ins_costs[j - 1]  # insertion
+            del_cost = d0[j] + elem_del_cost  # deletion
             sub_cost = d0[j - 1] + cost  # substitution/match
 
             # The match scores of each action
             ins_match = m1[j - 1]
             del_match = m0[j]
-            sub_match = m0[j - 1] + int(not cost)
+            sub_match = m0[j - 1] + (1 if equal else 0)
 
             # Ties break in favor of substitution, then insertion, then deletion.
             if sub_cost <= ins_cost and sub_cost <= del_cost:
                 d1[j] = sub_cost
                 m1[j] = sub_match
-                bp[i][j] = EQUAL if cost == 0 else REPLACE
+                bp[i][j] = EQUAL if equal else REPLACE
             elif ins_cost <= del_cost:
                 d1[j] = ins_cost
                 m1[j] = ins_match

edit_distance/sequence_matcher.py

@@ -21,9 +21,18 @@
 import operator
 from collections.abc import Sequence
 
-from edit_distance.edit_distance import EQUAL, edit_distance, edit_distance_backpointer
-
-
+from edit_distance.edit_distance import (
+    EQUAL,
+    GapCostFunction,
+    SubstitutionCostFunction,
+    TestFunction,
+    edit_distance,
+    edit_distance_backpointer,
+)
+
+
+# Two sequences, four comparison/cost hooks, and three cached results.
+# pylint: disable-next=too-many-instance-attributes
 class SequenceMatcher:
     """
     Similar to the :py:mod:`difflib` :py:class:`~difflib.SequenceMatcher`, but
@@ -34,12 +43,20 @@ def __init__(
         self,
         a: Sequence | None = None,
         b: Sequence | None = None,
-        test=operator.eq,
+        test: TestFunction = operator.eq,
+        *,
+        substitution_cost: SubstitutionCostFunction | None = None,
+        insertion_cost: GapCostFunction | None = None,
+        deletion_cost: GapCostFunction | None = None,
     ):
         """
         Initialize the object with sequences a and b.  Optionally, one can
         specify a test function that is used to compare sequence elements. This
         defaults to the built in ``eq`` operator (i.e. :py:func:`operator.eq`).
+
+        The three cost functions are passed straight through to
+        :py:func:`~edit_distance.edit_distance`; see that function for what
+        they mean and what they default to.
         """
         if a is None:
             a = []
@@ -49,8 +66,11 @@ def __init__(
         self.seq2 = b
         self._reset_object()
         self.test = test
-        self.dist = None
-        self._matches = None
+        self.substitution_cost = substitution_cost
+        self.insertion_cost = insertion_cost
+        self.deletion_cost = deletion_cost
+        self.dist: float | None = None
+        self._matches: int | None = None
         self.opcodes = None
 
     def set_seqs(self, a: Sequence, b: Sequence) -> None:
@@ -97,6 +117,9 @@ def get_opcodes(self):
                 self.seq1,
                 self.seq2,
                 test=self.test,
+                substitution_cost=self.substitution_cost,
+                insertion_cost=self.insertion_cost,
+                deletion_cost=self.deletion_cost,
             )
             if self.dist is not None:
                 assert d == self.dist
@@ -128,7 +151,14 @@ def real_quick_ratio(self) -> float:
     def _compute_distance_fast(self) -> None:
         """Calls edit_distance, and asserts that if we already have values for
         matches and distance, that they match."""
-        d, m = edit_distance(self.seq1, self.seq2, test=self.test)
+        d, m = edit_distance(
+            self.seq1,
+            self.seq2,
+            test=self.test,
+            substitution_cost=self.substitution_cost,
+            insertion_cost=self.insertion_cost,
+            deletion_cost=self.deletion_cost,
+        )
         if self.dist is not None:
             assert d == self.dist
         if self._matches is not None: