def measure(self, protocol: Protocol) -> float:
disentanglement_scores = []
non_constant_positions = 0
for j in range(self.max_message_length):
symbols_j = [message[j] for message in protocol.values()]
symbol_mutual_info = []
symbol_entropy = compute_entropy(symbols_j)
for i in range(self.num_concept_slots):
concepts_i = [
flatten_derivation(derivation)[i]
for derivation in protocol.keys()
]
mutual_info = compute_mutual_information(concepts_i, symbols_j)
symbol_mutual_info.append(mutual_info)
symbol_mutual_info.sort(reverse=True)
if symbol_entropy > 0:
disentanglement_score = (
symbol_mutual_info[0] -
symbol_mutual_info[1]) / symbol_entropy
disentanglement_scores.append(disentanglement_score)
non_constant_positions += 1
if non_constant_positions > 0:
return sum(disentanglement_scores) / non_constant_positions
else:
return np.nan
def _protocol_to_tensor(
self, protocol: Protocol
) -> Dict[Tuple[torch.LongTensor, torch.LongTensor], torch.LongTensor]:
vocab = get_vocab_from_protocol(protocol)
concept_set = set(concept for derivation in protocol.keys()
for concept in flatten_derivation(derivation))
concepts = {concept: idx for idx, concept in enumerate(concept_set)}
tensorized_protocol = {}
for derivation, message in protocol.items():
derivation = derivation_to_tensor(derivation, concepts)
message = torch.LongTensor([vocab[char] for char in message])
tensorized_protocol[derivation] = torch.nn.functional.one_hot(
message, num_classes=len(vocab)).reshape(-1)
return tensorized_protocol