<jupyter_start><jupyter_text>Structure Function Turbulent Data<jupyter_code>%load_ext autoreload
%autoreload 2
import numpy as np

import torch
from turboflow.models.rff import RFFNet
from turboflow.dataloaders import Turbo2D_simple, load_turbo2D_simple_numpy
from torch.utils.data import Dataset, DataLoader

import matplotlib.pyplot as plt

# set up cuda 
torch.cuda.is_available()
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
# ds=downsampling factor
Xlr, Ulr = load_turbo2D_simple_numpy(ds=4) # low resolution (64x64)
Xmr, Umr = load_turbo2D_simple_numpy(ds=2) # mid resolution (128x128)
Xhr, Uhr = load_turbo2D_simple_numpy(ds=1) # high (max) resolution (256x256)

L = int(Xlr.shape[0]**0.5)
M = int(Xmr.shape[0]**0.5)
H = int(Xhr.shape[0]**0.5)
from turboflow.models.rff import RFFNet

# model (RFF + MLP)
fft_scale = 10
fft_nfeat = 256
nin = 2  # = x and y coordinates
nout = 2 # = ux and uy components
mlp_layers = [nin*fft_nfeat] + [256]*3 + [nout]

rffnet = RFFNet('RFFNet', mlp_layers, fft_nfeat, fft_scale)

# training data
path_to_data = '../data/2021-Turb2D_velocities.npy'
trainset = Turbo2D_simple(path_to_data, device, ds=4, img=42)
trainloader = DataLoader(trainset, batch_size=6666, shuffle=False)

# TRAIN!
rffnet.to(device)
rffnet.fit(trainloader, epochs=1000) # we are not afraid to overfit the data, this is coord-based MLP!
from turboflow.models.phyrff_soft import PiRFFNet

# model (RFF + MLP)
fft_scale = 10
fft_nfeat = 256
nin = 2
nout = 2
mlp_layers = [nin*fft_nfeat] + [256]*3 + [nout]

pirffntet = PiRFFNet('PIRFFNet', mlp_layers, fft_nfeat, fft_scale)

# training data
path_to_data = '../data/2021-Turb2D_velocities.npy'
trainset = Turbo2D_simple(path_to_data, device, img=42)
trainloader = DataLoader(trainset, batch_size=6666, shuffle=False)

# TRAIN!
pirffntet.to(device)
pirffntet.fit(trainloader, epochs=2000)
# training data
path_to_data = '../data/2021-Turb2D_velocities.npy'
trainset = Turbo2D_simple(path_to_data, device, ds=4, img=42)
trainloader = DataLoader(trainset, batch_size=6666, shuffle=False)

X, y = trainset[:5]
delta_x = (X[2,1] - X[1,1]).item()
smallest_increment = delta_x/4<jupyter_output>Dataset Turbo2D, img # 42
(64, 64, 2)
(64, 64, 2)
Original size 64
Y shape (64, 64, 2)
Y min, max: -2.365612 2.941536
after normalization, Y min, max: -0.8042097733972999 1.0<jupyter_text>Structure Function<jupyter_code>print(Uhr)
nx, ny = Uhr.reshape(H,H,2)[:,:,0].shape

U2h = np.concatenate([Uhr.reshape(H,H,2)]*2, axis=1)
print(U2h.shape)
U2v = np.concatenate([Uhr.reshape(H,H,2)]*2, axis=0)
print(U2v.shape)
U2d = np.concatenate([U2h, U2h], axis=0)
print(U2d.shape)

increments = np.arange(H)
sfun = np.zeros((H,3))

Uhr_sq = Uhr.reshape(H,H,2)

for d in increments:
    
    # horizontal axis 
    diff_h = (U2h[:,d:H+d,:] - Uhr_sq)**2
    diff_h = diff_h[1]
    
    # vertical axis 
    diff_v = (U2v[d:H+d,:,:] - Uhr_sq)**2
    diff_v = diff_v[0]
    
    # diagonal axis
    diff_d = (U2d[d:H+d,d:H+d,:] - Uhr_sq)**2
    diff_d = (diff_d[0]+diff_d[1])/2
    
    sfun[d,0] = diff_h.mean()
    sfun[d,1] = diff_v.mean()
    sfun[d,2] = diff_d.mean()

c = 0.25
plt.loglog(sfun[:,0], label='horz')
plt.loglog(sfun[:,1], label='vert')
plt.loglog(sfun[:,2], label='diag')
plt.loglog(np.arange(nx//4), c*0.00275*np.arange(nx//4)**2, '--', label='model')
plt.legend()
plt.show()
del pirffntet
del rffnet
torch.cuda.empty_cache()
from turboflow.models.phyrff_sfun import SfunRFFNet

# model (RFF + MLP)
fft_scale = 10
fft_nfeat = 256
nin = 2  # = x and y coordinates
nout = 2 # = ux and uy components
mlp_layers = [nin*fft_nfeat] + [256]*4 + [nout]

lam_pde = 1e-4
lam_sfun = .5e-4

# model
n_centers = 20 # 100 too much
n_increments = 10
sfunrffnet = SfunRFFNet('SfunRFFNet', mlp_layers, fft_nfeat, fft_scale,
                        lam_pde, lam_sfun,
                        smallest_increment = smallest_increment,
                        n_centers=n_centers, n_increments=n_increments)

sfunrffnet.to(device)
sfunrffnet.fit(trainloader, epochs=3000) # we are not afraid to overfit the data, this is coord-based MLP!
x = torch.from_numpy(Xlr).float().contiguous().to(device)

sfunrffnet.eval().to(device)
pred = sfunrffnet(x)
Udfphy = pred.cpu().detach().numpy().squeeze().reshape(L,L,2)

plt.figure(figsize=(5,5))
plt.subplot(121)
plt.title('RFF: U upsampled: u_x')
plt.imshow(Udfphy[:,:,0])
plt.subplot(122)
plt.title('RFF: U upsampled: u_y')
plt.imshow(Udfphy[:,:,1])
plt.tight_layout()
plt.show()


plt.figure(figsize=(5,5))
plt.subplot(121)
plt.title('Error Ux')
plt.imshow(Udfphy[:,:,0] - Ulr.reshape(L,L,2)[:,:,0])
plt.subplot(122)
plt.title('Error Uy')
plt.imshow(Udfphy[:,:,1] - Ulr.reshape(L,L,2)[:,:,1])
plt.tight_layout()
plt.show()
x = torch.from_numpy(Xhr).float().contiguous().to(device)

sfunrffnet.eval().to(device)
pred = sfunrffnet(x)
Usfun = pred.cpu().detach().numpy().squeeze().reshape(H,H,2)

plt.figure(figsize=(10,10))
plt.subplot(121)
plt.title('RFF: U upsampled: u_x')
plt.imshow(Usfun[:,:,0])
plt.subplot(122)
plt.title('RFF: U upsampled: u_y')
plt.imshow(Usfun[:,:,1])
plt.tight_layout()
plt.show()


plt.figure(figsize=(10,10))
plt.subplot(121)
plt.title('Error Ux')
plt.imshow(Usfun[:,:,0] - Uhr.reshape(H,H,2)[:,:,0])
plt.subplot(122)
plt.title('Error Uy')
plt.imshow(Usfun[:,:,1] - Uhr.reshape(H,H,2)[:,:,1])
plt.tight_layout()
plt.show()
x = torch.from_numpy(Xhr).float().contiguous().to(device)

Upde = pirffntet(x).cpu().detach().numpy().squeeze()
Usfun = sfunrffnet(x).cpu().detach().numpy().squeeze()
Urff = rffnet(x).cpu().detach().numpy().squeeze()
from turboflow.utils import phy_utils as phy

fig, axs = plt.subplots(1,1,figsize=(8,6))
axs.loglog(np.arange(len(phy.powerspec(Ulr.reshape(L,L,2)[:,:,1]))), 
                         phy.powerspec(Ulr.reshape(L,L,2)[:,:,1]), label='Ulr 64x64')
axs.loglog(np.arange(len(phy.powerspec(Uhr.reshape(H,H,2)[:,:,0]))), 
                         phy.powerspec(Uhr.reshape(H,H,2)[:,:,0]), label='Uhr 256x256')
# axs.loglog(np.arange(len(phy.powerspec(Urff.reshape(H,H,2)[:,:,0]))), 
#                          phy.powerspec(Urff.reshape(H,H,2)[:,:,0]), label='Urff 256x256')
# axs.loglog(np.arange(len(phy.powerspec(Upde.reshape(H,H,2)[:,:,0]))), 
#                          phy.powerspec(Upde.reshape(H,H,2)[:,:,0]), label='Upde 256x256')
axs.loglog(np.arange(len(phy.powerspec(Usfun.reshape(H,H,2)[:,:,0]))), 
                         phy.powerspec(Usfun.reshape(H,H,2)[:,:,0]), label='Usfun 256x256')
axs.set_ylim(10**(-13), 10)
axs.set_xlim(1, 512)
axs.legend()
plt.ylabel('Energy specturm')
plt.xlabel('Frequency')
plt.tight_layout()
plt.show()
# Let s zoom in it!
ngrid = 256
coords = np.linspace(0.1, 0.4, ngrid)
coords = np.stack(np.meshgrid(coords, coords), -1) # X x Y x 2
coords = coords.reshape(ngrid*ngrid,2)
print(coords.shape)

x = torch.from_numpy(coords).float().contiguous().to(device)

sfunrffnet.eval().to(device)
pred = sfunrffnet(x)
pred = pred.cpu().detach().numpy().squeeze().reshape(ngrid,ngrid,2)

plt.figure(figsize=(5,5))
plt.imshow(pred[:,:,0])
plt.show()

s = 6
plt.figure(figsize=((5,5))
plt.quiver(coords.reshape(ngrid,ngrid,2)[::s,::s,0],
           coords.reshape(ngrid,ngrid,2)[::s,::s,1],
           pred[::s,::s,0],pred[::s,::s,1])
nx, ny = Uhr.reshape(H,H,2)[:,:,0].shape

U2h = np.concatenate([Uhr.reshape(H,H,2)]*2, axis=1)
print(U2h.shape)
U2v = np.concatenate([Uhr.reshape(H,H,2)]*2, axis=0)
print(U2v.shape)
U2d = np.concatenate([U2h, U2h], axis=0)
print(U2d.shape)

increments = np.arange(H)
sfun = np.zeros((H,3))

Uhr_sq = Uhr.reshape(H,H,2)

for d in increments:
    
    # horizontal axis 
    diff_h = (U2h[:,d:H+d,:] - Uhr_sq)**2
    diff_h = diff_h[1]
    
    # vertical axis 
    diff_v = (U2v[d:H+d,:,:] - Uhr_sq)**2
    diff_v = diff_v[0]
    
    # diagonal axis
    diff_d = (U2d[d:H+d,d:H+d,:] - Uhr_sq)**2
    diff_d = (diff_d[0]+diff_d[1])/2
    
    sfun[d,0] = diff_h.mean()
    sfun[d,1] = diff_v.mean()
    sfun[d,2] = diff_d.mean()

c = 0.25
plt.loglog(sfun[:,0], label='horz')
plt.loglog(sfun[:,1], label='vert')
plt.loglog(sfun[:,2], label='diag')
plt.loglog(np.arange(nx//4), c*0.00275*np.arange(nx//4)**2, '--', label='model')
plt.legend()
plt.show()<jupyter_output>(256, 512, 2)
(512, 256, 2)
(512, 512, 2)
{'fraction_non_alphanumeric': 0.19646317829457363, 'fraction_numerical': 0.05377906976744186, 'mean_word_length': 4.754006968641115, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '359299', 'license': 'Apache-2.0', 'n_tokens_mistral': 4071, 'n_tokens_neox': 3783, 'n_words': 588, 'path': 'notebooks/[Sandbox] RFFNet and StructureFun - Vectorilazed.ipynb', 'repo_name': 'Chutlhu/TurboSuperResultion', 'size': '8256'}