# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import tensorflow as tf
from IMNN import IMNN
from IMNN.ABC import ABC, priors

print("TensorFlow version =", tf.__version__)
print("IMNN version =", IMNN.__version__)

details = np.load("data/small_details.npz")
input_shape = tuple(details["input_shape"])
n_params = int(details["n_params"])
n_summaries = int(details["n_summaries"])
n_s = int(details["n_s"])
n_d = int(details["n_d"])
fiducial = details["fiducial"]
delta = details["delta"]


# +
class fiducial_generator:
    def __init__(self, file):
        self.file = file
    def __call__(self, i):
        data = np.load(self.file)
        yield data[i], i
        
class derivative_generator:
    def __init__(self, file):
        self.file = file
    def __call__(self, i, j, k):
        data = np.load(self.file)
        yield data[i, j, k], (i, j, k)


# -

print("n_params =", n_params)
print("n_summaries =", n_summaries)
print("n_s =", n_s)
print("n_d =", n_d)
print("input_shape =", input_shape)
print("θ_fid =", fiducial)
print("δθ =", delta)

model = tf.keras.Sequential(
    [tf.keras.Input(shape=input_shape),
     tf.keras.layers.Dense(128),
     tf.keras.layers.LeakyReLU(0.01),
     tf.keras.layers.Dense(128),
     tf.keras.layers.LeakyReLU(0.01),
     tf.keras.layers.Dense(n_summaries),
    ])
opt = tf.keras.optimizers.Adam()

imnn = IMNN.IMNN(n_s=n_s, n_d=n_d, n_params=n_params, n_summaries=n_summaries,
                 model=model, optimiser=opt, θ_fid=fiducial, δθ=delta,
                 fiducial_loader=fiducial_generator("data/fiducial_data.npy"), 
                 derivative_loader=derivative_generator("data/small_derivative_data.npy"),
                 validation_fiducial_loader=fiducial_generator("data/fiducial_validation_data.npy"), 
                 validation_derivative_loader=derivative_generator("data/small_derivative_validation_data.npy"),
                 at_once=1000,
                 map_fn=None, check_shape=True,
                 verbose=True, filename="model", save=True)

imnn.fit(None, λ=10., ϵ=0.01, patience=10, min_iterations=1000)

plt.plot(imnn.history["det_F"])
plt.plot(imnn.history["val_det_F"])



StarcoderdataJupyterScriptsDedupFiltered