Merge branch 'adversarial'

Merge adversarial branch into main

Author: daniel-habermann

.gitignore

@@ -170,3 +170,5 @@ cython_debug/
 experiments/normal_dim_2/notebooks/self-consistency-abi/
 checkpoints/
 experiments/normal_dim_2/stan/normal_dim_2
+data/
+plots/

experiments/hodgkin_huxley/__init__.py

@@ -0,0 +1,123 @@
+import tensorflow as tf
+from bayesflow.amortizers import AmortizedPosterior
+
+from src.self_consistency_real.schedules import (
+    ConstantSchedule,
+)
+
+
+class AmortizedPosteriorSC(AmortizedPosterior):
+    def __init__(
+        self,
+        prior,
+        simulator,
+        real_data,
+        lambda_schedule=ConstantSchedule(1.0),
+        n_consistency_samples=32,
+        theta_clip_value_min=-float("inf"),
+        theta_clip_value_max=float("inf"),
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.prior = prior
+        self.simulator = simulator
+        self.real_data = real_data  # tf.convert_to_tensor(real_data, dtype=tf.float32)
+        self.step = tf.Variable(0, trainable=False, dtype=tf.int32)
+        self.lambda_schedule = lambda_schedule
+        self.n_consistency_samples = n_consistency_samples
+        self.theta_clip_value_min = theta_clip_value_min
+        self.theta_clip_value_max = theta_clip_value_max
+
+    def compute_loss(self, input_dict, **kwargs):
+        self.step.assign_add(1)
+        lambda_ = self.lambda_schedule(self.step)
+
+        # Get amortizer outputs
+        net_out, sum_out = self(input_dict, return_summary=True, **kwargs)
+        z, log_det_J = net_out
+
+        # Case summary loss should be computed
+        if self.summary_loss is not None:
+            sum_loss = self.summary_loss(sum_out)
+        # Case no summary loss, simply add 0 for convenience
+        else:
+            sum_loss = 0.0
+
+        # Case dynamic latent space - function of summary conditions
+        if self.latent_is_dynamic:
+            logpdf = self.latent_dist(sum_out).log_prob(z)
+        # Case _static latent space
+        else:
+            logpdf = self.latent_dist.log_prob(z)
+
+        # Compute and return total posterior loss
+        posterior_loss = tf.reduce_mean(-logpdf - log_det_J) + sum_loss
+
+        # SELF CONSISTENCY LOSS
+
+        if tf.greater(lambda_, 0.0):
+            # x has shape (n_datasets, data_dim)
+
+            # indices = tf.stop_gradient(tf.range(tf.shape(self.real_data)[0]))
+            # prior_draw = tf.random.normal((64, 7), 1, 2)
+            # x = tf.stop_gradient(self.simulator(prior_draw))
+            if callable(self.real_data):
+                x = self.real_data()
+            else:
+                x = self.real_data
+
+            n_datasets = tf.shape(x)[0]
+            #
+            # z shape: n_consistency_samples, n_datasets, data_dim
+            z = self.latent_dist.sample(
+                (self.n_consistency_samples, n_datasets), to_numpy=False
+            )
+
+            # add a n_consistency_samples dimension as first (0th) index to x
+            # conditions shape: n_consistency_samples, n_datasets, summary_dim
+            data_summary = self.summary_net(x)
+            data_summary = tf.expand_dims(data_summary, axis=0)
+            conditions = tf.tile(data_summary, [self.n_consistency_samples, 1, 1])
+
+            # x_repeated shape: n_consistency_samples, n_datasets, data_dim
+            x_reshaped = tf.expand_dims(x, axis=0)
+            x_repeated = tf.tile(x_reshaped, [self.n_consistency_samples, 1, 1])
+
+            # theta shape: n_consistency_samples, n_datasets, n_params
+            theta = tf.stop_gradient(
+                self.inference_net.inverse(z, conditions, training=False)
+            )
+
+            # log_prior is log(p(theta)) with shape n_consistency_samples, n_datasets
+            log_prior = self.prior.log_prob(theta)
+
+            # log_lik is log(p(y | theta)) with shape n_consistency_samples, n_datasets
+            log_lik = tf.stop_gradient(self.simulator.log_prob(theta, x_repeated))
+
+            # log_post is log(p(theta | y)) with shape n_consistency_samples, n_datasets
+            sc_input_dict = {
+                "parameters": tf.reshape(theta, (-1, tf.shape(theta)[-1])),
+                "summary_conditions": tf.reshape(
+                    x_repeated, (-1, tf.shape(x_repeated)[-1])
+                ),
+            }
+            log_post = self.log_posterior(sc_input_dict, to_numpy=False)
+            log_post = tf.reshape(log_post, (tf.shape(theta)[:-1]))
+
+            # marginal likelihood p(y) = p(theta) * p(y | theta) / p(theta | y)
+            # shape: n_consistency_samples, n_datasets
+            log_ml = log_prior + log_lik - log_post
+
+            # shape: data_size
+            log_ml_var = tf.math.reduce_variance(log_ml, axis=-2)
+
+            # shape: 1
+            sc_loss = tf.math.reduce_mean(log_ml_var, axis=-1)
+        else:
+            sc_loss = tf.constant(0.0)
+
+        return {
+            "Post.Loss": posterior_loss,
+            "SC.Loss": tf.multiply(lambda_, sc_loss),
+        }

experiments/hodgkin_huxley/amortized_posterior_sc.py

@@ -0,0 +1,82 @@
+import tensorflow as tf
+import tensorflow_probability as tfp
+from bayesflow.simulation import GenerativeModel
+from .ode import HodgkinHuxleyODE
+
+
+def get_generative_model():
+    prior = PriorWithLogProb()
+    simulator = SimulatorWithLogProb()
+
+    model = GenerativeModel(
+        prior=prior,
+        simulator=simulator,
+        prior_is_batched=False,
+        simulator_is_batched=True,
+    )
+
+    return model
+
+
+class SimulatorWithLogProb:
+    def __init__(self):
+        self.ode = HodgkinHuxleyODE()
+
+    def __call__(self, z):
+        theta = z_to_theta(z)
+        result = self.ode.solve_ode(theta)
+        dist = tfp.distributions.StudentT(loc=result, scale=0.1, df=10)
+
+        y = dist.sample()
+
+        return y
+
+    def log_prob(self, z, x):
+        theta = z_to_theta(z)
+        y = self.ode.solve_ode(theta)
+        dist = tfp.distributions.StudentT(loc=y, scale=0.1, df=10)
+
+        x_flat = tf.reshape(x, (-1, tf.shape(x)[-1]))
+        pointwise_log_prob = dist.log_prob(x_flat)
+        log_prob = tf.reduce_mean(pointwise_log_prob, axis=-1)
+        log_prob = tf.reshape(log_prob, (tf.shape(x)[:-1]))
+
+        return log_prob
+
+
+class PriorWithLogProb:
+    def __call__(self):
+        z = tfp.distributions.Normal(loc=0, scale=1).sample(7)
+
+        return z
+
+    def log_prob(self, z):
+        return tf.reduce_sum(
+            tfp.distributions.Normal(loc=0, scale=1).log_prob(z), axis=-1
+        )
+
+
+def theta_to_z(theta):
+    z_1 = (tf.math.log(theta[..., 0]) - tf.math.log(110.0)) / 0.1
+    z_2 = (tf.math.log(theta[..., 1]) - tf.math.log(36.0)) / 0.1
+    z_3 = (tf.math.log(theta[..., 2]) - tf.math.log(0.2)) / 0.5
+    z_4 = (theta[..., 3] - 1.0) / 0.05
+    z_5 = (theta[..., 4] + 55.0) / 5.0
+    z_6 = (theta[..., 5] - 50.0) / 5.0
+    z_7 = (theta[..., 6] + 77.0) / 5.0
+
+    return tf.stack([z_1, z_2, z_3, z_4, z_5, z_6, z_7], axis=-1)
+
+
+def z_to_theta(z):
+    theta_1 = tf.exp(tf.math.log(110.0) + 0.1 * z[..., 0])
+    theta_2 = tf.exp(tf.math.log(36.0) + 0.1 * z[..., 1])
+    theta_3 = tf.exp(tf.math.log(0.2) + 0.5 * z[..., 2])
+    theta_4 = z[..., 3] * 0.05 + 1.0
+    theta_5 = z[..., 4] * 5.0 - 55.0
+    theta_6 = z[..., 5] * 5.0 + 50.0
+    theta_7 = z[..., 6] * 5.0 - 77.0
+
+    return tf.stack(
+        [theta_1, theta_2, theta_3, theta_4, theta_5, theta_6, theta_7], axis=-1
+    )

experiments/hodgkin_huxley/generative_model.py

@@ -0,0 +1,17 @@
+import numpy as np
+import tensorflow as tf
+
+
+# mean(abs(mean(y_pred) - y))
+def mean_absolute_bias(trainer, y, n_samples=5000):
+    posterior_draws = trainer.amortizer.sample(
+        {"summary_conditions": y}, n_samples=n_samples
+    )
+
+    y_pred = trainer.generative_model.simulator(posterior_draws)["sim_data"]
+    y_pred = tf.reshape(y_pred, [*posterior_draws.shape[:-1], 200]).numpy()
+
+    absolute_bias_i = np.abs(np.mean(y_pred, axis=-2) - y)
+    mean_absolute_bias = np.mean(absolute_bias_i, axis=-1)
+
+    return mean_absolute_bias

experiments/hodgkin_huxley/metrics.py

@@ -0,0 +1,135 @@
+import bayesflow as bf
+import tensorflow as tf
+import pickle
+from pathlib import Path
+from .generative_model import get_generative_model
+from bayesflow.amortizers import AmortizedPosterior
+from .amortized_posterior_sc import AmortizedPosteriorSC
+
+
+def get_real_data():
+    file_path = Path(__file__).parents[0] / "data" / "real_data.pkl"
+
+    if not file_path.exists():
+        model = get_generative_model()
+        prior = tf.random.normal((1024, 7), 0, 2)
+        data = model.simulator(prior)["sim_data"]
+        real_data = data  # + tf.random.uniform(data.shape, minval=-2.0, maxval=2.0)
+
+        with open(file_path, "wb") as file:
+            pickle.dump(real_data, file)
+
+    with open(file_path, "rb") as file:
+        real_data = pickle.load(file)
+
+    return real_data
+
+
+def get_real_data_subset(n=32):
+    x = get_real_data()
+    indices = tf.random.shuffle(tf.range(tf.shape(x)[0]))[:n]
+
+    subset = tf.gather(x, indices)
+
+    return subset
+
+
+def get_training_data():
+    file_path = Path(__file__).parents[0] / "data" / "training_data.pkl"
+
+    if not file_path.exists():
+        model = get_generative_model()
+        forward_dict = model(2**15)
+
+        with open(file_path, "wb") as file:
+            pickle.dump(forward_dict, file)
+
+    with open(file_path, "rb") as file:
+        forward_dict = pickle.load(file)
+
+    return forward_dict
+
+
+def get_summary_network():
+    return tf.keras.Sequential(
+        [
+            tf.keras.layers.Lambda(lambda x: tf.expand_dims(x, -1)),
+            tf.keras.layers.LSTM(100),
+            tf.keras.layers.Dense(400, activation="relu"),
+            tf.keras.layers.Dense(200, activation="relu"),
+            tf.keras.layers.Dense(100, activation="relu"),
+            tf.keras.layers.Dense(50, activation="relu"),
+        ]
+    )
+
+
+def get_inference_network():
+    return bf.networks.InvertibleNetwork(
+        num_params=7,
+        num_coupling_layers=10,
+        coupling_design="spline",
+        coupling_settings={
+            "dense_args": {"units": 256},
+            "kernel_regularizer": tf.keras.regularizers.l2(1e-3),
+        },
+    )
+
+
+def configurator(forward_dict):
+    input_dict = {}
+
+    # expand dims so summary network works on 4-dimensional inputs
+    input_dict["parameters"] = forward_dict["prior_draws"]
+    input_dict["summary_conditions"] = forward_dict["sim_data"]
+
+    return input_dict
+
+
+def get_amortizer():
+    model = get_generative_model()
+    summary_net = get_summary_network()
+    inference_net = get_inference_network()
+
+    simulator = model.simulator.simulator
+    prior = model.prior.prior
+
+    amortizer = AmortizedPosteriorSC(
+        prior=prior,
+        simulator=simulator,
+        real_data=get_real_data_subset,
+        inference_net=inference_net,
+        summary_net=summary_net,
+        n_consistency_samples=8,
+    )
+
+    return amortizer
+
+
+def get_trainer(**kwargs):
+    generative_model = get_generative_model()
+    amortizer = get_amortizer()
+
+    trainer = bf.trainers.Trainer(
+        amortizer=amortizer,
+        generative_model=generative_model,
+        configurator=configurator,
+        **kwargs,
+    )
+
+    return trainer
+
+
+def get_trainer_no_sc(**kwargs):
+    generative_model = get_generative_model()
+    amortizer = AmortizedPosterior(
+        inference_net=get_inference_network(), summary_net=get_summary_network()
+    )
+
+    trainer = bf.trainers.Trainer(
+        amortizer=amortizer,
+        generative_model=generative_model,
+        configurator=configurator,
+        **kwargs,
+    )
+
+    return trainer