return loss_tot, grad_tot.ravel()
if __name__ == "__main__":
## Build simulated noisy map
# Noise map
if cplx:
n = np.random.normal(0.0, 1 / SNR, s.shape) + 1j*np.random.normal(0.0, 1 / SNR, s.shape)
else:
n = np.random.normal(0.0, 1 / SNR, s.shape)
# Noisy map
d = s + n
print("Building operator...")
start_time = time.time()
wph_op = pw.WPHOp(M, N, J, L=L, dn=dn, device=devices[0], cplx=cplx)
print(f"Done! (in {time.time() - start_time}s)")
print("Computing stats of target image...")
start_time = time.time()
coeffs = wph_op.apply(d, norm=norm).to("cpu")
wph_op.to("cpu") # Move back to CPU before the actual denoising algorithm
if torch.cuda.is_available():
torch.cuda.empty_cache() # Empty the memory cache to clear devices[0] memory
print(f"Done! (in {time.time() - start_time}s)")
## Minimization
eval_cnt = 0
total_start_time = time.time()
target_hist_imag = softhist(data_torch.imag.flatten())
target_hist_real = target_hist_real.to(device)
target_hist_imag = target_hist_imag.to(device)
target_hist_real = target_hist_real / torch.sum(target_hist_real)
target_hist_imag = target_hist_imag / torch.sum(target_hist_imag)
print(target_hist_real, torch.sum(target_hist_real))
print(target_hist_imag, torch.sum(target_hist_imag))
cplx = np.iscomplexobj(data)
print("Building operator...", flush=True)
start_time = time.time()
wph_op = pw.WPHOp(M, N, J, L=L, dn=dn, device=device)
print("Load model...", flush=True)
wph_op.load_model(p_list=p_list, dl=dl, dj=dj)
print(f"Done! (in {time.time() - start_time}s)", flush=True)
print("Computing stats of target image...", flush=True)
start_time = time.time()
if cplx:
coeffs = wph_op.apply(np.stack(
(data, data.real + 1j * 1e-9, data.imag + 1j * 1e-9), axis=0),
norm=norm,
padding=not pbc)
else:
coeffs = wph_op.apply(data, norm=norm, padding=not pbc)
print(f"Done! (in {time.time() - start_time}s)")
import pywph as pw
import torch
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
M, N = 512, 512
J = 8
L = 8
dn = 0
cplx = True
norm = "auto"
data = torch.from_numpy(np.load('data/I_1.npy'))
start = time.time()
wph_op = pw.WPHOp(M, N, J, L=L, dn=dn, cplx=cplx)
print(time.time() - start)
wph_op.to(0)
start = time.time()
data, nb_chunks = wph_op.preconfigure(data, requires_grad=True)
for i in range(nb_chunks):
print(i, f"/{nb_chunks}")
coeffs = wph_op.apply(data, i, norm=norm)
loss = (torch.abs(coeffs)**2).mean()
loss.backward(retain_graph=True)
del coeffs, loss
print(time.time() - start)
wph_op.to("cpu")
os.chdir('../pywph/')
import pywph as pw
from wph_quijote.wph_syntheses.wph_operator_wrapper import WPHOp_quijote
from wph_quijote_legacy_model import wph_quijote_legacy_model
M, N = 256, 256
J = 6
L = 8
dn = 2
norm = "auto"
data = np.load(
'data/Q_1.npy')[::2, ::2] + 1j * np.load('data/U_1.npy')[::2, ::2]
for norm in [None, "auto"]:
wph_op = pw.WPHOp(M, N, J, L=L, dn=dn, device=0, cplx=True)
wph_moments_indices, scaling_moments_indices = wph_quijote_legacy_model(
J, L, dn=dn)
wph_op.load_model([],
extra_wph_moments=wph_moments_indices,
extra_scaling_moments=scaling_moments_indices)
wph_pywph = wph_op(data, norm=norm, ret_wph_obj=True)
wph_op.to("cpu")
stat_params = {
"J": J,
"L": L,
"delta_j": J - 1,
"delta_l": 4,
"delta_n": dn,
"scaling_function_moments": [0, 1, 2, 3],