actually add changes to informativity

src/ultk/effcomm/informativity.py

@@ -15,27 +15,53 @@
 )
 
 
-def build_utility_matrix(
-    universe: Universe, utility: Callable[[Referent, Referent], float]
+def build_pairwise_matrix(
+    universe: Universe, func: Callable[[Referent, Referent], float]
 ) -> np.ndarray:
-    """Construct the square matrix specifying the utility function defined for pairs of meanings, used for computing communicative success."""
+    """Construct the square matrix specifying the utility/cost function defined for pairs of meanings, used for computing communicative success."""
     return np.array(
         [
-            [utility(ref, ref_) for ref_ in universe.referents]
+            [func(ref, ref_) for ref_ in universe.referents]
             for ref in universe.referents
         ]
     )
 
 
-def indicator_utility(ref1: Referent, ref2: Referent) -> float:
-    """Indicator utility function, i.e. delta.  Returns 1.0 iff ref1 equals ref2."""
-    return float(ref1 == ref2)
+def expected_communication_score(
+    speaker: Speaker,
+    listener: Listener,
+    prior: np.ndarray,
+    score_matrix: np.ndarray,
+) -> float:
+    """Compute an expected utility or cost function w.r.t  speaker, listener, and prior.
+
+
+    $I(L) = \sum_{m, \hat{m}} P(m, \hat{m}) \cdot u(m, \hat{m})$
+
+    $ = \sum_{m \in M} p(m) \sum_{i \in L} p(i|m) \sum_{\hat{m} \in i} p(\hat{m} |i) \cdot u(m, m')$
+
+    $ = \sum \\text{diag}(p)SR \odot U $
+
+    For more details, see [docs/vectorized_informativity](https://github.com/CLMBRs/altk/blob/main/docs/vectorized_informativity.pdf).
+
+    Args:
+        speaker: a literal or pragmatic speaker, containing a matrix S for P(e | m)
+
+        listener: a literal or pragmatic listener, containing a matrix R for P(m | e)
+
+        prior: p(m), distribution over meanings representing communicative need
+
+        utility: a matrix encoding u(m, m'), e.g. similarity of meanings
+    """
+    S = speaker.normalized_weights()
+    R = listener.normalized_weights()
+    return float(np.sum(np.diag(prior) @ S @ R * score_matrix))
 
 
 def informativity(
     language: Language,
     prior: np.ndarray,
-    utility: Callable[[Referent, Referent], float] = indicator_utility,
+    score: np.ndarray = None,
     agent_type: str = "literal",
 ) -> float:
     """The informativity of a language is identified with the successful communication between a speaker and a listener.
@@ -47,7 +73,9 @@ def informativity(
 
         prior: a probability distribution representing communicative need (frequency) for Referents.
 
-        utility: a function representing the usefulness of listener guesses about speaker Referents, e.g. Referent similarity. To reward only exact recovery of meanings, use the indicator function (default).
+        score: a 2D matrix representing the usefulness of listener guesses about speaker Referents, 
+        where utility[i][j] specifies the utility of guessing Referent j when the true Referent is i. 
+        To reward only exact recovery of meanings, use an identity matrix as the utility matrix (default).
 
         kind: {"literal, pragmatic"} Whether to measure informativity using literal or pragmatic agents, as canonically described in the Rational Speech Act framework. The default is "literal".
 
@@ -70,6 +98,9 @@ def informativity(
     speaker = LiteralSpeaker(language)
     listener = LiteralListener(language)
 
+    if score is None:
+        score = np.eye(len(language.universe))
+
     if agent_type == "literal":
         pass
     elif agent_type == "pragmatic":
@@ -80,47 +111,4 @@ def informativity(
             f"agent_type must be either 'literal' or 'pragmatic'. Received: {agent_type}."
         )
 
-    inf = communicative_success(speaker, listener, prior, utility)
-
-    # Check informativity > 0
-    utility_matrix = build_utility_matrix(speaker.language.universe, utility)
-    m, _ = utility_matrix.shape  # square matrix
-    if np.array_equal(utility_matrix, np.eye(m)):
-        if isclose(inf, 0.0):
-            raise ValueError(
-                f"Informativity must be nonzero for indicator utility reward function, but was: {inf}"
-            )
-
-    return inf
-
-
-def communicative_success(
-    speaker: Speaker,
-    listener: Listener,
-    prior: np.ndarray,
-    utility: Callable[[Referent, Referent], float],
-) -> float:
-    """Helper function to compute the literal informativity of a language.
-
-
-    $I(L) = \sum_{m, \hat{m}} P(m, \hat{m}) \cdot u(m, \hat{m})$
-
-    $ = \sum_{m \in M} p(m) \sum_{i \in L} p(i|m) \sum_{\hat{m} \in i} p(\hat{m} |i) \cdot u(m, m')$
-
-    $ = \sum \\text{diag}(p)SR \odot U $
-
-    For more details, see [docs/vectorized_informativity](https://github.com/CLMBRs/altk/blob/main/docs/vectorized_informativity.pdf).
-
-    Args:
-        speaker: a literal or pragmatic speaker, containing a matrix S for P(e | m)
-
-        listener: a literal or pragmatic listener, containing a matrix R for P(m | e)
-
-        prior: p(m), distribution over meanings representing communicative need
-
-        utility: a function u(m, m') representing similarity of meanings, or pair-wise usefulness of listener guesses about speaker meanings.
-    """
-    S = speaker.normalized_weights()
-    R = listener.normalized_weights()
-    U = build_utility_matrix(speaker.language.universe, utility)
-    return float(np.sum(np.diag(prior) @ S @ R * U))
+    return expected_communication_score(speaker, listener, prior, score)