import json
from pathlib import Path
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import yaml
def cfg(r):return yaml.safe_load((Path(r)/"conf/config.yaml").read_text())
def truth(i):
 y,x=torch.meshgrid(torch.linspace(-1,1,128),torch.linspace(-1,1,128),indexing='ij');return torch.stack([torch.sin((v+1)*x+.1*i)*torch.cos((v+1)*y) for v in range(12)]).float()
def observe(x):return F.avg_pool2d(x.reshape(3,4,128,128),16).reshape(12,8,8)
class ScoreUNet(nn.Module):
 def __init__(self,channels=12,hidden=16):super().__init__();self.net=nn.Sequential(nn.Conv2d(channels+1,hidden,3,padding=1),nn.SiLU(),nn.Conv2d(hidden,hidden,3,padding=1),nn.SiLU(),nn.Conv2d(hidden,channels,3,padding=1));self.model_config={'channels':channels,'hidden':hidden}
 def forward(self,x,sigma):return self.net(torch.cat((x,torch.ones_like(x[:,:1])*sigma),1))
def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n')