|
|
|
|
|
|
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
import numpy as np |
|
|
import math |
|
|
import matplotlib.pyplot as plt |
|
|
from typing import List, Dict, Tuple, Optional |
|
|
|
|
|
SAFE_MIN = -1e6 |
|
|
SAFE_MAX = 1e6 |
|
|
EPS = 1e-8 |
|
|
|
|
|
|
|
|
|
|
|
def make_safe(tensor, min_val=SAFE_MIN, max_val=SAFE_MAX): |
|
|
zero = torch.tensor(0.0, device=tensor.device, dtype=tensor.dtype) |
|
|
maxv = torch.tensor(max_val, device=tensor.device, dtype=tensor.dtype) |
|
|
tensor = torch.where(torch.isnan(tensor), zero, tensor) |
|
|
tensor = torch.where(torch.isinf(tensor), maxv, tensor) |
|
|
return torch.clamp(tensor, min_val, max_val) |
|
|
|
|
|
|
|
|
def safe_complex_division(numerator, denominator, eps=EPS): |
|
|
denominator_conj = torch.conj(denominator) |
|
|
norm_sq = torch.real(denominator * denominator_conj) |
|
|
norm_sq = torch.clamp(norm_sq, min=eps) |
|
|
return (numerator * denominator_conj) / norm_sq |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class MobiusTransform(nn.Module): |
|
|
def __init__(self, learnable=True, init_identity=True): |
|
|
super().__init__() |
|
|
self.learnable = learnable |
|
|
|
|
|
if init_identity: |
|
|
a_init, b_init, c_init, d_init = 1.0, 0.0, 0.0, 1.0 |
|
|
else: |
|
|
a_init, d_init = 1.0, 1.0 |
|
|
b_init, c_init = 0.1, 0.1 |
|
|
|
|
|
if learnable: |
|
|
self.a = nn.Parameter(torch.tensor([a_init, 0.0])) |
|
|
self.b = nn.Parameter(torch.tensor([b_init, 0.0])) |
|
|
self.c = nn.Parameter(torch.tensor([c_init, 0.0])) |
|
|
self.d = nn.Parameter(torch.tensor([d_init, 0.0])) |
|
|
else: |
|
|
self.register_buffer('a', torch.tensor([a_init, 0.0])) |
|
|
self.register_buffer('b', torch.tensor([b_init, 0.0])) |
|
|
self.register_buffer('c', torch.tensor([c_init, 0.0])) |
|
|
self.register_buffer('d', torch.tensor([d_init, 0.0])) |
|
|
|
|
|
def to_complex(self, param): |
|
|
return torch.complex(param[0], param[1]) |
|
|
|
|
|
def get_determinant(self): |
|
|
a_complex = self.to_complex(self.a) |
|
|
b_complex = self.to_complex(self.b) |
|
|
c_complex = self.to_complex(self.c) |
|
|
d_complex = self.to_complex(self.d) |
|
|
|
|
|
det = a_complex * d_complex - b_complex * c_complex |
|
|
return det |
|
|
|
|
|
def normalize_parameters(self): |
|
|
if self.learnable: |
|
|
with torch.no_grad(): |
|
|
det = torch.abs(self.get_determinant()) |
|
|
if det < EPS: |
|
|
one = torch.tensor([1.0, 0.0], device=self.a.device, dtype=self.a.dtype) |
|
|
self.a.copy_(one) |
|
|
self.d.copy_(one) |
|
|
self.b.mul_(0.1) |
|
|
self.c.mul_(0.1) |
|
|
for p in (self.a, self.b, self.c, self.d): |
|
|
p.clamp_(-10.0, 10.0) |
|
|
|
|
|
|
|
|
def transform(self, z): |
|
|
self.normalize_parameters() |
|
|
|
|
|
a_complex = self.to_complex(self.a) |
|
|
b_complex = self.to_complex(self.b) |
|
|
c_complex = self.to_complex(self.c) |
|
|
d_complex = self.to_complex(self.d) |
|
|
|
|
|
numerator = a_complex * z + b_complex |
|
|
denominator = c_complex * z + d_complex |
|
|
transformed = safe_complex_division(numerator, denominator) |
|
|
|
|
|
return transformed |
|
|
|
|
|
def inverse_transform(self, w): |
|
|
self.normalize_parameters() |
|
|
|
|
|
a_complex = self.to_complex(self.a) |
|
|
b_complex = self.to_complex(self.b) |
|
|
c_complex = self.to_complex(self.c) |
|
|
d_complex = self.to_complex(self.d) |
|
|
|
|
|
numerator = d_complex * w - b_complex |
|
|
denominator = -c_complex * w + a_complex |
|
|
|
|
|
return safe_complex_division(numerator, denominator) |
|
|
|
|
|
def get_transform_info(self): |
|
|
det = self.get_determinant() |
|
|
one = torch.tensor(1.0, device=det.device, dtype=det.real.dtype) |
|
|
return { |
|
|
'determinant': det, |
|
|
'is_identity': torch.allclose(torch.abs(det), one, atol=1e-6), |
|
|
'parameters': {'a': self.to_complex(self.a), 'b': self.to_complex(self.b), |
|
|
'c': self.to_complex(self.c), 'd': self.to_complex(self.d)} |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ComplexStateMarkovChain(nn.Module): |
|
|
def __init__(self, num_states, state_embedding_dim=64, distance_kernel='gaussian'): |
|
|
super().__init__() |
|
|
self.num_states = num_states |
|
|
self.state_embedding_dim = state_embedding_dim |
|
|
self.distance_kernel = distance_kernel |
|
|
|
|
|
self.state_positions = nn.Parameter( |
|
|
torch.complex( |
|
|
torch.randn(num_states) * 2.0, |
|
|
torch.randn(num_states) * 2.0 |
|
|
) |
|
|
) |
|
|
|
|
|
self.state_embeddings = nn.Parameter(torch.randn(num_states, state_embedding_dim) * 0.1) |
|
|
|
|
|
self.base_transition_logits = nn.Parameter(torch.randn(num_states, num_states) * 0.1) |
|
|
self.distance_scale = nn.Parameter(torch.tensor(1.0)) |
|
|
self.distance_bias = nn.Parameter(torch.tensor(0.0)) |
|
|
|
|
|
if distance_kernel == 'gaussian': |
|
|
self.kernel_width = nn.Parameter(torch.tensor(1.0)) |
|
|
elif distance_kernel == 'inverse': |
|
|
self.kernel_power = nn.Parameter(torch.tensor(1.0)) |
|
|
|
|
|
def compute_transformed_distances(self, mobius_transform): |
|
|
transformed_positions = mobius_transform.transform(self.state_positions) |
|
|
|
|
|
pos_i = transformed_positions.unsqueeze(0) |
|
|
pos_j = transformed_positions.unsqueeze(1) |
|
|
|
|
|
complex_diff = pos_i - pos_j |
|
|
distances = torch.abs(complex_diff) |
|
|
|
|
|
return distances, transformed_positions |
|
|
|
|
|
def distance_to_probability(self, distances): |
|
|
distances = torch.clamp(distances, min=EPS) |
|
|
|
|
|
if self.distance_kernel == 'gaussian': |
|
|
width = torch.clamp(self.kernel_width, min=0.1, max=10.0) |
|
|
prob_contrib = torch.exp(-distances**2 / (2 * width**2)) |
|
|
elif self.distance_kernel == 'inverse': |
|
|
power = torch.clamp(self.kernel_power, min=0.5, max=3.0) |
|
|
prob_contrib = 1.0 / (distances**power + EPS) |
|
|
else: |
|
|
prob_contrib = torch.clamp(1.0 - distances, min=0.0) |
|
|
|
|
|
return prob_contrib |
|
|
|
|
|
def compute_transition_matrix(self, mobius_transform): |
|
|
distances, transformed_positions = self.compute_transformed_distances(mobius_transform) |
|
|
|
|
|
distance_contrib = self.distance_to_probability(distances) |
|
|
|
|
|
scale = torch.clamp(self.distance_scale, min=0.1, max=10.0) |
|
|
bias = torch.clamp(self.distance_bias, min=-5.0, max=5.0) |
|
|
scaled_distance = scale * distance_contrib + bias |
|
|
|
|
|
transition_logits = self.base_transition_logits + scaled_distance |
|
|
transition_logits = transition_logits + torch.eye(self.num_states, device=transition_logits.device)*0.05 |
|
|
|
|
|
transition_matrix = F.softmax(transition_logits, dim=1) |
|
|
|
|
|
return transition_matrix, transformed_positions |
|
|
|
|
|
def forward(self, initial_state, num_steps, mobius_transform): |
|
|
batch_size = initial_state.shape[0] if initial_state.dim() > 1 else 1 |
|
|
|
|
|
if initial_state.dim() == 1: |
|
|
current_state = initial_state.unsqueeze(0) |
|
|
else: |
|
|
current_state = initial_state |
|
|
|
|
|
transition_matrix, transformed_positions = self.compute_transition_matrix(mobius_transform) |
|
|
|
|
|
trajectory = [current_state.clone()] |
|
|
state_positions = [transformed_positions[current_state.argmax(dim=-1)]] |
|
|
|
|
|
for step in range(num_steps): |
|
|
current_state = torch.matmul(current_state, transition_matrix) |
|
|
trajectory.append(current_state.clone()) |
|
|
|
|
|
most_likely_states = current_state.argmax(dim=-1) |
|
|
state_positions.append(transformed_positions[most_likely_states]) |
|
|
|
|
|
return { |
|
|
'trajectory': torch.stack(trajectory), |
|
|
'final_state': current_state, |
|
|
'state_positions': torch.stack(state_positions), |
|
|
'transition_matrix': transition_matrix, |
|
|
'transformed_positions': transformed_positions |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class MobiusMarkovSystem(nn.Module): |
|
|
def __init__(self, num_states, state_embedding_dim=64, evolution_steps=10): |
|
|
super().__init__() |
|
|
self.num_states = num_states |
|
|
self.evolution_steps = evolution_steps |
|
|
|
|
|
self.mobius_transform = MobiusTransform(learnable=True, init_identity=True) |
|
|
self.markov_chain = ComplexStateMarkovChain(num_states, state_embedding_dim) |
|
|
|
|
|
self.mobius_evolution = nn.Sequential( |
|
|
nn.Linear(state_embedding_dim, state_embedding_dim), |
|
|
nn.Tanh(), |
|
|
nn.Linear(state_embedding_dim, 8), |
|
|
) |
|
|
|
|
|
self.state_encoder = nn.Sequential( |
|
|
nn.Linear(num_states, state_embedding_dim), |
|
|
nn.LayerNorm(state_embedding_dim), |
|
|
nn.ReLU(), |
|
|
nn.Linear(state_embedding_dim, state_embedding_dim) |
|
|
) |
|
|
|
|
|
self.state_decoder = nn.Sequential( |
|
|
nn.Linear(state_embedding_dim, state_embedding_dim), |
|
|
nn.ReLU(), |
|
|
nn.Linear(state_embedding_dim, num_states), |
|
|
nn.Softmax(dim=-1) |
|
|
) |
|
|
|
|
|
self.geometry_controller = nn.Parameter(torch.tensor(0.1)) |
|
|
|
|
|
def evolve_mobius_parameters(self, state_embedding): |
|
|
evolution_signal = self.mobius_evolution(state_embedding) |
|
|
evolution_rate = torch.clamp(self.geometry_controller, 0.01, 1.0) |
|
|
if self.mobius_transform.learnable: |
|
|
with torch.no_grad(): |
|
|
updates = (evolution_signal.view(4, 2) * evolution_rate * 0.01)\ |
|
|
.to(device=self.mobius_transform.a.device, dtype=self.mobius_transform.a.dtype) |
|
|
self.mobius_transform.a.add_(updates[0]) |
|
|
self.mobius_transform.b.add_(updates[1]) |
|
|
self.mobius_transform.c.add_(updates[2]) |
|
|
self.mobius_transform.d.add_(updates[3]) |
|
|
self.mobius_transform.normalize_parameters() |
|
|
|
|
|
|
|
|
def forward(self, initial_state, return_full_trajectory=False): |
|
|
state_embedding = self.state_encoder(initial_state) |
|
|
|
|
|
evolution_history = { |
|
|
'states': [], |
|
|
'geometries': [], |
|
|
'transition_matrices': [], |
|
|
'transformed_positions': [] |
|
|
} |
|
|
|
|
|
current_state = initial_state |
|
|
|
|
|
for step in range(self.evolution_steps): |
|
|
state_embedding = self.state_encoder(current_state) |
|
|
|
|
|
self.evolve_mobius_parameters(state_embedding.mean(dim=0)) |
|
|
|
|
|
markov_output = self.markov_chain.forward( |
|
|
current_state, |
|
|
num_steps=1, |
|
|
mobius_transform=self.mobius_transform |
|
|
) |
|
|
|
|
|
current_state = markov_output['final_state'] |
|
|
|
|
|
if return_full_trajectory: |
|
|
evolution_history['states'].append(current_state.clone()) |
|
|
evolution_history['geometries'].append(self.mobius_transform.get_transform_info()) |
|
|
evolution_history['transition_matrices'].append(markov_output['transition_matrix']) |
|
|
evolution_history['transformed_positions'].append(markov_output['transformed_positions']) |
|
|
|
|
|
final_embedding = self.state_encoder(current_state) |
|
|
final_prediction = self.state_decoder(final_embedding) |
|
|
|
|
|
output = { |
|
|
'final_state': current_state, |
|
|
'final_prediction': final_prediction, |
|
|
'final_embedding': final_embedding, |
|
|
'final_geometry': self.mobius_transform.get_transform_info() |
|
|
} |
|
|
|
|
|
if return_full_trajectory: |
|
|
output['evolution_history'] = evolution_history |
|
|
|
|
|
return output |
|
|
|
|
|
def predict_sequence(self, initial_state, sequence_length): |
|
|
predictions = [] |
|
|
current_state = initial_state |
|
|
|
|
|
for _ in range(sequence_length): |
|
|
output = self.forward(current_state) |
|
|
predictions.append(output['final_prediction']) |
|
|
current_state = output['final_state'] |
|
|
|
|
|
return torch.stack(predictions) |
|
|
|
|
|
def get_system_info(self): |
|
|
return { |
|
|
'num_states': self.num_states, |
|
|
'evolution_steps': self.evolution_steps, |
|
|
'current_geometry': self.mobius_transform.get_transform_info(), |
|
|
'state_positions': self.markov_chain.state_positions, |
|
|
'geometry_evolution_rate': self.geometry_controller.item() |
|
|
} |
|
|
|
|
|
|
