latents.py

#!/usr/bin/env python3
import numpy as np
import einops
import math
from collections.abc import Callable
import sys

def ber(p: float, val: int) -> float:
  assert 0 <= p <= 1 and val in [0, 1]
  return [1 - p, p][val]

def is_distribution(prob: np.ndarray) -> bool:
  return bool(np.all(prob >= 0)) and math.isclose(np.sum(prob), 1)

"""
Given distribution over X, Y, Z, extend to variables:
- A = X ⊕ Y
- B = Y ⊕ Z
- A ⊕ B
- X ⊕ Y ⊕ Z
"""
ALL_VARS    : list[str] = ['X', 'Y', 'Z', 'A', 'B', 'A_xor_B', 'A_and_B', 'A_or_B', 'X_xor_Y_xor_Z']
PRETTY_VARS : list[str] = ['X', 'Y', 'Z', 'A', 'B', 'A^B',     'A&B',     'A|B',    'X^Y^Z']
def get_dist_with_fns(prob: np.ndarray) -> np.ndarray:
  prob_ext : np.ndarray = np.zeros([2] * len(ALL_VARS))
  for x, y, z in np.ndindex(prob.shape):
    a = x ^ y
    b = z ^ y
    prob_ext[x, y, z,
                 a,
                 b,
                 a ^ b,
                 a & b,
                 a | b,
                 x ^ y ^ z,
               ] \
      = prob[x, y, z]
  
  return prob_ext
  
"""
Distribution for 0 <= p, q, r <= 1:
- X ~ Bernoulli(p)
- Z: With probability r, Z := X; otherwise Z ~ Bernoulli(p)
- Y ~ Bernoulli(q)
"""
def get_dist(p: float, q: float, r: float) -> np.ndarray:
  assert 0 <= p <= 1 and 0 <= q <= 1 and 0 <= r <= 1
  prob : np.ndarray = np.zeros((2, 2, 2))
  
  for x in [0, 1]:
    for y in [0, 1]:
        p_x = ber(p, x)
        p_y = ber(q, y)

        # case 1: z copies x
        prob[x, y, x] += p_x * p_y * r

        # case 2: z does not copy x
        for z in [0, 1]:
          prob[x, y, z] += p_x * p_y * (1 - r) * ber(p, z)  
  assert is_distribution(prob)
  return prob

def all_unique(l: list) -> bool:
  return np.unique(l).size == len(l)

def marginal(prob: np.ndarray, vars: list[str]) -> np.ndarray:
  assert all_unique(vars)
  return einops.einsum(prob, f"{' '.join(ALL_VARS)} -> {' '.join(vars)}")

def H(prob: np.ndarray, vars: list[str]) -> float:
  # assert is_distribution(prob)

  unique_vars = list(set(vars))
  prob = marginal(prob, unique_vars)
  pos_prob = prob[prob != 0]
  return np.sum(-pos_prob * np.log2(pos_prob))

def H_cond(prob: np.ndarray, vars: list[str], vars_cond: list[str]) -> float:
  return H(prob, vars + vars_cond) - H(prob, vars_cond)

def I(prob: np.ndarray, vars_x: list[str], vars_y: list[str]) -> float:
  return H(prob, vars_x) + H(prob, vars_y) - H(prob, vars_x + vars_y)

def I_cond(prob: np.ndarray, vars_x: list[str], vars_y: list[str], vars_cond: list[str]) -> float:
  return H_cond(prob, vars_x, vars_cond) + H_cond(prob, vars_y, vars_cond) \
         - H_cond(prob, vars_x + vars_y, vars_cond)

def det_error(prob: np.ndarray, vars_x: list[str], vars_y: list[str], latent: list[str], latent_label: str='', log=False) -> float:
  err = max(
    I_cond(prob, vars_x, vars_y, latent),
    H_cond(prob, latent, vars_x),
    H_cond(prob, latent, vars_y)
  )
  if log:
    print(f"{latent_label:<8} {abs(I_cond(prob, vars_x, vars_y, latent)):2.2f}     {abs(H_cond(prob, latent, vars_x)):2.2f}    {abs(H_cond(prob, latent, vars_y)):2.2f}   {err:2.2f}")
  return err

def stoch_error(prob: np.ndarray, vars_x: list[str], vars_y: list[str], latent: list[str]) -> float:
  return max(
    I_cond(prob, vars_x, vars_y, latent),
    I_cond(prob, vars_x, latent, vars_y),
    I_cond(prob, vars_y, latent, vars_x)
  )

def get_ratio(prob: np.ndarray, log=False) -> float:
  prob = get_dist_with_fns(prob)
  stoch_y = stoch_error(prob, ['A'], ['B'], ['Y'])
  latents : list[list[str]] = [[var] for var in ALL_VARS] + [[]]
  labels  : list[str]       = PRETTY_VARS                 + ['⊥']
  if log:
    print(f'ε = {stoch_y:2.2f}')
    print(f"{'Λ':<6} I(A;B|Λ)  H(Λ|A)  H(Λ|B)  max")
  det_errors = [det_error(prob, ['A'], ['B'], latent, label, log) \
                   for latent, label in zip(latents, labels)]
  return min(det_errors) / stoch_y

rng = np.random.default_rng()
def get_rand_dist() -> np.ndarray:
  prob = rng.dirichlet(np.ones(8))
  prob = prob.reshape((2, 2, 2))
  assert is_distribution(prob)
  return prob

def hill_climb(prob : np.ndarray, obj : Callable[[np.ndarray], float], alpha=0.02) -> tuple[float, np.ndarray]:
  verts = [np.zeros(prob.shape) for i in range(prob.size)]
  for idx, vert in enumerate(verts):
    vert.flat[idx] = 1
  assert all(is_distribution(vert) for vert in verts)

  best = obj(prob)
  while True:
    cands = [alpha*vert + (1 - alpha)*prob for vert in verts]
    # assert all(is_distribution(cand) for cand in cands)
    cur_best, prob = max((obj(cand), cand) for cand in cands)
    if cur_best < best:
      return best, prob
    alpha = min(alpha, cur_best - best)
    if alpha < 0.0001:
      return cur_best, prob
    best = cur_best

def main():
  try:
    while True:
      best, prob = hill_climb(get_rand_dist(), get_ratio, alpha=0.01)
      if best < 5:
        print('#', end='')
        sys.stdout.flush()
      else:
        print(f'\nratio={best:1.1f}')
        print(prob)
        get_ratio(prob, log=True)
  except KeyboardInterrupt:
    print('\nAborted.')
    sys.exit(0)

if __name__ == "__main__":
    main()