flipcosmo 1.0.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

flip/covariance/emulators/generator.py

@@ -2,8 +2,6 @@ import importlib
 
 import numpy as np
 
-import flip.covariance.emulators.gpmatrix as gpmatrix
-import flip.covariance.emulators.nnmatrix as nnmatrix
 from flip.covariance import cov_utils
 from flip.utils import create_log
 
@@ -12,6 +10,7 @@ log = create_log()
 _avail_emulator_models = [
     "gpmatrix",
     "nnmatrix",
+    "skgpmatrix",
 ]
 
 
@@ -24,6 +23,23 @@ def compute_cov(
     test_emulator=False,
     **kwargs,
 ):
+    """Train an emulator for a given covariance type and return evaluators.
+
+    Args:
+        emulator_model_name: Backend name, one of ``gpmatrix``, ``nnmatrix``, ``skgpmatrix``.
+        covariance_list: List of covariance objects with precomputed matrices.
+        parameter_values: Array ``(n_samples, n_params)`` training inputs.
+        emulator_parameter_names: Names of parameters to pass at evaluation time.
+        covariance_type: One of ``"gg"``, ``"vv"``, ``"gv"`` indicating block type.
+        test_emulator: If True, runs a quick evaluation vs. training data and logs errors.
+        **kwargs: Extra keyword arguments forwarded to backend ``train``.
+
+    Returns:
+        List of callables, one per covariance term, taking a parameter dict and returning a matrix.
+
+    Raises:
+        ValueError: If list lengths mismatch or model name is invalid.
+    """
     if len(covariance_list) != len(parameter_values):
         raise ValueError(
             "covariance_list and parameter_values must have the same length."
@@ -38,7 +54,7 @@ def compute_cov(
         f"flip.covariance.emulators.{emulator_model_name}"
     )
 
-    emulator_type = emulator_module._emulator_type
+    # emulator_type = emulator_module._emulator_type
     # Only implemented for matrix emulators, so far.
 
     square_covariance, emulator_output_variance, emulator_output_non_diagonal = (
@@ -86,6 +102,19 @@ def prepare_covariance_matrices(
     covariance_list,
     covariance_type,
 ):
+    """Extract and stack covariance training targets for emulator backends.
+
+    Converts each covariance object into arrays for diagonal and non-diagonal
+    parts depending on the covariance type.
+
+    Args:
+        covariance_list: List of covariance objects.
+        covariance_type: Block type ``("gg"|"vv"|"gv")``.
+
+    Returns:
+        Tuple ``(square_covariance, output_variance, output_non_diagonal)`` where
+        arrays have shapes compatible with backend trainers.
+    """
     if covariance_type[0] == covariance_type[1]:
         square_covariance = True
     else:
@@ -132,6 +161,21 @@ def return_evaluation_functions(
     covariance_type,
     **kwargs,
 ):
+    """Wrap backend models into per-term evaluation functions.
+
+    Each returned function accepts a parameter dict mapping names to values and
+    returns either a square matrix (gg/vv) or a flattened vector (gv).
+
+    Args:
+        models: Tuple of backend-trained models and evaluation dictionaries.
+        emulator_module: Module providing ``evaluate(model, x, dict)`` function.
+        emulator_parameter_names: Names used to extract parameter order from dict.
+        covariance_type: Block type ``("gg"|"vv"|"gv")``.
+        **kwargs: Extra keyword arguments (unused).
+
+    Returns:
+        List of evaluation callables for each covariance term.
+    """
 
     if covariance_type[0] == covariance_type[1]:
         square_covariance = True
@@ -222,6 +266,19 @@ def test_training(
     verbose=True,
     **kwargs,
 ):
+    """Compute emulator errors against training covariances for sanity checks.
+
+    Args:
+        evaluation_functions: List of term evaluators returned by ``return_evaluation_functions``.
+        covariance_list: List of covariance objects (ground truth).
+        parameter_values_dict: List of parameter dicts matching training inputs.
+        covariance_type: Block type ``("gg"|"vv"|"gv")``.
+        verbose: If True, logs per-term maximal relative errors.
+        **kwargs: Extra keyword arguments (unused).
+
+    Returns:
+        Nested list of errors per sample and term: ``[ [err_diag, rel_err_diag, err_all, rel_err_all], ... ]``.
+    """
     if covariance_type[0] == covariance_type[1]:
         square_covariance = True
     else:
@@ -281,6 +338,19 @@ def generate_covariance(
     emulator_parameter_names,
     **kwargs,
 ):
+    """Generate emulator evaluators for requested covariance blocks.
+
+    Args:
+        emulator_model_name: Backend name.
+        model_kind: One of ``"density"``, ``"velocity"``, ``"full"``, ``"density_velocity"``.
+        covariance_list: List of covariance objects for training.
+        parameter_values: Array ``(n_samples, n_params)`` training inputs.
+        emulator_parameter_names: Parameter names used at evaluation.
+        **kwargs: Forwarded to backend trainers.
+
+    Returns:
+        Dict with keys among ``{"gg","vv","gv"}`` mapping to lists of evaluator functions.
+    """
 
     emulator_covariance_dict = {}
 

flip/covariance/emulators/gpmatrix.py

@@ -8,7 +8,7 @@ try:
     import GPy
 
     gpy_installed = True
-except:
+except ImportError:
     gpy_installed = False
     log.add(
         "Install GPy to use the gpmatrix emulator",
@@ -38,6 +38,35 @@ def train(
     num_restarts_non_diagonal=None,
     **kwargs,
 ):
+    """Train Gaussian Process emulators for covariance matrices.
+
+    Trains one GP per covariance term. For square covariances (gg or vv), it
+    trains both a GP for the diagonal element (variance) and one for the
+    non-diagonal elements. For rectangular covariances (gv), only the
+    non-diagonal GP is trained.
+
+    Args:
+        square_covariance: Whether the covariance is square (gg or vv) or not (gv).
+        output_variance: Array of shape ``(n_terms, n_samples)`` containing diagonal values per sample. Ignored if not square.
+        output_non_diagonal: Array of shape ``(n_terms, n_samples, n_nd)`` containing flattened non-diagonal values per sample.
+        parameter_values: Array of shape ``(n_samples, n_params)`` with emulator input parameters.
+        kernel_variance: Optional GPy kernel to use for variance models. Defaults to ``Exponential + Poly`` if GPy is available.
+        kernel_non_diagonal: Optional GPy kernel to use for non-diagonal models. Defaults to ``Exponential + Poly`` if GPy is available.
+        verbose: If True, prints optimization messages.
+        num_restarts_variance: Optional number of optimizer restarts for variance GPs.
+        num_restarts_non_diagonal: Optional number of optimizer restarts for non-diagonal GPs.
+        **kwargs: Extra keyword arguments (unused).
+
+    Returns:
+        Tuple containing:
+        - ``gp_models_variance``: List of variance GP models or ``None``.
+        - ``gp_models_non_diagonal``: List of non-diagonal GP models.
+        - ``gp_evaluation_dictionary_variance``: Per-term evaluation dictionaries or ``None``.
+        - ``gp_evaluation_dictionary_non_diagonal``: Per-term evaluation dictionaries.
+
+    Raises:
+        ImportError: If GPy is not installed when training is requested.
+    """
 
     if square_covariance:
         if kernel_variance is None:
@@ -91,5 +120,15 @@ def evaluate(
     evaluation_value,
     evaluation_dictionary,
 ):
+    """Evaluate a trained GP model at given parameter values.
+
+    Args:
+        model: Trained GPy GPRegression model.
+        evaluation_value: Array of shape ``(1, n_params)`` with input parameters.
+        evaluation_dictionary: Placeholder for API consistency (unused for GPy).
+
+    Returns:
+        Tuple ``(mean, std)`` as returned by ``GPy.models.GPRegression.predict``.
+    """
     output = model.predict(evaluation_value)
     return output

flip/covariance/emulators/nnmatrix.py

@@ -9,7 +9,7 @@ try:
 
     torch_install = True
 
-except:
+except ImportError:
     torch_install = False
     log.add(
         "Install pytorch to use the nnmatrix emulator",
@@ -36,6 +36,15 @@ class RegressionNet(default_regression_object):
         output_dimension=1,
         activation_function=default_activation_function,
     ):
+        """Simple fully connected regression network.
+
+        Args:
+            input_dimension: Number of input features.
+            dimension_hidden_layers: Width of hidden layers.
+            number_hidden_layers: Number of hidden layers.
+            output_dimension: Number of output targets.
+            activation_function: Torch activation class to use.
+        """
         super().__init__()
         layers = []
 
@@ -52,6 +61,7 @@ class RegressionNet(default_regression_object):
         self.model = torch.nn.Sequential(*layers)
 
     def forward(self, x):
+        """Forward pass returning network output."""
         return self.model(x)
 
 
@@ -65,6 +75,18 @@ def train_torch_model(
     verbose,
     model_name,
 ):
+    """Train a torch model on normalized data.
+
+    Args:
+        number_epochs: Number of training epochs.
+        model: Torch model to train.
+        normalized_input: Normalized input array ``(n_samples, n_features)``.
+        normalized_output: Normalized target array ``(n_samples, n_targets)``.
+        optimizer: Torch optimizer instance.
+        loss_function: Torch loss function.
+        verbose: If True, logs periodic losses.
+        model_name: Label used in logging.
+    """
     normalized_input = torch.tensor(normalized_input, dtype=torch.float32)
     normalized_output = torch.tensor(normalized_output, dtype=torch.float32)
     for epoch in range(1, number_epochs + 1):
@@ -92,7 +114,31 @@ def train(
     activation_function=default_activation_function,
     loss_function=default_loss_function,
     tolerance_optimizer=1e-3,
+    **kwargs,
 ):
+    """Train neural network emulators for covariance matrices.
+
+    Normalizes inputs and outputs; trains separate nets per covariance term.
+    For square covariances (gg/vv) also trains a variance net.
+
+    Args:
+        square_covariance: Whether the covariance is square (gg/vv) or rectangular (gv).
+        output_variance: Array ``(n_terms, n_samples)`` diagonal entries; ignored if not square.
+        output_non_diagonal: Array ``(n_terms, n_samples, n_nd)`` flattened non-diagonals.
+        parameter_values: Array ``(n_samples, n_params)`` emulator inputs.
+        verbose: If True, prints training progress.
+        dimension_hidden_layers: Hidden layer width.
+        number_hidden_layers: Hidden layer count.
+        number_epochs: Epochs to train.
+        activation_function: Torch activation class.
+        loss_function: Torch loss function.
+        tolerance_optimizer: Learning rate for Adam optimizer.
+        **kwargs: Extra keyword arguments (unused).
+
+    Returns:
+        Tuple of lists/dicts: ``(nn_models_variance, nn_models_non_diagonal,
+        nn_evaluation_dictionary_variance, nn_evaluation_dictionary_non_diagonal)``.
+    """
 
     parameter_values_mean, parameter_values_std = (
         parameter_values.mean(axis=0),
@@ -202,6 +248,16 @@ def evaluate(
     evaluation_value,
     evaluation_dictionary,
 ):
+    """Evaluate a trained NN emulator, denormalizing output.
+
+    Args:
+        model: Trained torch model.
+        evaluation_value: Array ``(1, n_params)`` input values.
+        evaluation_dictionary: Dict with normalization stats (input_mean/std, output_mean/std).
+
+    Returns:
+        Tuple ``(output, None)`` where output is the denormalized numpy array.
+    """
 
     normalized_evaluation_value = (
         evaluation_value - evaluation_dictionary["input_mean"]

flip/covariance/emulators/skgpmatrix.py

@@ -0,0 +1,125 @@
+import numpy as np
+
+from flip.utils import create_log
+
+log = create_log()
+
+try:
+    from sklearn.gaussian_process import GaussianProcessRegressor
+    from sklearn.gaussian_process.kernels import RBF, ConstantKernel
+
+    sklearngp_installed = True
+
+except ImportError:
+    sklearngp_installed = False
+    log.add(
+        "Install scikit-learn to use the sklearnmatrix emulator",
+        level="warning",
+    )
+
+
+_emulator_type = "matrix"
+
+if sklearngp_installed:
+    default_kernel_variance = RBF(length_scale=1.0) + ConstantKernel(constant_value=1.0)
+    default_kernel_non_diagonal = RBF(length_scale=1.0) + ConstantKernel(
+        constant_value=1.0
+    )
+else:
+    default_kernel_variance = None
+    default_kernel_non_diagonal = None
+
+
+def train(
+    square_covariance,
+    output_variance,
+    output_non_diagonal,
+    parameter_values,
+    kernel_variance=None,
+    kernel_non_diagonal=None,
+    num_restarts_variance=0,
+    num_restarts_non_diagonal=0,
+    **kwargs,
+):
+    """Train scikit-learn GP emulators for covariance matrices.
+
+    Args:
+        square_covariance: Whether covariance is square (gg/vv) or rectangular (gv).
+        output_variance: Array ``(n_terms, n_samples)`` of diagonal entries; ignored if not square.
+        output_non_diagonal: Array ``(n_terms, n_samples, n_nd)`` of flattened non-diagonals.
+        parameter_values: Array ``(n_samples, n_params)`` emulator inputs.
+        kernel_variance: Kernel for variance GP; defaults to ``RBF + Constant``.
+        kernel_non_diagonal: Kernel for non-diagonal GP; defaults to ``RBF + Constant``.
+        num_restarts_variance: Optimizer restarts for variance GP.
+        num_restarts_non_diagonal: Optimizer restarts for non-diagonal GP.
+        **kwargs: Extra keyword arguments (unused).
+
+    Returns:
+        Tuple of models and evaluation dictionaries analogous to other emulator backends.
+    """
+
+    if square_covariance:
+        if kernel_variance is None:
+            kernel_variance = default_kernel_variance
+        gp_models_variance = []
+        gp_evaluation_dictionary_variance = [None for j in range(len(output_variance))]
+    else:
+        gp_models_variance = None
+        gp_evaluation_dictionary_variance = None
+    if kernel_non_diagonal is None:
+        kernel_non_diagonal = default_kernel_non_diagonal
+
+    gp_models_non_diagonal = []
+    gp_evaluation_dictionary_non_diagonal = [
+        None for j in range(len(output_non_diagonal))
+    ]
+
+    for j in range(len(output_non_diagonal)):
+
+        if square_covariance:
+
+            model_variance = GaussianProcessRegressor(
+                kernel=kernel_variance,
+                n_restarts_optimizer=num_restarts_variance,
+            )
+            model_variance.fit(
+                parameter_values,
+                output_variance[j][:, np.newaxis],
+            )
+            gp_models_variance.append(model_variance)
+
+        model_non_diagonal = GaussianProcessRegressor(
+            kernel=kernel_non_diagonal,
+            n_restarts_optimizer=num_restarts_non_diagonal,
+        )
+        model_non_diagonal.fit(
+            parameter_values,
+            output_non_diagonal[j],
+        )
+        gp_models_non_diagonal.append(model_non_diagonal)
+
+    return (
+        gp_models_variance,
+        gp_models_non_diagonal,
+        gp_evaluation_dictionary_variance,
+        gp_evaluation_dictionary_non_diagonal,
+    )
+
+
+def evaluate(
+    model,
+    evaluation_value,
+    evaluation_dictionary,
+):
+    """Evaluate a scikit-learn GP model.
+
+    Args:
+        model: Trained ``GaussianProcessRegressor``.
+        evaluation_value: Array ``(1, n_params)`` input values.
+        evaluation_dictionary: Placeholder for API consistency (unused).
+
+    Returns:
+        Tuple ``(mean, std)`` from ``GaussianProcessRegressor.predict``.
+    """
+    output = model.predict(evaluation_value, return_std=True)
+    return output