<jupyter_start><jupyter_text>set up model<jupyter_code>n_params = 2
n_summaries = 2
θ_fid = np.array([1.0, 0.5], dtype=np.float32)
N = Lside
input_shape = (1,1, N,N)

# IMNN params
n_s = 5000
n_d = 5000
λ = 100.0
ϵ = 0.1
LN.rsq.shape
rng, initial_model_key = jax.random.split(rng)
rng, fitting_key = jax.random.split(rng)
def InceptBlock2(filters, strides, do_5x5=True, do_3x3=True):
    """InceptNet convolutional striding block.
    filters: tuple: (f1,f2,f3)
    filters1: for conv1x1
    filters2: for conv1x1,conv3x3
    filters3L for conv1x1,conv5x5"""
    
    filters1, filters2, filters3 = filters
    
    conv1x1 = stax.serial(stax.Conv(filters1, (1,1), strides, padding="SAME"))
    
    filters4 = filters2
    conv3x3 = stax.serial(stax.Conv(filters2, (1,1), strides=None, padding="SAME"),
                        stax.Conv(filters4, (3,3), strides, padding="SAME"))
                        
    filters5 = filters3
    conv5x5 = stax.serial(stax.Conv(filters3, (1,1), strides=None, padding="SAME"),
                         stax.Conv(filters5, (5,5), strides, padding="SAME"))

    maxpool = stax.serial(stax.MaxPool((3,3), padding="SAME"),
                         stax.Conv(filters4, (1,1), strides, padding="SAME"))
                            
    if do_3x3:
        if do_5x5:
            return stax.serial(
                  stax.FanOut(4),  # should num=3 or 2 here ?
                  stax.parallel(conv1x1, conv3x3, conv5x5, maxpool),
                  stax.FanInConcat(), 
                  stax.LeakyRelu)
        else:
            return stax.serial(
                  stax.FanOut(3),  # should num=3 or 2 here ?
                  stax.parallel(conv1x1, conv3x3, maxpool),
                  stax.FanInConcat(), 
                  stax.LeakyRelu)
            
    else:
        return stax.serial(
              stax.FanOut(2),  # should num=3 or 2 here ?
              stax.parallel(conv1x1, maxpool),
              stax.FanInConcat(), 
              stax.LeakyRelu)
fs = 32
model = stax.serial(
        #InceptBlock2((fs,fs,fs), strides=(4,4)),
        #InceptBlock2((fs,fs,fs), strides=(4,4)),
        InceptBlock2((fs,fs,fs), strides=(4,4)),
        InceptBlock2((fs,fs,fs), strides=(2,2), do_5x5=False, do_3x3=False),
        InceptBlock2((fs,fs,fs), strides=(2,2), do_5x5=False, do_3x3=False),
        stax.Conv(n_summaries, (1,1), strides=(1,1), padding="SAME"),
        stax.Flatten
)
optimiser = optimizers.adam(step_size=1e-3)
from IMNN.experimental.jax.imnn._imnn import _IMNN
from IMNN.experimental.jax.utils import check_simulator, value_and_jacrev


class SimIMNN(_IMNN):
    def __init__(self, n_s, n_d, n_params, n_summaries, input_shape, θ_fid,
                 model, optimiser, key_or_state, simulator, verbose=True):
        super().__init__(
            n_s=n_s,
            n_d=n_d,
            n_params=n_params,
            n_summaries=n_summaries,
            input_shape=input_shape,
            θ_fid=θ_fid,
            model=model,
            key_or_state=key_or_state,
            optimiser=optimiser,
            verbose=verbose)
        self.simulator = check_simulator(simulator)
        self.simulate = True

    def get_fitting_keys(self, rng):
        return jax.random.split(rng, num=3)

    def get_summaries(self, w, key, validate=False):
        def get_summary(key, θ):
            return self.model(w, self.simulator(key, θ))

        def get_derivatives(key):
            return value_and_jacrev(get_summary, argnums=1)(key, self.θ_fid)

        keys = np.array(jax.random.split(key, num=self.n_s))
        summaries, derivatives = jax.vmap(get_derivatives)(keys[:self.n_d])
        if self.n_s > self.n_d:
            summaries = np.vstack([
                summaries,
                jax.vmap(partial(get_summary, θ=self.θ_fid))(keys[self.n_d:])])
        return np.squeeze(summaries), np.squeeze(derivatives)
    
    
import jax
import jax.numpy as np
from IMNN.experimental.jax.imnn import SimulatorIMNN
from IMNN.experimental.jax.utils import value_and_jacrev, check_devices, \
    check_type, check_splitting


class AggregatedSimulatorIMNN(SimulatorIMNN):
    def __init__(self, n_s, n_d, n_params, n_summaries, input_shape, θ_fid,
                 model, optimiser, key_or_state, simulator, devices,
                 n_per_device, verbose=True):
        super().__init__(
            n_s=n_s,
            n_d=n_d,
            n_params=n_params,
            n_summaries=n_summaries,
            input_shape=input_shape,
            θ_fid=θ_fid,
            model=model,
            key_or_state=key_or_state,
            optimiser=optimiser,
            simulator=simulator,
            verbose=verbose)
        self.devices = check_devices(devices)
        self.n_devices = len(self.devices)
        self.n_per_device = check_type(n_per_device, int, "n_per_device")
        if self.n_s == self.n_d:
            check_splitting(self.n_s, "n_s and n_d", self.n_devices,
                            self.n_per_device)
        else:
            check_splitting(self.n_s, "n_s", self.n_devices, self.n_per_device)
            check_splitting(self.n_d, "n_d", self.n_devices, self.n_per_device)

    def get_summaries(self, w, key=None, validate=False):
        def derivative_scan(counter, rng):
            def get_device_summaries(rng):
                def get_summary(key, θ):
                    return self.model(w, self.simulator(key, θ))

                def get_derivatives(rng):
                    return value_and_jacrev(get_summary, argnums=1)(
                        rng, self.θ_fid)

                keys = np.array(jax.random.split(rng, num=self.n_per_device))
                return jax.vmap(get_derivatives)(keys)

            keys = np.array(jax.random.split(rng, num=self.n_devices))
            summaries, derivatives = jax.pmap(
                get_device_summaries, devices=self.devices)(keys)
            return counter, (summaries, derivatives)

        def summary_scan(counter, rng):
            def get_device_summaries(rng):
                def get_summary(key):
                    return self.model(w, self.simulator(key, self.θ_fid))

                keys = np.array(jax.random.split(rng, num=self.n_per_device))
                return jax.vmap(get_summary)(keys)

            keys = np.array(jax.random.split(rng, num=self.n_devices))
            summaries = jax.pmap(
                get_device_summaries, devices=self.devices)(keys)
            return counter, summaries

        n = self.n_d // (self.n_devices * self.n_per_device)
        if self.n_s > self.n_d:
            n_r = (self.n_s - self.n_d) // (self.n_devices * self.n_per_device)
            key, *keys = jax.random.split(key, num=n_r + 1)
            counter, remaining_summaries = jax.lax.scan(
                summary_scan, n_r, np.array(keys))
        keys = np.array(jax.random.split(key, num=n))
        counter, results = jax.lax.scan(
            derivative_scan, 0, keys)
        summaries, derivatives = results
        if self.n_s > self.n_d:
            summaries = np.vstack([summaries, remaining_summaries])
        return (summaries.reshape((-1, self.n_summaries)),
                derivatives.reshape((-1, self.n_summaries, self.n_params)))


IMNN = SimIMNN(
        n_s=5000,
        n_d=5000,
        n_params=n_params,
        n_summaries=n_summaries,
        input_shape=input_shape,
        θ_fid=θ_fid,
        model=model,
        optimiser=optimiser,
        key_or_state=initial_model_key,
        simulator=lambda rng, θ: simulator(rng, θ, simulator_args=simulator_args),
#         devices=[jax.devices()[0]],
#         n_per_device=1000
        )
IMNN_rngs = 1 * [fitting_key] #+ 12 * [None]
labels = [
    "Simulator, InceptNet\n"
]
%%time
for i in range(1):
    rng,fit_rng = jax.random.split(rng)
    IMNN.fit(λ=10., ϵ=ϵ, rng=fit_rng, min_iterations=500) #for IMNN, IMNN_rng in zip(IMNNs, IMNN_rngs);
IMNNs = [IMNN]
for i, (IMNN, label) in enumerate(zip(IMNNs, labels)):
    if i == 0:
        ax = IMNN.training_plot(expected_detF=None, colour="C{}".format(i), label=label)
    elif i == 10:
        other_ax = IMNN.training_plot(
            expected_detF=None, colour="C{}".format(i), label=label
        )
    elif i == 11:
        IMNN.training_plot(
            ax=other_ax,
            expected_detF=50, colour="C{}".format(i), label=label
        )
        other_ax[0].set_yscale("log")
        other_ax[2].set_yscale("log")
    else:
        IMNN.training_plot(
            ax=ax, expected_detF=None, colour="C{}".format(i), label=label, ncol=5
        );
ax[0].set_yscale("log")
IMNN.F<jupyter_output><empty_output><jupyter_text>ABC inference<jupyter_code>class uniform:
    def __init__(self, low, high):
        self.low = np.array(low)
        self.high = np.array(high)
        self.event_shape = [[] for i in range(self.low.shape[0])]
    def sample(self, n=None, seed=None):
        if n is None:
            n = 1
        keys = np.array(jax.random.split(
            seed, 
            num=len(self.event_shape)))
        return jax.vmap(
            lambda key, low, high : jax.random.uniform(
                key, 
                shape=(n,), 
                minval=low, 
                maxval=high))(
            keys, self.low, self.high)


prior = uniform([0.1, 0.1], [1.6, 1.6])
#prior = uniform([0.1, 0.1], [5.0, 3.0])
rng, key = jax.random.split(rng)
θ_target = np.array([0.9, 0.6])
target_data = simulator(
    key,
    θ_target,
    simulator_args={**simulator_args, **{'squeeze':False}})
@jit #partial(jax.jit, static_argnums=0)
def get_estimate(d):
    if len(d.shape) == 1:
        return IMNN.θ_fid + np.einsum(
            "ij,kj,kl,l->i",
            IMNN.invF,
            IMNN.dμ_dθ,
            IMNN.invC,
            IMNN.model(IMNN.best_w, d, rng=rng) - IMNN.μ)
    else:
        return IMNN.θ_fid + np.einsum(
            "ij,kj,kl,ml->mi",
            IMNN.invF,
            IMNN.dμ_dθ,
            IMNN.invC,
            IMNN.model(IMNN.best_w, d, rng=rng) - IMNN.μ)
estimates = get_estimate(target_data) #[i.get_estimate(target_data) for i in IMNNs];
GAs = [GaussianApproximation(get_estimate(target_data), IMNN.invF, prior)]
%matplotlib inline
for i, (GA, label) in enumerate(zip(GAs, labels)):
    if i == 0:
        ax = GA.marginal_plot(
            axis_labels=[r"$\alpha$", r"$\beta$"], label='on-the-fly IMNN', colours="C{}".format(i)
        )
    else:
        GA.marginal_plot(ax=ax, label=label, colours="C{}".format(i), ncol=8)
ABC = ApproximateBayesianComputation(
        target_data, prior, 
     lambda A,B : simulator(A,B, simulator_args={**simulator_args, **{'squeeze':True}}), 
                                       get_estimate, F=IMNN.F, gridsize=50
    )
target_data.shape
%%time
rng,abc_key = jax.random.split(rng)
ABC(rng=abc_key,
        n_samples=int(1e3),
        min_accepted=15000,
        max_iterations=20000,
        ϵ=0.05,
        smoothing=0.);
ABC.parameters.accepted[0].shape
α
%matplotlib inline
#plt.style.use('default')
new_colors = [ '#2c0342', '#286d87', '#4fb49d', '#9af486']
fig,ax = plt.subplots(nrows=2, ncols=2, figsize=(3.37*2, 3.37*2))
#latexify(fig_width=3.37, fig_height=3.37)

ABC.scatter_plot(ax=ax,
                      colours=new_colors[0], 
                      axis_labels=[r"$\alpha$", r"$\beta$"], 
                      s=8,
                      label='ABC estimate')
# ABC.marginal_plot(ax=ax,
#             axis_labels=[r"$A$", r"$B$"], colours='green',
#             label='ABC marginal plot')


GAs[0].marginal_plot(ax=ax, colours=new_colors[2], axis_labels=[r"$\alpha$", r"$\beta$"], label=None, ncol=1)
ax[0,1].imshow(target_data[0, 0])
ax[0,1].set_yticks([])

#ax[0,1].set_title(r'$\theta_{\rm target} = A,B = (%.2f,%.2f)$'%(θ_target[0], θ_target[1]))

ax[0,0].axvline(θ_target[0], linestyle='--', c='k')
ax[1,0].axvline(θ_target[0], linestyle='--', c='k')
ax[1,0].axhline(θ_target[1], linestyle='--', c='k')
ax[1,1].axhline(θ_target[1], linestyle='--', c='k', label=r'$\theta_{\rm target}$')

ax[1,0].set_xlabel(r'$\alpha$')
ax[1,0].set_ylabel(r'$\beta$')

ax[0,0].axvline(θ_fid[0], linestyle='--', c='k', alpha=0.4)
ax[0,0].set_yticks([])
#ax[1,0].contourf(A_range, B_range, L1.reshape((size, size)))
#ax[0, 0].plot(A_range, np.real(loglikeA), color='g', label='loglikeA')

ax[1,0].axvline(θ_fid[0], linestyle='--', c='k', alpha=0.4)
ax[1,0].axhline(θ_fid[1], linestyle='--', c='k', alpha=0.4)
ax[1,1].axhline(θ_fid[1], linestyle='--', c='k', alpha=0.4, label=r'$\theta_{\rm fid}$')
ax[1,1].set_yticks([])
#ax[1,1].plot(np.real(loglikeB), B_range, color='g', label='loglikeB')
ax[1,1].legend(framealpha=0.)


# add in the likelihood estimate
# ax[0, 0].plot(A_range, likelihoodA, color='#FF8D33', label=None)
# ax[0, 1].axis("off")
# ax[1, 0].contour(A_range, B_range, np.real(likelihood), levels=value, colors='#FF8D33')
# ax[1, 1].plot(likelihoodB, B_range, color='#FF8D33', label='loglike')

ax[0,0].legend(framealpha=0.)


#plt.savefig('/mnt/home/tmakinen/repositories/field-plots/128x128-contours.png', dpi=400)
plt.subplots_adjust(wspace=0, hspace=0)
plt.show()
import powerbox as pbox
α = 1.0
β = 0.5
#mf = hmf.MassFunction(z=0)

# Generate a callable function that returns the cosmological power spectrum.
power = lambda k : ((1 / α**2) * np.exp(α**2 * np.exp(-0.25 * (k**2 / β**2)) - 1))

# Create the power-box instance. The boxlength is in inverse units of the k of which pk is a function, i.e.
# Mpc/h in this case.
pb = pbox.LogNormalPowerBox(N=128, dim=2, pk = power, boxlength= 128., seed=123)
plt.imshow(pb.delta_x(), extent=(0,100,0,100))
plt.colorbar()
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15043274167522372, 'fraction_numerical': 0.028018189819568724, 'mean_word_length': 2.548009367681499, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '320400', 'license': 'Apache-2.0', 'n_tokens_mistral': 5354, 'n_tokens_neox': 5094, 'n_words': 857, 'path': 'scratch/fields/.ipynb_checkpoints/florent-LN-checkpoint.ipynb', 'repo_name': 'tlmakinen/kosmo-kompress', 'size': '13634'}