Initial import

extra/python/example/2mpp-chain.ini.txt

@@ -0,0 +1,89 @@
+[system]
+console_output=logares.txt
+mask_precision=0.9
+VERBOSE_LEVEL = 2
+N0 = 16 
+N1 = 16
+N2 = 16
+
+#Ndata0=64
+#Ndata1=64
+#Ndata2=64
+
+L0 = 200
+L1 = 200
+L2 = 200
+
+corner0 = 0
+corner1 = 0
+corner2 = 0
+
+NUM_MODES=100
+test_mode=true
+
+# If true, the initial power spectrum of the chain is set to the cosmological one
+seed_cpower=true
+
+# Indicate which samplers should be blocked for testing purposes
+
+[block_loop]
+# Indicate which samplers should be blocked for testing purposes
+bias_sampler_blocked=true
+nmean_sampler_blocked=true
+hades_sampler_blocked=false
+
+ares_heat=1.0
+
+[gravity]
+model=CHAIN
+models=PRIMORDIAL,TRANSFER_EHU
+
+[gravity_chain_0]
+model=PRIMORDIAL
+a_final=1.0
+
+[mcmc]
+number_to_generate=1
+random_ic=false
+init_random_scaling=1.0
+
+[likelihood]
+
+[hades]
+algorithm=HMC
+max_epsilon=0.01
+max_timesteps=50
+mixing=1
+
+[run]
+NCAT = 1
+SIMULATION=true
+
+[cosmology]
+omega_r = 0
+fnl = 0
+omega_k = 0
+omega_m = 0.3175
+omega_b = 0.049
+omega_q = 0.6825
+w = -1
+wprime = 0
+n_s = 0.9624
+sigma8 = 0.8344
+h100 = 0.6711
+beta = 1.5 
+z0 = 0
+
+# 11.5 mag cut
+
+[catalog_0]
+datafile = 2MPP.txt
+bias=8,1,0.1,1
+halo_selection=none
+refbias = false
+nmean=10
+
+
+[python]
+likelihood_path=test_likelihood.py
+bias_sampler_type=slice

extra/python/example/model/example_jax.py

@@ -0,0 +1,73 @@
+#
+# This python script gives an example on using JAX and PyBORG together
+#
+import jax
+import numpy as np
+import borg
+
+cons = borg.console()
+
+myprint=lambda x: cons.print_std(x) if type(x) == str else cons.print_std(repr(x))
+myprint("Hello!")
+
+@jax.jit
+def jax_func(rho):
+   return rho**2
+
+class MyModel(borg.forward.BaseForwardModel):
+   def __init__(self, box):
+     myprint("Start forward model")
+     super().__init__(box, box)
+
+   def getPreferredInput(self):
+     return borg.forward.PREFERRED_REAL
+
+   def getPreferredOutput(self):
+     return borg.forward.PREFERRED_REAL
+
+   def forwardModel_v2_impl(self, input_array):
+     # Save the input data in a jax array (i.e. upload to accelerator)
+     self.save = jax.numpy.array(input_array)
+ 
+   def getDensityFinal_impl(self, output_array):
+     # Run forward, and save the AG function
+     fwd, self.ag_fun = jax.vjp(jax_func, self.save)
+     output_array[:] = fwd
+
+   def adjointModel_v2_impl(self, input_ag):
+     # Save the ag vector
+     self.ag = input_ag
+   
+   def getAdjointModel_impl(self, output_ag):
+     # Evaluate the new function with ag_fun
+     out_ag, = self.ag_fun(self.ag)
+     output_ag[:] = out_ag
+
+def build_gravity_model(box):
+  chain = borg.forward.ChainForwardModel(box)
+  chain.addModel(borg.forward.models.Primordial(box, 1.0))
+  chain.addModel(borg.forward.models.EisensteinHu(box))
+  chain.addModel(borg.forward.models.BorgLpt(box, box, ai=1.0))
+  chain.addModel(MyModel(box))
+  return chain
+
+
+if __name__ == "__main__":
+  box = borg.forward.BoxModel()
+  cpar = borg.cosmo.CosmologicalParameters()
+
+  chain = build_gravity_model(box)
+  chain.setCosmoParams(cpar)
+
+  s_hat = np.fft.rfftn(np.random.randn(*box.N)/np.sqrt(box.N[0]**3))
+
+  myprint(np.var(s_hat))
+
+  chain.forwardModel_v2(s_hat)
+  rho = np.zeros(chain.getOutputBoxModel().N)
+  chain.getDensityFinal(rho)
+
+
+  ag = 2*rho
+  chain.adjointModel_v2(ag)
+  chain.getAdjointModel(ag)

extra/python/example/run_altair.py

@@ -0,0 +1,73 @@
+#
+# Example to run a Altair forward model with PyBORG
+#
+import borg
+import numpy as np
+
+
+# Setup resolution of the initial mesh
+Ng = 64
+# Box size in Mpc/h
+L = 8000.0
+
+# setup the box
+bb = borg.forward.BoxModel()
+bb.L = L, L, L
+bb.N = Ng, Ng, Ng
+bb.xmin = -L/2,-L/2,-L/2
+
+print(bb)
+
+# Initialize some default cosmology
+cosmo = borg.cosmo.CosmologicalParameters()
+
+# Fiducial scale factor to express initial conditions
+a0 = 0.1
+
+chain = borg.forward.ChainForwardModel(bb)
+# Add primordial fluctuations
+chain.addModel(borg.forward.models.Primordial(bb, a0))
+# Add E&Hu transfer function
+chain.addModel(borg.forward.models.EisensteinHu(bb))
+# Run an LPT model from a=0.0 to af. The ai=a0 is the scale factor at which the IC are expressed
+lpt = borg.forward.models.BorgLpt(bb, bb, ai=a0, af=1.0)
+chain.addModel(lpt)
+
+Lz=10000
+altair = borg.forward.models.newModel(
+    "ALTAIR_AP",
+    bb,
+    dict(
+        corner0_z=-Lz/2,
+        corner1_z=-Lz/2,
+        corner2_z=-Lz/2,
+        L0_z=Lz,
+        L1_z=Lz,
+        L2_z=Lz,
+        N0_z=64,
+        N1_z=64,
+        N2_z=64,
+        is_contrast=True
+    ),
+)
+chain.addModel(altair)
+
+# Set the cosmology
+chain.setCosmoParams(cosmo)
+
+
+# Generate white noise: it has to be scaled by 1/N**(3./2) to be one in Fourier
+ic = np.random.randn(Ng, Ng, Ng) / np.sqrt(Ng ** 3)
+delta_m = np.zeros(lpt.getOutputBoxModel().N)
+delta_m_ap = np.zeros(chain.getOutputBoxModel().N)
+
+# RUN!
+chain.forwardModel_v2(ic)
+lpt.getDensityFinal(delta_m)
+chain.getDensityFinal(delta_m_ap)
+
+# Obtain the particles
+pp = np.zeros((lpt.getNumberOfParticles(), 3))
+lpt.getParticlePositions(pp)
+
+print(pp)

extra/python/example/run_forward.py

@@ -0,0 +1,49 @@
+#
+# Example to run a forward model with PyBORG
+#
+import borg
+import numpy as np
+
+
+# Setup resolution of the initial mesh
+Ng=64
+# Box size in Mpc/h
+L=100.
+
+# setup the box
+bb = borg.forward.BoxModel()
+bb.L = L,L,L
+bb.N = Ng,Ng,Ng
+
+# Initialize some default cosmology
+cosmo = borg.cosmo.CosmologicalParameters()
+
+# Fiducial scale factor to express initial conditions
+a0=0.1
+
+chain = borg.forward.ChainForwardModel(bb)
+# Add primordial fluctuations
+chain.addModel(borg.forward.models.Primordial(bb, a0))
+# Add E&Hu transfer function
+chain.addModel(borg.forward.models.EisensteinHu(bb))
+# Run an LPT model from a=0.0 to af. The ai=a0 is the scale factor at which the IC are expressed
+lpt = borg.forward.models.BorgLpt(bb, bb, ai=a0, af=1.0)
+chain.addModel(lpt)
+
+# Set the cosmology
+chain.setCosmoParams(cosmo)
+
+
+# Generate white noise: it has to be scaled by 1/N**(3./2) to be one in Fourier
+ic = np.random.randn(Ng, Ng, Ng)/np.sqrt(Ng**3)
+delta_m = np.zeros((Ng,Ng,Ng))
+
+# RUN!
+chain.forwardModel_v2(ic)
+chain.getDensityFinal(delta_m)
+
+# Obtain the particles
+pp = np.zeros((lpt.getNumberOfParticles(),3))
+lpt.getParticlePositions(pp)
+
+print(pp)

extra/python/example/test_gradient_cic.py

@@ -0,0 +1,87 @@
+import itertools
+from tqdm import tqdm
+import borg
+import numpy as np
+
+
+def build_gravity_model(state, box, cpar):
+    print(cpar)
+    lpt = borg.forward.models.BorgLpt(box, box, ai=1.0)
+    chain = borg.buildDefaultChain(box, cpar, 1.0, lpt)
+    return chain, lpt
+
+
+def build_likelihood(state, info):
+    return borg.likelihood.GaussianLinear(info)
+
+
+if not borg.EMBEDDED and __name__ == "__main__":
+    from tqdm import tqdm
+
+    Ng =  8
+
+    cpar = borg.cosmo.CosmologicalParameters()
+    state = borg.likelihood.MarkovState()
+    info = borg.likelihood.LikelihoodInfo()
+    info["GRID_LENGTH"] = np.array([0, 100., 0, 100., 0, 100.])
+    info["GRID"] = np.array([Ng, Ng, Ng], dtype=np.uint32)
+    box = borg.forward.BoxModel(100., Ng)
+
+    state.newArray3d("galaxy_data_0", Ng, Ng, Ng)
+
+    fwd, lpt = build_gravity_model(state, box, cpar)
+    state.newForwardModel("BORG_model", fwd)
+    state.newScalar("corner0", 0.0)
+    state.newScalar("corner1", 0.0)
+    state.newScalar("corner2", 0.0)
+    state.newScalar("localNdata0", 0, False, "L")
+    state.newScalar("localNdata1", Ng, False, "L")
+    state.newScalar("localNdata2", 0, False, "L")
+    state.newScalar("localNdata3", Ng, False, "L")
+    state.newScalar("localNdata4", 0, False, "L")
+    state.newScalar("localNdata5", Ng, False, "L")
+    state.newScalar("ares_heat", 1.0)
+    state.newScalar("NCAT", 1, False, "L")
+    state.newScalar("galaxy_bias_ref_0", False)
+
+    state.newScalar("cosmology", cpar)
+    state.newScalar("galaxy_nmean_0", 0.0)
+    state.newArray1d("galaxy_bias_0", 2)
+    state.newArray3d("galaxy_synthetic_sel_window_0", Ng, Ng, Ng)
+    state["galaxy_bias_0"][:] = np.array([1.0, 1.0])
+    state["galaxy_synthetic_sel_window_0"][:] = 1
+
+    like = build_likelihood(state, info)
+
+    like.initializeLikelihood(state)
+    like.updateMetaParameters(state)
+
+    s_hat = np.fft.rfftn(np.random.randn(Ng, Ng, Ng) / Ng**(1.5))
+    like.generateMockData(s_hat, state)
+
+    like.logLikelihood(s_hat)
+    analytic_gradient = like.gradientLikelihood(s_hat)
+
+    num_gradient = np.zeros((Ng, Ng, Ng // 2 + 1), dtype=np.complex128)
+    s_hat_epsilon = s_hat.copy()
+
+    epsilon = 0.001
+    for i, j, k in tqdm(itertools.product(*map(range, [Ng, Ng, Ng // 2 + 1])),
+                        total=Ng * Ng * (Ng // 2 + 1)):
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k] + epsilon
+        L = like.logLikelihood(s_hat_epsilon)
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k] - epsilon
+        L -= like.logLikelihood(s_hat_epsilon)
+        QQ = L / (2.0 * epsilon)
+
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k] + 1j * epsilon
+        L = like.logLikelihood(s_hat_epsilon)
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k] - 1j * epsilon
+        L -= like.logLikelihood(s_hat_epsilon)
+        QQ = QQ + L * 1j / (2.0 * epsilon)
+
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k]
+
+        num_gradient[i, j, k] = QQ
+
+    np.savez("gradients.npz", num=num_gradient, ana=analytic_gradient)

extra/python/example/test_likelihood.py

@@ -0,0 +1,89 @@
+import numpy as np
+import borg
+
+cons = borg.console()
+
+myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(
+    repr(x))
+
+sigma_noise = 0.1
+
+
+class MyLikelihood(borg.likelihood.BaseLikelihood):
+    def __init__(self, fwd, N, L):
+        myprint(f" Init {N}, {L}")
+        super().__init__(fwd, N, L)
+        self.comm = fwd.getCommunicator()
+
+    def initializeLikelihood(self, state):
+        myprint("Init likelihood")
+        self.data = state['galaxy_data_0']
+        state.newArray3d("my_density_field", self.data.shape[0],
+                         self.data.shape[1], self.data.shape[2], True)
+
+    def updateMetaParameters(self, state):
+        cpar = state['cosmology']
+        myprint(f"Cosmology is {cpar}")
+        self.getForwardModel().setCosmoParams(cpar)
+
+    def generateMockData(self, s_hat, state):
+
+        fwd = self.getForwardModel()
+        output = np.zeros(fwd.getOutputBoxModel().N)
+        fwd.forwardModel_v2(s_hat)
+        fwd.getDensityFinal(output)
+
+        state['galaxy_data_0'][:] = output + np.random.normal(
+            size=output.shape) * sigma_noise
+        state['my_density_field'][:] = output
+        like = ((state['galaxy_data_0'][:] - output)**2).sum() / sigma_noise**2
+        myprint(
+            f"Initial log_likelihood: {like}, var(s_hat) = {np.var(s_hat)}")
+
+    def logLikelihoodComplex(self, s_hat, gradientIsNext):
+        fwd = self.getForwardModel()
+
+        output = np.zeros(fwd.getBoxModel().N)
+        fwd.forwardModel_v2(s_hat)
+        fwd.getDensityFinal(output)
+        L = 0.5 * ((output - self.data)**2).sum() / sigma_noise**2
+        myprint(f"var(s_hat): {np.var(s_hat)}, Call to logLike: {L}")
+        return self.comm.allreduce(L)
+
+    def gradientLikelihoodComplex(self, x_hat):
+        fwd = self.getForwardModel()
+        output = np.zeros(fwd.getOutputBoxModel().N)
+        fwd.forwardModel_v2(x_hat)
+        fwd.getDensityFinal(output)
+        mygradient = (output - self.data) / sigma_noise**2
+        fwd.adjointModel_v2(mygradient)
+        mygrad_hat = np.zeros(s_hat.shape, dtype=np.complex128)
+        fwd.getAdjointModel(mygrad_hat)
+        return mygrad_hat
+
+
+model = None
+
+
+@borg.registerGravityBuilder
+def build_gravity_model(state, box):
+    global model
+    chain = borg.forward.ChainForwardModel(box)
+    chain.addModel(borg.forward.models.HermiticEnforcer(box))
+    chain.addModel(borg.forward.models.Primordial(box, 1.0))
+    chain.addModel(borg.forward.models.EisensteinHu(box))
+    model = chain
+    return chain
+
+
+@borg.registerLikelihoodBuilder
+def build_likelihood(state, info):
+    boxm = model.getBoxModel()
+    like = MyLikelihood(model, boxm.N, boxm.L)
+    #like.initializeLikelihood(state)
+    return like
+
+
+@borg.registerSamplerBuilder
+def build_samplers(state, info):
+    return []

extra/python/example/velocity/test_gradient_cic.py

@@ -0,0 +1,84 @@
+import itertools
+from tqdm import tqdm
+import borg
+import numpy as np
+
+cons = borg.console()
+
+myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(
+    repr(x))
+
+
+def build_gravity_model(box, cpar):
+    lpt = borg.forward.models.BorgLpt(box, box, ai=1.0)
+    chain = borg.buildDefaultChain(box, cpar, 1.0, lpt)
+    vfield = borg.forward.velocity.CICModel(box, lpt)
+    return chain, lpt, vfield
+
+
+if not borg.EMBEDDED and __name__ == "__main__":
+    from tqdm import tqdm
+
+    Ng = 8
+
+    # Create the auxiliary objects
+    cpar = borg.cosmo.CosmologicalParameters()
+    box = borg.forward.BoxModel(100., Ng)
+
+    fwd, lpt, vfield = build_gravity_model(box, cpar)
+
+    s_hat = np.fft.rfftn(np.random.randn(Ng, Ng, Ng) / Ng**(1.5))
+
+    def vfield_like(s_hat):
+        fwd.forwardModel_v2(s_hat)
+        rho = np.zeros((Ng, Ng, Ng))
+        fwd.getDensityFinal(rho)
+        vgrid = vfield.getVelocityField()
+        return (vgrid**2).sum()
+
+    def vfield_ag(s_hat):
+        fwd.forwardModel_v2(s_hat)
+        rho = np.zeros((Ng, Ng, Ng))
+        fwd.getDensityFinal(rho)
+        # The derivative of square is 2 * vector
+        vgrid = 2 * vfield.getVelocityField()
+
+        fwd.clearAdjointGradient()
+        vfield.computeAdjointModel(vgrid)
+        myprint("Calling adjoint")
+        # We have to trigger the adjoint computation in any case
+        fwd.adjointModel_v2(None)
+        myprint("Done with adjoint")
+        analytic_gradient = np.zeros((Ng, Ng, Ng // 2 + 1),
+                                     dtype=np.complex128)
+        fwd.getAdjointModel(analytic_gradient)
+        return analytic_gradient
+
+    myprint("Running adjoint")
+
+    num_gradient = np.zeros((Ng, Ng, Ng // 2 + 1), dtype=np.complex128)
+    s_hat_epsilon = s_hat.copy()
+    cons.setVerboseLevel(5)
+    analytic_gradient = vfield_ag(s_hat)
+    cons.setVerboseLevel(1)
+
+    epsilon = 0.001
+    for i, j, k in tqdm(itertools.product(*map(range, [Ng, Ng, Ng // 2 + 1])),
+                        total=Ng * Ng * (Ng // 2 + 1)):
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k] + epsilon
+        L = vfield_like(s_hat_epsilon)
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k] - epsilon
+        L -= vfield_like(s_hat_epsilon)
+        QQ = L / (2.0 * epsilon)
+
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k] + 1j * epsilon
+        L = vfield_like(s_hat_epsilon)
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k] - 1j * epsilon
+        L -= vfield_like(s_hat_epsilon)
+        QQ = QQ + L * 1j / (2.0 * epsilon)
+
+        s_hat_epsilon[i, j, k] = s_hat[i, j, k]
+
+        num_gradient[i, j, k] = QQ
+
+    np.savez("gradients.npz", num=num_gradient, ana=analytic_gradient)