ximinf 0.0.2__tar.gz → 0.0.8__tar.gz

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

@@ -1,10 +1,11 @@
 Metadata-Version: 2.4
 Name: ximinf
-Version: 0.0.2
+Version: 0.0.8
 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.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ximinf"
-version = "0.0.2"
+version = "0.0.8"
 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.8/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.8/src/ximinf/generate_sim.py

@@ -0,0 +1,120 @@
+# 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={
+            "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.8/src/ximinf/nn_inference.py

@@ -0,0 +1,136 @@
+# 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_size_embed=shared["N_size_embed"],
+                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