ximinf 0.0.2__tar.gz → 0.0.16__tar.gz

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

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: ximinf
-Version: 0.0.2
+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
 Project-URL: Homepage, https://github.com/a-trigui/ximinf
+Project-URL: Documentation, https://ximinf.readthedocs.io
 Keywords: cosmology,supernovae,simulation based inference
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ximinf"
-version = "0.0.2"
+version = "0.0.16"
 description = "Simulation Based Inference of Cosmological parameters in Jax using type Ia supernovae. "
 readme = "README.md"
 requires-python = ">=3.10"
@@ -35,7 +35,7 @@ classifiers = [
 
 [project.urls]
 Homepage = "https://github.com/a-trigui/ximinf"
-# Documentation = "https://my_package.readthedocs.io"
+Documentation = "https://ximinf.readthedocs.io"
 
 [project.optional-dependencies]
 notebooks = ["jupyter", "matplotlib"]

ximinf-0.0.16/src/ximinf/__init__.py

@@ -0,0 +1,7 @@
+# src/ximinf/__init__.py
+
+# from .generate_sim import *
+# from .nn_inference import *
+# from .nn_train import *
+# from .nn_test import *
+# from .selection_effects import *

ximinf-0.0.16/src/ximinf/generate_sim.py

@@ -0,0 +1,117 @@
+# Simulation libraries
+import skysurvey
+import numpy as np
+from pyDOE import lhs  # LHS sampler
+
+def scan_params(ranges, N, dtype=np.float32):
+    """
+    Generate sampled parameter sets using Latin Hypercube Sampling (LHS).
+
+    Parameters
+    ----------
+    ranges : dict
+        Mapping parameter names to (min, max) tuples.
+    N : int
+        Number of samples.
+    dtype : data-type, optional
+        Numeric type for the sampled arrays (default is np.float32).
+
+    Returns
+    -------
+    params_dict : dict
+        Dictionary of parameter arrays of shape (N,).
+    """
+    param_names = list(ranges.keys())
+    n_params = len(param_names)
+
+    # LHS unit samples in [0,1]
+    unit_samples = lhs(n_params, samples=N)
+
+    # Scale unit samples to parameter ranges
+    params_dict = {}
+    for i, p in enumerate(param_names):
+        low, high = ranges[p]
+        params_dict[p] = (unit_samples[:, i] * (high - low) + low).astype(dtype)
+
+    return params_dict
+    
+def simulate_one(params_dict, z_max, M, cols, N=None, i=None):
+    """
+    Simulate a single dataset of SNe Ia.
+
+    Parameters
+    ----------
+    params_dict : dict
+        Dictionary of model parameters (alpha, beta, mabs, gamma, sigma_int, etc.).
+    z_max : float
+        Maximum redshift.
+    M : int
+        Number of SNe to simulate.
+    cols : list of str
+        List of columns to include in the output.
+    N : int, optional
+        Total number of simulations (for progress printing).
+    i : int, optional
+        Current simulation index (for progress printing).
+
+    Returns
+    -------
+    data_dict : dict
+        Dictionary of lists (one per column) containing the simulated data.
+    """
+    import ztfidr.simulation as sim
+    import skysurvey_sniapop
+
+    # Print progress
+    if N is not None and i is not None:
+        if (i+1) % max(1, N//10) == 0 or i == N-1:
+            print(f"Simulation {i+1}/{N}", end="\r", flush=True)
+
+    # Define default parameters including sigma_int
+    default_params = {
+        "alpha": 0.0,
+        "beta": 0.0,
+        "mabs": -19.3,
+        # "gamma": 0.0,
+        "sigma_int": 0.0,  # default intrinsic scatter
+    }
+
+    # Merge defaults with provided params (params_dict takes priority)
+    params = {**default_params, **params_dict}
+
+    # Ensure all are floats
+    alpha_ = float(params["alpha"])
+    beta_  = float(params["beta"])
+    mabs_  = float(params["mabs"])
+    # gamma_ = float(params["gamma"])
+    sigma_int_ = float(params["sigma_int"])
+
+    # brokenalpha_model = skysurvey_sniapop.brokenalpha_model
+
+    # Generate SNe sample
+    snia = skysurvey.SNeIa.from_draw(
+        size=M,
+        zmax=z_max,
+        # model=brokenalpha_model,
+        magabs={
+            "mabs": mabs_,
+            "sigmaint": sigma_int_,
+            "alpha": alpha_,
+            "beta": beta_,
+        }
+    )
+
+    # 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
+    df = snia.data
+
+    # Collect requested columns as lists
+    data_dict = {col: list(df[col]) for col in cols if col in df}
+
+    return data_dict

ximinf-0.0.16/src/ximinf/nn_inference.py

@@ -0,0 +1,135 @@
+# Standard
+import os
+import json
+
+# Jax
+from flax import nnx
+
+# Checkpointing
+import orbax.checkpoint as ocp  # Checkpointing library
+ckpt_dir = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')
+import pathlib  # File path handling library
+
+# Modules
+import ximinf.nn_train as nntr
+
+# def load_nn(path):
+#     """
+#     Load a neural network model from a checkpoint.
+
+#     Parameters
+#     ----------
+#     path : str
+#         Path to the checkpoint directory.
+
+#     Returns
+#     -------
+#     model : nnx.Module
+#         The loaded neural network model.
+
+#     Raises
+#     ------
+#     ValueError
+#         If the checkpoint directory or config file does not exist.
+#     """
+#     # Define the checkpoint directory
+#     ckpt_dir = os.path.abspath(path)
+#     ckpt_dir = pathlib.Path(ckpt_dir).resolve()
+
+#     # Ensure the folder is removed before saving
+#     if ckpt_dir.exists()==False:
+#         # Make an error
+#         raise ValueError(f"Checkpoint directory {ckpt_dir} does not exist. Please check the path.")
+    
+#     # Load model configuration
+#     config_path = ckpt_dir / 'config.json'
+#     if not config_path.exists():
+#         raise ValueError("Model config file not found in checkpoint directory.")
+    
+#     with open(config_path, 'r') as f:
+#         model_config = json.load(f)
+
+#     Nsize_p = model_config['Nsize_p']
+#     Nsize_r = model_config['Nsize_r']
+#     n_cols = model_config['n_cols']
+#     n_params = model_config['n_params']
+#     N_size_embed = model_config['N_size_embed']
+
+#     # 1. Re-create the checkpointer
+#     checkpointer = ocp.StandardCheckpointer()
+
+#     # Split the model into GraphDef (structure) and State (parameters + buffers)
+#     abstract_model = nnx.eval_shape(lambda: nntr.DeepSetClassifier(0.0, Nsize_p, Nsize_r, N_size_embed, n_cols, n_params, rngs=nnx.Rngs(0)))
+#     abs_graphdef, abs_rngkey, abs_rngcount, _ = nnx.split(abstract_model, nnx.RngKey, nnx.RngCount, ...)
+
+#     # 3. Restore
+#     state_restored = checkpointer.restore(ckpt_dir / 'state')
+#     print('NNX State restored: ')
+
+#     model = nnx.merge(abs_graphdef, abs_rngkey, abs_rngcount, state_restored)
+
+#     nnx.display(model)
+
+#     return model
+
+def load_autoregressive_nn(path):
+    """
+    Load an autoregressive stack of NNX models.
+
+    Parameters
+    ----------
+    path : str
+        Checkpoint directory.
+
+    Returns
+    -------
+    models_per_group : list[nnx.Module]
+        Reconstructed models, one per group.
+    model_config : dict
+        Loaded configuration dictionary.
+    """
+    ckpt_dir = pathlib.Path(path).resolve()
+    if not ckpt_dir.exists():
+        raise ValueError(f"Checkpoint directory {ckpt_dir} does not exist.")
+
+    config_path = ckpt_dir / "config.json"
+    if not config_path.exists():
+        raise ValueError("Model config file not found.")
+
+    with open(config_path, "r") as f:
+        model_config = json.load(f)
+
+    shared = model_config["shared"]
+    group_configs = model_config["groups"]
+
+    checkpointer = ocp.StandardCheckpointer()
+    models_per_group = []
+
+    for gconf in group_configs:
+        n_params_visible = gconf["n_params_visible"]
+
+        # Recreate abstract model (shape-only)
+        abstract_model = nnx.eval_shape(
+            lambda: nntr.DeepSetClassifier( # It should not work, there is no class DeepSetClassifier defined in nntr, check how this should be properly done
+                dropout_rate=0.0,
+                Nsize_p=shared["Nsize_p"],
+                Nsize_r=shared["Nsize_r"],
+                n_cols=shared["n_cols"],
+                n_params=n_params_visible,
+                rngs=nnx.Rngs(0),
+            )
+        )
+
+        graphdef, rngkey, rngcount, _ = nnx.split(
+            abstract_model, nnx.RngKey, nnx.RngCount, ...
+        )
+
+        # Restore parameters
+        state = checkpointer.restore(
+            ckpt_dir / f"state_group_{gconf['group_id']}"
+        )
+
+        model = nnx.merge(graphdef, rngkey, rngcount, state)
+        models_per_group.append(model)
+
+    return models_per_group, model_config

{ximinf-0.0.2 → ximinf-0.0.16}/src/ximinf/nn_test.py

@@ -24,56 +24,43 @@ def distance(theta1, 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.
-    """
+# 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])
+    )
 
-    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):
     """
@@ -126,6 +113,50 @@ def inference_loop(rng_key, kernel, initial_state, 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)
@@ -137,42 +168,17 @@ def sample_posterior(log_prob, n_warmup, n_samples, init_position, rng_key):
     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):
+def one_sample_step_groups(rng_key, xi, theta_star, n_warmup, n_samples, 
+                           models_per_group, bounds, param_groups, param_names):
     """
-    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.
+    Sample from posterior using sum of log-likelihoods over all groups.
     """
-
-    # 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)
+        return log_prob_fn_groups(theta, models_per_group, xi, bounds, param_groups, param_names)
 
-    # 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)
@@ -180,29 +186,28 @@ def one_sample_step(rng_key, xi, theta_star, n_warmup, n_samples, model, bounds)
     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.
-    Returns proper f_vals for ECP computation.
-    """
+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(rng, x[None, :], theta, n_warmup, n_samples, model, bounds),
+        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):
+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):
     """
-    Compute expected coverage probabilities (ECP) using vectorized sampling.
-    Returns proper f_vals for ECP computation.
+    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(
-        rng_keys, x_list, theta_star_list, n_warmup, n_samples, model, bounds
+    _, 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
@@ -210,14 +215,9 @@ def compute_ecp_tarp_jitted(model, x_list, theta_star_list, alpha_list, n_warmup
 
     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.
-    Returns correct f_vals for all simulations.
-    """
+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 = []
@@ -229,19 +229,18 @@ def compute_ecp_tarp_jitted_with_progress(model, x_list, theta_star_list, alpha_
         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
+        _, 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)
 
-    # 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
+    return ecp_vals, posterior_uns, rng_key

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

@@ -3,6 +3,8 @@ import os
 import json
 import numpy as np  # Numerical Python
 import scipy as sp
+import matplotlib.pyplot as plt
+from IPython.display import clear_output
 
 # JAX and Flax (new NNX API)
 import jax  # Automatic differentiation library
@@ -19,7 +21,7 @@ ckpt_dir = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')
 # Cosmology
 from astropy.cosmology import Planck18
 
-def rm_cosmo(z, magobs, magabs, ref_mag=19.3, z_max=0.1, n_grid=100_000):
+def rm_cosmo(z, magobs, ref_mag=19.3, z_max=0.1, n_grid=100_000):
     """
     Interpolate Planck18 distance modulus and compute residuals to the cosmology
     
@@ -59,14 +61,27 @@ def rm_cosmo(z, magobs, magabs, ref_mag=19.3, z_max=0.1, n_grid=100_000):
     print('... done')
 
     magobs_corr = magobs - mu_planck18 + ref_mag
-    magabs_corr = magabs + ref_mag
 
-    return mu_planck18, magobs_corr, magabs_corr
+    return mu_planck18, magobs_corr
 
 
 def gaussian(x, mu, sigma):
     """
     Compute the normalized Gaussian function.
+
+    Parameters
+    ----------
+    x : array-like
+        Input values.
+    mu : float
+        Mean of the Gaussian.
+    sigma : float
+        Standard deviation of the Gaussian.
+
+    Returns
+    -------
+    array-like
+        The values of the Gaussian function evaluated at x.
     """
     prefactor = 1 / (np.sqrt(2 * np.pi * sigma**2))
     exponent = np.exp(-((x - mu)**2) / (2 * sigma**2))
@@ -167,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):
     """
@@ -290,30 +345,50 @@ def pred_step(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)
+        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, x):
-        h = nnx.relu(self.linear1(x))
+    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):
-    def __init__(self, Nsize_p, Nsize_r, n_params, *, rngs):
-        self.linear1 = nnx.Linear(Nsize_p + n_params, Nsize_r, rngs=rngs)
-        self.linear2 = nnx.Linear(Nsize_r, 1, rngs=rngs)
+    """
+    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)
 
-    def __call__(self, dropout, pooled_features, theta):
-        x = jnp.concatenate([pooled_features, theta], axis=-1)
         x = nnx.relu(self.linear1(x))
         x = dropout(x)
-        return self.linear2(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):
 
@@ -325,7 +400,14 @@ class DeepSetClassifier(nnx.Module):
         self.rho = Rho(Nsize_p, Nsize_r, n_params, rngs=rngs)
 
     def __call__(self, input_data):
-        N, input_dim = input_data.shape
+        # ----------------------------------------------------
+        # 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)
@@ -340,40 +422,166 @@ class DeepSetClassifier(nnx.Module):
         # Parameters
         theta = input_data[:, -self.n_params:]  # shape (N, n_params)
 
-        # print(theta)
-
         # Apply Phi
-        h = self.phi(self.dropout, data)
+        h = self.phi(data)
 
         # Apply mask
         h_masked = h * mask[..., None]
 
-        # Pool (masked average)mask_sum = jnp.sum(mask, axis=1, keepdims=True)
+        # 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 # Try jnp.sqrt(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,
+               train_labels,
+               test_data,
+               test_labels,
+               key,
+               epochs,
+               batch_size,
+               patience,
+               metrics_history,
+               M,
+               N,
+               cpu,
+               gpu,
+               group_id,
+               group_params,
+               plot_flag=False):
+    """
+    Train loop with early stopping and optional plotting.
+    """
 
+    # 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):
+        
+        epoch_train_loss = 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 = 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_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(current_train_accuracy)
+
+        epoch_test_loss = 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 = 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_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(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:
+            strikes += 1
+        elif current_train_loss < best_train_loss:
+            best_train_loss = current_train_loss # Update best train loss
+
+        if strikes >= patience:
+            print(f"\n Early stopping at epoch {epoch+1} due to {patience} consecutive increases in loss gap \n")
+            break
+
+        # Plotting (optional)
+        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
+            ax1.set_title(f'Loss for M:{M} and N:{N}')
+            for dataset in ('train', 'test'):
+                ax1.plot(metrics_history[f'{dataset}_loss'], label=f'{dataset}_loss')
+            ax1.legend()
+            ax1.set_yscale("log")
+
+            # Accuracy subplot
+            ax2.set_title('Accuracy')
+            for dataset in ('train', 'test'):
+                ax2.plot(metrics_history[f'{dataset}_accuracy'], label=f'{dataset}_accuracy')
+            ax2.legend()
+
+            plt.show()
+
+        if epoch == epochs-1:
+            print(f"\n Reached maximum epochs: {epochs} \n")
+
+    return model, metrics_history, key
+
+def save_autoregressive_nn(models_per_group, path, model_config):
+    """
+    Save an autoregressive stack of NNX models.
 
-def save_nn(model, path, model_config):
+    Parameters
+    ----------
+    models_per_group : list[nnx.Module]
+        One model per autoregressive group.
+    path : str
+        Checkpoint directory.
+    model_config : dict
+        Full model configuration (shared + per-group).
+    """
     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)
+    for g, model in enumerate(models_per_group):
+        # Split model into graph-independent state
+        _, _, _, state = nnx.split(model, nnx.RngKey, nnx.RngCount, ...)
+        checkpointer.save(ckpt_dir / f"state_group_{g}", state)
 
-    # Save model configuration for later loading
-    with open(ckpt_dir / 'config.json', 'w') as f:
-        json.dump(model_config, f)
+    # Save configuration
+    with open(ckpt_dir / "config.json", "w") as f:
+        json.dump(model_config, f, indent=2)

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

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: ximinf
-Version: 0.0.2
+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
 Project-URL: Homepage, https://github.com/a-trigui/ximinf
+Project-URL: Documentation, https://ximinf.readthedocs.io
 Keywords: cosmology,supernovae,simulation based inference
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)

ximinf-0.0.2/src/ximinf/__init__.py

@@ -1 +0,0 @@
-# src/ximinf/__init__.py

ximinf-0.0.2/src/ximinf/generate_sim.py

@@ -1,132 +0,0 @@
-import numpy as np
-import pandas as pd
-
-# Simulation libraries
-import skysurvey
-import skysurvey_sniapop
-import ztfidr.simulation as sim
-
-
-# def flatten_df(df: pd.DataFrame, columns: list, params: list = None) -> np.ndarray:
-#     """
-#     Flatten selected columns from a DataFrame into a single 1D numpy array.
-    
-#     Parameters
-#     ----------
-#     df : pd.DataFrame
-#         Input dataframe containing the data.
-#     columns : list of str
-#         Column names to extract and flatten.
-#     prepend_params : list or None
-#         Optional list of parameters to prepend to the flattened array.
-        
-#     Returns
-#     -------
-#     np.ndarray
-#         1D array containing [prepend_params..., col1..., col2..., ...]
-#     """
-#     arrays = [df[col].to_numpy(dtype=np.float32) for col in columns]
-#     flat = np.concatenate(arrays)
-    
-#     if params is not None:
-#         flat = np.concatenate([np.array(params, dtype=np.float32), flat])
-        
-#     return flat
-
-# def unflatten_array(flat_array: np.ndarray, columns: list, n_points: int = 0):
-#     """
-#     Convert a flattened array back into its original columns and optional prepended parameters.
-
-#     Parameters
-#     ----------
-#     flat_array : np.ndarray
-#         1D array containing the prepended parameters (optional) and column data.
-#     columns : list of str
-#         Original column names in the same order as they were flattened.
-#     n_points : int
-#         Number of rows (SNe) in the data. If > 0, the function will deduce
-#         the number of prepended parameters automatically.
-
-#     Returns
-#     -------
-#     tuple
-#         If prepended_params exist: (prepended_params, df)  
-#         Else: df
-#     """
-#     flat_array = flat_array.astype(np.float32)
-    
-#     if n_points > 0:
-#         # Deduce number of prepended parameters
-#         n_params = flat_array.size - n_points * len(columns)
-#         if n_params < 0:
-#             raise ValueError("Number of points incompatible with flat array size")
-#         prepended_params = flat_array[:n_params] if n_params > 0 else None
-#         data_array = flat_array[n_params:]
-#     else:
-#         prepended_params = None
-#         data_array = flat_array
-
-#     n_rows = data_array.size // len(columns)
-#     if n_rows * len(columns) != data_array.size:
-#         raise ValueError("Flat array size is not compatible with number of columns")
-    
-#     # Split array into columns
-#     split_arrays = np.split(data_array, len(columns))
-#     df = pd.DataFrame({col: arr for col, arr in zip(columns, split_arrays)})
-
-#     if prepended_params is not None:
-#         return prepended_params, df
-#     else:
-#         return df
-    
-def simulate_one(params_dict, sigma_int, z_max, M, cols, N=None, i=None):
-    """
-    params_dict: dict of model parameters (alpha, beta, mabs, gamma, etc.)
-    cols: list of columns to include in the output
-    Returns a dict with:
-        'data': dict of lists (one per column)
-        'params': dict of parameter values
-    """
-    # Print progress
-    if N is not None and i is not None:
-        if (i+1) % max(1, N//10) == 0 or i == N-1:
-            print(f"Simulation {i+1}/{N}", end="\r", flush=True)
-
-    # Unpack parameters
-    alpha_ = float(params_dict.get("alpha", 0))
-    beta_  = float(params_dict.get("beta", 0))
-    mabs_  = float(params_dict.get("mabs", 0))
-    gamma_ = float(params_dict.get("gamma", 0))
-
-    brokenalpha_model = skysurvey_sniapop.brokenalpha_model
-
-    # Generate SNe sample
-    snia = skysurvey.SNeIa.from_draw(
-        size=M,
-        zmax=z_max,
-        model=brokenalpha_model,
-        magabs={
-            "x1": "@x1",
-            "c": "@c",
-            "mabs": mabs_,
-            "sigmaint": sigma_int,
-            "alpha_low": alpha_,
-            "alpha_high": 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
-
-    # Collect requested columns as lists
-    data_dict = {col: list(df[col]) for col in cols if col in df}
-
-    return data_dict

ximinf-0.0.2/src/ximinf/nn_inference.py

@@ -1,56 +0,0 @@
-# Standard
-import os
-import json
-
-# Jax
-from flax import nnx
-
-# Checkpointing
-import orbax.checkpoint as ocp  # Checkpointing library
-ckpt_dir = ocp.test_utils.erase_and_create_empty('/tmp/my-checkpoints/')
-import pathlib  # File path handling library
-
-# Modules
-import ximinf.nn_train as nntr
-
-def load_nn(path):
-    # Define the checkpoint directory
-    ckpt_dir = os.path.abspath(path)
-    ckpt_dir = pathlib.Path(ckpt_dir).resolve()
-
-    # Ensure the folder is removed before saving
-    if ckpt_dir.exists()==False:
-        # Make an error
-        raise ValueError(f"Checkpoint directory {ckpt_dir} does not exist. Please check the path.")
-    
-    # Load model configuration
-    config_path = ckpt_dir / 'config.json'
-    if not config_path.exists():
-        raise ValueError("Model config file not found in checkpoint directory.")
-    
-    with open(config_path, 'r') as f:
-        model_config = json.load(f)
-
-    Nsize_p = model_config['Nsize_p']
-    Nsize_r = model_config['Nsize_r']
-    n_cols = model_config['n_cols']
-    n_params = model_config['n_params']
-
-    # 1. Re-create the checkpointer
-    checkpointer = ocp.StandardCheckpointer()
-
-    # Split the model into GraphDef (structure) and State (parameters + buffers)
-    abstract_model = nnx.eval_shape(lambda: nntr.DeepSetClassifier(0.05, Nsize_p, Nsize_r, n_cols, n_params, rngs=nnx.Rngs(0)))
-    abs_graphdef, abs_rngkey, abs_rngcount, _ = nnx.split(abstract_model, nnx.RngKey, nnx.RngCount, ...)
-
-    # 3. Restore
-    state_restored = checkpointer.restore(ckpt_dir / 'state')
-    #jax.tree.map(np.testing.assert_array_equal, abstract_state, state_restored)
-    print('NNX State restored: ')
-    # nnx.display(state_restored)
-
-    model = nnx.merge(abs_graphdef, abs_rngkey, abs_rngcount, state_restored)
-
-    nnx.display(model)
-
-    return model