Spaces:
Sleeping
Sleeping
File size: 3,137 Bytes
d46bde8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | # model.py
# Updated: Modern torch.amp API (device_type='cuda') for all autocast blocks
# - Replaced torch.cuda.amp.autocast → torch.amp.autocast('cuda', ...)
import torch
import torch.nn as nn
from geoopt import ManifoldParameter
from geoopt.manifolds import Stiefel
from torch.amp import autocast # Add this import if not present
from .config import device, IDEAL_HARMONICS
from .utils import manifold
class StiefelDampedCoupledInharmGR(nn.Module):
def __init__(self, dim: int, k_modes: int, initial_basis: torch.Tensor):
super().__init__()
self.dim = dim
self.k_modes = k_modes
self.base = ManifoldParameter(initial_basis, manifold=manifold)
self.vel_dir_raw = nn.Parameter(torch.randn(dim, k_modes, device=device))
self.log_speed = nn.Parameter(torch.log(torch.ones(k_modes, device=device) * 2.0))
self.log_base_rate = nn.Parameter(torch.log(torch.tensor(0.06, device=device)))
self.log_slope = nn.Parameter(torch.log(torch.tensor(0.064, device=device)))
self.raw_lin_b = nn.Parameter(torch.tensor(-16.0, device=device))
self.raw_quad_b = nn.Parameter(torch.tensor(-14.0, device=device))
self.raw_coupling_strength = nn.Parameter(torch.log(torch.tensor(0.30, device=device)))
self.coupling_raw = nn.Parameter(torch.randn(k_modes, k_modes, device=device) * 0.03)
def forward(self, t):
base = self.base
dim, k_modes = self.dim, self.k_modes
# Force full precision for manifold projection/norm
with torch.amp.autocast('cuda', enabled=False):
vel_dir = manifold.proju(base, self.vel_dir_raw)
vel_dir = vel_dir / (vel_dir.norm(dim=0, keepdim=True) + 1e-8)
speed_scalars = torch.exp(self.log_speed)
damping_rates = torch.exp(self.log_base_rate) + torch.exp(self.log_slope) * (IDEAL_HARMONICS - 1)
inharm_b = torch.exp(self.raw_lin_b) * IDEAL_HARMONICS + torch.exp(self.raw_quad_b) * IDEAL_HARMONICS.pow(2)
freq = IDEAL_HARMONICS * torch.sqrt(1 + inharm_b * IDEAL_HARMONICS.pow(2))
full_freq = speed_scalars * freq
vel = vel_dir * speed_scalars.unsqueeze(0) * freq.unsqueeze(0)
coupling_strength = torch.exp(self.raw_coupling_strength)
coupling_skew = self.coupling_raw.tril(diagonal=-1) - self.coupling_raw.triu(diagonal=1)
# Force full precision for coupling projection
with torch.amp.autocast('cuda', enabled=False):
coupling_vel = manifold.proju(base, base @ coupling_skew)
vel_total = vel + coupling_strength * coupling_vel
abs_t = torch.abs(t).view(-1, 1)
envelope = torch.exp(-damping_rates * abs_t)
vel_batch = t.view(-1, 1, 1) * vel_total.unsqueeze(0) * envelope.unsqueeze(1)
base_batch = base.unsqueeze(0).expand(t.shape[0], dim, k_modes)
# Force full precision for expmap (most sensitive manifold op)
with torch.amp.autocast('cuda', enabled=False):
preds = manifold.expmap(base_batch, vel_batch)
return preds, damping_rates, coupling_strength, inharm_b, speed_scalars, full_freq |