Added make_replicas from Validphys

colibri/mc_utils.py

@@ -16,6 +16,9 @@
 from colibri.export_results import write_exportgrid
 from colibri.core import MCPseudodata
 
+from validphys.pseudodata import make_replica
+from validphys.n3fit_data import replica_mcseed
+
 import logging
 
 log = logging.getLogger(__name__)
@@ -25,25 +28,35 @@ def mc_pseudodata(
     pseudodata_central_covmat_index,
     replica_index,
     trval_seed,
+    mcseed,
     shuffle_indices=True,
+    positive_pseudodata=False,
     mc_validation_fraction=0.2,
 ):
     """Produces Monte Carlo pseudodata for the replica with index replica_index.
     The pseudodata is returned with a set of training indices, which account for
     a fraction mc_validation_fraction of the data.
-    """
 
-    central_values = pseudodata_central_covmat_index.central_values
+    If positive_pseudodata is True, the pseudodata will be resampled until all values
+    are positive"""
+
+    central_values = [pseudodata_central_covmat_index.central_values]
     covmat = pseudodata_central_covmat_index.covmat
     all_indices = pseudodata_central_covmat_index.central_values_idx
-
-    # Generate pseudodata according to a multivariate Gaussian centred on
-    # central_values and with covariance matrix covmat.
-    key = jax.random.PRNGKey(replica_index)
-    pseudodata = jax.random.multivariate_normal(
-        key,
-        central_values,
-        covmat,
+    seed = replica_mcseed(replica_index, mcseed, genrep=True)
+
+    if positive_pseudodata:
+        group_positivity_mask = np.ones_like(central_values, dtype=bool)
+    else:
+        group_positivity_mask = None
+
+    pseudodata = jnp.array(
+        make_replica(
+            central_values,
+            seed,
+            covmat,
+            group_positivity_mask=group_positivity_mask,
+        ).squeeze()
     )
 
     # Now select a subset of 1 - mc_validation_fraction indices to be the

colibri/monte_carlo_fit.py

@@ -68,14 +68,14 @@ def monte_carlo_fit(
 
     @jax.jit
     def loss_training(parameters, batch):
-        return -2 * mc_log_likelihood[0](parameters, batch) / len_tr_idx
+        return -2 * mc_log_likelihood[0](parameters, batch) 
 
     @jax.jit
     def loss_validation(parameters):
 
         val = -2 * mc_log_likelihood[1](parameters)
 
-        return val / len_val_idx if len_val_idx > 0 else val
+        return val 
 
     log.info(f"Running fit with backend: {jbackend.get_backend().platform}")
     log.info("Starting Monte Carlo fit...")

colibri/param_initialisation.py

@@ -1,6 +1,7 @@
 import jax
 import jax.numpy as jnp
 import logging
+from jax.nn.initializers import glorot_normal, glorot_uniform, zeros, constant
 
 log = logging.getLogger(__name__)
 
@@ -26,7 +27,7 @@ def pdf_initial_parameters(pdf_model, param_initialiser_settings, replica_index=
     initial_values: jnp.array
         The initial values for the parameters.
     """
-    if param_initialiser_settings["type"] not in ("zeros", "normal", "uniform"):
+    if param_initialiser_settings["type"] not in ("zeros", "normal", "uniform", "glorot_norm"):
         log.warning(
             f"MC initialiser type {param_initialiser_settings['type']} not recognised, using default: 'zeros' instead."
         )
@@ -133,3 +134,34 @@ def expand(setting, default, name):
         )
 
         return initial_values
+
+    if param_initialiser_settings["type"] == "glorot_norm":
+        # Get layer shapes
+        if "layer_shapes" not in param_initialiser_settings:
+            raise ValueError("'layer_shapes' must be specified for Glorot initialization")
+
+        layer_shapes = param_initialiser_settings["layer_shapes"]
+
+        # For biases: zeros or constant
+        bias_init_type = param_initialiser_settings.get("init_biases", "zeros")
+        if bias_init_type == "zeros":
+            bias_init_fn = zeros
+        else:  # constant
+            bias_value = param_initialiser_settings.get("bias_init_value", 0.01)
+            bias_init_fn = lambda: constant(bias_value)
+
+        # For weights: glorot
+
+        weight_init_fn = glorot_normal()
+
+        subkeys = jax.random.split(random_seed, len(param_names))
+
+        initialized_params = []
+        for i, (shape, subkey) in enumerate(zip(layer_shapes, subkeys)):
+            if len(shape) == 1:  # Bias
+                init_val = bias_init_fn(subkey, shape) if callable(bias_init_fn) else bias_init_fn(subkey, shape)
+            else:  # Weight
+                init_val = weight_init_fn(subkey, shape)
+            initialized_params.append(init_val.flatten())
+
+        return jnp.concatenate(initialized_params)