diff --git a/README.md b/README.md index e76831d..c6299fd 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,3 @@ # mgborg_emulator +bash build.sh --c-compiler icx --cxx-compiler icpx --python=/home/bartlett//anaconda3/envs/borg_env/bin/python3 --hades-python --with-mpi --install-system-python --build-dir /data101/bartlett/build_borg \ No newline at end of file diff --git a/conf/basic_ini.ini b/conf/basic_ini.ini index 93096ef..7833f72 100644 --- a/conf/basic_ini.ini +++ b/conf/basic_ini.ini @@ -37,6 +37,7 @@ mixing = 1 [model] gravity = lpt logfR0 = -5.0 +gauss_sigma = 0.05 af = 1.0 ai = 0.05 @@ -65,11 +66,14 @@ z0 = 0 [emulator] use_emulator = True -model_weights_path = /home/bartlett/mgborg_emulator/weights/emulator_weights.file_type +model_weights_path = /home/bartlett/mgborg_emulator/weights/d2d_borg.pt architecture = StyledVNet use_float64 = False use_pad_and_NN = True requires_grad = False +style_size = 1 +in_chan = 3 +out_chan = 3 [run] run_type = mock diff --git a/mgborg_emulator/forwards.py b/mgborg_emulator/forwards.py index ec12248..44fd2af 100644 --- a/mgborg_emulator/forwards.py +++ b/mgborg_emulator/forwards.py @@ -4,7 +4,8 @@ import numpy as np # Emulator imports import torch -# from .model_architecture.cosmology import * +torch.cuda.empty_cache() + from networks import StyledVNet, StyledVNet_distill # JAX set-up @@ -16,6 +17,7 @@ from jax.config import config as jax_config jax_config.update("jax_enable_x64", True) from utils import myprint +import utils class Emulator(borg.forward.BaseForwardModel): @@ -85,7 +87,7 @@ class Emulator(borg.forward.BaseForwardModel): # Step 2 - correct for particles that moved over the periodic boundary if self.debug: start_time = time.time() - disp_temp = correct_displacement_over_periodic_boundaries(disp,L=self.box.L[0],max_disp_1d=self.box.L[0]//2) + disp_temp = utils.correct_displacement_over_periodic_boundaries(disp,L=self.box.L[0],max_disp_1d=self.box.L[0]//2) if self.debug: myprint("Step 2 of forward pass took %s seconds" % (time.time() - start_time)) @@ -102,7 +104,7 @@ class Emulator(borg.forward.BaseForwardModel): # Step 4 - normalize if self.debug: start_time = time.time() - dis(dis_in) + utils.dis(dis_in) if self.debug: myprint("Step 4 of forward pass took %s seconds" % (time.time() - start_time)) @@ -162,7 +164,7 @@ class Emulator(borg.forward.BaseForwardModel): # Step 9 - undo the normalization if self.debug: start_time = time.time() - dis(dis_out,undo=True) + utils.dis(dis_out,undo=True) if self.debug: self.dis_out = dis_out myprint("Step 9 of forward pass took %s seconds" % (time.time() - start_time)) @@ -227,7 +229,7 @@ class Emulator(borg.forward.BaseForwardModel): # reverse step 9 if self.debug: start_time = time.time() - dis(ag,undo=False) + utils.dis(ag,undo=False) if self.debug: myprint("Reverse step 9 took %s seconds" % (time.time() - start_time)) @@ -274,7 +276,7 @@ class Emulator(borg.forward.BaseForwardModel): # reverse step 4 if self.debug: start_time = time.time() - dis(ag,undo=True) + utils.dis(ag,undo=True) if self.debug: myprint("Reverse step 4 took %s seconds" % (time.time() - start_time)) @@ -314,9 +316,9 @@ class Emulator(borg.forward.BaseForwardModel): y = y * Ny / Ly z = z * Nz / Lz - qx, ix = get_cell_coord(x) - qy, iy = get_cell_coord(y) - qz, iz = get_cell_coord(z) + qx, ix = utils.get_cell_coord(x) + qy, iy = utils.get_cell_coord(y) + qz, iz = utils.get_cell_coord(z) ix = ix.astype(int) iy = iy.astype(int) @@ -352,9 +354,9 @@ class Emulator(borg.forward.BaseForwardModel): y = y * Ny / Ly z = z * Nz / Lz - qx, ix = get_cell_coord(x) - qy, iy = get_cell_coord(y) - qz, iz = get_cell_coord(z) + qx, ix = utils.get_cell_coord(x) + qy, iy = utils.get_cell_coord(y) + qz, iz = utils.get_cell_coord(z) ix = ix.astype(int) iy = iy.astype(int) diff --git a/mgborg_emulator/likelihoods.py b/mgborg_emulator/likelihood.py similarity index 95% rename from mgborg_emulator/likelihoods.py rename to mgborg_emulator/likelihood.py index 6f4b09f..15c33e9 100644 --- a/mgborg_emulator/likelihoods.py +++ b/mgborg_emulator/likelihood.py @@ -7,9 +7,10 @@ import symbolic_pofk.linear import jax from functools import partial import ast +import torch import utils as utils -from utils import myprint +from utils import myprint, initial_pos import forwards import networks @@ -56,7 +57,7 @@ class GaussianLikelihood(borg.likelihood.BaseLikelihood): # Initialise model parameters model_params = { - 'logfR0':float(config['model']['logfr0']) + 'logfR0':float(config['model']['logfr0']), 'gauss_sigma':float(config['model']['gauss_sigma']) } self.fwd_param.setModelParams(model_params) @@ -134,7 +135,12 @@ class GaussianLikelihood(borg.likelihood.BaseLikelihood): if self.run_type == 'data': raise NotImplementedError elif self.run_type == 'mock': - raise NotImplementedError + output_density = np.zeros(self.fwd.getOutputBoxModel().N) + self.fwd.forwardModel_v2(s_hat) + self.fwd.getDensityFinal(output_density) + state["BORG_final_density"][:] = output_density + output_density += np.random.normal(size=self.fwd.getOutputBoxModel().N) * self.getModelParam('nullforward', 'gauss_sigma') + state["mock"][:] = output_density else: raise NotImplementedError @@ -179,9 +185,6 @@ class GaussianLikelihood(borg.likelihood.BaseLikelihood): output_density = np.zeros((N,N,N)) self.fwd.forwardModel_v2(s_hat) self.fwd.getDensityFinal(output_density) - - # Get velocity field - output_velocity = self.fwd_vel.getVelocityField() self.delta = output_density self.vel = output_velocity @@ -219,7 +222,7 @@ class GaussianLikelihood(borg.likelihood.BaseLikelihood): self.fwd.adjointModel_v2(mygradient) mygrad_hat = np.zeros(s_hat.shape, dtype=np.complex128) self.fwd.getAdjointModel(mygrad_hat) - elf.fwd.clearAdjointGradient() + self.fwd.clearAdjointGradient() return mygrad_hat @@ -263,6 +266,10 @@ def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.Bo config.read(ini_file) ai = float(config['model']['ai']) af = float(config['model']['af']) + + # Setup forward model + chain = borg.forward.ChainForwardModel(box) + chain.addModel(borg.forward.models.HermiticEnforcer(box)) # Cosmological parameters if ini_file is None: @@ -270,10 +277,6 @@ def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.Bo else: cpar = utils.get_cosmopar(ini_file) chain.setCosmoParams(cpar) - - # Setup forward model - chain = borg.forward.ChainForwardModel(box) - chain.addModel(borg.forward.models.HermiticEnforcer(box)) # CLASS transfer function chain @= borg.forward.model_lib.M_PRIMORDIAL_AS(box) @@ -322,10 +325,7 @@ def build_gravity_model(state: borg.likelihood.MarkovState, box: borg.forward.Bo # Initialize model model = getattr(networks, config['emulator']['architecture']) - # if use_distilled: - # model = networks.StyledVNet_distill(1,3,3,num_filt=num_filt) - # else: - # model = networks.StyledVNet(1,3,3) + model = model(int(config['emulator']['style_size']), int(config['emulator']['in_chan']), int(config['emulator']['out_chan'])) model.load_state_dict(emu_weights['model']) use_float64 = config['emulator']['use_float64'].lower().strip() == 'true' @@ -505,7 +505,7 @@ def build_likelihood(state: borg.likelihood.MarkovState, info: borg.likelihood.L myprint("Building likelihood") myprint(chain.getCosmoParams()) boxm = chain.getBoxModel() - likelihood = VelocityBORGLikelihood(chain, fwd_param, fwd_vel, borg.getIniConfigurationFilename()) + likelihood = VelocityBORGLikelihood(chain, fwd_param, borg.getIniConfigurationFilename()) return likelihood diff --git a/mgborg_emulator/networks.py b/mgborg_emulator/networks.py index 64833f5..60a704a 100644 --- a/mgborg_emulator/networks.py +++ b/mgborg_emulator/networks.py @@ -1,8 +1,8 @@ import torch import torch.nn as nn -from .styled_conv import ConvStyledBlock, ResStyledBlock -from .narrow import narrow_by +from map2map_tools.styled_conv import ConvStyledBlock, ResStyledBlock +from map2map_tools.narrow import narrow_by class StyledVNet(nn.Module): def __init__(self, style_size, in_chan, out_chan, bypass=None, **kwargs): diff --git a/mgborg_emulator/utils.py b/mgborg_emulator/utils.py index f59cd55..582fd5a 100644 --- a/mgborg_emulator/utils.py +++ b/mgborg_emulator/utils.py @@ -1,5 +1,9 @@ import aquila_borg as borg import configparser +import numpy as np +import jax.numpy as jnp +from scipy.special import hyp2f1 +import symbolic_pofk.linear cons = borg.console() myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(repr(x)) @@ -34,3 +38,83 @@ def get_cosmopar(ini_file): cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s) return cpar + +def initial_pos(L,N,order="F"): + values = np.linspace(0,L,N+1)[:-1] #ensure LLC + xx,yy,zz = np.meshgrid(values,values,values) + + if order=="F": + pos_mesh = np.vstack((yy.flatten(),xx.flatten(),zz.flatten())).T + if order=="C": + pos_mesh = np.vstack((zz.flatten(),xx.flatten(),yy.flatten())).T + + return pos_mesh + + +def correct_displacement_over_periodic_boundaries(disp,L,max_disp_1d=125): + # Need to correct for positions moving over the periodic boundary + + moved_over_bound = L - max_disp_1d + axis = ['x','y','z'] + + for i in [0,1,2]: + idx_sup, idx_sub = check(disp,L,moved_over_bound,max_disp_1d,i,axis) + + # Correct positions + disp[:,i][idx_sup] -= L + disp[:,i][idx_sub] += L + + check(disp,L,moved_over_bound,max_disp_1d,i,axis) + + assert np.amin(disp[:,i]) >= -max_disp_1d and np.amax(disp[:,i]) <= max_disp_1d, "Particles outside allowed region" + + return disp + + +def check(disp,L,moved_over_bound,max_disp_1d,i,axis): + idxsup = disp[:,i]>moved_over_bound + idx = np.abs(disp[:,i])<=max_disp_1d + idxsub = disp[:,i]<-moved_over_bound + + sup = len(disp[:,i][idxsup]) + did_not_cross_boundary = len(disp[:,i][idx]) + sub = len(disp[:,i][idxsub]) + + if not sub+did_not_cross_boundary+sup == len(disp[:,i]): + myprint(f'Disp in {axis[i]} direction under -{moved_over_bound} Mpc/h is = '+str(sub)) + myprint(f'|Disp| in {axis[i]} direction under {max_disp_1d} Mpc/h is = '+str(did_not_cross_boundary)) + myprint(f'Disp in {axis[i]} direction over {moved_over_bound} Mpc/h is = '+str(sup)) + myprint('These add up to: '+str(sub+did_not_cross_boundary+sup)) + myprint(f"Should add up to: len(disp[:,i]) {len(disp[:,i])}") + myprint('\n') + + assert sub+did_not_cross_boundary+sup == len(disp[:,i]), "Incorrect summation" # cannot lose/gain particles + + return idxsup, idxsub + + +def dis(x, undo=False, z=0.0, dis_std=6.0, **kwargs): + dis_norm = dis_std * linear_D(z) # [Mpc/h] + + if not undo: + dis_norm = 1 / dis_norm + + x *= dis_norm + + +def linear_D(z, Om=0.31): + """linear growth function for flat LambdaCDM, normalized to 1 at redshift zero + """ + OL = 1 - Om + a = 1 / (1+z) + return a * hyp2f1(1, 1/3, 11/6, - OL * a**3 / Om) \ + / hyp2f1(1, 1/3, 11/6, - OL / Om) + + +def get_cell_coord(x): + ix = jnp.floor(x) + return x - ix, ix + +def get_cell_coord_np(x): + ix = np.floor(x) + return x - ix, ix \ No newline at end of file diff --git a/tests/allocation_stats_0.txt b/tests/allocation_stats_0.txt new file mode 100644 index 0000000..9195239 --- /dev/null +++ b/tests/allocation_stats_0.txt @@ -0,0 +1,19 @@ +Memory still allocated at the end: 479.03 MB + +Statistics per context (name, allocated, freed, peak) +====================== + +*none* 64.0001 122.07 312.32 +BORG LPT MODEL 113.12 0 553.44 +BORGForwardModel::setup 0.000205994 0 216.82 +BorgLptModel::BorgLptModel 144 0 216.82 +[/home/bartlett/borg/libLSS/physics/chain_forward_model.cpp]virtual void LibLSS::ChainForwardModel::forwardModel_v2(ModelInput<3>) 160 16.25 617.35 +[/home/bartlett/borg/libLSS/physics/forward_model.cpp]void LibLSS::BORGForwardModel::setupDefault() 64 0 72.8201 +[/home/bartlett/borg/libLSS/physics/forwards/borg_lpt.cpp]std::shared_ptr build_borg_lpt(std::shared_ptr, const BoxModel &, const PropertyProxy &) [Grid = LibLSS::ClassicCloudInCell] 0 0 8.82007 +[/home/bartlett/borg/libLSS/physics/forwards/particle_balancer/balanceinfo.hpp]void LibLSS::BalanceInfo::allocate(MPI_Communication *, size_t) 16.16 0 569.6 +[/home/bartlett/borg/libLSS/physics/forwards/primordial_as.cpp]std::shared_ptr build_primordial_as(std::shared_ptr, const BoxModel &, const PropertyProxy &) 4.41 7.62939e-06 4.41006 +[/home/bartlett/borg/libLSS/physics/forwards/transfer_class.cpp]std::shared_ptr build_class(std::shared_ptr, const BoxModel &, const PropertyProxy &) 4.41 7.62939e-06 8.82008 +[/home/bartlett/borg/libLSS/samplers/core/gridLikelihoodBase.cpp]LibLSS::GridDensityLikelihoodBase<3>::GridDensityLikelihoodBase(MPI_Communication *, const GridSizes &, const GridLengths &) [Dims = 3] 64 32.5 280.82 +[/home/bartlett/borg/libLSS/tools/mpi/ghost_planes.hpp]void LibLSS::GhostPlanes, 2>::setup(MPI_Communication *, PlaneList &&, PlaneSet &&, DimList &&, size_t) [T = std::complex, Nd = 2, PlaneList = std::set &, PlaneSet = std::set &, DimList = std::array &] 7.62939e-06 0.000495911 5.34058e-05 +dispatch_plane_map 0.000988007 0.000499725 0.00104141 +lpt_ic 32 16.25 601.6 diff --git a/tests/fft_wisdom b/tests/fft_wisdom new file mode 100644 index 0000000..71f83aa --- /dev/null +++ b/tests/fft_wisdom @@ -0,0 +1,63 @@ +(fftw-3.3.10 fftw_wisdom #x029d284e #xc5dfc84e #xce4ac490 #x98c4e7bb + (fftw_dft_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x8c3bd361 #x74c82e49 #xd9d08cec #x490fb3c3) + (fftw_dft_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #xa79e8a60 #xe87bfd8a #x902920f9 #x2ec9719b) + (fftw_mpi_rdft2_serial_register 0 #x10bdd #x10bdd #x0 #x4ee85eaa #x68518300 #x99d7298f #x8e84c365) + (fftw_dft_r2hc_register 0 #x10bdd #x10bdd #x0 #xbe8bb67c #x0147d50a #xd6a7d883 #xaba02422) + (fftw_rdft2_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x851f9a4b #x19b2c27b #xed08cfa9 #x4bb3370b) + (fftw_rdft2_rank_geq2_register 0 #x10bdd #x10bdd #x0 #x658ee741 #x13268820 #x59089f57 #x487d6156) + (fftw_codelet_t1_16 0 #x10bdd #x10bdd #x0 #x8d3a0404 #x4ced2878 #xd7b0010b #x343f82db) + (fftw_rdft_rank0_register 3 #x10bdd #x10bdd #x0 #xb076fe66 #xe428602a #xff6bd3c3 #xf198ea92) + (fftw_dft_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #xb26542a9 #x6f11da3d #xcc2f3b11 #xf58ffaf3) + (fftw_dft_r2hc_register 0 #x10bdd #x10bdd #x0 #x19986bf9 #x1b1c9f9a #xbfa313e4 #x8089c52d) + (fftw_codelet_n1_8 0 #x10bdd #x10bdd #x0 #xd47121ba #x8509d296 #xda71f68f #x32819f4c) + (fftw_dft_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x676d367b #x9d081048 #x84d35015 #xadb7f14f) + (fftw_codelet_n1_8 0 #x10bdd #x10bdd #x0 #x96b8529e #xb560f141 #x58923d7d #x29e83836) + (fftw_rdft2_rank_geq2_register 0 #x10bdd #x10bdd #x0 #xafd851da #x99c0aebd #x0230e6a1 #x4dc1a835) + (fftw_dft_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x9d85b3d1 #x9fe51a73 #x90554ee0 #xae9d884a) + (fftw_rdft2_rank_geq2_register 0 #x10bdd #x10bdd #x0 #x432cb311 #x3851b42e #xe2a123c6 #xe40063b1) + (fftw_dft_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x3b11e766 #xcac44a9a #xebe02560 #xf4741891) + (fftw_rdft2_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x1e5b754b #xc1ee385c #x32bb8ea2 #x03c0f524) + (fftw_codelet_r2cb_128 2 #x10bdd #x10bdd #x0 #x3537951e #x3fde8ce2 #x944269f6 #xe658bdf2) + (fftw_rdft2_rank_geq2_register 0 #x10bdd #x10bdd #x0 #x87a70733 #x168e4628 #x5985aa94 #x15260cc8) + (fftw_codelet_n1_8 0 #x10bdd #x10bdd #x0 #x808c2c0b #xbc5616e2 #xc7cd6770 #x7d6b57c8) + (fftw_dft_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #xc5735095 #x158eb2be #x520352f6 #x7c5ad18a) + (fftw_dft_buffered_register 1 #x10bdd #x10bdd #x0 #x1bb138e5 #x126b4895 #x84639765 #xc6e16ed2) + (fftw_mpi_rdft2_serial_register 0 #x10bdd #x10bdd #x0 #xb92bad75 #xe367ef0e #x7081a83c #x6bca5b81) + (fftw_rdft2_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x85be544a #x6ec428e5 #x464a9e59 #x2d009947) + (fftw_rdft2_rank_geq2_register 0 #x10bdd #x10bdd #x0 #xd82cc4f0 #x264fda99 #x5da09515 #xaf7ad82a) + (fftw_codelet_r2cf_8 2 #x10bdd #x10bdd #x0 #x2b8efd4b #xf9271562 #x58810e69 #x2a843b48) + (fftw_mpi_rdft2_serial_register 0 #x10bdd #x10bdd #x0 #x64e50173 #xb27777d0 #xab348092 #xbcf5cef7) + (fftw_dft_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #xdfbe0f99 #x9b0e292d #x4a5c0b38 #x807ed5d9) + (fftw_codelet_n1_8 0 #x10bdd #x10bdd #x0 #x786bb47e #xf1901331 #x7f050641 #xf40ac495) + (fftw_dft_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #xe414074e #xb5f6c556 #x331a5df5 #x42b72bfe) + (fftw_dft_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x1c43a133 #x3a02cc9a #x06061ea2 #x8bd59400) + (fftw_codelet_n1_8 0 #x10bdd #x10bdd #x0 #x8780ffa9 #x5043b17e #xf6b31f7c #x09308bf3) + (fftw_rdft_rank0_register 3 #x10bdd #x10bdd #x0 #x0a7949d7 #x8e617643 #x1dc37829 #xb9c99559) + (fftw_codelet_n1_8 0 #x10bdd #x10bdd #x0 #x6cd039b7 #x416fd217 #x9a1f1005 #xae8b516f) + (fftw_codelet_r2cf_128 2 #x10bdd #x10bdd #x0 #x01ed1f43 #x9eb2a8be #x54a463f4 #x4065bd8c) + (fftw_dft_buffered_register 1 #x10bdd #x10bdd #x0 #x8cb4c207 #xbbe7f9d8 #x6142bd03 #x1fef6cbd) + (fftw_codelet_t1_16 0 #x10bdd #x10bdd #x0 #x07493f2d #xdbff69bf #xc99ae69c #x99aebe87) + (fftw_mpi_rdft2_serial_register 0 #x10bdd #x10bdd #x0 #x12333882 #x58ed4c26 #x2573fe73 #x58166349) + (fftw_dft_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x55dc834f #xe6778f48 #x5fb9338e #xb26105c2) + (fftw_codelet_n1_16 0 #x10bdd #x10bdd #x0 #x31243853 #xcd3ffac1 #x216a5369 #xbed877e1) + (fftw_codelet_n1_8 0 #x10bdd #x10bdd #x0 #x07538465 #x05c872d2 #x2e3a8015 #x2c1d6fb3) + (fftw_codelet_t1_8 0 #x10bdd #x10bdd #x0 #xe1c49a16 #x6e50ae21 #x6a60b5f7 #x922c9861) + (fftw_rdft2_rank_geq2_register 0 #x10bdd #x10bdd #x0 #xd747290f #x47e022d8 #x5f66ed62 #x0703deff) + (fftw_rdft2_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x6e5782da #xff8f0cb7 #xb7b3553d #xff227114) + (fftw_dft_nop_register 0 #x10bdd #x10bdd #x0 #x54333caa #x860f0b6a #xeeb1b4b0 #xbc1bf503) + (fftw_dft_nop_register 0 #x10bdd #x10bdd #x0 #x329f88d1 #x8b794130 #x1d6cd482 #x226d9d8c) + (fftw_dft_nop_register 0 #x10bdd #x10bdd #x0 #xcaa1aff4 #x76b5ab1f #xa1c4e63a #x795533ee) + (fftw_dft_r2hc_register 0 #x10bdd #x10bdd #x0 #xed14e8ff #x6a3fa4af #xcf6bc976 #x9ed6862b) + (fftw_rdft2_rank_geq2_register 0 #x10bdd #x10bdd #x0 #xf68eeb4b #x9707c500 #xea967ca2 #x94317695) + (fftw_dft_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x8e7540ba #xdf13a27a #xc5dfd29d #xcc18c450) + (fftw_dft_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #x4d5128a3 #x2c886ce4 #xae533803 #xd036a3e3) + (fftw_codelet_t1_16 0 #x10bdd #x10bdd #x0 #xec8d8b00 #x73468233 #xfd711f80 #x8216811f) + (fftw_rdft2_rank_geq2_register 0 #x10bdd #x10bdd #x0 #x69620706 #x1630ea0a #x7845ac91 #x107c1ea6) + (fftw_dft_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #xfa2f8f06 #x1e602e94 #x45491c98 #xb72edccf) + (fftw_dft_buffered_register 1 #x10bdd #x10bdd #x0 #x9d8339bb #x21e617a3 #x19bcf6ec #x843e58c6) + (fftw_dft_r2hc_register 0 #x10bdd #x10bdd #x0 #x2b1407ee #x9181dd05 #x70f5ba55 #x9d03e54b) + (fftw_dft_buffered_register 1 #x10bdd #x10bdd #x0 #xb5a35887 #x56e26ed1 #x50fc7926 #x21aa2465) + (fftw_rdft2_thr_vrank_geq1_register 0 #x10bdd #x10bdd #x0 #xde80c147 #x2f2b1766 #x646a051f #xbce0e7ea) + (fftw_dft_nop_register 0 #x10bdd #x10bdd #x0 #xa211000f #x16096f47 #x8977fda7 #x5859feab) + (fftw_codelet_r2cb_8 2 #x10bdd #x10bdd #x0 #xcc33ac01 #x3a093b70 #xce721838 #xe388f482) +) diff --git a/tests/test_gradient.py b/tests/test_gradient.py new file mode 100644 index 0000000..5d47e5b --- /dev/null +++ b/tests/test_gradient.py @@ -0,0 +1,198 @@ +import aquila_borg as borg +import configparser +import numpy as np +import matplotlib.pyplot as plt +import matplotlib +import itertools +import h5py as h5 +import os +import sys +import contextlib +from tqdm import tqdm + +sys.path.insert(0, '../mgborg_emulator/') +import likelihood +import forwards +import utils + +run_test = True +# run_test = False +epsilon = 1e-2 + + +# Create a context manager to suppress stdout +@contextlib.contextmanager +def suppress_stdout(): + with open(os.devnull, 'w') as devnull: + old_stdout = sys.stdout + sys.stdout = devnull + try: + yield + finally: + sys.stdout = old_stdout + +def compare_gradients( + ag_lh_auto_real: np.ndarray, + ag_lh_auto_imag: np.ndarray, + ag_lh_num_real: np.ndarray, + ag_lh_num_imag: np.ndarray, + plot_step: int, + filename: str="gradients.png", +) -> None: + """ + Comparison of an autodiff adjoint gradient of the likelihood against a + numerical one evaluated with finite differences. + + Args: + - ag_lh_auto_real (np.ndarray): Real part of the adjoint gradient (autodiff) + - ag_lh_auto_imag (np.ndarray): Imaginary part of the adjoint gradient (autodiff) + - ag_lh_num_real (np.ndarray): Real part of the adjoint gradient (numerical) + - ag_lh_num_imag (np.ndarray): Imaginary part of the adjoint gradient (numerical) + - plot_step (int): How frequently to sample the arrays + - filename (str): Name of the file to save the figure + + """ + # Plot colors + colors = { + "red": "#ba3d3b", + "blue": "#3d5792", + } + + fig, axs = plt.subplots(2, 2, figsize=(10, 7)) + + # Real part + axs[0,0].axhline(0.0, color="black", linestyle=":") + axs[0,0].plot(ag_lh_num_real[::plot_step], c=colors["blue"], label="Finite differences") + axs[0,0].plot(ag_lh_auto_real[::plot_step], "o", c=colors["red"], ms=3, label="Autodiff") + axs[0,0].yaxis.get_major_formatter().set_powerlimits((-2, 2)) + axs[0,0].set_ylabel("Real part") + axs[0,0].legend() + axs[0,1].plot(ag_lh_num_real[::plot_step], + ag_lh_auto_real[::plot_step] - ag_lh_num_real[::plot_step], + "o", + c=colors["red"], + ms=3 + ) + x = axs[0,1].get_xlim() + axs[0,1].axhline(y=0, color='k') + axs[0,1].set_xlabel("Numerical") + axs[0,1].set_ylabel("Autodiff - Numerical (real)") + + # Imaginary part + axs[1,0].axhline(0.0, color="black", linestyle=":") + axs[1,0].plot(ag_lh_num_imag[::plot_step][4:], c=colors["blue"], label="Finite differences") + axs[1,0].plot(ag_lh_auto_imag[::plot_step][4:], "o", c=colors["red"], ms=3, label="Autodiff") + axs[1,0].yaxis.get_major_formatter().set_powerlimits((-2, 2)) + axs[1,0].set_ylabel("Imaginary part") + axs[1,0].set_xlabel("Voxel ID") + axs[1,1].plot(ag_lh_num_imag[::plot_step], + ag_lh_auto_imag[::plot_step] - ag_lh_num_imag[::plot_step], + ".", + c=colors["red"] + ) + x = axs[1,1].get_xlim() + axs[1,1].axhline(y=0, color='k') + axs[1,1].set_xlabel("Numerical") + axs[1,1].set_ylabel("Autodiff - Numerical (imag)") + + fig.suptitle("Adjoint gradient of the likelihood w.r.t. initial conditions") + fig.tight_layout() + fig.subplots_adjust(hspace=0.) + + path = "../figs/" + fig.savefig(path + filename, bbox_inches="tight") + + +ini_file = '../conf/basic_ini.ini' + +# Input box +box_in = borg.forward.BoxModel() +config = configparser.ConfigParser() +config.read(ini_file) +box_in.L = (float(config['system']['L0']), float(config['system']['L1']), float(config['system']['L2'])) +box_in.N = (int(config['system']['N0']), int(config['system']['N1']), int(config['system']['N2'])) +box_in.xmin = (float(config['system']['corner0']), float(config['system']['corner1']), float(config['system']['corner2'])) + +# Setup BORG forward model and likelihood +model = likelihood.build_gravity_model(None, box_in, ini_file=ini_file) +cosmo = utils.get_cosmopar(ini_file) +model.setCosmoParams(cosmo) +fwd_param = forwards.NullForward(box_in) +mylike = likelihood.GaussianLikelihood(model, fwd_param, ini_file) + +# Create mock data +state = borg.likelihood.MarkovState() +mylike.initializeLikelihood(state) +mylike.updateCosmology(cosmo) +s_hat = np.fft.rfftn(np.random.randn(*box_in.N)) / box_in.Ntot ** (0.5) +mylike.generateMockData(s_hat, state) + +# Compute density field +output_density = np.zeros(box_in.N) +mylike.fwd.forwardModel_v2(s_hat) +print('SUM START', output_density.sum()) +mylike.fwd.getDensityFinal(output_density) +print('SUM NOW', output_density.sum()) +L = mylike.logLikelihoodComplex(s_hat, None) +print(L) +quit() + +# Autodiff +autodiff_gradient = mylike.gradientLikelihoodComplex(s_hat) +print(autodiff_gradient.min(), autodiff_gradient.max(), np.sum(np.isfinite(autodiff_gradient)), np.prod(autodiff_gradient.shape)) + +# Finite differences +if run_test: + s_hat_epsilon = s_hat.copy() + num_gradient = np.zeros(s_hat.shape, dtype=np.complex128) + for i, j, k in tqdm( + itertools.product(*map(range, [box_in.N[0], box_in.N[1], box_in.N[2] // 2 + 1])), + total=box_in.N[0] * box_in.N[1] * (box_in.N[2] // 2 + 1), + mininterval=1, + ): + + # +/- epsilon + s_hat_epsilon[i, j, k] = s_hat[i, j, k] + epsilon + with suppress_stdout(): + L = mylike.logLikelihoodComplex(s_hat_epsilon, None) + s_hat_epsilon[i, j, k] = s_hat[i, j, k] - epsilon + with suppress_stdout(): + L -= mylike.logLikelihoodComplex(s_hat_epsilon, None) + QQ = L / (2.0 * epsilon) + + # +/- i * epsilon + s_hat_epsilon[i, j, k] = s_hat[i, j, k] + 1j * epsilon + with suppress_stdout(): + L = mylike.logLikelihoodComplex(s_hat_epsilon, None) + s_hat_epsilon[i, j, k] = s_hat[i, j, k] - 1j * epsilon + with suppress_stdout(): + L -= mylike.logLikelihoodComplex(s_hat_epsilon, None) + QQ = QQ + L * 1j / (2.0 * epsilon) + + s_hat_epsilon[i, j, k] = s_hat[i, j, k] + num_gradient[i, j, k] = QQ + + + with h5.File(f"gradients_{box_in.N}.h5", mode="w") as ff: + ff["scalars/gradient_array_lh"] = autodiff_gradient + ff["scalars/gradient_array_lh_ref"] = num_gradient + ff["scalars/gradient_array_prior"] = np.zeros_like(autodiff_gradient) + ff["scalars/gradient_array_prior_ref"] = np.zeros_like(autodiff_gradient) + +slice_step = 2 +plot_step = 2 + +with h5.File(f'gradients_{box_in.N}.h5', 'r') as f: + ag_lh_auto_real = f["scalars"]["gradient_array_lh"][::slice_step, :, :].flatten().real + ag_lh_auto_imag = f["scalars"]["gradient_array_lh"][::slice_step, :, :].flatten().imag + ag_lh_num_real = f["scalars"]["gradient_array_lh_ref"][::slice_step, :, :].flatten().real + ag_lh_num_imag = f["scalars"]["gradient_array_lh_ref"][::slice_step, :, :].flatten().imag + + compare_gradients( + ag_lh_auto_real, + ag_lh_auto_imag, + ag_lh_num_real, + ag_lh_num_imag, + plot_step, + f'gradient_test_{box_in.N[0]}.png', + ) diff --git a/tests/timing_stats_0.txt b/tests/timing_stats_0.txt new file mode 100644 index 0000000..45c0f98 --- /dev/null +++ b/tests/timing_stats_0.txt @@ -0,0 +1,58 @@ +ARES version a9c88f3e8f91e04ec7c84eaea3b11ac5546aa6a5 modules + +Cumulative timing spent in different context +-------------------------------------------- +Context, Total time (seconds) + + [/home/bartlett/borg/libLSS/physics/forward_model.cpp]void LibLSS::ForwardModel::setCosmoParams(const CosmologicalParameters &) 39 2.96271 + [/home/bartlett/borg/libLSS/physics/chain_forward_model.cpp]virtual void LibLSS::ChainForwardModel::forwardModel_v2(ModelInput<3>) 1 2.54629 + [/home/bartlett/borg/libLSS/physics/forwards/borg_lpt.cpp]virtual void LibLSS::BorgLptModel<>::updateCosmo() [CIC = LibLSS::ClassicCloudInCell] 4 0.882557 + lightcone computation 1 0.878502 + [/home/bartlett/borg/libLSS/physics/cosmo.cpp]void LibLSS::Cosmology::precompute_com2a() 1 0.678918 + [/home/bartlett/borg/libLSS/physics/forwards/transfer_class.cpp]virtual void LibLSS::ForwardClass::updateCosmo() 4 0.593989 + [/home/bartlett/borg/libLSS/physics/class_cosmo.cpp]LibLSS::ClassCosmo::ClassCosmo(const CosmologicalParameters &, unsigned int, double, std::string, unsigned int, const std::map &) 1 0.482735 + [/home/bartlett/borg/libLSS/physics/forwards/primordial_as.cpp]std::shared_ptr build_primordial_as(std::shared_ptr, const BoxModel &, const PropertyProxy &) 1 0.222871 + [/home/bartlett/borg/libLSS/physics/forwards/transfer_class.cpp]std::shared_ptr build_class(std::shared_ptr, const BoxModel &, const PropertyProxy &) 1 0.202131 + BORG LPT MODEL 1 0.193024 + BORG forward model 1 0.191335 + lpt_ic 1 0.186872 + [/home/bartlett/borg/libLSS/physics/cosmo.cpp]void LibLSS::Cosmology::precompute_d_plus() 1 0.174369 + [/home/bartlett/borg/libLSS/physics/forwards/lpt/borg_fwd_lpt.cpp]virtual void LibLSS::BorgLptModel<>::getDensityFinal(ModelOutput<3>) [CIC = LibLSS::ClassicCloudInCell] 1 0.0468955 + Classic CIC projection 1 0.0429608 + FFTW_Manager::execute_c2r 3 0.0177956 + [/home/bartlett/borg/libLSS/physics/forwards/borg_lpt.cpp]std::shared_ptr build_borg_lpt(std::shared_ptr, const BoxModel &, const PropertyProxy &) [Grid = LibLSS::ClassicCloudInCell] 1 0.0122158 + BorgLptModel::BorgLptModel 1 0.0117337 + [/home/bartlett/borg/libLSS/tools/hermiticity_fixup.cpp]LibLSS::Hermiticity_fixer::Hermiticity_fixer(Mgr_p) [T = double, Nd = 3] 1 0.00922349 + [/home/bartlett/borg/libLSS/tools/mpi/ghost_planes.hpp]void LibLSS::GhostPlanes, 2>::setup(MPI_Communication *, PlaneList &&, PlaneSet &&, DimList &&, size_t) [T = std::complex, Nd = 2, PlaneList = std::set &, PlaneSet = std::set &, DimList = std::array &] 1 0.00915645 + FFTW_Manager::create_r2c_plan 3 0.00870736 + [/home/bartlett/borg/python/pyforward.cpp]void transfer_in(std::shared_ptr &, T &, U &, bool) [T = boost::multi_array_ref, 3>, U = pybind11::detail::unchecked_reference, 3>] 1 0.00718587 + dispatch_plane_map 1 0.00579799 + FFTW_Manager::create_c2r_plan 2 0.00561102 + [/home/bartlett/borg/libLSS/physics/forward_model.cpp]void LibLSS::BORGForwardModel::setupDefault() 1 0.00491136 + [/home/bartlett/borg/libLSS/physics/forwards/primordial_as.cpp]virtual void LibLSS::ForwardPrimordial_As::updateCosmo() 8 0.00461581 + [/home/bartlett/borg/libLSS/physics/forwards/primordial_as.cpp]void LibLSS::ForwardPrimordial_As::updatePower() 4 0.00431693 + [/home/bartlett/borg/libLSS/samplers/core/gridLikelihoodBase.cpp]LibLSS::GridDensityLikelihoodBase<3>::GridDensityLikelihoodBase(MPI_Communication *, const GridSizes &, const GridLengths &) [Dims = 3] 1 0.0029058 + [/home/bartlett/borg/libLSS/physics/hermitic.hpp]virtual void LibLSS::ForwardHermiticOperation::getDensityFinal(ModelOutput<3>) 1 0.00258788 + [/home/bartlett/borg/libLSS/physics/forwards/particle_balancer/balanceinfo.hpp]void LibLSS::BalanceInfo::allocate(MPI_Communication *, size_t) 1 0.00160495 + BORGForwardModel::setup 9 0.00148339 + Initializing peer system 15 0.00125319 + [/home/bartlett/borg/libLSS/physics/forwards/adapt_generic_bias.cpp]void (anonymous namespace)::bias_registrator() 1 0.0011465 + [/home/bartlett/borg/libLSS/physics/class_cosmo.cpp]void LibLSS::ClassCosmo::retrieve_Tk(const double) 2 0.00064302 + [/home/bartlett/borg/libLSS/physics/class_cosmo.cpp]void LibLSS::ClassCosmo::reinterpolate(const array_ref_1d &, const array_ref_1d &, LibLSS::auto_interpolator &) 6 0.000562291 + [/home/bartlett/borg/libLSS/tools/hermiticity_fixup.cpp]void LibLSS::Hermiticity_fixer::forward(CArrayRef &) [T = double, Nd = 3] 1 0.00051675 + [/home/bartlett/borg/libLSS/physics/model_io.cpp]virtual LibLSS::detail_output::ModelOutputBase<3>::~ModelOutputBase() [Nd = 3, Super = LibLSS::detail_model::ModelIO<3>] 14 0.000265913 + [/home/bartlett/borg/libLSS/tools/hermiticity_fixup.cpp]typename std::enable_if::type fix_plane(Mgr &, Ghosts &&, CArray &&, size_t *) [rank = 0UL, Mgr = LibLSS::FFTW_Manager, Ghosts = (lambda at /home/bartlett/borg/libLSS/tools/hermiticity_fixup.cpp:212:7), CArray = boost::detail::multi_array::multi_array_view, 2>, Dim = 2UL] 1 0.000250649 + gather_peer_by_plane 1 0.00023947 + [/home/bartlett/borg/libLSS/tools/hermiticity_fixup.cpp]typename std::enable_if::type fix_plane(Mgr &, Ghosts &&, CArray &&, size_t *) [rank = 0UL, Mgr = LibLSS::FFTW_Manager, Ghosts = (lambda at /home/bartlett/borg/libLSS/tools/hermiticity_fixup.cpp:218:7), CArray = boost::detail::multi_array::multi_array_view, 2>, Dim = 2UL] 1 0.000165529 + [/home/bartlett/borg/libLSS/physics/model_io.cpp]void LibLSS::detail_output::ModelOutputBase<3>::transfer(ModelOutputBase &&) [Nd = 3, Super = LibLSS::detail_model::ModelIO<3>] 9 7.3155e-05 + ghost synchronize 1 3.1071e-05 + [/home/bartlett/borg/libLSS/physics/model_io/base.hpp]void LibLSS::detail_model::ModelIO<3>::transfer(ModelIO &&) [Nd = 3] 30 2.4868e-05 + [/home/bartlett/borg/libLSS/physics/forwards/transfer_class.cpp]virtual void LibLSS::ForwardClass::forwardModel_v2(ModelInput<3>) 1 1.2496e-05 + [/home/bartlett/borg/libLSS/physics/model_io.cpp]virtual void LibLSS::detail_output::ModelOutputBase<3>::close() [Nd = 3, Super = LibLSS::detail_model::ModelIO<3>] 14 1.2394e-05 + [/home/bartlett/borg/libLSS/physics/forwards/primordial_as.cpp]virtual void LibLSS::ForwardPrimordial_As::forwardModel_v2(ModelInput<3>) 1 1.231e-05 + [/home/bartlett/borg/libLSS/physics/forwards/transfer_class.cpp]virtual void LibLSS::ForwardClass::setModelParams(const ModelDictionnary &) 1 9.429e-06 + [/home/bartlett/borg/libLSS/physics/model_io.cpp]void LibLSS::detail_input::ModelInputBase<3>::setRequestedIO(PreferredIO) [Nd = 3, Super = LibLSS::detail_model::ModelIO<3>] 5 3.766e-06 + [/home/bartlett/borg/libLSS/physics/model_io.cpp]void LibLSS::detail_output::ModelOutputBase<3>::setRequestedIO(PreferredIO) [Nd = 3, Super = LibLSS::detail_model::ModelIO<3>] 4 3.162e-06 + [/home/bartlett/borg/libLSS/physics/model_io.cpp]void LibLSS::detail_input::ModelInputBase<3>::needDestroyInput() [Nd = 3, Super = LibLSS::detail_model::ModelIO<3>] 1 2.173e-06 + particle distribution 1 1.5e-06 + [/home/bartlett/borg/libLSS/physics/forward_model.cpp]virtual void LibLSS::ForwardModel::setModelParams(const ModelDictionnary &) 1 1.258e-06 diff --git a/weights/d2d_borg.pt b/weights/d2d_borg.pt new file mode 100644 index 0000000..e69e5a7 Binary files /dev/null and b/weights/d2d_borg.pt differ