ximinf 0.0.8__tar.gz → 0.0.16__tar.gz

{ximinf-0.0.8/src/ximinf.egg-info → ximinf-0.0.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ximinf
-Version: 0.0.8
+Version: 0.0.16
 Summary: Simulation Based Inference of Cosmological parameters in Jax using type Ia supernovae. 
 Author-email: Adam Trigui <a.trigui@ip2i.in2p3.fr>
 License: GPL-3.0-or-later

{ximinf-0.0.8 → ximinf-0.0.16}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ximinf"
-version = "0.0.8"
+version = "0.0.16"
 description = "Simulation Based Inference of Cosmological parameters in Jax using type Ia supernovae. "
 readme = "README.md"
 requires-python = ">=3.10"

{ximinf-0.0.8 → ximinf-0.0.16}/src/ximinf/generate_sim.py

@@ -72,7 +72,7 @@ def simulate_one(params_dict, z_max, M, cols, N=None, i=None):
         "alpha": 0.0,
         "beta": 0.0,
         "mabs": -19.3,
-        "gamma": 0.0,
+        # "gamma": 0.0,
         "sigma_int": 0.0,  # default intrinsic scatter
     }
 
@@ -83,36 +83,33 @@ def simulate_one(params_dict, z_max, M, cols, N=None, i=None):
     alpha_ = float(params["alpha"])
     beta_  = float(params["beta"])
     mabs_  = float(params["mabs"])
-    gamma_ = float(params["gamma"])
+    # gamma_ = float(params["gamma"])
     sigma_int_ = float(params["sigma_int"])
 
-    brokenalpha_model = skysurvey_sniapop.brokenalpha_model
+    # brokenalpha_model = skysurvey_sniapop.brokenalpha_model
 
     # Generate SNe sample
     snia = skysurvey.SNeIa.from_draw(
         size=M,
         zmax=z_max,
-        model=brokenalpha_model,
+        # model=brokenalpha_model,
         magabs={
-            "x1": "@x1",
-            "c": "@c",
             "mabs": mabs_,
             "sigmaint": sigma_int_,
-            "alpha_low": alpha_,
-            "alpha_high": alpha_,
+            "alpha": alpha_,
             "beta": beta_,
-            "gamma": gamma_
         }
     )
 
     # Apply noise
-    errormodel = sim.noise_model
-    errormodel["localcolor"]["kwargs"]["a"] = 2
-    errormodel["localcolor"]["kwargs"]["loc"] = 0.005
-    errormodel["localcolor"]["kwargs"]["scale"] = 0.05
-    noisy_snia = snia.apply_gaussian_noise(errormodel)
-
-    df = noisy_snia.data
+    # errormodel = sim.noise_model
+    # errormodel["localcolor"]["kwargs"]["a"] = 2
+    # errormodel["localcolor"]["kwargs"]["loc"] = 0.005
+    # errormodel["localcolor"]["kwargs"]["scale"] = 0.05
+    # noisy_snia = snia.apply_gaussian_noise(errormodel)
+
+    # df = noisy_snia.data
+    df = snia.data
 
     # Collect requested columns as lists
     data_dict = {col: list(df[col]) for col in cols if col in df}

{ximinf-0.0.8 → ximinf-0.0.16}/src/ximinf/nn_inference.py

@@ -114,7 +114,6 @@ def load_autoregressive_nn(path):
                 dropout_rate=0.0,
                 Nsize_p=shared["Nsize_p"],
                 Nsize_r=shared["Nsize_r"],
-                N_size_embed=shared["N_size_embed"],
                 n_cols=shared["n_cols"],
                 n_params=n_params_visible,
                 rngs=nnx.Rngs(0),

ximinf-0.0.16/src/ximinf/nn_test.py

@@ -0,0 +1,246 @@
+# Import libraries
+import jax
+import jax.numpy as jnp
+import blackjax
+from functools import partial
+from tqdm.notebook import tqdm
+
+def distance(theta1, theta2):
+    """
+    Compute the Euclidean distance between two points in NDIM space.
+
+    Parameters
+    ----------
+    theta1 : array-like
+        First point in NDIM-dimensional space.
+    theta2 : array-like
+        Second point in NDIM-dimensional space.
+
+    Returns
+    -------
+    float
+        The Euclidean distance between `theta1` and `theta2`.
+    """
+    diff = theta1 - theta2
+    return jnp.linalg.norm(diff)
+
+# def log_prior(theta, bounds):
+#     """
+#     Compute the log-prior probability for the parameter `theta`, 
+#     assuming uniform prior within given bounds.
+
+#     Parameters
+#     ----------
+#     theta : array-like
+#         The parameter values for which the prior is to be calculated.
+#     bounds : jnp.ndarray, optional
+#         The bounds on each parameter (default is the global `BOUNDS`).
+
+#     Returns
+#     -------
+#     float
+#         The log-prior of `theta`, or negative infinity if `theta` is out of bounds.
+#     """
+
+#     in_bounds = jnp.all((theta >= bounds[:, 0]) & (theta <= bounds[:, 1]))
+#     return jnp.where(in_bounds, 0.0, -jnp.inf)
+
+def log_group_prior(theta, bounds, group_indices):
+    """
+    Log prior for a single parameter group.
+    Uniform within bounds, -inf otherwise.
+    """
+    theta_g = theta[group_indices]
+    bounds_g = bounds[group_indices]
+
+    in_bounds = jnp.all(
+        (theta_g >= bounds_g[:, 0]) &
+        (theta_g <= bounds_g[:, 1])
+    )
+
+    return jnp.where(in_bounds, 0.0, -jnp.inf)
+
+
+
+def sample_reference_point(rng_key, bounds):
+    """
+    Sample a reference point within the given bounds uniformly.
+
+    Parameters
+    ----------
+    rng_key : jax.random.PRNGKey
+        The random key used for sampling.
+    bounds : jnp.ndarray, optional
+        The bounds for each parameter (default is the global `BOUNDS`).
+
+    Returns
+    -------
+    tuple
+        A tuple containing the updated `rng_key` and the sampled reference point `theta`.
+    """
+    ndim = bounds.shape[0]
+    rng_key, subkey = jax.random.split(rng_key)
+    u = jax.random.uniform(subkey, shape=(ndim,))
+    span = bounds[:, 1] - bounds[:, 0]
+    theta = bounds[:, 0] + u * span
+    return rng_key, theta
+
+def inference_loop(rng_key, kernel, initial_state, num_samples):
+    """
+    Perform an inference loop using a Markov Chain Monte Carlo (MCMC) kernel.
+
+    Parameters
+    ----------
+    rng_key : jax.random.PRNGKey
+        The random key used for sampling.
+    kernel : callable
+        The MCMC kernel (e.g., NUTS) used for updating the state.
+    initial_state : object
+        The initial state of the MCMC chain.
+    num_samples : int
+        The number of samples to generate in the chain.
+
+    Returns
+    -------
+    jax.numpy.ndarray
+        The sampled states from the inference loop.
+    """
+
+    def one_step(state, rng):
+        state, _ = kernel(rng, state)
+        return state, state
+    keys = jax.random.split(rng_key, num_samples)
+    _, states = jax.lax.scan(one_step, initial_state, keys)
+    return states
+
+def log_prob_fn_groups(theta, models_per_group, data, bounds,
+                       param_groups, global_param_names):
+
+    log_r_sum = 0.0
+    log_p_group_sum = 0.0
+
+    data = data.reshape(1, -1)
+
+    for g, group in enumerate(param_groups):
+
+        # --- parameter bookkeeping (unchanged) ---
+        prev_groups = [
+            p
+            for i in range(g)
+            for p in (param_groups[i] if isinstance(param_groups[i], list)
+                      else [param_groups[i]])
+        ]
+
+        group_list = [group] if isinstance(group, str) else group
+        visible_param_names = prev_groups + group_list
+
+        visible_idx = jnp.array(
+            [global_param_names.index(name) for name in visible_param_names]
+        )
+
+        theta_visible = theta[visible_idx].reshape(1, -1)
+        input_g = jnp.concatenate([data, theta_visible], axis=-1)
+
+        # --- ratio estimator ---
+        logits = models_per_group[g](input_g)
+        p = jax.nn.sigmoid(logits)
+        log_r_sum += jnp.log(p) - jnp.log1p(-p)
+
+        # --- marginal prior for this group ---
+        group_idx = jnp.array(
+            [global_param_names.index(name) for name in group_list]
+        )
+
+        log_p_group_sum += log_group_prior(theta, bounds, group_idx)
+
+    return jnp.squeeze(log_r_sum + log_p_group_sum)
+
+
+
+@partial(jax.jit, static_argnums=(0, 1, 2))
+def sample_posterior(log_prob, n_warmup, n_samples, init_position, rng_key):
+    warmup = blackjax.window_adaptation(blackjax.nuts, log_prob)
+    rng_key, warmup_key, sample_key = jax.random.split(rng_key, 3)
+    (warmup_state, params), _ = warmup.run(warmup_key, init_position, num_steps=n_warmup)
+    kernel = blackjax.nuts(log_prob, **params).step
+    rng_key, sample_key = jax.random.split(rng_key)
+    states = inference_loop(sample_key, kernel, warmup_state, n_samples)
+    return rng_key, states.position
+
+
+def one_sample_step_groups(rng_key, xi, theta_star, n_warmup, n_samples, 
+                           models_per_group, bounds, param_groups, param_names):
+    """
+    Sample from posterior using sum of log-likelihoods over all groups.
+    """
+    rng_key, theta_r0 = sample_reference_point(rng_key, bounds)
+
+    def log_post(theta):
+        return log_prob_fn_groups(theta, models_per_group, xi, bounds, param_groups, param_names)
+
+    rng_key, posterior = sample_posterior(log_post, n_warmup, n_samples, theta_star, rng_key)
+    d_star = distance(theta_star, theta_r0)
+    d_samples = jnp.linalg.norm(posterior - theta_r0, axis=1)
+    f_val = jnp.mean(d_samples < d_star)
+
+    return rng_key, f_val, posterior
+
+
+def batched_one_sample_step_groups(rng_keys, x_batch, theta_star_batch,
+                                   n_warmup, n_samples, models_per_group, bounds, param_groups, param_names):
+    return jax.vmap(
+        lambda rng, x, theta: one_sample_step_groups(rng, x[None, :], theta, n_warmup, n_samples,
+                                                     models_per_group, bounds, param_groups, param_names),
+        in_axes=(0, 0, 0)
+    )(rng_keys, x_batch, theta_star_batch)
+
+def compute_ecp_tarp_jitted_groups(models_per_group, x_list, theta_star_list, alpha_list,
+                                   n_warmup, n_samples, rng_key, bounds,
+                                   param_groups, param_names):
+    """
+    Batched ECP computation using multiple group models.
+    """
+    N = x_list.shape[0]
+    rng_key, split_key = jax.random.split(rng_key)
+    rng_keys = jax.random.split(split_key, N)
+
+    # Batched MCMC and distance evaluation
+    _, f_vals, posterior_uns = batched_one_sample_step_groups(
+        rng_keys, x_list, theta_star_list, n_warmup, n_samples,
+        models_per_group, bounds, param_groups, param_names
+    )
+
+    # Compute ECP values for each alpha
+    ecp_vals = [jnp.mean(f_vals < (1 - alpha)) for alpha in alpha_list]
+
+    return ecp_vals, f_vals, posterior_uns, rng_key
+
+def compute_ecp_tarp_jitted_with_progress_groups(models_per_group, x_list, theta_star_list, alpha_list,
+                                                 n_warmup, n_samples, rng_key, bounds,
+                                                 param_groups, param_names, batch_size=20):
+    N = x_list.shape[0]
+
+    posterior_list = []
+    f_vals_list = []
+
+    for start in tqdm(range(0, N, batch_size), desc="Computing ECP batches"):
+        end = min(start + batch_size, N)
+        x_batch = x_list[start:end]
+        theta_batch = theta_star_list[start:end]
+
+        # Compute ECP and posterior for batch
+        _, f_vals_batch, posterior_batch, rng_key = compute_ecp_tarp_jitted_groups(
+            models_per_group, x_batch, theta_batch, alpha_list,
+            n_warmup, n_samples, rng_key, bounds,
+            param_groups, param_names
+        )
+
+        posterior_list.append(posterior_batch)
+        f_vals_list.append(f_vals_batch)
+
+    posterior_uns = jnp.concatenate(posterior_list, axis=0)
+    f_vals_all = jnp.concatenate(f_vals_list, axis=0)
+
+    ecp_vals = [jnp.mean(f_vals_all < (1 - alpha)) for alpha in alpha_list]
+
+    return ecp_vals, posterior_uns, rng_key

{ximinf-0.0.8 → ximinf-0.0.16}/src/ximinf/nn_train.py

@@ -182,6 +182,46 @@ def train_test_split_jax(X, y, test_size=0.3, shuffle=False, key=None):
 
     return X[:N_train], X[N_train:], y[:N_train], y[N_train:]
 
+def train_test_split_indices_jax(N, test_size=0.3, shuffle=False, key=None, fixed_test_idx=None):
+    """
+    Generate train/test indices in JAX, optionally using a fixed test set.
+
+    Parameters
+    ----------
+    N : int
+        Total number of samples.
+    test_size : float
+        Fraction of the dataset to use as test data.
+    shuffle : bool
+        Whether to shuffle before splitting (ignored if fixed_test_idx is provided).
+    key : jax.random.PRNGKey
+        Random key used for shuffling (required if shuffle=True and fixed_test_idx is None).
+    fixed_test_idx : jax.numpy.ndarray, optional
+        Predefined indices to use as test set (persistent across rounds).
+
+    Returns
+    -------
+    train_idx : jax.numpy.ndarray
+        Indices for the training set.
+    test_idx : jax.numpy.ndarray
+        Indices for the test set.
+    """
+
+    N_test = int(jnp.floor(test_size * N))
+
+    if fixed_test_idx is None:
+        if shuffle:
+            perm = jax.random.permutation(key, N)
+        else:
+            perm = jnp.arange(N)
+        test_idx = perm[:N_test]
+    else:
+        test_idx = fixed_test_idx
+
+    train_idx = jnp.setdiff1d(jnp.arange(N), test_idx)
+    return train_idx, test_idx
+
+
 @nnx.jit
 def l2_loss(model, alpha):
     """
@@ -205,7 +245,7 @@ def l2_loss(model, alpha):
     return alpha * sum((param ** 2).sum() for param in params)
 
 @nnx.jit
-def loss_fn(model, batch, l2_reg=1e-7):
+def loss_fn(model, batch, l2_reg=1e-5):
     """
     Compute the total loss, which is the sum of the data loss and L2 regularization.
 
@@ -304,6 +344,100 @@ def pred_step(model, x_batch):
     logits = model(x_batch)
     return logits
 
+class Phi(nnx.Module):
+    """
+    Neural network module for the Phi network in a Deep Set architecture.
+    """
+    def __init__(self, Nsize, n_cols, *, rngs):
+        self.linear1 = nnx.Linear(n_cols, Nsize, rngs=rngs) #+n_params
+        self.linear2 = nnx.Linear(Nsize, Nsize, rngs=rngs)
+        self.linear3 = nnx.Linear(Nsize, Nsize, rngs=rngs)
+
+    def __call__(self, data):
+        h = data
+        
+        h = nnx.relu(self.linear1(h))
+        h = nnx.relu(self.linear2(h))
+        h = nnx.relu(self.linear3(h))
+        return h
+
+
+class Rho(nnx.Module):
+    """
+    Neural network module for the Rho network in a Deep Set architecture
+    with separate LayerNorm for pooled features and theta.
+    """
+    def __init__(self, Nsize_p, Nsize_r, N_size_params, *, rngs):
+        self.linear1 = nnx.Linear(Nsize_p + N_size_params, Nsize_r, rngs=rngs) #
+        self.linear2 = nnx.Linear(Nsize_r, Nsize_r, rngs=rngs)
+        self.linear3 = nnx.Linear(Nsize_r, 1, rngs=rngs)
+
+    def __call__(self, dropout, pooled_features, params):
+        # Concatenate pooled features and embedding
+        x = jnp.concatenate([pooled_features, params], axis=-1)
+
+        x = nnx.relu(self.linear1(x))
+        x = dropout(x)
+
+        x = nnx.relu(self.linear2(x)) #leaky_relu
+        x = dropout(x)
+
+        return self.linear3(x)
+
+
+class DeepSetClassifier(nnx.Module):
+    """
+    Deep Set Classifier model combining Phi and Rho networks.
+    """
+    def __init__(self, dropout_rate, Nsize_p, Nsize_r,
+                 n_cols, n_params, *, rngs):
+
+        self.dropout = nnx.Dropout(rate=dropout_rate, rngs=rngs)
+        self.n_cols   = n_cols
+        self.n_params = n_params
+
+        self.phi = Phi(Nsize_p, n_cols, rngs=rngs)
+        self.rho = Rho(Nsize_p, Nsize_r, n_params, rngs=rngs)
+
+    def __call__(self, input_data):
+        # ----------------------------------------------------
+        # Accept both shape (N, D) and (D,) without failing
+        # ----------------------------------------------------
+        if input_data.ndim == 1:
+            input_data = input_data[None, :]
+
+        N = input_data.shape[0]
+        input_dim = input_data.shape[1]
+
+        # Compute M first from input size
+        # Total input columns = M*n_cols + n_params + M (mask)
+        M = (input_dim - self.n_params) // (self.n_cols + 1)
+
+        # Reshape data columns
+        data = input_data[:, :M*self.n_cols].reshape(N, M, self.n_cols)
+
+        # Slice mask (last M columns)
+        mask = input_data[:, -M-self.n_params:-self.n_params]         # shape (N, M)
+
+        # Parameters
+        theta = input_data[:, -self.n_params:]  # shape (N, n_params)
+
+        # Apply Phi
+        h = self.phi(data)
+
+        # Apply mask
+        h_masked = h * mask[..., None]
+
+        # Pool (masked average)
+        mask_sum = jnp.sum(mask, axis=1, keepdims=True)
+        mask_sum = jnp.where(mask_sum == 0, 1.0, mask_sum)
+        pooled = jnp.sum(h_masked, axis=1) / mask_sum
+
+        # pooled_N = jnp.concatenate([pooled, mask_sum], axis=-1)
+
+        # Apply Rho
+        return self.rho(self.dropout, pooled, theta)
+
 def train_loop(model,
                optimizer,
                train_data,
@@ -317,6 +451,10 @@ def train_loop(model,
                metrics_history,
                M,
                N,
+               cpu,
+               gpu,
+               group_id,
+               group_params,
                plot_flag=False):
     """
     Train loop with early stopping and optional plotting.
@@ -325,66 +463,63 @@ def train_loop(model,
     # Initialise stopping criteria
     best_train_loss = jnp.inf
     best_test_loss = jnp.inf
+    best_train_accuracy = 0.0
+    best_test_accuracy = 0.0
     strikes = 0
 
     model.train()
 
     for epoch in range(epochs):
-        # Shuffle the training data using JAX.
-        # key, subkey = jax.random.split(key)
-        # perm = jax.random.permutation(subkey, len(train_data))
-        # train_data = train_data[perm]
-        # train_labels = train_labels[perm]
-        # del perm
         
         epoch_train_loss = 0
-        epoch_train_correct = 0
-        epoch_train_total = 0
+        epoch_train_accuracy = 0
         
         for i in range(0, len(train_data), batch_size):
             # Get the current batch of data and labels
-            batch_data = train_data[i:i+batch_size]
-            batch_labels = train_labels[i:i+batch_size]
+            batch_data = jax.device_put(train_data[i:i+batch_size], gpu)
+            batch_labels = jax.device_put(train_labels[i:i+batch_size], gpu)
             
             # Perform a training step
             loss, _ = loss_fn(model, (batch_data, batch_labels))
             accuracy = accuracy_fn(model, (batch_data, batch_labels))
             epoch_train_loss += loss
             # Multiply batch accuracy by batch size to get number of correct predictions
-            epoch_train_correct += accuracy * len(batch_data)
-            epoch_train_total += len(batch_data)
+            epoch_train_accuracy += accuracy * len(batch_data)
             train_step(model, optimizer, (batch_data, batch_labels))
         
         # Log the training metrics.
         current_train_loss = epoch_train_loss / (len(train_data) / batch_size)
+        current_train_accuracy = epoch_train_accuracy / len(train_data)
         metrics_history['train_loss'].append(current_train_loss)
         # Compute overall epoch accuracy
-        metrics_history['train_accuracy'].append(epoch_train_correct / epoch_train_total)
+        metrics_history['train_accuracy'].append(current_train_accuracy)
 
         epoch_test_loss = 0
-        epoch_test_correct = 0
-        epoch_test_total = 0
+        epoch_test_accuracy = 0
 
         # Compute the metrics on the test set using the same batching as training
         for i in range(0, len(test_data), batch_size):
-            batch_data = test_data[i:i+batch_size]
-            batch_labels = test_labels[i:i+batch_size]
+            batch_data = jax.device_put(test_data[i:i+batch_size], gpu)
+            batch_labels = jax.device_put(test_labels[i:i+batch_size], gpu)
 
             loss, _ = loss_fn(model, (batch_data, batch_labels))
             accuracy = accuracy_fn(model, (batch_data, batch_labels))
             epoch_test_loss += loss
-            epoch_test_correct += accuracy * len(batch_data)
-            epoch_test_total += len(batch_data)
+            epoch_test_accuracy += accuracy * len(batch_data)
 
         # Log the test metrics.
         current_test_loss = epoch_test_loss / (len(test_data) / batch_size)
+        current_test_accuracy = epoch_test_accuracy / len(test_data)
         metrics_history['test_loss'].append(current_test_loss)
-        metrics_history['test_accuracy'].append(epoch_test_correct / epoch_test_total)
+        metrics_history['test_accuracy'].append(current_test_accuracy)
         
         # Early Stopping Check
         if current_test_loss < best_test_loss:
             best_test_loss = current_test_loss  # Update best test loss
             strikes = 0
+        # elif current_test_accuracy > best_test_accuracy:
+        #     best_test_accuracy = current_test_accuracy  # Update best test accuracy
+        #     strikes = 0
         elif current_train_loss >= best_train_loss:
             strikes = 0
         elif current_test_loss > best_test_loss and current_train_loss < best_train_loss:
@@ -400,6 +535,8 @@ def train_loop(model,
         if plot_flag and epoch % 1 == 0:
             clear_output(wait=True)
 
+            print(f"=== Training model for group {group_id}: {group_params} ===")
+
             fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
 
             # Loss subplot
@@ -417,40 +554,10 @@ def train_loop(model,
 
             plt.show()
 
-    return model, metrics_history, key
-
-
-# def save_nn(model, path, model_config):
-#     """
-#     Save a neural network model to a checkpoint.
+        if epoch == epochs-1:
+            print(f"\n Reached maximum epochs: {epochs} \n")
 
-#     Parameters
-#     ----------
-#     model : nnx.Module
-#         The model to save.
-#     path : str
-#         Path to the checkpoint directory.
-#     model_config : dict
-#         Configuration dictionary for the model.
-#     """
-#     ckpt_dir = os.path.abspath(path)
-#     ckpt_dir = ocp.test_utils.erase_and_create_empty(ckpt_dir)
-
-#     # Split the model into GraphDef (structure) and State (parameters + buffers)
-#     _, _, _, state = nnx.split(model, nnx.RngKey, nnx.RngCount, ...)
-
-#     # Display for debugging (optional)
-#     # nnx.display(state)
-
-#     # Initialize the checkpointer
-#     checkpointer = ocp.StandardCheckpointer()
-
-#     # Save State (parameters & non-trainable variables)
-#     checkpointer.save(ckpt_dir / 'state', state)
-
-#     # Save model configuration for later loading
-#     with open(ckpt_dir / 'config.json', 'w') as f:
-#         json.dump(model_config, f)
+    return model, metrics_history, key
 
 def save_autoregressive_nn(models_per_group, path, model_config):
     """

{ximinf-0.0.8 → ximinf-0.0.16/src/ximinf.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ximinf
-Version: 0.0.8
+Version: 0.0.16
 Summary: Simulation Based Inference of Cosmological parameters in Jax using type Ia supernovae. 
 Author-email: Adam Trigui <a.trigui@ip2i.in2p3.fr>
 License: GPL-3.0-or-later

ximinf-0.0.8/src/ximinf/nn_test.py

@@ -1,453 +0,0 @@
-# Import libraries
-import jax
-import jax.numpy as jnp
-import blackjax
-from functools import partial
-from tqdm.notebook import tqdm
-
-def distance(theta1, theta2):
-    """
-    Compute the Euclidean distance between two points in NDIM space.
-
-    Parameters
-    ----------
-    theta1 : array-like
-        First point in NDIM-dimensional space.
-    theta2 : array-like
-        Second point in NDIM-dimensional space.
-
-    Returns
-    -------
-    float
-        The Euclidean distance between `theta1` and `theta2`.
-    """
-    diff = theta1 - theta2
-    return jnp.linalg.norm(diff)
-
-def log_prior(theta, bounds):
-    """
-    Compute the log-prior probability for the parameter `theta`, 
-    assuming uniform prior within given bounds.
-
-    Parameters
-    ----------
-    theta : array-like
-        The parameter values for which the prior is to be calculated.
-    bounds : jnp.ndarray, optional
-        The bounds on each parameter (default is the global `BOUNDS`).
-
-    Returns
-    -------
-    float
-        The log-prior of `theta`, or negative infinity if `theta` is out of bounds.
-    """
-
-    in_bounds = jnp.all((theta >= bounds[:, 0]) & (theta <= bounds[:, 1]))
-    return jnp.where(in_bounds, 0.0, -jnp.inf)
-
-# def log_prob_fn(theta, model, xy_noise, bounds):
-#     """
-#     Compute the log-probability for the parameter `theta` using a 
-#     log-prior and the log-likelihood from the neural likelihood ratio approximation.
-
-#     Parameters
-#     ----------
-#     theta : array-like
-#         The parameter values for which the log-probability is computed.
-#     model : callable
-#         A function that takes `theta` and produces model logits for computing the likelihood.
-#     xy_noise : array-like
-#         Input data with added noise for evaluating the likelihood.
-
-#     Returns
-#     -------
-#     float
-#         The log-probability, which is the sum of the log-prior and the log-likelihood.
-#     """
-
-#     lp = log_prior(theta, bounds)
-#     lp = jnp.where(jnp.isfinite(lp), lp, -jnp.inf)
-#     xy_flat = xy_noise.squeeze()
-#     inp = jnp.concatenate([xy_flat, theta])[None, :]
-#     logits = model(inp)
-#     p = jax.nn.sigmoid(logits).squeeze()
-#     p = jnp.clip(p, 1e-6, 1 - 1e-6)
-#     log_like = jnp.log(p) - jnp.log1p(-p)
-#     return lp + log_like
-
-def sample_reference_point(rng_key, bounds):
-    """
-    Sample a reference point within the given bounds uniformly.
-
-    Parameters
-    ----------
-    rng_key : jax.random.PRNGKey
-        The random key used for sampling.
-    bounds : jnp.ndarray, optional
-        The bounds for each parameter (default is the global `BOUNDS`).
-
-    Returns
-    -------
-    tuple
-        A tuple containing the updated `rng_key` and the sampled reference point `theta`.
-    """
-    ndim = bounds.shape[0]
-    rng_key, subkey = jax.random.split(rng_key)
-    u = jax.random.uniform(subkey, shape=(ndim,))
-    span = bounds[:, 1] - bounds[:, 0]
-    theta = bounds[:, 0] + u * span
-    return rng_key, theta
-
-def inference_loop(rng_key, kernel, initial_state, num_samples):
-    """
-    Perform an inference loop using a Markov Chain Monte Carlo (MCMC) kernel.
-
-    Parameters
-    ----------
-    rng_key : jax.random.PRNGKey
-        The random key used for sampling.
-    kernel : callable
-        The MCMC kernel (e.g., NUTS) used for updating the state.
-    initial_state : object
-        The initial state of the MCMC chain.
-    num_samples : int
-        The number of samples to generate in the chain.
-
-    Returns
-    -------
-    jax.numpy.ndarray
-        The sampled states from the inference loop.
-    """
-
-    def one_step(state, rng):
-        state, _ = kernel(rng, state)
-        return state, state
-    keys = jax.random.split(rng_key, num_samples)
-    _, states = jax.lax.scan(one_step, initial_state, keys)
-    return states
-
-def log_prob_fn_groups(theta, models_per_group, x, bounds, param_groups, param_names):
-    """
-    Compute the sum of log-likelihoods for all groups given full theta.
-
-    Parameters
-    ----------
-    theta : jnp.ndarray, shape (n_params,)
-        Full parameter vector.
-    models_per_group : list
-        List of DeepSetClassifier models, one per group.
-    x : jnp.ndarray
-        Input data sample (shape: (data_features + ... + n_params))
-    bounds : jnp.ndarray
-        Parameter bounds.
-    param_groups : list
-        List of parameter groups.
-    param_names : list
-        List of all parameter names in order.
-
-    Returns
-    -------
-    float
-        Sum of log-likelihoods over all groups.
-    """
-    log_lik_sum = 0.0
-
-    n_params = len(param_names)
-
-    # Use everything except the last n_params entries as data
-    data_part = x[:-n_params].reshape(1, -1)   # 2D
-    # If mask is required, you can extract it similarly here
-
-    for g, group in enumerate(param_groups):
-        # Determine visible parameters for this group
-        prev_groups = [
-            p
-            for i in range(g)
-            for p in (param_groups[i] if isinstance(param_groups[i], list) else [param_groups[i]])
-        ]
-        group_list = [group] if isinstance(group, str) else group
-        visible_param_names = prev_groups + group_list
-
-        # Get visible theta values
-        visible_idx = jnp.array([param_names.index(name) for name in visible_param_names])
-        theta_visible = theta[visible_idx].reshape(1, -1)  # make 2D
-
-        # Concatenate data with visible parameters
-        input_g = jnp.concatenate([data_part, theta_visible], axis=-1)
-
-        # Forward pass through the model
-        logits = models_per_group[g](input_g)
-        p = jax.nn.sigmoid(logits)
-
-        log_lik_sum += jnp.log(p) + jnp.log(1 - p)
-
-    return jnp.squeeze(log_lik_sum) + log_prior(theta, bounds)
-
-
-
-
-@partial(jax.jit, static_argnums=(0, 1, 2))
-def sample_posterior(log_prob, n_warmup, n_samples, init_position, rng_key):
-    warmup = blackjax.window_adaptation(blackjax.nuts, log_prob)
-    rng_key, warmup_key, sample_key = jax.random.split(rng_key, 3)
-    (warmup_state, params), _ = warmup.run(warmup_key, init_position, num_steps=n_warmup)
-    kernel = blackjax.nuts(log_prob, **params).step
-    rng_key, sample_key = jax.random.split(rng_key)
-    states = inference_loop(sample_key, kernel, warmup_state, n_samples)
-    return rng_key, states.position
-
-
-# # ========== JIT‐compiled per‐sample step ==========
-# @partial(jax.jit, static_argnums=(3, 4, 5))
-# def one_sample_step(rng_key, xi, theta_star, n_warmup, n_samples, model, bounds):
-#     """
-#     Sample from the posterior distribution using Hamiltonian Monte Carlo (HMC) 
-#     with NUTS (No-U-Turn Sampler) for a given `log_prob`.
-
-#     Parameters
-#     ----------
-#     log_prob : callable
-#         The log-probability function for the model and parameters.
-#     n_warmup : int
-#         The number of warmup steps to adapt the sampler.
-#     n_samples : int
-#         The number of samples to generate after warmup.
-#     init_position : array-like
-#         The initial position for the chain (parameter values).
-#     rng_key : jax.random.PRNGKey
-#         The random key used for sampling.
-
-#     Returns
-#     -------
-#     jax.numpy.ndarray
-#         The sampled positions (parameters) from the posterior distribution.
-#     """
-
-#     # Draw a random reference
-#     rng_key, theta_r0 = sample_reference_point(rng_key, bounds)
-
-#     def log_post(theta):
-#         return log_prob_fn(theta, model, xi, bounds)
-
-#     # Run MCMC
-#     rng_key, posterior = sample_posterior(log_post, n_warmup, n_samples, theta_star, rng_key)
-
-#     # Compute e-c-p distances
-#     d_star = distance(theta_star, theta_r0)
-#     d_samples = jnp.linalg.norm(posterior - theta_r0, axis=1)
-#     f_val = jnp.mean(d_samples < d_star)
-
-#     return rng_key, f_val, posterior
-
-def one_sample_step_groups(rng_key, xi, theta_star, n_warmup, n_samples, 
-                           models_per_group, bounds, param_groups, param_names):
-    """
-    Sample from posterior using sum of log-likelihoods over all groups.
-    """
-    rng_key, theta_r0 = sample_reference_point(rng_key, bounds)
-
-    def log_post(theta):
-        return log_prob_fn_groups(theta, models_per_group, xi, bounds, param_groups, param_names)
-
-    rng_key, posterior = sample_posterior(log_post, n_warmup, n_samples, theta_star, rng_key)
-    d_star = distance(theta_star, theta_r0)
-    d_samples = jnp.linalg.norm(posterior - theta_r0, axis=1)
-    f_val = jnp.mean(d_samples < d_star)
-
-    return rng_key, f_val, posterior
-
-
-# def batched_one_sample_step(rng_keys, x_batch, theta_star_batch, n_warmup, n_samples, model, bounds):
-#     """
-#     Vectorized wrapper over `one_sample_step` using jax.vmap.
-
-#     Parameters
-#     ----------
-#     rng_keys : jax.random.PRNGKey
-#         Batch of random keys.
-#     x_batch : array-like
-#         Batch of input data.
-#     theta_star_batch : array-like
-#         Batch of true parameter values.
-#     n_warmup : int
-#         Number of warmup steps.
-#     n_samples : int
-#         Number of samples.
-#     model : callable
-#         The model function.
-#     bounds : array-like
-#         Parameter bounds.
-
-#     Returns
-#     -------
-#     tuple
-#         (rng_keys, f_vals, posterior_samples)
-#     """
-#     return jax.vmap(
-#         lambda rng, x, theta: one_sample_step(rng, x[None, :], theta, n_warmup, n_samples, model, bounds),
-#         in_axes=(0, 0, 0)
-#     )(rng_keys, x_batch, theta_star_batch)
-
-def batched_one_sample_step_groups(rng_keys, x_batch, theta_star_batch,
-                                   n_warmup, n_samples, models_per_group, bounds, param_groups, param_names):
-    return jax.vmap(
-        lambda rng, x, theta: one_sample_step_groups(rng, x[None, :], theta, n_warmup, n_samples,
-                                                     models_per_group, bounds, param_groups, param_names),
-        in_axes=(0, 0, 0)
-    )(rng_keys, x_batch, theta_star_batch)
-
-
-
-# def compute_ecp_tarp_jitted(model, x_list, theta_star_list, alpha_list, n_warmup, n_samples, rng_key, bounds):
-#     """
-#     Compute expected coverage probabilities (ECP) using vectorized sampling.
-
-#     Parameters
-#     ----------
-#     model : callable
-#         The model function.
-#     x_list : array-like
-#         List of input data.
-#     theta_star_list : array-like
-#         List of true parameter values.
-#     alpha_list : list of float
-#         List of alpha values for ECP computation.
-#     n_warmup : int
-#         Number of warmup steps.
-#     n_samples : int
-#         Number of samples.
-#     rng_key : jax.random.PRNGKey
-#         Random key.
-#     bounds : array-like
-#         Parameter bounds.
-
-#     Returns
-#     -------
-#     tuple
-#         (ecp_vals, f_vals, posterior_uns, rng_key)
-#     """
-#     N = x_list.shape[0]
-#     rng_key, split_key = jax.random.split(rng_key)
-#     rng_keys = jax.random.split(split_key, N)
-
-#     # Batched MCMC and distance evaluation
-#     _, f_vals, posterior_uns = batched_one_sample_step_groups(
-#         rng_keys, x_list, theta_star_list, n_warmup, n_samples, model, bounds
-#     )
-
-#     # Compute ECP values for each alpha
-#     ecp_vals = [jnp.mean(f_vals < (1 - alpha)) for alpha in alpha_list]
-
-#     return ecp_vals, f_vals, posterior_uns, rng_key
-
-
-# def compute_ecp_tarp_jitted_with_progress(model, x_list, theta_star_list, alpha_list,
-#                                           n_warmup, n_samples, rng_key, bounds,
-#                                           batch_size=20):
-#     """
-#     Compute ECP using JITed MCMC in batches with progress reporting via tqdm.
-
-#     Parameters
-#     ----------
-#     model : callable
-#         The model function.
-#     x_list : array-like
-#         List of input data.
-#     theta_star_list : array-like
-#         List of true parameter values.
-#     alpha_list : list of float
-#         List of alpha values for ECP computation.
-#     n_warmup : int
-#         Number of warmup steps.
-#     n_samples : int
-#         Number of samples.
-#     rng_key : jax.random.PRNGKey
-#         Random key.
-#     bounds : array-like
-#         Parameter bounds.
-#     batch_size : int, optional
-#         Batch size for processing (default is 20).
-
-#     Returns
-#     -------
-#     tuple
-#         (ecp_vals, posterior_uns, rng_key)
-#     """
-#     N = x_list.shape[0]
-
-#     posterior_list = []
-#     f_vals_list = []
-
-#     for start in tqdm(range(0, N, batch_size), desc="Computing ECP batches"):
-#         end = min(start + batch_size, N)
-#         x_batch = x_list[start:end]
-#         theta_batch = theta_star_list[start:end]
-
-#         # Compute ECP and posterior for batch
-#         _, f_vals_batch, posterior_batch, rng_key = compute_ecp_tarp_jitted(
-#             model, x_batch, theta_batch, alpha_list,
-#             n_warmup, n_samples, rng_key, bounds
-#         )
-
-#         posterior_list.append(posterior_batch)
-#         f_vals_list.append(f_vals_batch)
-
-#     # Concatenate across batches
-#     posterior_uns = jnp.concatenate(posterior_list, axis=0)
-#     f_vals_all = jnp.concatenate(f_vals_list, axis=0)
-
-#     # Compute final ECP for each alpha
-#     ecp_vals = [jnp.mean(f_vals_all < (1 - alpha)) for alpha in alpha_list]
-
-#     return ecp_vals, posterior_uns, rng_key
-
-def compute_ecp_tarp_jitted_groups(models_per_group, x_list, theta_star_list, alpha_list,
-                                   n_warmup, n_samples, rng_key, bounds,
-                                   param_groups, param_names):
-    """
-    Batched ECP computation using multiple group models.
-    """
-    N = x_list.shape[0]
-    rng_key, split_key = jax.random.split(rng_key)
-    rng_keys = jax.random.split(split_key, N)
-
-    # Batched MCMC and distance evaluation
-    _, f_vals, posterior_uns = batched_one_sample_step_groups(
-        rng_keys, x_list, theta_star_list, n_warmup, n_samples,
-        models_per_group, bounds, param_groups, param_names
-    )
-
-    # Compute ECP values for each alpha
-    ecp_vals = [jnp.mean(f_vals < (1 - alpha)) for alpha in alpha_list]
-
-    return ecp_vals, f_vals, posterior_uns, rng_key
-
-def compute_ecp_tarp_jitted_with_progress_groups(models_per_group, x_list, theta_star_list, alpha_list,
-                                                 n_warmup, n_samples, rng_key, bounds,
-                                                 param_groups, param_names, batch_size=20):
-    N = x_list.shape[0]
-
-    posterior_list = []
-    f_vals_list = []
-
-    for start in tqdm(range(0, N, batch_size), desc="Computing ECP batches"):
-        end = min(start + batch_size, N)
-        x_batch = x_list[start:end]
-        theta_batch = theta_star_list[start:end]
-
-        # Compute ECP and posterior for batch
-        _, f_vals_batch, posterior_batch, rng_key = compute_ecp_tarp_jitted_groups(
-            models_per_group, x_batch, theta_batch, alpha_list,
-            n_warmup, n_samples, rng_key, bounds,
-            param_groups, param_names
-        )
-
-        posterior_list.append(posterior_batch)
-        f_vals_list.append(f_vals_batch)
-
-    posterior_uns = jnp.concatenate(posterior_list, axis=0)
-    f_vals_all = jnp.concatenate(f_vals_list, axis=0)
-
-    ecp_vals = [jnp.mean(f_vals_all < (1 - alpha)) for alpha in alpha_list]
-
-    return ecp_vals, posterior_uns, rng_key