import math
import numpy as np
from typing import Tuple
class Tonnetz:
def __init__(self, grid_size: int = 12):
self.grid_size = grid_size
@property
def total_positions(self) -> int:
return self.grid_size * self.grid_size
def coordinates(self, index: int) -> Tuple[int, int]:
n = self.grid_size
return (index % n, (index // n) % n)
def to_index(self, row: int, col: int) -> int:
n = self.grid_size
return (row % n) + (col % n) * n
def distance(self, i: int, j: int) -> int:
n = self.grid_size
xi, yi = self.coordinates(i)
xj, yj = self.coordinates(j)
dx = abs(xi - xj)
dy = abs(yi - yj)
return min(dx, n - dx) + min(dy, n - dy)
def distance_coords(self, a: Tuple[int, int], b: Tuple[int, int]) -> int:
n = self.grid_size
dx = abs(a[0] - b[0])
dy = abs(a[1] - b[1])
return min(dx, n - dx) + min(dy, n - dy)
def spectral_gap(self) -> float:
return 2.0 - 2.0 * math.cos(2.0 * math.pi / self.grid_size)
def decay_rate(self, t: float) -> float:
return math.exp(-self.spectral_gap() * t)
def distance_matrix(n_tokens: int, grid_size: int = 12) -> np.ndarray:
idx = np.arange(n_tokens)
x = idx % grid_size
y = (idx // grid_size) % grid_size
dx = np.abs(x[:, None] - x[None, :])
dy = np.abs(y[:, None] - y[None, :])
dx = np.minimum(dx, grid_size - dx)
dy = np.minimum(dy, grid_size - dy)
return dx + dy