Core SNM Functionality

This section covers the core set of functions and classes that implement the main spectral normative modeling algorithms and concepts. This includes classes for fitting univariate and multivariate (spectral) normative models, as well as classes used for defining covariates and how they influence the model.

---

spectranorm.snm

snm.py

Core implementation of spectral normative modeling (SNM).

This module provides the code base for using spectral normative models. It can be used to fit direct and spectral normative models to data and also to predict normative centiles using pre-trained models.

See full documentation at: https://sina-mansour.github.io/spectranorm

CovarianceModelSpec dataclass

General specification of a normative model for covariance.

This model aims to learn the relationships between two variables of interest for both of which a normative model is specified. This can capture patterns where the normative trends in two variables are related.

Attributes:

Name	Type	Description
variable_of_interest_1	str	str Name of the first variable of interest.
variable_of_interest_2	str	str Name of the second variable of interest.
covariates	list[CovariateSpec]	list[CovariateSpec] Listing all model covariates and specifying how each covariate is modeled.
influencing_covariance	list[str]	list[str] List of covariate names that influence the covariance between the two variables of interest.

Source code in src/spectranorm/snm.py

@dataclass
class CovarianceModelSpec:
    """
    General specification of a normative model for covariance.

    This model aims to learn the relationships between two variables of interest for
    both of which a normative model is specified. This can capture patterns where the
    normative trends in two variables are related.

    Attributes:
        variable_of_interest_1: str
            Name of the first variable of interest.
        variable_of_interest_2: str
            Name of the second variable of interest.
        covariates: list[CovariateSpec]
            Listing all model covariates and specifying how each covariate is modeled.
        influencing_covariance: list[str]
            List of covariate names that influence the covariance between the two
            variables of interest.
    """

    variable_of_interest_1: str
    variable_of_interest_2: str
    covariates: list[CovariateSpec]
    influencing_covariance: list[str]

    def __post_init__(self) -> None:
        if not isinstance(self.variable_of_interest_1, str):
            err = "variable_of_interest_1 must be a string."
            raise TypeError(err)
        if not isinstance(self.variable_of_interest_2, str):
            err = "variable_of_interest_2 must be a string."
            raise TypeError(err)
        if not isinstance(self.covariates, list):
            err = "covariates must be a list of CovariateSpec instances."
            raise TypeError(err)
        if not all(isinstance(cov, CovariateSpec) for cov in self.covariates):
            err = "All items in covariates must be CovariateSpec instances."
            raise TypeError(err)
        if not isinstance(self.influencing_covariance, list):
            err = "influencing_covariance must be a list of covariate names."
            raise TypeError(err)

CovarianceNormativeModel dataclass

Covariance normative model implementation.

This class implements covariance normative modeling, which models the covariance structure between a pair of variables as a normative random variable.

Attributes:

Name	Type	Description
spec	CovarianceModelSpec	CovarianceModelSpec Specification of the covariance model including variables of interest, and list of covariates.
defaults	dict[str, Any]	dict Default parameters for the model, including spline specifications, ADVI iterations, convergence tolerance, random seed, and Adam optimizer learning rates.

Source code in src/spectranorm/snm.py

@dataclass
class CovarianceNormativeModel:
    """
    Covariance normative model implementation.

    This class implements covariance normative modeling, which models the covariance
    structure between a pair of variables as a normative random variable.

    Attributes:
        spec: CovarianceModelSpec
            Specification of the covariance model including variables of interest,
            and list of covariates.
        defaults: dict
            Default parameters for the model, including spline specifications,
            ADVI iterations, convergence tolerance, random seed, and Adam optimizer
            learning rates.
    """

    spec: CovarianceModelSpec
    defaults: dict[str, Any] = field(
        default_factory=lambda: {
            "spline_df": DEFAULT_SPLINE_DF,
            "spline_degree": DEFAULT_SPLINE_DEGREE,
            "spline_extrapolation_factor": DEFAULT_SPLINE_EXTRAPOLATION_FACTOR,
            "advi_iterations": DEFAULT_ADVI_ITERATIONS,
            "advi_convergence_tolerance": DEFAULT_ADVI_CONVERGENCE_TOLERANCE,
            "random_seed": DEFAULT_RANDOM_SEED,
            "adam_learning_rate": DEFAULT_ADAM_LEARNING_RATE,
            "adam_learning_rate_decay": DEFAULT_ADAM_LEARNING_RATE_DECAY,
        },
    )

    def __repr__(self) -> str:
        """
        String representation of the CovarianceNormativeModel instance.
        """
        return f"CovarianceNormativeModel(\n\tspec={self.spec}\n)"

    @classmethod
    def from_direct_model(
        cls,
        direct_model: DirectNormativeModel,
        variable_of_interest_1: str,
        variable_of_interest_2: str,
        influencing_covariance: list[str] | None = None,
        defaults_overwrite: dict[str, Any] | None = None,
    ) -> CovarianceNormativeModel:
        """
        Initialize the model from a direct model instance, and two variable names.

        Args:
            direct_model: DirectNormativeModel
                This model will be used to instantiate a similar covariance model.
            variable_of_interest_1: str
                Name of the first target variable to model.
            variable_of_interest_2: str
                Name of the second target variable to model.
            influencing_covariance: list[str] | None
                List of covariates that influence the covariance structure. If not
                provided, this will be copied from the direct model's
                `influencing_variance`.

        Returns:
            CovarianceNormativeModel
                An instance of CovarianceNormativeModel initialized with the provided
                data.
        """
        # Validity checks for input parameters
        if not isinstance(direct_model, DirectNormativeModel):
            err = "direct_model must be an instance of DirectNormativeModel."
            raise TypeError(err)
        if not (
            isinstance(variable_of_interest_1, str)
            and isinstance(variable_of_interest_2, str)
        ):
            err = "Variables of interest must be strings."
            raise TypeError(err)

        # Substitute influencing_covariance if not provided
        if influencing_covariance is None:
            influencing_covariance = direct_model.spec.influencing_variance

        # Use the same setup as the direct model
        model = cls(
            spec=CovarianceModelSpec(
                variable_of_interest_1=variable_of_interest_1,
                variable_of_interest_2=variable_of_interest_2,
                covariates=direct_model.spec.covariates,
                influencing_covariance=influencing_covariance,
            ),
        )

        # update defaults
        model.defaults.update(direct_model.defaults)
        model.defaults.update(defaults_overwrite or {})

        return model

    def _validate_model(self) -> None:
        """
        Validate the covariance model instance.

        This method checks if the model instance is complete and valid.
        It raises errors if any required fields are missing or if there are
        inconsistencies in the model instance.
        """
        if self.spec is None:
            err = (
                "Model specification is not set. Please initialize the model,"
                " e.g., with 'from_dataframe'."
            )
            raise ValueError(err)
        if len(self.spec.covariates) == 0:
            err = (
                "No covariates specified in the model. "
                "Please add covariates to the specification."
            )
            raise ValueError(err)
        if len(self.spec.influencing_covariance) == 0:
            err = (
                "No covariates specified to influence the covariance "
                "between the variables of interest."
            )
            raise ValueError(err)

    def save_model(self, directory: Path, *, save_posterior: bool = False) -> None:
        """
        Save the fitted model and it's posterior to a directory.
        The model will be saved in a subdirectory named 'saved_model'.
        If this directory is not empty, an error is raised.

        Args:
            directory: Path
                Path to a directory to save the model.
            save_posterior: bool (default=False)
                If True, save the model's posterior trace inference data.
        """
        # Prepare the save directory
        directory = Path(directory)
        saved_model_dir = utils.general.prepare_save_directory(directory, "saved_model")

        model_dict = {
            "spec": self.spec,
            "defaults": self.defaults,
        }
        if hasattr(self, "model_params"):
            model_dict["model_params"] = self.model_params
            if hasattr(self, "model_inference_data") and save_posterior:
                self.model_inference_data.to_netcdf(
                    saved_model_dir / "model_inference_data.nc",
                )
        joblib.dump(model_dict, saved_model_dir / "model_dict.joblib")

    @classmethod
    def load_model(
        cls,
        directory: Path,
        *,
        load_posterior: bool = False,
    ) -> CovarianceNormativeModel:
        """
        Load the model and its posterior from a directory.
        The model will be loaded from a subdirectory named 'saved_model'.

        Args:
            directory: Path
                Path to the directory containing the model.
            load_posterior: bool (default=False)
                If True, load the model's posterior trace from the saved inference data.
        """
        # Validate the load directory
        directory = Path(directory)
        saved_model_dir = utils.general.validate_load_directory(
            directory,
            "saved_model",
        )

        # Load the saved model dict
        model_dict = joblib.load(saved_model_dir / "model_dict.joblib")

        # Create an instance of the class
        instance = cls(
            spec=model_dict["spec"],
        )

        # Set the attributes from the loaded model dictionary
        instance.defaults.update(model_dict["defaults"])
        if "model_params" in model_dict:
            instance.model_params = model_dict["model_params"]
            if load_posterior:
                instance.model_inference_data = az.from_netcdf(  # type: ignore[no-untyped-call]
                    saved_model_dir / "model_inference_data.nc",
                )

        return instance

    def _validate_dataframe_for_fitting(self, train_data: pd.DataFrame) -> None:
        """
        Validate the training DataFrame for fitting.
        """
        utils.general.validate_dataframe(
            train_data,
            (
                [cov.name for cov in self.spec.covariates]
                + [
                    self.spec.variable_of_interest_1,
                    self.spec.variable_of_interest_2,
                    f"{self.spec.variable_of_interest_1}_mu_estimate",
                    f"{self.spec.variable_of_interest_2}_mu_estimate",
                    f"{self.spec.variable_of_interest_1}_std_estimate",
                    f"{self.spec.variable_of_interest_2}_std_estimate",
                ]
            ),
        )

    def _build_model_coordinates(
        self,
        observations: npt.NDArray[np.integer[Any]],
    ) -> dict[str, Any]:
        """
        Build the model coordinates for the training DataFrame.
        """
        # Data coordinates
        model_coords = {"observations": observations, "scalar": [0]}

        # Additional coordinates for covariates
        for cov in self.spec.covariates:
            if cov.cov_type == "numerical":
                if cov.effect == "spline":
                    if cov.spline_spec is not None:  # to satisfy type checker
                        model_coords[f"{cov.name}_splines"] = np.arange(
                            cov.spline_spec.df,
                        )
                elif cov.effect == "linear":
                    model_coords[f"{cov.name}_linear"] = np.arange(1)
            elif cov.cov_type == "categorical":
                model_coords[cov.name] = cov.categories
            else:
                err = f"Invalid covariate type '{cov.cov_type}' for '{cov.name}'."
                raise ValueError(err)
        return model_coords

    def _model_linear_correlation_effect(
        self,
        train_data: pd.DataFrame,
        cov: CovariateSpec,
        effects_list: list[TensorVariable],
        sigma_prior: float = 0.1,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Model a linear effect for a numerical covariate.
        """
        # Linear effect
        if adapt is None:  # Model fitting
            linear_beta = pm.Normal(
                f"linear_beta_{cov.name}",
                mu=0,
                sigma=sigma_prior,
                size=1,
                dims=(f"{cov.name}_linear",),
            )
            # Increment parameter count for linear effect
            self.model_params["n_params"] += 1
        else:  # Freeze during adaptation/fine-tuning
            linear_beta = pm.Deterministic(
                f"linear_beta_{cov.name}",
                pt.as_tensor_variable(
                    adapt["pretrained_model_params"]["posterior_means"][
                        f"linear_beta_{cov.name}"
                    ],
                ),
                dims=(f"{cov.name}_linear",),
            )
        if cov.moments is not None:  # to satisfy type checker
            effects_list.append(
                (
                    (
                        (cast("npt.NDArray[Any]", train_data[cov.name].to_numpy()))
                        - cov.moments[0]
                    )
                    / cov.moments[1]
                )
                * linear_beta,
            )

    def _model_spline_correlation_effect(
        self,
        train_data: pd.DataFrame,
        cov: CovariateSpec,
        effects_list: list[TensorVariable],
        spline_bases: dict[str, npt.NDArray[np.floating[Any]]],
        sigma_prior: float = 1,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Model a spline effect for a numerical covariate.
        """
        # Spline effect
        spline_bases[cov.name] = spline_bases.get(
            cov.name,
            cov.make_spline_bases(
                cast("npt.NDArray[Any]", train_data[cov.name].to_numpy()),
            ),
        )
        if adapt is None:  # Model fitting
            spline_betas = pm.ZeroSumNormal(
                f"spline_betas_{cov.name}",
                sigma=sigma_prior,
                shape=spline_bases[cov.name].shape[1],
                dims=(f"{cov.name}_splines",),
            )
            # Note ZeroSumNormal imposes a centering constraint
            # (ensuring identifiability)
            # Increment parameter count for spline effects
            if cov.spline_spec is not None:  # to satisfy type checker
                self.model_params["n_params"] += cov.spline_spec.df - 1
        else:  # Freeze during adaptation/fine-tuning
            spline_betas = pm.Deterministic(
                f"spline_betas_{cov.name}",
                pt.as_tensor_variable(
                    adapt["pretrained_model_params"]["posterior_means"][
                        f"spline_betas_{cov.name}"
                    ],
                ),
                dims=(f"{cov.name}_splines",),
            )
        effects_list.append(pt.dot(spline_bases[cov.name], spline_betas.T))

    def _model_categorical_correlation_effect(
        self,
        train_data: pd.DataFrame,
        cov: CovariateSpec,
        effects_list: list[TensorVariable],
        category_indices: dict[str, npt.NDArray[np.integer[Any]]],
        sigma_prior: float = 1,
        hierarchical_sigma_prior: float = 0.1,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Model the effect of a categorical covariate.
        """
        # Factorize categories
        category_indices[cov.name] = category_indices.get(
            cov.name,
            cov.factorize_categories(
                cast("npt.NDArray[Any]", train_data[cov.name].to_numpy()),
            ),
        )
        if adapt is None:  # Model fitting
            if cov.hierarchical:
                # Hierarchical categorical effect
                # Hyperpriors for category (Bayesian equivalent of random effects)
                sigma_intercept_category = pm.HalfNormal(
                    f"sigma_intercept_{cov.name}",
                    sigma=sigma_prior,
                    dims=("scalar",),
                )

                # Hierarchical intercepts for each category (using reparameterized form)
                categorical_intercept_offset = pm.ZeroSumNormal(
                    f"intercept_offset_{cov.name}",
                    sigma=hierarchical_sigma_prior,
                    dims=(cov.name,),
                )
                # Note ZeroSumNormal imposes a centering constraint
                # (ensuring identifiability)
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    (
                        categorical_intercept_offset
                        * pt.reshape(sigma_intercept_category, (1,))  # pyright: ignore[reportPrivateImportUsage]
                    ),
                    dims=(cov.name,),
                )

                # Increment parameter count for hierarchical intercept
                self.model_params["n_params"] += 1

            else:
                # Non-hierarchical (linear) categorical effect
                categorical_intercept = pm.ZeroSumNormal(
                    f"intercept_{cov.name}",
                    sigma=sigma_prior,
                    dims=(cov.name,),
                )
                # Note ZeroSumNormal imposes a centering constraint
                # (ensuring identifiability)
            # Increment parameter count for categorical effects
            if cov.categories is not None:  # to satisfy type checker
                self.model_params["n_params"] += len(cov.categories) - 1
        elif cov.name != adapt["covariate_to_adapt"]:
            # Freeze during adaptation/fine-tuning
            if cov.hierarchical:
                # Hierarchical categorical effect
                # Hyperpriors for category (Bayesian equivalent of random effects)
                sigma_intercept_category = pm.Deterministic(
                    f"sigma_intercept_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"sigma_intercept_{cov.name}"
                        ],
                    ),
                    dims=("scalar",),
                )
                # Hierarchical intercepts for each category (using reparameterized form)
                categorical_intercept_offset = pm.Deterministic(
                    f"intercept_offset_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"intercept_offset_{cov.name}"
                        ],
                    ),
                    dims=(cov.name,),
                )
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    (
                        categorical_intercept_offset
                        * pt.reshape(sigma_intercept_category, (1,))  # pyright: ignore[reportPrivateImportUsage]
                    ),
                    dims=(cov.name,),
                )
            else:
                # Non-hierarchical (linear) categorical effect
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"intercept_{cov.name}"
                        ],
                    ),
                    dims=(cov.name,),
                )
        else:
            if cov.hierarchical:
                # Hierarchical categorical effect
                # Hyperpriors for category (Bayesian equivalent of random effects)
                # Hyperpriors are fixed during adaptation
                sigma_intercept_category = pm.Deterministic(
                    f"sigma_intercept_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"sigma_intercept_{cov.name}"
                        ],
                    ),
                    dims=("scalar",),
                )
                # Hierarchical intercepts for each category (using reparameterized form)
                # New categories get new parameters, old categories are fixed
                # Freeze old category parameters during adaptation
                fixed_categorical_intercept_offset = pm.Deterministic(
                    f"intercept_offset_{cov.name}_fixed",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"intercept_offset_{cov.name}"
                        ],
                    ),
                )
                # Create new parameters for new categories
                new_category_count = len(adapt["new_category_names"])
                new_categorical_intercept_offset = pm.Normal(
                    f"intercept_offset_{cov.name}_adapt",
                    mu=0,
                    sigma=hierarchical_sigma_prior,
                    size=new_category_count,
                )
                # Combine fixed and new offsets
                categorical_intercept_offset = pm.Deterministic(
                    f"intercept_offset_{cov.name}",
                    pt.concatenate(
                        [
                            fixed_categorical_intercept_offset,
                            new_categorical_intercept_offset,
                        ],
                    ),
                    dims=(cov.name,),
                )
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    (
                        categorical_intercept_offset
                        * pt.reshape(sigma_intercept_category, (1,))  # pyright: ignore[reportPrivateImportUsage]
                    ),
                    dims=(cov.name,),
                )
            else:
                # Non-hierarchical (linear) categorical effect
                # New categories get new parameters, old categories are fixed
                # Freeze old category parameters during adaptation
                fixed_categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}_fixed",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"intercept_{cov.name}"
                        ],
                    ),
                )
                # Create new parameters for new categories
                new_category_count = len(adapt["new_category_names"])
                new_categorical_intercept = pm.Normal(
                    f"intercept_{cov.name}_adapt",
                    mu=0,
                    sigma=sigma_prior,
                    size=new_category_count,
                )
                # Combine fixed and new intercepts
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    pt.concatenate(
                        [
                            fixed_categorical_intercept,
                            new_categorical_intercept,
                        ],
                    ),
                    dims=(cov.name,),
                )
            self.model_params["n_params"] += new_category_count
        effects_list.append(
            categorical_intercept[category_indices[cov.name]],
        )

    def _model_all_correlation_effects(
        self,
        train_data: pd.DataFrame,
        spline_bases: dict[str, npt.NDArray[np.floating[Any]]],
        category_indices: dict[str, npt.NDArray[np.integer[Any]]],
        adapt: dict[str, Any] | None = None,
    ) -> list[TensorVariable]:
        """
        Model all covariate correlation effects.
        """
        # Create a list to contain the effects of covariates on
        # the z-transformed correlation
        z_transformed_correlation_effects = []

        # Model the z-transformed correlation between the variables of interest
        # Model the global intercept for z
        if adapt is None:
            global_intercept_z = pm.Normal(
                "global_intercept_z",
                mu=0,
                sigma=5,
                dims=("scalar",),
            )
            # Increment parameter count for global intercept
            self.model_params["n_params"] += 1
        else:
            # Use the pretrained global intercept
            global_intercept_z = pm.Deterministic(
                "global_intercept_z",
                pt.as_tensor_variable(
                    adapt["pretrained_model_params"]["posterior_means"][
                        "global_intercept_z"
                    ],
                ),
                dims=("scalar",),
            )
        z_transformed_correlation_effects.append(global_intercept_z)
        # Model additional covariate effects on the z estimate
        for cov in self.spec.covariates:
            if cov.name in self.spec.influencing_covariance:
                if cov.cov_type == "numerical":
                    if cov.effect == "linear":
                        self._model_linear_correlation_effect(
                            train_data=train_data,
                            cov=cov,
                            effects_list=z_transformed_correlation_effects,
                            adapt=adapt,
                        )
                    elif cov.effect == "spline":
                        self._model_spline_correlation_effect(
                            train_data=train_data,
                            cov=cov,
                            effects_list=z_transformed_correlation_effects,
                            spline_bases=spline_bases,
                            adapt=adapt,
                        )
                elif cov.cov_type == "categorical":
                    self._model_categorical_correlation_effect(
                        train_data=train_data,
                        cov=cov,
                        effects_list=z_transformed_correlation_effects,
                        category_indices=category_indices,
                        adapt=adapt,
                    )

                else:
                    err = f"Invalid covariate type '{cov.cov_type}' for '{cov.name}'."
                    raise ValueError(err)
        return z_transformed_correlation_effects

    def _combine_all_correlation_effects(
        self,
        z_transformed_correlation_effects: list[TensorVariable],
        combination_indices: npt.NDArray[np.integer[Any]],
        combination_weights: npt.NDArray[np.floating[Any]],
        standardized_vois: npt.NDArray[np.floating[Any]],
        standardized_vois_mu_estimate: npt.NDArray[np.floating[Any]],
        standardized_vois_std_estimate: npt.NDArray[np.floating[Any]],
    ) -> None:
        """
        Combine all effects to model the observed data likelihood from the list of
        correlation effects.
        """
        # Combine all covariance effects
        z_transformed_correlation_estimate = sum(z_transformed_correlation_effects)

        # Convert z-transformed score to correlation
        correlation_estimate = pt.tanh(z_transformed_correlation_estimate)

        # Now apply the random combinations to get final distribution estimates
        # Apply combination weights for mu estimate
        combined_mu_estimate = pt.sum(
            pt.mul(
                standardized_vois_mu_estimate[combination_indices, :],
                combination_weights,
            ),
            axis=1,
        )
        # Apply combination weights for sigma estimate
        combined_sigma_estimate = pt.mul(
            standardized_vois_std_estimate[combination_indices, :],
            combination_weights,
        )
        # Apply combination weights for correlation estimate
        combined_correlation_estimate = correlation_estimate[combination_indices]  # pyright: ignore[reportOptionalSubscript]
        # Now build the std estimate
        combined_std_estimate = pt.sqrt(
            combined_sigma_estimate[:, 0] ** 2  # pyright: ignore[reportOptionalSubscript]
            + combined_sigma_estimate[:, 1] ** 2  # pyright: ignore[reportOptionalSubscript]
            + (
                2
                * combined_sigma_estimate[:, 0]  # pyright: ignore[reportOptionalSubscript]
                * combined_sigma_estimate[:, 1]  # pyright: ignore[reportOptionalSubscript]
                * combined_correlation_estimate
            ),
        )

        # Apply combination to the variables of interest
        combined_variable_of_interest = pt.sum(
            pt.mul(
                standardized_vois[combination_indices, :],
                combination_weights,
            ),
            axis=1,
        )

        effective_sample_size = self.model_params["sample_size"]

        # Model likelihood estimation for covariance model
        _likelihood = pm.Normal(
            f"likelihood_cov_{self.spec.variable_of_interest_1}_{self.spec.variable_of_interest_2}",
            mu=combined_mu_estimate,
            sigma=combined_std_estimate,
            observed=combined_variable_of_interest,
            total_size=(effective_sample_size),
        )

    def _fit_model_with_advi(self, *, progress_bar: bool = True) -> None:
        """
        Fit the model using Automatic Differentiation Variational Inference (ADVI).
        """
        base_lr = self.defaults["adam_learning_rate"]
        decay = self.defaults["adam_learning_rate_decay"]
        lr = shared(base_lr)
        optimizer = pm.adam(learning_rate=cast("float", lr))

        # Adaptive learning rate schedule callback
        def update_learning_rate(_approx: Any, _loss: Any, iteration: int) -> None:
            lr.set_value(base_lr * (decay**iteration))

        # Run automatic differential variational inference to fit the model
        self._trace = pm.fit(
            method="advi",
            n=self.defaults["advi_iterations"],
            random_seed=self.defaults["random_seed"],  # For reproducibility
            obj_optimizer=optimizer,
            callbacks=[
                update_learning_rate,
                pm.callbacks.CheckParametersConvergence(
                    tolerance=self.defaults["advi_convergence_tolerance"],
                    diff="relative",
                ),
            ],
            progressbar=progress_bar,
        )

        # Sample from the posterior distribution and store the results
        self.model_inference_data = self._trace.sample(
            2000,
            random_seed=self.defaults["random_seed"],
        )

        # Compute posterior means and standard deviations
        posterior_means = self.model_inference_data.posterior.mean(
            dim=("chain", "draw"),
        )
        posterior_stds = self.model_inference_data.posterior.std(
            dim=("chain", "draw"),
        )

        # Store posterior means and stds as a dictionary in model parameters
        self.model_params["posterior_means"] = {
            x: posterior_means.data_vars[x].to_numpy()
            for x in posterior_means.data_vars
        }
        self.model_params["posterior_stds"] = {
            x: posterior_stds.data_vars[x].to_numpy() for x in posterior_stds.data_vars
        }

    def fit(
        self,
        train_data: pd.DataFrame,
        *,
        save_directory: Path | None = None,
        progress_bar: bool = True,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Fit the normative model to the training data.

        This method implements the fitting logic for the normative model
        based on the provided training data and model specification.

        Args:
            train_data: pd.DataFrame
                DataFrame containing the training data. It must include the variable
                of interest, their predicted moments, and all specified covariates.
            save_directory: Path | None
                A path to a directory to save the model. If provided, the fitted model
                will be saved to this path.
            progress_bar: bool
                If True, display a progress bar during fitting. Defaults to True.
            adapt: dict[str, Any] | None
                If provided, adapt a pre-trained model to a new covariate.
                Note: We recommended using the `adapt_fit` method, and not directly
                changing this argument, unless you know what you are doing.
        """
        # Validation checks
        self._validate_model()
        self._validate_dataframe_for_fitting(train_data)

        # Extract the variables of interest
        variables_of_interest = train_data[
            [self.spec.variable_of_interest_1, self.spec.variable_of_interest_2]
        ].to_numpy()

        # A dictionary to hold the model parameters after fitting
        if adapt is None:
            self.model_params = {}
            self.model_params["mean_vois"] = variables_of_interest.mean(axis=0)
            self.model_params["std_vois"] = variables_of_interest.std(axis=0)
            self.model_params["sample_size"] = variables_of_interest.shape[0]
            # Initialize parameter count
            self.model_params["n_params"] = 0
        else:
            # Update the pretrained model parameters
            if not hasattr(self, "model_params") or self.model_params is None:
                self.model_params = copy.deepcopy(adapt["pretrained_model_params"])
            self.model_params["sample_size"] += variables_of_interest.shape[0]

        # Data preparation
        # Combination weights
        combination_weights = np.ones(shape=(train_data.shape[0], 2))

        # Data coordinates
        combination_indices = np.arange(train_data.shape[0])
        model_coords = self._build_model_coordinates(
            observations=combination_indices,
        )

        # Fitting logic
        with pm.Model(coords=model_coords) as self._model:
            # Standardize the variable of interest, and store mean and std
            # This is done to ensure that the model is not sensitive to
            # the scale of the variable
            standardized_vois = (
                variables_of_interest - self.model_params["mean_vois"]
            ) / self.model_params["std_vois"]
            variables_of_interest_mu_estimate = train_data[
                [
                    f"{self.spec.variable_of_interest_1}_mu_estimate",
                    f"{self.spec.variable_of_interest_2}_mu_estimate",
                ]
            ].to_numpy()
            standardized_vois_mu_estimate = (
                variables_of_interest_mu_estimate - self.model_params["mean_vois"]
            ) / self.model_params["std_vois"]
            variables_of_interest_std_estimate = train_data[
                [
                    f"{self.spec.variable_of_interest_1}_std_estimate",
                    f"{self.spec.variable_of_interest_2}_std_estimate",
                ]
            ].to_numpy()
            standardized_vois_std_estimate = (
                variables_of_interest_std_estimate / self.model_params["std_vois"]
            )

            # A dictionary for precomputed bspline basis functions
            spline_bases: dict[str, npt.NDArray[np.floating[Any]]] = {}

            # A dictionary for factorized categories
            category_indices: dict[str, npt.NDArray[np.integer[Any]]] = {}

            # Model the covariance between the variables of interest
            z_transformed_correlation_effects = self._model_all_correlation_effects(
                train_data,
                spline_bases,
                category_indices,
                adapt=adapt,
            )

            # Combine all covariance effects
            self._combine_all_correlation_effects(
                z_transformed_correlation_effects=z_transformed_correlation_effects,
                combination_indices=combination_indices,
                combination_weights=combination_weights,
                standardized_vois=standardized_vois,
                standardized_vois_mu_estimate=standardized_vois_mu_estimate,
                standardized_vois_std_estimate=standardized_vois_std_estimate,
            )

            # Fit the model using ADVI
            self._fit_model_with_advi(progress_bar=progress_bar)

        # Save the model if a save path is provided
        if save_directory is not None:
            self.save_model(Path(save_directory))

    def adapt_fit(
        self,
        covariate_to_adapt: str,
        new_category_names: npt.NDArray[np.str_],
        train_data: pd.DataFrame,
        *,
        pretrained_model_params: dict[str, Any] | None = None,
        save_directory: Path | None = None,
        progress_bar: bool = True,
    ) -> None:
        """
        Using a previously fitted model, adapt the model to a new batch.
        This method enables adaptation of the model to data from a new
        batch/site by freezing all fitted parameters, and only estimating
        new parameters for the new batch/site category.

        Args:
            covariate_to_adapt: str
                Name of the categorical covariate representing the batch/site
                to which the model should be adapted.
                Note: This covariate must have been specified in the original
                model.
            new_category_names: list[str]
                Names of the new categories in the covariate_to_adapt representing
                the new batch/site labels (e.g. names of the new site).
                Note: These names must not have been present in the original
                fitted model.
            train_data: pd.DataFrame
                DataFrame containing the training data for adaptation.
                It must include the variable of interest and all specified covariates.
                Note: The covariate_to_adapt column must only contain the
                new_category_names (no new data from previously trained batches).
            pretrained_model_params: dict[str, Any] | None
                The model parameters from a previously fitted model to adapt.
                If None, the model parameters from the current instance will be used
                (assuming fitting was done).
            save_directory: Path | None
                A path to a directory to save the adapted model. If provided,
                the fitted model will be saved to this path.
            progress_bar: bool
                If True, display a progress bar during fitting. Defaults to True.
        """
        # Validation checks
        self._validate_model()
        self._validate_dataframe_for_fitting(train_data)

        # Locate the covariate to adapt
        cov_to_adapt_index = [cov.name for cov in self.spec.covariates].index(
            covariate_to_adapt,
        )

        # Extend the covariate categories to include the new categories
        self.spec.covariates[cov_to_adapt_index].extend_categories(new_category_names)

        # Extract the pre-trained model parameters
        if pretrained_model_params is None:
            if not hasattr(self, "model_params") or self.model_params is None:
                err = (
                    "No pretrained model parameters found. "
                    "Please provide pretrained_model_params or fit the model first."
                )
                raise ValueError(err)
            pretrained_model_params = copy.deepcopy(self.model_params)

        # Fit the adapted model
        self.fit(
            train_data,
            save_directory=save_directory,
            progress_bar=progress_bar,
            adapt={
                "covariate_to_adapt": covariate_to_adapt,
                "new_category_names": new_category_names,
                "pretrained_model_params": pretrained_model_params,
            },
        )

    def predict(
        self,
        test_covariates: pd.DataFrame,
        model_params: dict[str, Any] | None = None,
        predict_without: list[str] | None = None,
    ) -> NormativePredictions:
        """
        Predict correlation for new data (from covariates) using the fitted model.

        Args:
            test_covariates: pd.DataFrame
                DataFrame containing the new covariate data to predict.
                This must include all specified covariates.
                Note: covariates listed in predict_without will be ignored and are
                hence not required.
            model_params: dict | None
                Optional dictionary of model parameters to use. If not provided,
                the stored parameters from model.fit() will be used.
            predict_without: list[str] | None
                Optional list of covariate names to ignore during prediction.
                This can be used to check the effect of removing certain covariates
                from the model.

        Returns:
            NormativePredictions: Object containing the predicted pairwise correlations
                for the variables of interest.
        """
        # Validate the new data
        validation_columns = [
            cov.name
            for cov in self.spec.covariates
            if cov.name not in (predict_without or [])
        ]
        utils.general.validate_dataframe(test_covariates, validation_columns)

        # Parameters
        model_params = model_params or self.model_params
        if model_params is None:
            err = "No model parameters found. Please provide model_params."
            raise ValueError(err)

        # Posterior means
        posterior_means = model_params["posterior_means"]

        # Calculate mean and variance effects
        z_transformed_correlation_estimate = np.zeros(test_covariates.shape[0]) + float(
            posterior_means["global_intercept_z"],
        )

        for cov in self.spec.covariates:
            if (cov.name in self.spec.influencing_covariance) and (
                cov.name not in (predict_without or [])
            ):
                if cov.cov_type == "numerical":
                    if cov.effect == "linear":
                        if cov.moments is None:
                            err = (
                                f"Covariate '{cov.name}' is missing moments for"
                                " standardization."
                            )
                            raise ValueError(err)
                        z_transformed_correlation_estimate += (
                            (
                                cast(
                                    "npt.NDArray[Any]",
                                    test_covariates[cov.name].to_numpy(),
                                )
                                - cov.moments[0]
                            )
                            / cov.moments[1]
                        ) * posterior_means[f"linear_beta_{cov.name}"]
                    elif cov.effect == "spline":
                        spline_bases = cov.make_spline_bases(
                            cast(
                                "npt.NDArray[Any]",
                                test_covariates[cov.name].to_numpy(),
                            ),
                        )
                        spline_betas = posterior_means[f"spline_betas_{cov.name}"]
                        z_transformed_correlation_estimate += (
                            spline_bases @ spline_betas
                        )
                elif cov.cov_type == "categorical":
                    category_indices = cov.factorize_categories(
                        cast("npt.NDArray[Any]", test_covariates[cov.name].to_numpy()),
                    )
                    categorical_intercept = None
                    if cov.hierarchical:
                        categorical_intercept = (
                            posterior_means[f"intercept_offset_{cov.name}"]
                            * posterior_means[f"sigma_intercept_{cov.name}"]
                        )
                    else:
                        categorical_intercept = posterior_means[f"intercept_{cov.name}"]
                    z_transformed_correlation_estimate += categorical_intercept[
                        category_indices
                    ]

        # Convert z-transformed score to correlation
        correlation_estimate = np.tanh(z_transformed_correlation_estimate)

        # Create a the predictions object and return
        return NormativePredictions({"correlation_estimate": correlation_estimate})

adapt_fit(covariate_to_adapt: str, new_category_names: npt.NDArray[np.str_], train_data: pd.DataFrame, *, pretrained_model_params: dict[str, Any] | None = None, save_directory: Path | None = None, progress_bar: bool = True) -> None

Using a previously fitted model, adapt the model to a new batch. This method enables adaptation of the model to data from a new batch/site by freezing all fitted parameters, and only estimating new parameters for the new batch/site category.

Parameters:

Name	Type	Description	Default
covariate_to_adapt	str	str Name of the categorical covariate representing the batch/site to which the model should be adapted. Note: This covariate must have been specified in the original model.	required
new_category_names	NDArray[str_]	list[str] Names of the new categories in the covariate_to_adapt representing the new batch/site labels (e.g. names of the new site). Note: These names must not have been present in the original fitted model.	required
train_data	DataFrame	pd.DataFrame DataFrame containing the training data for adaptation. It must include the variable of interest and all specified covariates. Note: The covariate_to_adapt column must only contain the new_category_names (no new data from previously trained batches).	required
pretrained_model_params	dict[str, Any] | None	dict[str, Any] | None The model parameters from a previously fitted model to adapt. If None, the model parameters from the current instance will be used (assuming fitting was done).	None
save_directory	Path | None	Path | None A path to a directory to save the adapted model. If provided, the fitted model will be saved to this path.	None
progress_bar	bool	bool If True, display a progress bar during fitting. Defaults to True.	True

Source code in src/spectranorm/snm.py

def adapt_fit(
    self,
    covariate_to_adapt: str,
    new_category_names: npt.NDArray[np.str_],
    train_data: pd.DataFrame,
    *,
    pretrained_model_params: dict[str, Any] | None = None,
    save_directory: Path | None = None,
    progress_bar: bool = True,
) -> None:
    """
    Using a previously fitted model, adapt the model to a new batch.
    This method enables adaptation of the model to data from a new
    batch/site by freezing all fitted parameters, and only estimating
    new parameters for the new batch/site category.

    Args:
        covariate_to_adapt: str
            Name of the categorical covariate representing the batch/site
            to which the model should be adapted.
            Note: This covariate must have been specified in the original
            model.
        new_category_names: list[str]
            Names of the new categories in the covariate_to_adapt representing
            the new batch/site labels (e.g. names of the new site).
            Note: These names must not have been present in the original
            fitted model.
        train_data: pd.DataFrame
            DataFrame containing the training data for adaptation.
            It must include the variable of interest and all specified covariates.
            Note: The covariate_to_adapt column must only contain the
            new_category_names (no new data from previously trained batches).
        pretrained_model_params: dict[str, Any] | None
            The model parameters from a previously fitted model to adapt.
            If None, the model parameters from the current instance will be used
            (assuming fitting was done).
        save_directory: Path | None
            A path to a directory to save the adapted model. If provided,
            the fitted model will be saved to this path.
        progress_bar: bool
            If True, display a progress bar during fitting. Defaults to True.
    """
    # Validation checks
    self._validate_model()
    self._validate_dataframe_for_fitting(train_data)

    # Locate the covariate to adapt
    cov_to_adapt_index = [cov.name for cov in self.spec.covariates].index(
        covariate_to_adapt,
    )

    # Extend the covariate categories to include the new categories
    self.spec.covariates[cov_to_adapt_index].extend_categories(new_category_names)

    # Extract the pre-trained model parameters
    if pretrained_model_params is None:
        if not hasattr(self, "model_params") or self.model_params is None:
            err = (
                "No pretrained model parameters found. "
                "Please provide pretrained_model_params or fit the model first."
            )
            raise ValueError(err)
        pretrained_model_params = copy.deepcopy(self.model_params)

    # Fit the adapted model
    self.fit(
        train_data,
        save_directory=save_directory,
        progress_bar=progress_bar,
        adapt={
            "covariate_to_adapt": covariate_to_adapt,
            "new_category_names": new_category_names,
            "pretrained_model_params": pretrained_model_params,
        },
    )

fit(train_data: pd.DataFrame, *, save_directory: Path | None = None, progress_bar: bool = True, adapt: dict[str, Any] | None = None) -> None

Fit the normative model to the training data.

This method implements the fitting logic for the normative model based on the provided training data and model specification.

Parameters:

Name	Type	Description	Default
train_data	DataFrame	pd.DataFrame DataFrame containing the training data. It must include the variable of interest, their predicted moments, and all specified covariates.	required
save_directory	Path | None	Path | None A path to a directory to save the model. If provided, the fitted model will be saved to this path.	None
progress_bar	bool	bool If True, display a progress bar during fitting. Defaults to True.	True
adapt	dict[str, Any] | None	dict[str, Any] | None If provided, adapt a pre-trained model to a new covariate. Note: We recommended using the adapt_fit method, and not directly changing this argument, unless you know what you are doing.	None

Source code in src/spectranorm/snm.py

def fit(
    self,
    train_data: pd.DataFrame,
    *,
    save_directory: Path | None = None,
    progress_bar: bool = True,
    adapt: dict[str, Any] | None = None,
) -> None:
    """
    Fit the normative model to the training data.

    This method implements the fitting logic for the normative model
    based on the provided training data and model specification.

    Args:
        train_data: pd.DataFrame
            DataFrame containing the training data. It must include the variable
            of interest, their predicted moments, and all specified covariates.
        save_directory: Path | None
            A path to a directory to save the model. If provided, the fitted model
            will be saved to this path.
        progress_bar: bool
            If True, display a progress bar during fitting. Defaults to True.
        adapt: dict[str, Any] | None
            If provided, adapt a pre-trained model to a new covariate.
            Note: We recommended using the `adapt_fit` method, and not directly
            changing this argument, unless you know what you are doing.
    """
    # Validation checks
    self._validate_model()
    self._validate_dataframe_for_fitting(train_data)

    # Extract the variables of interest
    variables_of_interest = train_data[
        [self.spec.variable_of_interest_1, self.spec.variable_of_interest_2]
    ].to_numpy()

    # A dictionary to hold the model parameters after fitting
    if adapt is None:
        self.model_params = {}
        self.model_params["mean_vois"] = variables_of_interest.mean(axis=0)
        self.model_params["std_vois"] = variables_of_interest.std(axis=0)
        self.model_params["sample_size"] = variables_of_interest.shape[0]
        # Initialize parameter count
        self.model_params["n_params"] = 0
    else:
        # Update the pretrained model parameters
        if not hasattr(self, "model_params") or self.model_params is None:
            self.model_params = copy.deepcopy(adapt["pretrained_model_params"])
        self.model_params["sample_size"] += variables_of_interest.shape[0]

    # Data preparation
    # Combination weights
    combination_weights = np.ones(shape=(train_data.shape[0], 2))

    # Data coordinates
    combination_indices = np.arange(train_data.shape[0])
    model_coords = self._build_model_coordinates(
        observations=combination_indices,
    )

    # Fitting logic
    with pm.Model(coords=model_coords) as self._model:
        # Standardize the variable of interest, and store mean and std
        # This is done to ensure that the model is not sensitive to
        # the scale of the variable
        standardized_vois = (
            variables_of_interest - self.model_params["mean_vois"]
        ) / self.model_params["std_vois"]
        variables_of_interest_mu_estimate = train_data[
            [
                f"{self.spec.variable_of_interest_1}_mu_estimate",
                f"{self.spec.variable_of_interest_2}_mu_estimate",
            ]
        ].to_numpy()
        standardized_vois_mu_estimate = (
            variables_of_interest_mu_estimate - self.model_params["mean_vois"]
        ) / self.model_params["std_vois"]
        variables_of_interest_std_estimate = train_data[
            [
                f"{self.spec.variable_of_interest_1}_std_estimate",
                f"{self.spec.variable_of_interest_2}_std_estimate",
            ]
        ].to_numpy()
        standardized_vois_std_estimate = (
            variables_of_interest_std_estimate / self.model_params["std_vois"]
        )

        # A dictionary for precomputed bspline basis functions
        spline_bases: dict[str, npt.NDArray[np.floating[Any]]] = {}

        # A dictionary for factorized categories
        category_indices: dict[str, npt.NDArray[np.integer[Any]]] = {}

        # Model the covariance between the variables of interest
        z_transformed_correlation_effects = self._model_all_correlation_effects(
            train_data,
            spline_bases,
            category_indices,
            adapt=adapt,
        )

        # Combine all covariance effects
        self._combine_all_correlation_effects(
            z_transformed_correlation_effects=z_transformed_correlation_effects,
            combination_indices=combination_indices,
            combination_weights=combination_weights,
            standardized_vois=standardized_vois,
            standardized_vois_mu_estimate=standardized_vois_mu_estimate,
            standardized_vois_std_estimate=standardized_vois_std_estimate,
        )

        # Fit the model using ADVI
        self._fit_model_with_advi(progress_bar=progress_bar)

    # Save the model if a save path is provided
    if save_directory is not None:
        self.save_model(Path(save_directory))

from_direct_model(direct_model: DirectNormativeModel, variable_of_interest_1: str, variable_of_interest_2: str, influencing_covariance: list[str] | None = None, defaults_overwrite: dict[str, Any] | None = None) -> CovarianceNormativeModel classmethod

Initialize the model from a direct model instance, and two variable names.

Parameters:

Name	Type	Description	Default
direct_model	DirectNormativeModel	DirectNormativeModel This model will be used to instantiate a similar covariance model.	required
variable_of_interest_1	str	str Name of the first target variable to model.	required
variable_of_interest_2	str	str Name of the second target variable to model.	required
influencing_covariance	list[str] | None	list[str] | None List of covariates that influence the covariance structure. If not provided, this will be copied from the direct model's influencing_variance.	None

Returns:

Type	Description
CovarianceNormativeModel	CovarianceNormativeModel An instance of CovarianceNormativeModel initialized with the provided data.

Source code in src/spectranorm/snm.py

@classmethod
def from_direct_model(
    cls,
    direct_model: DirectNormativeModel,
    variable_of_interest_1: str,
    variable_of_interest_2: str,
    influencing_covariance: list[str] | None = None,
    defaults_overwrite: dict[str, Any] | None = None,
) -> CovarianceNormativeModel:
    """
    Initialize the model from a direct model instance, and two variable names.

    Args:
        direct_model: DirectNormativeModel
            This model will be used to instantiate a similar covariance model.
        variable_of_interest_1: str
            Name of the first target variable to model.
        variable_of_interest_2: str
            Name of the second target variable to model.
        influencing_covariance: list[str] | None
            List of covariates that influence the covariance structure. If not
            provided, this will be copied from the direct model's
            `influencing_variance`.

    Returns:
        CovarianceNormativeModel
            An instance of CovarianceNormativeModel initialized with the provided
            data.
    """
    # Validity checks for input parameters
    if not isinstance(direct_model, DirectNormativeModel):
        err = "direct_model must be an instance of DirectNormativeModel."
        raise TypeError(err)
    if not (
        isinstance(variable_of_interest_1, str)
        and isinstance(variable_of_interest_2, str)
    ):
        err = "Variables of interest must be strings."
        raise TypeError(err)

    # Substitute influencing_covariance if not provided
    if influencing_covariance is None:
        influencing_covariance = direct_model.spec.influencing_variance

    # Use the same setup as the direct model
    model = cls(
        spec=CovarianceModelSpec(
            variable_of_interest_1=variable_of_interest_1,
            variable_of_interest_2=variable_of_interest_2,
            covariates=direct_model.spec.covariates,
            influencing_covariance=influencing_covariance,
        ),
    )

    # update defaults
    model.defaults.update(direct_model.defaults)
    model.defaults.update(defaults_overwrite or {})

    return model

load_model(directory: Path, *, load_posterior: bool = False) -> CovarianceNormativeModel classmethod

Load the model and its posterior from a directory. The model will be loaded from a subdirectory named 'saved_model'.

Parameters:

Name	Type	Description	Default
directory	Path	Path Path to the directory containing the model.	required
load_posterior	bool	bool (default=False) If True, load the model's posterior trace from the saved inference data.	False

Source code in src/spectranorm/snm.py

@classmethod
def load_model(
    cls,
    directory: Path,
    *,
    load_posterior: bool = False,
) -> CovarianceNormativeModel:
    """
    Load the model and its posterior from a directory.
    The model will be loaded from a subdirectory named 'saved_model'.

    Args:
        directory: Path
            Path to the directory containing the model.
        load_posterior: bool (default=False)
            If True, load the model's posterior trace from the saved inference data.
    """
    # Validate the load directory
    directory = Path(directory)
    saved_model_dir = utils.general.validate_load_directory(
        directory,
        "saved_model",
    )

    # Load the saved model dict
    model_dict = joblib.load(saved_model_dir / "model_dict.joblib")

    # Create an instance of the class
    instance = cls(
        spec=model_dict["spec"],
    )

    # Set the attributes from the loaded model dictionary
    instance.defaults.update(model_dict["defaults"])
    if "model_params" in model_dict:
        instance.model_params = model_dict["model_params"]
        if load_posterior:
            instance.model_inference_data = az.from_netcdf(  # type: ignore[no-untyped-call]
                saved_model_dir / "model_inference_data.nc",
            )

    return instance

predict(test_covariates: pd.DataFrame, model_params: dict[str, Any] | None = None, predict_without: list[str] | None = None) -> NormativePredictions

Predict correlation for new data (from covariates) using the fitted model.

Parameters:

Name	Type	Description	Default
test_covariates	DataFrame	pd.DataFrame DataFrame containing the new covariate data to predict. This must include all specified covariates. Note: covariates listed in predict_without will be ignored and are hence not required.	required
model_params	dict[str, Any] | None	dict | None Optional dictionary of model parameters to use. If not provided, the stored parameters from model.fit() will be used.	None
predict_without	list[str] | None	list[str] | None Optional list of covariate names to ignore during prediction. This can be used to check the effect of removing certain covariates from the model.	None

Returns:

Name	Type	Description
NormativePredictions	NormativePredictions	Object containing the predicted pairwise correlations for the variables of interest.

Source code in src/spectranorm/snm.py

def predict(
    self,
    test_covariates: pd.DataFrame,
    model_params: dict[str, Any] | None = None,
    predict_without: list[str] | None = None,
) -> NormativePredictions:
    """
    Predict correlation for new data (from covariates) using the fitted model.

    Args:
        test_covariates: pd.DataFrame
            DataFrame containing the new covariate data to predict.
            This must include all specified covariates.
            Note: covariates listed in predict_without will be ignored and are
            hence not required.
        model_params: dict | None
            Optional dictionary of model parameters to use. If not provided,
            the stored parameters from model.fit() will be used.
        predict_without: list[str] | None
            Optional list of covariate names to ignore during prediction.
            This can be used to check the effect of removing certain covariates
            from the model.

    Returns:
        NormativePredictions: Object containing the predicted pairwise correlations
            for the variables of interest.
    """
    # Validate the new data
    validation_columns = [
        cov.name
        for cov in self.spec.covariates
        if cov.name not in (predict_without or [])
    ]
    utils.general.validate_dataframe(test_covariates, validation_columns)

    # Parameters
    model_params = model_params or self.model_params
    if model_params is None:
        err = "No model parameters found. Please provide model_params."
        raise ValueError(err)

    # Posterior means
    posterior_means = model_params["posterior_means"]

    # Calculate mean and variance effects
    z_transformed_correlation_estimate = np.zeros(test_covariates.shape[0]) + float(
        posterior_means["global_intercept_z"],
    )

    for cov in self.spec.covariates:
        if (cov.name in self.spec.influencing_covariance) and (
            cov.name not in (predict_without or [])
        ):
            if cov.cov_type == "numerical":
                if cov.effect == "linear":
                    if cov.moments is None:
                        err = (
                            f"Covariate '{cov.name}' is missing moments for"
                            " standardization."
                        )
                        raise ValueError(err)
                    z_transformed_correlation_estimate += (
                        (
                            cast(
                                "npt.NDArray[Any]",
                                test_covariates[cov.name].to_numpy(),
                            )
                            - cov.moments[0]
                        )
                        / cov.moments[1]
                    ) * posterior_means[f"linear_beta_{cov.name}"]
                elif cov.effect == "spline":
                    spline_bases = cov.make_spline_bases(
                        cast(
                            "npt.NDArray[Any]",
                            test_covariates[cov.name].to_numpy(),
                        ),
                    )
                    spline_betas = posterior_means[f"spline_betas_{cov.name}"]
                    z_transformed_correlation_estimate += (
                        spline_bases @ spline_betas
                    )
            elif cov.cov_type == "categorical":
                category_indices = cov.factorize_categories(
                    cast("npt.NDArray[Any]", test_covariates[cov.name].to_numpy()),
                )
                categorical_intercept = None
                if cov.hierarchical:
                    categorical_intercept = (
                        posterior_means[f"intercept_offset_{cov.name}"]
                        * posterior_means[f"sigma_intercept_{cov.name}"]
                    )
                else:
                    categorical_intercept = posterior_means[f"intercept_{cov.name}"]
                z_transformed_correlation_estimate += categorical_intercept[
                    category_indices
                ]

    # Convert z-transformed score to correlation
    correlation_estimate = np.tanh(z_transformed_correlation_estimate)

    # Create a the predictions object and return
    return NormativePredictions({"correlation_estimate": correlation_estimate})

save_model(directory: Path, *, save_posterior: bool = False) -> None

Save the fitted model and it's posterior to a directory. The model will be saved in a subdirectory named 'saved_model'. If this directory is not empty, an error is raised.

Parameters:

Name	Type	Description	Default
directory	Path	Path Path to a directory to save the model.	required
save_posterior	bool	bool (default=False) If True, save the model's posterior trace inference data.	False

Source code in src/spectranorm/snm.py

def save_model(self, directory: Path, *, save_posterior: bool = False) -> None:
    """
    Save the fitted model and it's posterior to a directory.
    The model will be saved in a subdirectory named 'saved_model'.
    If this directory is not empty, an error is raised.

    Args:
        directory: Path
            Path to a directory to save the model.
        save_posterior: bool (default=False)
            If True, save the model's posterior trace inference data.
    """
    # Prepare the save directory
    directory = Path(directory)
    saved_model_dir = utils.general.prepare_save_directory(directory, "saved_model")

    model_dict = {
        "spec": self.spec,
        "defaults": self.defaults,
    }
    if hasattr(self, "model_params"):
        model_dict["model_params"] = self.model_params
        if hasattr(self, "model_inference_data") and save_posterior:
            self.model_inference_data.to_netcdf(
                saved_model_dir / "model_inference_data.nc",
            )
    joblib.dump(model_dict, saved_model_dir / "model_dict.joblib")

CovariateSpec dataclass

Specification of a single covariate and how it should be modeled.

Attributes:

Name	Type	Description
name	str	str Name of the covariate (e.g., 'age', 'site').
cov_type	CovariateType	str Type of the covariate ('numerical' or 'categorical').
effect	NumericalEffect | None	str For numerical covariates, how the effect is modeled ('linear' or 'spline').
categories	NDArray[str_] | None	np.ndarray | None For categorical covariates, the category labels stored as a NumPy array.
hierarchical	bool | None	bool For categorical covariates, whether to model with a hierarchical structure.
spline_spec	SplineSpec | None	SplineSpec | None Optional SplineSpec instance for spline modeling; required if effect is 'spline'.

Validation

• Numerical covariates must specify 'effect'.
• If 'effect' is 'spline', 'spline_spec' must be provided.
• Categorical covariates must specify 'hierarchical'.
• Categorical covariates cannot have 'effect' or 'spline_spec'.
• Categorical covariates must have categories listed.

Source code in src/spectranorm/snm.py

@dataclass
class CovariateSpec:
    """
    Specification of a single covariate and how it should be modeled.

    Attributes:
        name: str
            Name of the covariate (e.g., 'age', 'site').
        cov_type: str
            Type of the covariate ('numerical' or 'categorical').
        effect: str
            For numerical covariates, how the effect is modeled ('linear'
            or 'spline').
        categories: np.ndarray | None
            For categorical covariates, the category labels stored as a NumPy array.
        hierarchical: bool
            For categorical covariates, whether to model with a
            hierarchical structure.
        spline_spec: SplineSpec | None
            Optional SplineSpec instance for spline modeling;
            required if effect is 'spline'.

    Validation:
        - Numerical covariates must specify 'effect'.
        - If 'effect' is 'spline', 'spline_spec' must be provided.
        - Categorical covariates must specify 'hierarchical'.
        - Categorical covariates cannot have 'effect' or 'spline_spec'.
        - Categorical covariates must have categories listed.
    """

    name: str
    cov_type: CovariateType  # "categorical" or "numerical"
    effect: NumericalEffect | None = None  # Only if numerical
    categories: npt.NDArray[np.str_] | None = None  # Only if categorical
    hierarchical: bool | None = None  # Only if categorical
    spline_spec: SplineSpec | None = None  # Only for spline modeling
    moments: tuple[float, float] | None = None  # Only for linear effects

    def __repr__(self) -> str:
        """
        String representation of the CovariateSpec instance.
        """
        representation = f"CovariateSpec(name={self.name}, cov_type={self.cov_type}"
        if self.cov_type == "numerical":
            representation += f", effect={self.effect}"
        elif self.cov_type == "categorical":
            representation += f", hierarchical={self.hierarchical}"
            if self.categories is not None:
                representation += f", n_categories={len(self.categories.tolist())}"
        representation += ")"
        return representation

    # Validation checks for the covariate specification.
    def validate_numerical(self) -> None:
        if self.effect not in {"linear", "spline"}:
            err = (
                f"Numerical covariate '{self.name}' must specify effect as "
                "'linear' or 'spline'."
            )
            raise ValueError(err)
        if self.hierarchical is not None:
            err = (
                f"Numerical covariate '{self.name}' should not specify 'hierarchical'."
            )
            raise ValueError(err)
        if self.categories is not None:
            err = f"Numerical covariate '{self.name}' should not specify 'categories'."
            raise ValueError(err)
        if self.effect == "spline":
            if self.spline_spec is None:
                err = (
                    f"Numerical covariate '{self.name}' must have spline "
                    "specification if effect is 'spline'."
                )
                raise ValueError(err)
            if self.moments is not None:
                err = (
                    f"Numerical covariate '{self.name}' should not specify "
                    "moments if effect is 'spline'."
                )
                raise ValueError(err)
        if self.effect == "linear":
            if self.spline_spec is not None:
                err = (
                    f"Numerical covariate '{self.name}' should not have spline "
                    "specification unless effect is 'spline'."
                )
                raise ValueError(err)
            if self.moments is None:
                err = (
                    f"Numerical covariate '{self.name}' must specify moments "
                    "(mean and standard deviation) for linear effects."
                )
                raise ValueError(err)

    def validate_categorical(self) -> None:
        if self.effect is not None:
            err = (
                f"Categorical covariate '{self.name}' should not have a "
                "numerical effect type."
            )
            raise ValueError(err)
        if self.spline_spec is not None:
            err = (
                f"Categorical covariate '{self.name}' should not have spline "
                "specification."
            )
            raise ValueError(err)
        if self.hierarchical is None:
            err = (
                f"Categorical covariate '{self.name}' must specify whether "
                "it is hierarchical."
            )
            raise ValueError(err)
        if self.categories is None:
            err = f"Categorical covariate '{self.name}' must specify categories."
            raise ValueError(err)
        if not isinstance(self.categories, np.ndarray):
            err = (
                f"Categorical covariate '{self.name}' must specify categories "
                "as a NumPy array."
            )
            raise TypeError(err)

    def __post_init__(self) -> None:
        if self.cov_type == "numerical":
            self.validate_numerical()
        elif self.cov_type == "categorical":
            self.validate_categorical()
        else:
            err = f"Invalid covariate type '{self.cov_type}' for '{self.name}'."
            raise ValueError(err)

    def make_spline_bases(
        self,
        values: npt.NDArray[np.floating[Any]],
        *,
        include_intercept: bool = True,
    ) -> npt.NDArray[np.floating[Any]]:
        """
        Create B-spline basis expansion functions for a given covariate.

        Args:
            values (np.ndarray): The values to create the spline basis functions for.

        Returns:
            np.ndarray: The B-spline basis function expansion.
        """
        if self.effect != "spline" or self.spline_spec is None:
            err = f"Covariate '{self.name}' is not a spline covariate."
            raise ValueError(err)

        # Create fixed set of knots if not already defined
        if self.spline_spec.knots is None:
            knots = np.linspace(
                self.spline_spec.lower_bound,
                self.spline_spec.upper_bound,
                self.spline_spec.df - self.spline_spec.degree + 1,
            )[1:-1].tolist()
        else:
            knots = self.spline_spec.knots

        # Create B-spline basis functions
        return np.array(
            patsy.bs(  # pyright: ignore[reportAttributeAccessIssue]
                values,
                knots=knots,
                df=self.spline_spec.df,
                degree=self.spline_spec.degree,
                lower_bound=self.spline_spec.lower_bound,
                upper_bound=self.spline_spec.upper_bound,
                include_intercept=include_intercept,
            ),
        )

    def factorize_categories(
        self,
        values: npt.NDArray[np.str_],
    ) -> npt.NDArray[np.int_]:
        """
        Factorize categorical covariate values into numerical indices.

        Args:
            values (np.ndarray): The values to factorize.

        Returns:
            np.ndarray: The factorized numerical indices for the categories.
        """
        if self.cov_type != "categorical":
            err = (
                f"Covariate '{self.name}' is not a categorical "
                "covariate to be factorized."
            )
            raise ValueError(err)

        # Create a mapping from category values to indices
        if self.categories is None:  # to satisfy type checker
            err = f"Covariate '{self.name}' does not have categories defined."
            raise ValueError(err)
        category_mapping: dict[str, int] = {
            category: idx for idx, category in enumerate(self.categories)
        }
        # Factorize the values using the mapping
        return np.array([category_mapping[val] for val in values], dtype=int)

    def extend_categories(
        self,
        new_categories: npt.NDArray[np.str_],
    ) -> None:
        """
        Extend the categories of a categorical covariate with new categories.

        Args:
            new_categories (np.ndarray): The new categories to add.

        Returns:
            None
        """
        if self.cov_type != "categorical":
            err = f"Covariate '{self.name}' is not a categorical covariate to extend."
            raise ValueError(err)
        if self.categories is None:  # to satisfy type checker
            err = f"Covariate '{self.name}' does not have categories defined."
            raise ValueError(err)

        # Make sure categories are unique and new
        unique_new_categories = np.setdiff1d(new_categories, self.categories)
        if unique_new_categories.size < new_categories.size:
            err = (
                f"Some new categories are already present in the "
                f"covariate '{self.name}'."
            )
            raise ValueError(err)

        # Extend the categories array
        self.categories = np.concatenate((self.categories, unique_new_categories))

extend_categories(new_categories: npt.NDArray[np.str_]) -> None

Extend the categories of a categorical covariate with new categories.

Parameters:

Name	Type	Description	Default
new_categories	ndarray	The new categories to add.	required

Returns:

Type	Description
None	None

Source code in src/spectranorm/snm.py

def extend_categories(
    self,
    new_categories: npt.NDArray[np.str_],
) -> None:
    """
    Extend the categories of a categorical covariate with new categories.

    Args:
        new_categories (np.ndarray): The new categories to add.

    Returns:
        None
    """
    if self.cov_type != "categorical":
        err = f"Covariate '{self.name}' is not a categorical covariate to extend."
        raise ValueError(err)
    if self.categories is None:  # to satisfy type checker
        err = f"Covariate '{self.name}' does not have categories defined."
        raise ValueError(err)

    # Make sure categories are unique and new
    unique_new_categories = np.setdiff1d(new_categories, self.categories)
    if unique_new_categories.size < new_categories.size:
        err = (
            f"Some new categories are already present in the "
            f"covariate '{self.name}'."
        )
        raise ValueError(err)

    # Extend the categories array
    self.categories = np.concatenate((self.categories, unique_new_categories))

factorize_categories(values: npt.NDArray[np.str_]) -> npt.NDArray[np.int_]

Factorize categorical covariate values into numerical indices.

Parameters:

Name	Type	Description	Default
values	ndarray	The values to factorize.	required

Returns:

Type	Description
NDArray[int_]	np.ndarray: The factorized numerical indices for the categories.

Source code in src/spectranorm/snm.py

def factorize_categories(
    self,
    values: npt.NDArray[np.str_],
) -> npt.NDArray[np.int_]:
    """
    Factorize categorical covariate values into numerical indices.

    Args:
        values (np.ndarray): The values to factorize.

    Returns:
        np.ndarray: The factorized numerical indices for the categories.
    """
    if self.cov_type != "categorical":
        err = (
            f"Covariate '{self.name}' is not a categorical "
            "covariate to be factorized."
        )
        raise ValueError(err)

    # Create a mapping from category values to indices
    if self.categories is None:  # to satisfy type checker
        err = f"Covariate '{self.name}' does not have categories defined."
        raise ValueError(err)
    category_mapping: dict[str, int] = {
        category: idx for idx, category in enumerate(self.categories)
    }
    # Factorize the values using the mapping
    return np.array([category_mapping[val] for val in values], dtype=int)

make_spline_bases(values: npt.NDArray[np.floating[Any]], *, include_intercept: bool = True) -> npt.NDArray[np.floating[Any]]

Create B-spline basis expansion functions for a given covariate.

Parameters:

Name	Type	Description	Default
values	ndarray	The values to create the spline basis functions for.	required

Returns:

Type	Description
NDArray[floating[Any]]	np.ndarray: The B-spline basis function expansion.

Source code in src/spectranorm/snm.py

def make_spline_bases(
    self,
    values: npt.NDArray[np.floating[Any]],
    *,
    include_intercept: bool = True,
) -> npt.NDArray[np.floating[Any]]:
    """
    Create B-spline basis expansion functions for a given covariate.

    Args:
        values (np.ndarray): The values to create the spline basis functions for.

    Returns:
        np.ndarray: The B-spline basis function expansion.
    """
    if self.effect != "spline" or self.spline_spec is None:
        err = f"Covariate '{self.name}' is not a spline covariate."
        raise ValueError(err)

    # Create fixed set of knots if not already defined
    if self.spline_spec.knots is None:
        knots = np.linspace(
            self.spline_spec.lower_bound,
            self.spline_spec.upper_bound,
            self.spline_spec.df - self.spline_spec.degree + 1,
        )[1:-1].tolist()
    else:
        knots = self.spline_spec.knots

    # Create B-spline basis functions
    return np.array(
        patsy.bs(  # pyright: ignore[reportAttributeAccessIssue]
            values,
            knots=knots,
            df=self.spline_spec.df,
            degree=self.spline_spec.degree,
            lower_bound=self.spline_spec.lower_bound,
            upper_bound=self.spline_spec.upper_bound,
            include_intercept=include_intercept,
        ),
    )

DirectNormativeModel dataclass

Direct normative model implementation.

This class implements the direct normative modeling approach, which directly models the variable of interest using the specified covariates. It can be used to fit a model to data and predict normative centiles.

Attributes:

Name	Type	Description
spec	NormativeModelSpec	NormativeModelSpec Specification of the normative model including variable of interest, covariates, and data source.
defaults	dict[str, Any]	dict Default parameters for the model, including spline specifications, ADVI iterations, convergence tolerance, random seed, and Adam optimizer learning rates.

Source code in src/spectranorm/snm.py

@dataclass
class DirectNormativeModel:
    """
    Direct normative model implementation.

    This class implements the direct normative modeling approach, which
    directly models the variable of interest using the specified covariates.
    It can be used to fit a model to data and predict normative centiles.

    Attributes:
        spec: NormativeModelSpec
            Specification of the normative model including variable of interest,
            covariates, and data source.
        defaults: dict
            Default parameters for the model, including spline specifications,
            ADVI iterations, convergence tolerance, random seed, and Adam optimizer
            learning rates.
    """

    spec: NormativeModelSpec
    defaults: dict[str, Any] = field(
        default_factory=lambda: {
            "spline_df": DEFAULT_SPLINE_DF,
            "spline_degree": DEFAULT_SPLINE_DEGREE,
            "spline_extrapolation_factor": DEFAULT_SPLINE_EXTRAPOLATION_FACTOR,
            "advi_iterations": DEFAULT_ADVI_ITERATIONS,
            "advi_convergence_tolerance": DEFAULT_ADVI_CONVERGENCE_TOLERANCE,
            "random_seed": DEFAULT_RANDOM_SEED,
            "adam_learning_rate": DEFAULT_ADAM_LEARNING_RATE,
            "adam_learning_rate_decay": DEFAULT_ADAM_LEARNING_RATE_DECAY,
        },
    )

    def __repr__(self) -> str:
        """
        String representation of the DirectNormativeModel instance.
        """
        return f"DirectNormativeModel(spec={self.spec})"

    @staticmethod
    def _validate_init_args(
        model_type: ModelType,
        variable_of_interest: str,
        numerical_covariates: list[str],
        categorical_covariates: list[str],
        batch_covariates: list[str],
        nonlinear_covariates: list[str],
    ) -> None:
        # Validity checks for input parameters
        if model_type not in {"HBR", "BLR"}:
            err = f"Invalid model type '{model_type}'. Must be 'HBR' or 'BLR'."
            raise ValueError(err)
        for list_name, covariate_list in [
            ("numerical", numerical_covariates),
            ("categorical", categorical_covariates),
            ("batch", batch_covariates),
            ("nonlinear", nonlinear_covariates),
        ]:
            if not all(isinstance(item, str) for item in covariate_list):
                err = f"All covariate names must be strings: {list_name} covariates."
                raise TypeError(err)
        if not isinstance(variable_of_interest, str):
            err = "Variable of interest must be a string."
            raise TypeError(err)
        if not all(col in categorical_covariates for col in batch_covariates):
            err = "All batch covariates must be included in categorical covariates."
            raise ValueError(err)
        if not all(col in numerical_covariates for col in nonlinear_covariates):
            err = "All nonlinear covariates must be included in numerical covariates."
            raise ValueError(err)

    @classmethod
    def from_dataframe(
        cls,
        model_type: ModelType,
        dataframe: pd.DataFrame,
        variable_of_interest: str,
        numerical_covariates: list[str] | None = None,
        categorical_covariates: list[str] | None = None,
        batch_covariates: list[str] | None = None,
        nonlinear_covariates: list[str] | None = None,
        influencing_mean: list[str] | None = None,
        influencing_variance: list[str] | None = None,
        spline_kwargs: dict[str, Any] | None = None,
    ) -> DirectNormativeModel:
        """
        Initialize a normative model from a pandas DataFrame.

        Args:
            model_type: ModelType
                Type of the model to create, either "HBR" (Hierarchical Bayesian
                Regression) or "BLR" (Bayesian Linear Regression).
            dataframe: pd.DataFrame
                DataFrame containing the data.
            variable_of_interest: str
                Name of the target variable to model.
            numerical_covariates: list[str] | None
                List of numerical covariate names.
            categorical_covariates: list[str] | None
                List of categorical covariate names.
            batch_covariates: list[str] | None
                List of batch covariate names which should also be included in
                categorical_covariates.
            nonlinear_covariates: list[str] | None
                List of covariate names to be modeled as nonlinear effects.
                These should also be included in numerical_covariates.
            influencing_mean: list[str] | None
                List of covariate names that influence the mean of the variable
                of interest. These should be included in either numerical_covariates
                or categorical_covariates.
            influencing_variance: list[str] | None
                List of covariate names that influence the variance of the variable
                of interest. These should be included in either numerical_covariates
                or categorical_covariates.
            spline_kwargs: dict
                Additional keyword arguments for spline specification, such as
                `df`, `degree`, and `knots`. These are passed to the
                `create_spline_spec` method to create spline specifications for
                nonlinear covariates.

        Returns:
            DirectNormativeModel
                An instance of DirectNormativeModel initialized with the provided data.
        """
        # Set default values for optional parameters
        numerical_covariates = numerical_covariates or []
        categorical_covariates = categorical_covariates or []
        batch_covariates = batch_covariates or []
        nonlinear_covariates = nonlinear_covariates or []
        influencing_mean = influencing_mean or []
        influencing_variance = influencing_variance or []
        spline_kwargs = spline_kwargs or {}

        # Validity checks for input parameters
        cls._validate_init_args(
            model_type,
            variable_of_interest,
            numerical_covariates,
            categorical_covariates,
            batch_covariates,
            nonlinear_covariates,
        )
        utils.general.validate_dataframe(
            dataframe,
            [variable_of_interest, *numerical_covariates, *categorical_covariates],
        )

        # Create an instance of the class
        self = cls(
            spec=NormativeModelSpec(
                variable_of_interest=variable_of_interest,
                covariates=[],
                influencing_mean=influencing_mean,
                influencing_variance=influencing_variance,
            ),
        )

        # Populate the spline_kwargs with defaults if not provided
        spline_kwargs["df"] = spline_kwargs.get("df", self.defaults["spline_df"])
        spline_kwargs["degree"] = spline_kwargs.get(
            "degree",
            self.defaults["spline_degree"],
        )
        spline_kwargs["extrapolation_factor"] = spline_kwargs.get(
            "extrapolation_factor",
            self.defaults["spline_extrapolation_factor"],
        )

        # Start building the model specification
        # Add categorical covariates
        for cov_name in categorical_covariates:
            hierarchical = False
            if cov_name in batch_covariates and model_type == "HBR":
                hierarchical = True
            self.spec.covariates.append(
                CovariateSpec(
                    name=cov_name,
                    cov_type="categorical",
                    categories=dataframe[cov_name].unique(),
                    hierarchical=hierarchical,
                ),
            )
        for cov_name in numerical_covariates:
            if cov_name not in nonlinear_covariates:
                self.spec.covariates.append(
                    CovariateSpec(
                        name=cov_name,
                        cov_type="numerical",
                        effect="linear",
                        moments=(
                            dataframe[cov_name].mean(),
                            dataframe[cov_name].std(),
                        ),
                    ),
                )
            else:
                self.spec.covariates.append(
                    CovariateSpec(
                        name=cov_name,
                        cov_type="numerical",
                        effect="spline",
                        spline_spec=SplineSpec.create_spline_spec(
                            dataframe[cov_name],
                            **spline_kwargs,
                        ),
                    ),
                )
        return self

    def _validate_model(self) -> None:
        """
        Validate the model instance.

        This method checks if the model instance is complete and valid.
        It raises errors if any required fields are missing or if there are
        inconsistencies in the model specification.
        """
        if self.spec is None:
            err = (
                "Model specification is not set. "
                "Please initialize the model, e.g., with 'from_dataframe'."
            )
            raise ValueError(err)
        if len(self.spec.covariates) == 0:
            err = (
                "No covariates specified in the model. "
                "Please add covariates to the specification."
            )
            raise ValueError(err)
        if (len(self.spec.influencing_mean) == 0) and (
            len(self.spec.influencing_variance) == 0
        ):
            err = (
                "No covariates specified to influence the mean or "
                "variance of the variable of interest."
            )
            raise ValueError(err)

    def save_model(self, directory: Path, *, save_posterior: bool = False) -> None:
        """
        Save the fitted model and it's posterior to a directory.
        The model will be saved in a subdirectory named 'saved_model'.
        If this directory is not empty, an error is raised.

        Args:
            directory: Path
                Path to a directory to save the model.
            save_posterior: bool (default=False)
                If True, save the model's posterior trace inference data.
        """
        # Prepare the save directory
        directory = Path(directory)
        saved_model_dir = utils.general.prepare_save_directory(directory, "saved_model")

        model_dict = {
            "spec": self.spec,
            "defaults": self.defaults,
        }
        if hasattr(self, "model_params"):
            model_dict["model_params"] = self.model_params
            if hasattr(self, "model_inference_data") and save_posterior:
                self.model_inference_data.to_netcdf(
                    saved_model_dir / "model_inference_data.nc",
                )
        joblib.dump(model_dict, saved_model_dir / "model_dict.joblib")

    @classmethod
    def load_model(
        cls,
        directory: Path,
        *,
        load_posterior: bool = False,
    ) -> DirectNormativeModel:
        """
        Load the model and its posterior from a directory.
        The model will be loaded from a subdirectory named 'saved_model'.

        Args:
            directory: Path
                Path to the directory containing the model.
            load_posterior: bool (default=False)
                If True, load the model's posterior trace from the saved inference data.
        """
        # Validate the load directory
        directory = Path(directory)
        saved_model_dir = utils.general.validate_load_directory(
            directory,
            "saved_model",
        )

        # Load the saved model dict
        model_dict = joblib.load(saved_model_dir / "model_dict.joblib")

        # Create an instance of the class
        instance = cls(
            spec=model_dict["spec"],
        )

        # Set the attributes from the loaded model dictionary
        instance.defaults.update(model_dict["defaults"])
        if "model_params" in model_dict:
            instance.model_params = model_dict["model_params"]
            if load_posterior:
                instance.model_inference_data = az.from_netcdf(  # type: ignore[no-untyped-call]
                    saved_model_dir / "model_inference_data.nc",
                )

        return instance

    def _validate_dataframe_for_fitting(self, train_data: pd.DataFrame) -> None:
        """
        Validate the training DataFrame for fitting.
        """
        utils.general.validate_dataframe(
            train_data,
            (
                [cov.name for cov in self.spec.covariates]
                + [self.spec.variable_of_interest]
            ),
        )

    def _build_model_coordinates(
        self,
        observations: npt.NDArray[np.integer[Any]],
    ) -> dict[str, Any]:
        """
        Build the model coordinates for the training DataFrame.
        """
        # Data coordinates
        model_coords = {"observations": observations, "scalar": [0]}

        # Additional coordinates for covariates
        for cov in self.spec.covariates:
            if cov.cov_type == "numerical":
                if cov.effect == "spline":
                    if cov.spline_spec is not None:  # to satisfy type checker
                        model_coords[f"{cov.name}_splines"] = np.arange(
                            cov.spline_spec.df,
                        )
                elif cov.effect == "linear":
                    model_coords[f"{cov.name}_linear"] = np.arange(1)
            elif cov.cov_type == "categorical":
                model_coords[cov.name] = cov.categories
            else:
                err = f"Invalid covariate type '{cov.cov_type}' for '{cov.name}'."
                raise ValueError(err)
        return model_coords

    def _model_linear_mean_effect(
        self,
        train_data: pd.DataFrame,
        cov: CovariateSpec,
        effects_list: list[TensorVariable],
        sigma_prior: float = 10,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Model a linear effect for a numerical covariate on the mean estimate.
        """
        # Linear effect
        if adapt is None:  # Model fitting
            linear_beta = pm.Normal(
                f"linear_beta_{cov.name}",
                mu=0,
                sigma=sigma_prior,
                size=1,
                dims=(f"{cov.name}_linear",),
            )
            # Increment parameter count for linear effect
            self.model_params["n_params"] += 1
        else:  # Freeze during adaptation/fine-tuning
            linear_beta = pm.Deterministic(
                f"linear_beta_{cov.name}",
                pt.as_tensor_variable(
                    adapt["pretrained_model_params"]["posterior_means"][
                        f"linear_beta_{cov.name}"
                    ],
                ),
                dims=(f"{cov.name}_linear",),
            )
        if cov.moments is not None:  # to satisfy type checker
            effects_list.append(
                (
                    cast("npt.NDArray[Any]", train_data[cov.name].to_numpy())
                    - cov.moments[0]
                )
                / cov.moments[1]
                * linear_beta,
            )

    def _model_spline_mean_effect(
        self,
        train_data: pd.DataFrame,
        cov: CovariateSpec,
        effects_list: list[TensorVariable],
        spline_bases: dict[str, npt.NDArray[np.floating[Any]]],
        sigma_prior: float = 10,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Model a spline effect for a numerical covariate on the mean estimate.
        """
        # Spline effect
        spline_bases[cov.name] = spline_bases.get(
            cov.name,
            cov.make_spline_bases(
                cast("npt.NDArray[Any]", train_data[cov.name].to_numpy()),
            ),
        )
        if adapt is None:  # Model fitting
            # ZeroSumNormal imposes a centering constraint ensuring identifiability
            spline_betas = pm.ZeroSumNormal(
                f"spline_betas_{cov.name}",
                sigma=sigma_prior,
                shape=spline_bases[cov.name].shape[1],
                dims=(f"{cov.name}_splines",),
            )
            # Increment parameter count for spline effects
            if cov.spline_spec is not None:  # to satisfy type checker
                self.model_params["n_params"] += cov.spline_spec.df - 1
        else:  # Freeze during adaptation/fine-tuning
            spline_betas = pm.Deterministic(
                f"spline_betas_{cov.name}",
                pt.as_tensor_variable(
                    adapt["pretrained_model_params"]["posterior_means"][
                        f"spline_betas_{cov.name}"
                    ],
                ),
                dims=(f"{cov.name}_splines",),
            )
        effects_list.append(pt.dot(spline_bases[cov.name], spline_betas.T))

    def _model_categorical_mean_effect(
        self,
        train_data: pd.DataFrame,
        cov: CovariateSpec,
        effects_list: list[TensorVariable],
        category_indices: dict[str, npt.NDArray[np.integer[Any]]],
        sigma_prior: float = 10,
        hierarchical_sigma_prior: float = 1,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Model the effect of a categorical covariate on the mean estimate.
        """
        # Factorize categories
        category_indices[cov.name] = category_indices.get(
            cov.name,
            cov.factorize_categories(
                cast("npt.NDArray[Any]", train_data[cov.name].to_numpy()),
            ),
        )
        if adapt is None:  # Model fitting
            if cov.hierarchical:
                # Hierarchical categorical effect
                # Hyperpriors for category (Bayesian equivalent of random effects)
                sigma_intercept_category = pm.HalfNormal(
                    f"sigma_intercept_{cov.name}",
                    sigma=sigma_prior,
                    dims=("scalar",),
                )

                # Hierarchical intercepts for each category (using reparameterized form)
                categorical_intercept_offset = pm.ZeroSumNormal(
                    f"intercept_offset_{cov.name}",
                    sigma=hierarchical_sigma_prior,
                    dims=(cov.name,),
                )
                # Note ZeroSumNormal imposes a centering constraint
                # (ensuring identifiability)
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    (
                        categorical_intercept_offset
                        * pt.reshape(sigma_intercept_category, (1,))  # pyright: ignore[reportPrivateImportUsage]
                    ),
                    dims=(cov.name,),
                )

                # Increment parameter count for hierarchical intercept
                self.model_params["n_params"] += 1

            else:
                # Non-hierarchical (linear) categorical effect
                categorical_intercept = pm.ZeroSumNormal(
                    f"intercept_{cov.name}",
                    sigma=sigma_prior,
                    dims=(cov.name,),
                )
                # Note ZeroSumNormal imposes a centering constraint
                # (ensuring identifiability)
            # Increment parameter count for categorical effects
            if cov.categories is not None:  # to satisfy type checker
                self.model_params["n_params"] += len(cov.categories) - 1
        elif cov.name != adapt["covariate_to_adapt"]:
            # Freeze during adaptation/fine-tuning
            if cov.hierarchical:
                # Hierarchical categorical effect
                # Hyperpriors for category (Bayesian equivalent of random effects)
                sigma_intercept_category = pm.Deterministic(
                    f"sigma_intercept_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"sigma_intercept_{cov.name}"
                        ],
                    ),
                    dims=("scalar",),
                )
                # Hierarchical intercepts for each category (using reparameterized form)
                categorical_intercept_offset = pm.Deterministic(
                    f"intercept_offset_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"intercept_offset_{cov.name}"
                        ],
                    ),
                    dims=(cov.name,),
                )
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    (
                        categorical_intercept_offset
                        * pt.reshape(sigma_intercept_category, (1,))  # pyright: ignore[reportPrivateImportUsage]
                    ),
                    dims=(cov.name,),
                )
            else:
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"intercept_{cov.name}"
                        ],
                    ),
                    dims=(cov.name,),
                )
        else:  # Partial freezing (fit parameters for the new site only)
            if cov.hierarchical:
                # Hierarchical categorical effect
                # Hyperpriors for category (Bayesian equivalent of random effects)
                # Hyperpriors are fixed during adaptation
                sigma_intercept_category = pm.Deterministic(
                    f"sigma_intercept_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"sigma_intercept_{cov.name}"
                        ],
                    ),
                    dims=("scalar",),
                )
                # Hierarchical intercepts for each category (using reparameterized form)
                # New categories get new parameters, old categories are fixed
                # Freeze old category parameters during adaptation
                fixed_categorical_intercept_offset = pm.Deterministic(
                    f"intercept_offset_{cov.name}_fixed",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"intercept_offset_{cov.name}"
                        ],
                    ),
                )
                # Create new parameters for new categories
                new_category_count = len(adapt["new_category_names"])
                pretrain_sigma_prior = adapt["pretrained_model_params"][
                    "posterior_means"
                ][f"variance_intercept_offset_{cov.name}"].std()
                new_categorical_intercept_offset = pm.Normal(
                    f"intercept_offset_{cov.name}_adapt",
                    mu=0,
                    sigma=pretrain_sigma_prior,
                    size=new_category_count,
                )
                # Combine fixed and new offsets
                categorical_intercept_offset = pm.Deterministic(
                    f"intercept_offset_{cov.name}",
                    pt.concatenate(
                        [
                            fixed_categorical_intercept_offset,
                            new_categorical_intercept_offset,
                        ],
                    ),
                    dims=(cov.name,),
                )
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    (
                        categorical_intercept_offset
                        * pt.reshape(sigma_intercept_category, (1,))  # pyright: ignore[reportPrivateImportUsage]
                    ),
                    dims=(cov.name,),
                )
            else:
                # Non-hierarchical (linear) categorical effect
                # New categories get new parameters, old categories are fixed
                # Freeze old category parameters during adaptation
                fixed_categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}_fixed",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"intercept_{cov.name}"
                        ],
                    ),
                )
                # Create new parameters for new categories
                new_category_count = len(adapt["new_category_names"])
                new_categorical_intercept = pm.Normal(
                    f"intercept_{cov.name}_adapt",
                    mu=0,
                    sigma=sigma_prior,
                    size=new_category_count,
                )
                # Combine fixed and new offsets
                categorical_intercept = pm.Deterministic(
                    f"intercept_{cov.name}",
                    pt.concatenate(
                        [
                            fixed_categorical_intercept,
                            new_categorical_intercept,
                        ],
                    ),
                    dims=(cov.name,),
                )
            self.model_params["n_params"] += new_category_count
        effects_list.append(
            categorical_intercept[category_indices[cov.name]],
        )

    def _model_all_mean_effects(
        self,
        train_data: pd.DataFrame,
        spline_bases: dict[str, npt.NDArray[np.floating[Any]]],
        category_indices: dict[str, npt.NDArray[np.integer[Any]]],
        adapt: dict[str, Any] | None = None,
    ) -> list[TensorVariable]:
        """
        Model all covariate mean effects.
        """
        mean_effects = []
        # Model the global intercept
        if adapt is None:  # Model fitting
            global_intercept = pm.Normal(
                "global_intercept",
                mu=0,
                sigma=5,
                dims=("scalar",),
            )
            # Increment parameter count for global intercept
            self.model_params["n_params"] += 1
        else:  # Freeze during adaptation/fine-tuning
            global_intercept = pm.Deterministic(
                "global_intercept",
                pt.as_tensor_variable(
                    adapt["pretrained_model_params"]["posterior_means"][
                        "global_intercept"
                    ],
                ),
                dims=("scalar",),
            )
        mean_effects.append(global_intercept)
        # Model additional covariate effects on the mean
        for cov in self.spec.covariates:
            if cov.name in self.spec.influencing_mean:
                if cov.cov_type == "numerical":
                    if cov.effect == "linear":
                        self._model_linear_mean_effect(
                            train_data,
                            cov,
                            mean_effects,
                            sigma_prior=5,
                            adapt=adapt,
                        )
                    elif cov.effect == "spline":
                        self._model_spline_mean_effect(
                            train_data,
                            cov,
                            mean_effects,
                            spline_bases,
                            sigma_prior=5,
                            adapt=adapt,
                        )
                elif cov.cov_type == "categorical":
                    self._model_categorical_mean_effect(
                        train_data,
                        cov,
                        mean_effects,
                        category_indices,
                        sigma_prior=1,
                        hierarchical_sigma_prior=5,
                        adapt=adapt,
                    )
                else:
                    err = f"Invalid covariate type '{cov.cov_type}' for '{cov.name}'."
                    raise ValueError(err)
        return mean_effects

    def _model_linear_variance_effect(
        self,
        train_data: pd.DataFrame,
        cov: CovariateSpec,
        effects_list: list[TensorVariable],
        sigma_prior: float = 0.1,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Model a linear effect for a numerical covariate on the variance estimate.
        """
        # Linear effect
        if adapt is None:  # Model fitting
            linear_beta = pm.Normal(
                f"variance_linear_beta_{cov.name}",
                mu=0,
                sigma=sigma_prior,
                size=1,
                dims=(f"{cov.name}_linear",),
            )
            # Increment parameter count for linear effect
            self.model_params["n_params"] += 1
        else:  # Freeze during adaptation/fine-tuning
            linear_beta = pm.Deterministic(
                f"variance_linear_beta_{cov.name}",
                pt.as_tensor_variable(
                    adapt["pretrained_model_params"]["posterior_means"][
                        f"variance_linear_beta_{cov.name}"
                    ],
                ),
                dims=(f"{cov.name}_linear",),
            )
        if cov.moments is not None:  # to satisfy type checker
            effects_list.append(
                (
                    cast("npt.NDArray[Any]", train_data[cov.name].to_numpy())
                    - cov.moments[0]
                )
                / cov.moments[1]
                * linear_beta,
            )

    def _model_spline_variance_effect(
        self,
        train_data: pd.DataFrame,
        cov: CovariateSpec,
        effects_list: list[TensorVariable],
        spline_bases: dict[str, npt.NDArray[np.floating[Any]]],
        sigma_prior: float = 0.1,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Model a spline effect for a numerical covariate on the variance estimate.
        """
        # Spline effect
        spline_bases[cov.name] = spline_bases.get(
            cov.name,
            cov.make_spline_bases(
                cast("npt.NDArray[Any]", train_data[cov.name].to_numpy()),
            ),
        )
        if adapt is None:  # Model fitting
            spline_betas = pm.ZeroSumNormal(
                f"variance_spline_betas_{cov.name}",
                sigma=sigma_prior,
                shape=spline_bases[cov.name].shape[1],
                dims=(f"{cov.name}_splines",),
            )
            # Note ZeroSumNormal imposes a centering constraint
            # (ensuring identifiability)
            # Increment parameter count for spline effects
            if cov.spline_spec is not None:  # to satisfy type checker
                self.model_params["n_params"] += cov.spline_spec.df - 1
        else:  # Freeze during adaptation/fine-tuning
            spline_betas = pm.Deterministic(
                f"variance_spline_betas_{cov.name}",
                pt.as_tensor_variable(
                    adapt["pretrained_model_params"]["posterior_means"][
                        f"variance_spline_betas_{cov.name}"
                    ],
                ),
                dims=(f"{cov.name}_splines",),
            )
        effects_list.append(pt.dot(spline_bases[cov.name], spline_betas.T))

    def _model_categorical_variance_effect(
        self,
        train_data: pd.DataFrame,
        cov: CovariateSpec,
        effects_list: list[TensorVariable],
        category_indices: dict[str, npt.NDArray[np.integer[Any]]],
        sigma_prior: float = 0.1,
        hierarchical_sigma_prior: float = 0.1,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Model the effect of a categorical covariate on the variance estimate.
        """
        # Factorize categories
        category_indices[cov.name] = category_indices.get(
            cov.name,
            cov.factorize_categories(
                cast("npt.NDArray[Any]", train_data[cov.name].to_numpy()),
            ),
        )
        if adapt is None:  # Model fitting
            if cov.hierarchical:
                # Hierarchical categorical effect
                # Hyperpriors for category (Bayesian equivalent of random effects)
                sigma_intercept_category = pm.HalfNormal(
                    f"variance_sigma_intercept_{cov.name}",
                    sigma=sigma_prior,
                    dims=("scalar",),
                )

                # Hierarchical intercepts for each category (using reparameterized form)
                categorical_intercept_offset = pm.ZeroSumNormal(
                    f"variance_intercept_offset_{cov.name}",
                    sigma=hierarchical_sigma_prior,
                    dims=(cov.name,),
                )
                # Note ZeroSumNormal imposes a centering constraint
                # (ensuring identifiability)
                categorical_intercept = pm.Deterministic(
                    f"variance_intercept_{cov.name}",
                    (
                        categorical_intercept_offset
                        * pt.reshape(sigma_intercept_category, (1,))  # pyright: ignore[reportPrivateImportUsage]
                    ),
                    dims=(cov.name,),
                )

                # Increment parameter count for hierarchical intercept
                self.model_params["n_params"] += 1

            else:
                # Non-hierarchical (linear) categorical effect
                categorical_intercept = pm.ZeroSumNormal(
                    f"variance_intercept_{cov.name}",
                    sigma=sigma_prior,
                    dims=(cov.name,),
                )
                # Note ZeroSumNormal imposes a centering constraint
                # (ensuring identifiability)
            # Increment parameter count for categorical effects
            if cov.categories is not None:  # to satisfy type checker
                self.model_params["n_params"] += len(cov.categories) - 1
        elif cov.name != adapt["covariate_to_adapt"]:
            # Freeze during adaptation/fine-tuning
            if cov.hierarchical:
                # Hierarchical categorical effect
                # Hyperpriors for category (Bayesian equivalent of random effects)
                sigma_intercept_category = pm.Deterministic(
                    f"variance_sigma_intercept_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"variance_sigma_intercept_{cov.name}"
                        ],
                    ),
                    dims=("scalar",),
                )
                # Hierarchical intercepts for each category (using reparameterized form)
                categorical_intercept_offset = pm.Deterministic(
                    f"variance_intercept_offset_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"variance_intercept_offset_{cov.name}"
                        ],
                    ),
                    dims=(cov.name,),
                )
                categorical_intercept = pm.Deterministic(
                    f"variance_intercept_{cov.name}",
                    (
                        categorical_intercept_offset
                        * pt.reshape(sigma_intercept_category, (1,))  # pyright: ignore[reportPrivateImportUsage]
                    ),
                    dims=(cov.name,),
                )
            else:
                categorical_intercept = pm.Deterministic(
                    f"variance_intercept_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"variance_intercept_{cov.name}"
                        ],
                    ),
                    dims=(cov.name,),
                )
        else:  # Partial freezing (fit parameters for the new site only)
            if cov.hierarchical:
                # Hierarchical categorical effect
                # Hyperpriors for category (Bayesian equivalent of random effects)
                # Hyperpriors are fixed during adaptation
                sigma_intercept_category = pm.Deterministic(
                    f"variance_sigma_intercept_{cov.name}",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"variance_sigma_intercept_{cov.name}"
                        ],
                    ),
                    dims=("scalar",),
                )
                # Hierarchical intercepts for each category (using reparameterized form)
                # New categories get new parameters, old categories are fixed
                # Freeze old category parameters during adaptation
                fixed_categorical_intercept_offset = pm.Deterministic(
                    f"variance_intercept_offset_{cov.name}_fixed",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"variance_intercept_offset_{cov.name}"
                        ],
                    ),
                )
                # Create new parameters for new categories
                new_category_count = len(adapt["new_category_names"])
                pretrain_sigma_prior = adapt["pretrained_model_params"][
                    "posterior_means"
                ][f"variance_intercept_offset_{cov.name}"].std()
                new_categorical_intercept_offset = pm.Normal(
                    f"variance_intercept_offset_{cov.name}_adapt",
                    mu=0,
                    sigma=pretrain_sigma_prior,
                    size=new_category_count,
                )
                # Combine fixed and new offsets
                categorical_intercept_offset = pm.Deterministic(
                    f"variance_intercept_offset_{cov.name}",
                    pt.concatenate(
                        [
                            fixed_categorical_intercept_offset,
                            new_categorical_intercept_offset,
                        ],
                    ),
                    dims=(cov.name,),
                )
                categorical_intercept = pm.Deterministic(
                    f"variance_intercept_{cov.name}",
                    (
                        categorical_intercept_offset
                        * pt.reshape(sigma_intercept_category, (1,))  # pyright: ignore[reportPrivateImportUsage]
                    ),
                    dims=(cov.name,),
                )
            else:
                # Non-hierarchical (linear) categorical effect
                # New categories get new parameters, old categories are fixed
                # Freeze old category parameters during adaptation
                fixed_categorical_intercept = pm.Deterministic(
                    f"variance_intercept_{cov.name}_fixed",
                    pt.as_tensor_variable(
                        adapt["pretrained_model_params"]["posterior_means"][
                            f"variance_intercept_{cov.name}"
                        ],
                    ),
                )
                # Create new parameters for new categories
                new_category_count = len(adapt["new_category_names"])
                new_categorical_intercept = pm.Normal(
                    f"variance_intercept_{cov.name}_adapt",
                    mu=0,
                    sigma=sigma_prior,
                    size=new_category_count,
                )
                # Combine fixed and new offsets
                categorical_intercept = pm.Deterministic(
                    f"variance_intercept_{cov.name}",
                    pt.concatenate(
                        [
                            fixed_categorical_intercept,
                            new_categorical_intercept,
                        ],
                    ),
                    dims=(cov.name,),
                )
            self.model_params["n_params"] += new_category_count
        effects_list.append(
            categorical_intercept[category_indices[cov.name]],
        )

    def _model_all_variance_effects(
        self,
        train_data: pd.DataFrame,
        spline_bases: dict[str, npt.NDArray[np.floating[Any]]],
        category_indices: dict[str, npt.NDArray[np.integer[Any]]],
        adapt: dict[str, Any] | None = None,
    ) -> list[TensorVariable]:
        """
        Model all covariate variance effects.
        """
        variance_effects = []
        # Model the global variance
        if adapt is None:
            global_variance_baseline = pm.Normal(
                "global_variance_baseline",
                mu=-0.0,
                sigma=0.5,
                dims=("scalar",),
            )
            # Increment parameter count for global variance
            self.model_params["n_params"] += 1
        else:
            global_variance_baseline = pm.Deterministic(
                "global_variance_baseline",
                pt.as_tensor_variable(
                    adapt["pretrained_model_params"]["posterior_means"][
                        "global_variance_baseline"
                    ],
                ),
                dims=("scalar",),
            )
        variance_effects.append(global_variance_baseline)
        # Model additional covariate effects on the variance
        for cov in self.spec.covariates:
            if cov.name in self.spec.influencing_variance:
                if cov.cov_type == "numerical":
                    if cov.effect == "linear":
                        self._model_linear_variance_effect(
                            train_data,
                            cov,
                            variance_effects,
                            sigma_prior=0.1,
                            adapt=adapt,
                        )
                    elif cov.effect == "spline":
                        self._model_spline_variance_effect(
                            train_data,
                            cov,
                            variance_effects,
                            spline_bases,
                            sigma_prior=0.1,
                            adapt=adapt,
                        )
                elif cov.cov_type == "categorical":
                    self._model_categorical_variance_effect(
                        train_data,
                        cov,
                        variance_effects,
                        category_indices,
                        sigma_prior=0.1,
                        hierarchical_sigma_prior=0.1,
                        adapt=adapt,
                    )
                else:
                    err = f"Invalid covariate type '{cov.cov_type}' for '{cov.name}'."
                    raise ValueError(err)
        return variance_effects

    def _combine_all_effects(
        self,
        mean_effects: list[TensorVariable],
        variance_effects: list[TensorVariable],
        standardized_voi: npt.NDArray[np.floating[Any]],
    ) -> None:
        """
        Combine all effects to model the observed data likelihood.
        """
        # Combine all mean and variance effects
        mu_estimate = sum(mean_effects)
        log_sigma_estimate = sum(variance_effects)
        sigma_estimate = pt.exp(log_sigma_estimate)

        effective_sample_size = self.model_params["sample_size"]

        # Model likelihood of the variable of interest
        _likelihood = pm.Normal(
            f"likelihood_{self.spec.variable_of_interest}",
            mu=mu_estimate,
            sigma=sigma_estimate,
            observed=standardized_voi,
            total_size=effective_sample_size,
        )

    def _fit_model_with_advi(self, *, progress_bar: bool = True) -> None:
        """
        Fit the model using Automatic Differentiation Variational Inference (ADVI).
        """
        base_lr = self.defaults["adam_learning_rate"]
        decay = self.defaults["adam_learning_rate_decay"]
        lr = shared(base_lr)
        optimizer = pm.adam(learning_rate=cast("float", lr))

        # Adaptive learning rate schedule callback
        def update_learning_rate(_approx: Any, _loss: Any, iteration: int) -> None:
            lr.set_value(base_lr * (decay**iteration))

        # Run automatic differential variational inference to fit the model
        self._trace = pm.fit(
            method="advi",
            n=self.defaults["advi_iterations"],
            random_seed=self.defaults["random_seed"],  # For reproducibility
            obj_optimizer=optimizer,
            callbacks=[
                update_learning_rate,
                pm.callbacks.CheckParametersConvergence(
                    tolerance=self.defaults["advi_convergence_tolerance"],
                    diff="relative",
                ),
            ],
            progressbar=progress_bar,
        )

        # Sample from the posterior distribution and store the results
        self.model_inference_data = self._trace.sample(
            2000,
            random_seed=self.defaults["random_seed"],
        )

        # Compute posterior means and standard deviations
        posterior_means = self.model_inference_data.posterior.mean(
            dim=("chain", "draw"),
        )
        posterior_stds = self.model_inference_data.posterior.std(dim=("chain", "draw"))

        # Store posterior means and stds as a dictionary in model parameters
        self.model_params["posterior_means"] = {
            x: posterior_means.data_vars[x].to_numpy()
            for x in posterior_means.data_vars
        }
        self.model_params["posterior_stds"] = {
            x: posterior_stds.data_vars[x].to_numpy() for x in posterior_stds.data_vars
        }

    def fit(
        self,
        train_data: pd.DataFrame,
        *,
        save_directory: Path | None = None,
        progress_bar: bool = True,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Fit the normative model to the training data.

        This method implements the fitting logic for the normative model
        based on the provided training data and model specification.

        Args:
            train_data: pd.DataFrame
                DataFrame containing the training data. It must include the variable
                of interest and all specified covariates.
            save_directory: Path | None
                A path to a directory to save the model. If provided, the fitted model
                will be saved to this path.
            progress_bar: bool
                If True, display a progress bar during fitting. Defaults to True.
            adapt: dict[str, Any] | None
                If provided, adapt a pre-trained model to a new covariate.
                Note: We recommended using the `adapt_fit` method, and not directly
                changing this argument, unless you know what you are doing.
        """
        # Validation checks
        self._validate_model()
        self._validate_dataframe_for_fitting(train_data)

        # Extract the variable of interest
        variable_of_interest = train_data[self.spec.variable_of_interest].to_numpy()

        # A dictionary to hold the model parameters after fitting
        if adapt is None:
            self.model_params = {}
            self.model_params["mean_VOI"] = variable_of_interest.mean()
            self.model_params["std_VOI"] = variable_of_interest.std()
            self.model_params["sample_size"] = variable_of_interest.shape[0]
            # Initialize parameter count
            self.model_params["n_params"] = 0
        else:
            # Update the pretrained model parameters
            if not hasattr(self, "model_params") or self.model_params is None:
                self.model_params = copy.deepcopy(adapt["pretrained_model_params"])
            self.model_params["sample_size"] += variable_of_interest.shape[0]

        # Data preparation
        model_coords = self._build_model_coordinates(
            observations=np.arange(train_data.shape[0]),
        )

        # Fitting logic
        with pm.Model(coords=model_coords) as self._model:
            # Standardize the variable of interest, and store mean and std
            # This ensures that the model is not sensitive to the scale of the variable
            standardized_voi = (
                variable_of_interest - self.model_params["mean_VOI"]
            ) / self.model_params["std_VOI"]

            # A dictionary for precomputed bspline basis functions
            spline_bases: dict[str, npt.NDArray[np.floating[Any]]] = {}

            # A dictionary for factorized categories
            category_indices: dict[str, npt.NDArray[np.integer[Any]]] = {}

            # Model the mean of the variable of interest
            mean_effects = self._model_all_mean_effects(
                train_data,
                spline_bases,
                category_indices,
                adapt=adapt,
            )

            # Model the variance of the variable of interest
            variance_effects = self._model_all_variance_effects(
                train_data,
                spline_bases,
                category_indices,
                adapt=adapt,
            )

            # Combine all mean and variance effects
            self._combine_all_effects(
                mean_effects,
                variance_effects,
                standardized_voi,
            )

            # Fit the model using ADVI
            self._fit_model_with_advi(progress_bar=progress_bar)

        # Save the model if a save path is provided
        if save_directory is not None:
            self.save_model(Path(save_directory))

    def adapt_fit(
        self,
        covariate_to_adapt: str,
        new_category_names: npt.NDArray[np.str_],
        train_data: pd.DataFrame,
        *,
        pretrained_model_params: dict[str, Any] | None = None,
        save_directory: Path | None = None,
        progress_bar: bool = True,
    ) -> None:
        """
        Using a previously fitted model, adapt the model to a new batch.
        This method enables adaptation of the model to data from a new
        batch/site by freezing all fitted parameters, and only estimating
        new parameters for the new batch/site category.

        Args:
            covariate_to_adapt: str
                Name of the categorical covariate representing the batch/site
                to which the model should be adapted.
                Note: This covariate must have been specified in the original
                model.
            new_category_names: list[str]
                Names of the new categories in the covariate_to_adapt representing
                the new batch/site labels (e.g. names of the new site).
                Note: These names must not have been present in the original
                fitted model.
            train_data: pd.DataFrame
                DataFrame containing the training data for adaptation.
                It must include the variable of interest and all specified covariates.
                Note: The covariate_to_adapt column must only contain the
                new_category_names (no new data from previously trained batches).
            pretrained_model_params: dict[str, Any] | None
                The model parameters from a previously fitted model to adapt.
                If None, the model parameters from the current instance will be used
                (assuming fitting was done).
            save_directory: Path | None
                A path to a directory to save the adapted model. If provided,
                the fitted model will be saved to this path.
            progress_bar: bool
                If True, display a progress bar during fitting. Defaults to True.
        """
        # Validation checks
        self._validate_model()
        self._validate_dataframe_for_fitting(train_data)

        # Locate the covariate to adapt
        cov_to_adapt_index = [cov.name for cov in self.spec.covariates].index(
            covariate_to_adapt,
        )

        # Extend the covariate categories to include the new categories
        self.spec.covariates[cov_to_adapt_index].extend_categories(new_category_names)

        # Extract the pre-trained model parameters
        if pretrained_model_params is None:
            if not self.model_params:
                err = (
                    "No pretrained model parameters found. "
                    "Please provide pretrained_model_params or fit the model first."
                )
                raise ValueError(err)
            pretrained_model_params = copy.deepcopy(self.model_params)

        # Fit the adapted model
        self.fit(
            train_data,
            save_directory=save_directory,
            progress_bar=progress_bar,
            adapt={
                "covariate_to_adapt": covariate_to_adapt,
                "new_category_names": new_category_names,
                "pretrained_model_params": pretrained_model_params,
            },
        )

    def _predict_mu(
        self,
        test_covariates: pd.DataFrame,
        model_params: dict[str, Any],
        predict_without: list[str],
    ) -> npt.NDArray[np.floating[Any]]:
        """
        Internal method to predict the mean of the variable of interest.
        """
        # Calculate mean effect
        mu_estimate = np.full(
            test_covariates.shape[0],
            model_params["posterior_means"]["global_intercept"].item(),
        )

        for cov in self.spec.covariates:
            if (cov.name in self.spec.influencing_mean) and (
                cov.name not in predict_without
            ):
                if cov.cov_type == "numerical":
                    if cov.effect == "linear":
                        if cov.moments is not None:  # to satisfy type checker
                            mu_estimate += (
                                (
                                    cast(
                                        "npt.NDArray[Any]",
                                        test_covariates[cov.name].to_numpy(),
                                    )
                                    - cov.moments[0]
                                )
                                / cov.moments[1]
                            ) * model_params["posterior_means"][
                                f"linear_beta_{cov.name}"
                            ]
                    elif cov.effect == "spline":
                        spline_bases = cov.make_spline_bases(
                            cast(
                                "npt.NDArray[Any]",
                                test_covariates[cov.name].to_numpy(),
                            ),
                        )
                        spline_betas = model_params["posterior_means"][
                            f"spline_betas_{cov.name}"
                        ]
                        mu_estimate += np.dot(spline_bases, spline_betas)
                elif cov.cov_type == "categorical":
                    category_indices = cov.factorize_categories(
                        cast("npt.NDArray[Any]", test_covariates[cov.name].to_numpy()),
                    )
                    if cov.hierarchical:
                        categorical_intercept = (
                            model_params["posterior_means"][
                                f"intercept_offset_{cov.name}"
                            ]
                            * model_params["posterior_means"][
                                f"sigma_intercept_{cov.name}"
                            ]
                        )
                    else:
                        categorical_intercept = model_params["posterior_means"][
                            f"intercept_{cov.name}"
                        ]
                    mu_estimate += categorical_intercept[category_indices]

        return np.array(
            mu_estimate * model_params["std_VOI"] + model_params["mean_VOI"],
        )

    def _predict_std(
        self,
        test_covariates: pd.DataFrame,
        model_params: dict[str, Any],
        predict_without: list[str],
    ) -> npt.NDArray[np.floating[Any]]:
        """
        Internal method to predict the standard deviation of the variable of interest.
        """
        # Calculate deviation effect
        log_sigma_estimate = np.full(
            test_covariates.shape[0],
            model_params["posterior_means"]["global_variance_baseline"].item(),
        )

        for cov in self.spec.covariates:
            if (
                cov.name in self.spec.influencing_variance
                and cov.name not in predict_without
            ):
                if cov.cov_type == "numerical":
                    if cov.effect == "linear":
                        if cov.moments is not None:  # to satisfy type checker
                            log_sigma_estimate += (
                                (
                                    cast(
                                        "npt.NDArray[Any]",
                                        test_covariates[cov.name].to_numpy(),
                                    )
                                    - cov.moments[0]
                                )
                                / cov.moments[1]
                            ) * model_params["posterior_means"][
                                f"variance_linear_beta_{cov.name}"
                            ]
                    elif cov.effect == "spline":
                        spline_bases = cov.make_spline_bases(
                            cast(
                                "npt.NDArray[Any]",
                                test_covariates[cov.name].to_numpy(),
                            ),
                        )
                        variance_spline_betas = model_params["posterior_means"][
                            f"variance_spline_betas_{cov.name}"
                        ]
                        log_sigma_estimate += spline_bases @ variance_spline_betas
                elif cov.cov_type == "categorical":
                    category_indices = cov.factorize_categories(
                        cast("npt.NDArray[Any]", test_covariates[cov.name].to_numpy()),
                    )
                    if cov.hierarchical:
                        categorical_variance_intercept = (
                            model_params["posterior_means"][
                                f"variance_intercept_offset_{cov.name}"
                            ]
                            * model_params["posterior_means"][
                                f"variance_sigma_intercept_{cov.name}"
                            ]
                        )
                    else:
                        categorical_variance_intercept = model_params[
                            "posterior_means"
                        ][f"variance_intercept_{cov.name}"]
                    log_sigma_estimate += categorical_variance_intercept[
                        category_indices
                    ]

        return np.array(np.exp(log_sigma_estimate) * model_params["std_VOI"])

    def predict(
        self,
        test_covariates: pd.DataFrame,
        *,
        extended: bool = False,
        model_params: dict[str, Any] | None = None,
        predict_without: list[str] | None = None,
    ) -> NormativePredictions:
        """
        Predict normative moments (mean, std) for new data using the fitted model.

        Args:
            test_covariates: pd.DataFrame
                DataFrame containing the new covariate data to predict.
                This must include all specified covariates.
                Note: covariates listed in predict_without will be ignored and are
                hence not required.
            extended: bool
                If True, return additional stats such as log-likelihood, centiles, etc.
                Note that extended predictions require variable_of_interest to be
                provided in the test_covariates DataFrame.
            model_params: dict | None
                Optional dictionary of model parameters to use. If not provided,
                the stored parameters from model.fit() will be used.
            predict_without: list[str] | None
                Optional list of covariate names to ignore during prediction.
                This can be used to check the effect of removing certain covariates
                from the model.

        Returns:
            NormativePredictions: Object containing the predicted moments (mean, std)
                for the variable of interest.
        """
        # Validate the new data
        validation_columns = [
            cov.name
            for cov in self.spec.covariates
            if cov.name not in (predict_without or [])
        ]
        if extended:
            validation_columns.append(self.spec.variable_of_interest)
        utils.general.validate_dataframe(test_covariates, validation_columns)

        # Parameters
        model_params = model_params or self.model_params
        if model_params is None:
            err = "No model parameters found. Please provide model_params."
            raise ValueError(err)

        # Calculate mean and variance effects and store in the predictions object
        predictions = NormativePredictions(
            {
                "mu_estimate": self._predict_mu(
                    test_covariates,
                    model_params,
                    (predict_without or []),
                ),
                "std_estimate": self._predict_std(
                    test_covariates,
                    model_params,
                    (predict_without or []),
                ),
            },
        )

        # Check if extended predictions are requested
        if extended:
            # Add extended statistics to predictions (e.g. centiles, log loss, etc.)
            predictions.extend_predictions(
                variable_of_interest=(
                    cast(
                        "npt.NDArray[Any]",
                        test_covariates[self.spec.variable_of_interest].to_numpy(),
                    )
                ),
            )

        return predictions

    def evaluate(self, new_data: pd.DataFrame) -> NormativePredictions:
        """
        Evaluate the model on new data and return predictions.

        Args:
            new_data: pd.DataFrame
                DataFrame containing the new data to evaluate.
                It must include all specified covariates and the variable of interest.

        Returns:
            NormativePredictions: Object containing the predictions and evaluation
            metrics.
        """
        # Run extended predictions
        return self.predict(test_covariates=new_data).evaluate_predictions(
            variable_of_interest=(
                cast(
                    "npt.NDArray[Any]",
                    new_data[self.spec.variable_of_interest].to_numpy(),
                )
            ),
            train_mean=self.model_params["mean_VOI"],
            train_std=self.model_params["std_VOI"],
            n_params=self.model_params["n_params"],
        )

    def harmonize(
        self,
        data: pd.DataFrame,
        covariates_to_harmonize: list[str],
        *,
        model_params: dict[str, Any] | None = None,
    ) -> npt.NDArray[np.floating[Any]]:
        """
        Harmonize the variable of interest in the data to remove effects of
        certain covariates (e.g. batch).

        Args:
            data: pd.DataFrame
                DataFrame containing the data to harmonize.
                It must include all specified covariates and the variable of interest.
            covariates_to_harmonize: list[str]
                List of covariate names to harmonize.
                The partial effects of these covariates will be removed from the
                variable of interest, and the harmonized values will be returned.
            model_params: dict | None
                Optional dictionary of model parameters to use. If not provided,
                the stored parameters from model.fit() will be used.

        Returns:
            npt.NDArray[np.floating[Any]]: Array of harmonized values for the
                variable of interest.
        """
        # Validate the new data
        validation_columns = [cov.name for cov in self.spec.covariates]
        validation_columns.append(self.spec.variable_of_interest)
        utils.general.validate_dataframe(data, validation_columns)

        # Parameters
        if model_params is None:
            model_params = self.model_params

        # Predict the mean and std with all covariates
        full_predictions = self.predict(
            test_covariates=data,
            model_params=model_params,
            predict_without=[],
        )

        # Predict the mean and std without the covariates to harmonize
        reduced_predictions = self.predict(
            test_covariates=data,
            model_params=model_params,
            predict_without=covariates_to_harmonize,
        )

        # First standardize the variable of interest based on the full model
        voi_standardized = (
            (cast("npt.NDArray[Any]", data[self.spec.variable_of_interest].to_numpy()))
            - full_predictions.predictions["mu_estimate"]
        ) / full_predictions.predictions["std_estimate"]

        # Then return the harmonized values based on the reduced model
        return np.asarray(
            (
                voi_standardized * reduced_predictions.predictions["std_estimate"]
                + reduced_predictions.predictions["mu_estimate"]
            ),
            dtype=np.float64,
        )

adapt_fit(covariate_to_adapt: str, new_category_names: npt.NDArray[np.str_], train_data: pd.DataFrame, *, pretrained_model_params: dict[str, Any] | None = None, save_directory: Path | None = None, progress_bar: bool = True) -> None

Using a previously fitted model, adapt the model to a new batch. This method enables adaptation of the model to data from a new batch/site by freezing all fitted parameters, and only estimating new parameters for the new batch/site category.

Parameters:

Name	Type	Description	Default
covariate_to_adapt	str	str Name of the categorical covariate representing the batch/site to which the model should be adapted. Note: This covariate must have been specified in the original model.	required
new_category_names	NDArray[str_]	list[str] Names of the new categories in the covariate_to_adapt representing the new batch/site labels (e.g. names of the new site). Note: These names must not have been present in the original fitted model.	required
train_data	DataFrame	pd.DataFrame DataFrame containing the training data for adaptation. It must include the variable of interest and all specified covariates. Note: The covariate_to_adapt column must only contain the new_category_names (no new data from previously trained batches).	required
pretrained_model_params	dict[str, Any] | None	dict[str, Any] | None The model parameters from a previously fitted model to adapt. If None, the model parameters from the current instance will be used (assuming fitting was done).	None
save_directory	Path | None	Path | None A path to a directory to save the adapted model. If provided, the fitted model will be saved to this path.	None
progress_bar	bool	bool If True, display a progress bar during fitting. Defaults to True.	True

Source code in src/spectranorm/snm.py

def adapt_fit(
    self,
    covariate_to_adapt: str,
    new_category_names: npt.NDArray[np.str_],
    train_data: pd.DataFrame,
    *,
    pretrained_model_params: dict[str, Any] | None = None,
    save_directory: Path | None = None,
    progress_bar: bool = True,
) -> None:
    """
    Using a previously fitted model, adapt the model to a new batch.
    This method enables adaptation of the model to data from a new
    batch/site by freezing all fitted parameters, and only estimating
    new parameters for the new batch/site category.

    Args:
        covariate_to_adapt: str
            Name of the categorical covariate representing the batch/site
            to which the model should be adapted.
            Note: This covariate must have been specified in the original
            model.
        new_category_names: list[str]
            Names of the new categories in the covariate_to_adapt representing
            the new batch/site labels (e.g. names of the new site).
            Note: These names must not have been present in the original
            fitted model.
        train_data: pd.DataFrame
            DataFrame containing the training data for adaptation.
            It must include the variable of interest and all specified covariates.
            Note: The covariate_to_adapt column must only contain the
            new_category_names (no new data from previously trained batches).
        pretrained_model_params: dict[str, Any] | None
            The model parameters from a previously fitted model to adapt.
            If None, the model parameters from the current instance will be used
            (assuming fitting was done).
        save_directory: Path | None
            A path to a directory to save the adapted model. If provided,
            the fitted model will be saved to this path.
        progress_bar: bool
            If True, display a progress bar during fitting. Defaults to True.
    """
    # Validation checks
    self._validate_model()
    self._validate_dataframe_for_fitting(train_data)

    # Locate the covariate to adapt
    cov_to_adapt_index = [cov.name for cov in self.spec.covariates].index(
        covariate_to_adapt,
    )

    # Extend the covariate categories to include the new categories
    self.spec.covariates[cov_to_adapt_index].extend_categories(new_category_names)

    # Extract the pre-trained model parameters
    if pretrained_model_params is None:
        if not self.model_params:
            err = (
                "No pretrained model parameters found. "
                "Please provide pretrained_model_params or fit the model first."
            )
            raise ValueError(err)
        pretrained_model_params = copy.deepcopy(self.model_params)

    # Fit the adapted model
    self.fit(
        train_data,
        save_directory=save_directory,
        progress_bar=progress_bar,
        adapt={
            "covariate_to_adapt": covariate_to_adapt,
            "new_category_names": new_category_names,
            "pretrained_model_params": pretrained_model_params,
        },
    )

evaluate(new_data: pd.DataFrame) -> NormativePredictions

Evaluate the model on new data and return predictions.

Parameters:

Name	Type	Description	Default
new_data	DataFrame	pd.DataFrame DataFrame containing the new data to evaluate. It must include all specified covariates and the variable of interest.	required

Returns:

Name	Type	Description
NormativePredictions	NormativePredictions	Object containing the predictions and evaluation
	NormativePredictions	metrics.

Source code in src/spectranorm/snm.py

def evaluate(self, new_data: pd.DataFrame) -> NormativePredictions:
    """
    Evaluate the model on new data and return predictions.

    Args:
        new_data: pd.DataFrame
            DataFrame containing the new data to evaluate.
            It must include all specified covariates and the variable of interest.

    Returns:
        NormativePredictions: Object containing the predictions and evaluation
        metrics.
    """
    # Run extended predictions
    return self.predict(test_covariates=new_data).evaluate_predictions(
        variable_of_interest=(
            cast(
                "npt.NDArray[Any]",
                new_data[self.spec.variable_of_interest].to_numpy(),
            )
        ),
        train_mean=self.model_params["mean_VOI"],
        train_std=self.model_params["std_VOI"],
        n_params=self.model_params["n_params"],
    )

fit(train_data: pd.DataFrame, *, save_directory: Path | None = None, progress_bar: bool = True, adapt: dict[str, Any] | None = None) -> None

Fit the normative model to the training data.

This method implements the fitting logic for the normative model based on the provided training data and model specification.

Parameters:

Name	Type	Description	Default
train_data	DataFrame	pd.DataFrame DataFrame containing the training data. It must include the variable of interest and all specified covariates.	required
save_directory	Path | None	Path | None A path to a directory to save the model. If provided, the fitted model will be saved to this path.	None
progress_bar	bool	bool If True, display a progress bar during fitting. Defaults to True.	True
adapt	dict[str, Any] | None	dict[str, Any] | None If provided, adapt a pre-trained model to a new covariate. Note: We recommended using the adapt_fit method, and not directly changing this argument, unless you know what you are doing.	None

Source code in src/spectranorm/snm.py

def fit(
    self,
    train_data: pd.DataFrame,
    *,
    save_directory: Path | None = None,
    progress_bar: bool = True,
    adapt: dict[str, Any] | None = None,
) -> None:
    """
    Fit the normative model to the training data.

    This method implements the fitting logic for the normative model
    based on the provided training data and model specification.

    Args:
        train_data: pd.DataFrame
            DataFrame containing the training data. It must include the variable
            of interest and all specified covariates.
        save_directory: Path | None
            A path to a directory to save the model. If provided, the fitted model
            will be saved to this path.
        progress_bar: bool
            If True, display a progress bar during fitting. Defaults to True.
        adapt: dict[str, Any] | None
            If provided, adapt a pre-trained model to a new covariate.
            Note: We recommended using the `adapt_fit` method, and not directly
            changing this argument, unless you know what you are doing.
    """
    # Validation checks
    self._validate_model()
    self._validate_dataframe_for_fitting(train_data)

    # Extract the variable of interest
    variable_of_interest = train_data[self.spec.variable_of_interest].to_numpy()

    # A dictionary to hold the model parameters after fitting
    if adapt is None:
        self.model_params = {}
        self.model_params["mean_VOI"] = variable_of_interest.mean()
        self.model_params["std_VOI"] = variable_of_interest.std()
        self.model_params["sample_size"] = variable_of_interest.shape[0]
        # Initialize parameter count
        self.model_params["n_params"] = 0
    else:
        # Update the pretrained model parameters
        if not hasattr(self, "model_params") or self.model_params is None:
            self.model_params = copy.deepcopy(adapt["pretrained_model_params"])
        self.model_params["sample_size"] += variable_of_interest.shape[0]

    # Data preparation
    model_coords = self._build_model_coordinates(
        observations=np.arange(train_data.shape[0]),
    )

    # Fitting logic
    with pm.Model(coords=model_coords) as self._model:
        # Standardize the variable of interest, and store mean and std
        # This ensures that the model is not sensitive to the scale of the variable
        standardized_voi = (
            variable_of_interest - self.model_params["mean_VOI"]
        ) / self.model_params["std_VOI"]

        # A dictionary for precomputed bspline basis functions
        spline_bases: dict[str, npt.NDArray[np.floating[Any]]] = {}

        # A dictionary for factorized categories
        category_indices: dict[str, npt.NDArray[np.integer[Any]]] = {}

        # Model the mean of the variable of interest
        mean_effects = self._model_all_mean_effects(
            train_data,
            spline_bases,
            category_indices,
            adapt=adapt,
        )

        # Model the variance of the variable of interest
        variance_effects = self._model_all_variance_effects(
            train_data,
            spline_bases,
            category_indices,
            adapt=adapt,
        )

        # Combine all mean and variance effects
        self._combine_all_effects(
            mean_effects,
            variance_effects,
            standardized_voi,
        )

        # Fit the model using ADVI
        self._fit_model_with_advi(progress_bar=progress_bar)

    # Save the model if a save path is provided
    if save_directory is not None:
        self.save_model(Path(save_directory))

from_dataframe(model_type: ModelType, dataframe: pd.DataFrame, variable_of_interest: str, numerical_covariates: list[str] | None = None, categorical_covariates: list[str] | None = None, batch_covariates: list[str] | None = None, nonlinear_covariates: list[str] | None = None, influencing_mean: list[str] | None = None, influencing_variance: list[str] | None = None, spline_kwargs: dict[str, Any] | None = None) -> DirectNormativeModel classmethod

Initialize a normative model from a pandas DataFrame.

Parameters:

Name	Type	Description	Default
model_type	ModelType	ModelType Type of the model to create, either "HBR" (Hierarchical Bayesian Regression) or "BLR" (Bayesian Linear Regression).	required
dataframe	DataFrame	pd.DataFrame DataFrame containing the data.	required
variable_of_interest	str	str Name of the target variable to model.	required
numerical_covariates	list[str] | None	list[str] | None List of numerical covariate names.	None
categorical_covariates	list[str] | None	list[str] | None List of categorical covariate names.	None
batch_covariates	list[str] | None	list[str] | None List of batch covariate names which should also be included in categorical_covariates.	None
nonlinear_covariates	list[str] | None	list[str] | None List of covariate names to be modeled as nonlinear effects. These should also be included in numerical_covariates.	None
influencing_mean	list[str] | None	list[str] | None List of covariate names that influence the mean of the variable of interest. These should be included in either numerical_covariates or categorical_covariates.	None
influencing_variance	list[str] | None	list[str] | None List of covariate names that influence the variance of the variable of interest. These should be included in either numerical_covariates or categorical_covariates.	None
spline_kwargs	dict[str, Any] | None	dict Additional keyword arguments for spline specification, such as df, degree, and knots. These are passed to the create_spline_spec method to create spline specifications for nonlinear covariates.	None

Returns:

Type	Description
DirectNormativeModel	DirectNormativeModel An instance of DirectNormativeModel initialized with the provided data.

Source code in src/spectranorm/snm.py

@classmethod
def from_dataframe(
    cls,
    model_type: ModelType,
    dataframe: pd.DataFrame,
    variable_of_interest: str,
    numerical_covariates: list[str] | None = None,
    categorical_covariates: list[str] | None = None,
    batch_covariates: list[str] | None = None,
    nonlinear_covariates: list[str] | None = None,
    influencing_mean: list[str] | None = None,
    influencing_variance: list[str] | None = None,
    spline_kwargs: dict[str, Any] | None = None,
) -> DirectNormativeModel:
    """
    Initialize a normative model from a pandas DataFrame.

    Args:
        model_type: ModelType
            Type of the model to create, either "HBR" (Hierarchical Bayesian
            Regression) or "BLR" (Bayesian Linear Regression).
        dataframe: pd.DataFrame
            DataFrame containing the data.
        variable_of_interest: str
            Name of the target variable to model.
        numerical_covariates: list[str] | None
            List of numerical covariate names.
        categorical_covariates: list[str] | None
            List of categorical covariate names.
        batch_covariates: list[str] | None
            List of batch covariate names which should also be included in
            categorical_covariates.
        nonlinear_covariates: list[str] | None
            List of covariate names to be modeled as nonlinear effects.
            These should also be included in numerical_covariates.
        influencing_mean: list[str] | None
            List of covariate names that influence the mean of the variable
            of interest. These should be included in either numerical_covariates
            or categorical_covariates.
        influencing_variance: list[str] | None
            List of covariate names that influence the variance of the variable
            of interest. These should be included in either numerical_covariates
            or categorical_covariates.
        spline_kwargs: dict
            Additional keyword arguments for spline specification, such as
            `df`, `degree`, and `knots`. These are passed to the
            `create_spline_spec` method to create spline specifications for
            nonlinear covariates.

    Returns:
        DirectNormativeModel
            An instance of DirectNormativeModel initialized with the provided data.
    """
    # Set default values for optional parameters
    numerical_covariates = numerical_covariates or []
    categorical_covariates = categorical_covariates or []
    batch_covariates = batch_covariates or []
    nonlinear_covariates = nonlinear_covariates or []
    influencing_mean = influencing_mean or []
    influencing_variance = influencing_variance or []
    spline_kwargs = spline_kwargs or {}

    # Validity checks for input parameters
    cls._validate_init_args(
        model_type,
        variable_of_interest,
        numerical_covariates,
        categorical_covariates,
        batch_covariates,
        nonlinear_covariates,
    )
    utils.general.validate_dataframe(
        dataframe,
        [variable_of_interest, *numerical_covariates, *categorical_covariates],
    )

    # Create an instance of the class
    self = cls(
        spec=NormativeModelSpec(
            variable_of_interest=variable_of_interest,
            covariates=[],
            influencing_mean=influencing_mean,
            influencing_variance=influencing_variance,
        ),
    )

    # Populate the spline_kwargs with defaults if not provided
    spline_kwargs["df"] = spline_kwargs.get("df", self.defaults["spline_df"])
    spline_kwargs["degree"] = spline_kwargs.get(
        "degree",
        self.defaults["spline_degree"],
    )
    spline_kwargs["extrapolation_factor"] = spline_kwargs.get(
        "extrapolation_factor",
        self.defaults["spline_extrapolation_factor"],
    )

    # Start building the model specification
    # Add categorical covariates
    for cov_name in categorical_covariates:
        hierarchical = False
        if cov_name in batch_covariates and model_type == "HBR":
            hierarchical = True
        self.spec.covariates.append(
            CovariateSpec(
                name=cov_name,
                cov_type="categorical",
                categories=dataframe[cov_name].unique(),
                hierarchical=hierarchical,
            ),
        )
    for cov_name in numerical_covariates:
        if cov_name not in nonlinear_covariates:
            self.spec.covariates.append(
                CovariateSpec(
                    name=cov_name,
                    cov_type="numerical",
                    effect="linear",
                    moments=(
                        dataframe[cov_name].mean(),
                        dataframe[cov_name].std(),
                    ),
                ),
            )
        else:
            self.spec.covariates.append(
                CovariateSpec(
                    name=cov_name,
                    cov_type="numerical",
                    effect="spline",
                    spline_spec=SplineSpec.create_spline_spec(
                        dataframe[cov_name],
                        **spline_kwargs,
                    ),
                ),
            )
    return self

harmonize(data: pd.DataFrame, covariates_to_harmonize: list[str], *, model_params: dict[str, Any] | None = None) -> npt.NDArray[np.floating[Any]]

Harmonize the variable of interest in the data to remove effects of certain covariates (e.g. batch).

Parameters:

Name	Type	Description	Default
data	DataFrame	pd.DataFrame DataFrame containing the data to harmonize. It must include all specified covariates and the variable of interest.	required
covariates_to_harmonize	list[str]	list[str] List of covariate names to harmonize. The partial effects of these covariates will be removed from the variable of interest, and the harmonized values will be returned.	required
model_params	dict[str, Any] | None	dict | None Optional dictionary of model parameters to use. If not provided, the stored parameters from model.fit() will be used.	None

Returns:

Type	Description
NDArray[floating[Any]]	npt.NDArray[np.floating[Any]]: Array of harmonized values for the variable of interest.

Source code in src/spectranorm/snm.py

def harmonize(
    self,
    data: pd.DataFrame,
    covariates_to_harmonize: list[str],
    *,
    model_params: dict[str, Any] | None = None,
) -> npt.NDArray[np.floating[Any]]:
    """
    Harmonize the variable of interest in the data to remove effects of
    certain covariates (e.g. batch).

    Args:
        data: pd.DataFrame
            DataFrame containing the data to harmonize.
            It must include all specified covariates and the variable of interest.
        covariates_to_harmonize: list[str]
            List of covariate names to harmonize.
            The partial effects of these covariates will be removed from the
            variable of interest, and the harmonized values will be returned.
        model_params: dict | None
            Optional dictionary of model parameters to use. If not provided,
            the stored parameters from model.fit() will be used.

    Returns:
        npt.NDArray[np.floating[Any]]: Array of harmonized values for the
            variable of interest.
    """
    # Validate the new data
    validation_columns = [cov.name for cov in self.spec.covariates]
    validation_columns.append(self.spec.variable_of_interest)
    utils.general.validate_dataframe(data, validation_columns)

    # Parameters
    if model_params is None:
        model_params = self.model_params

    # Predict the mean and std with all covariates
    full_predictions = self.predict(
        test_covariates=data,
        model_params=model_params,
        predict_without=[],
    )

    # Predict the mean and std without the covariates to harmonize
    reduced_predictions = self.predict(
        test_covariates=data,
        model_params=model_params,
        predict_without=covariates_to_harmonize,
    )

    # First standardize the variable of interest based on the full model
    voi_standardized = (
        (cast("npt.NDArray[Any]", data[self.spec.variable_of_interest].to_numpy()))
        - full_predictions.predictions["mu_estimate"]
    ) / full_predictions.predictions["std_estimate"]

    # Then return the harmonized values based on the reduced model
    return np.asarray(
        (
            voi_standardized * reduced_predictions.predictions["std_estimate"]
            + reduced_predictions.predictions["mu_estimate"]
        ),
        dtype=np.float64,
    )

load_model(directory: Path, *, load_posterior: bool = False) -> DirectNormativeModel classmethod

Load the model and its posterior from a directory. The model will be loaded from a subdirectory named 'saved_model'.

Parameters:

Name	Type	Description	Default
directory	Path	Path Path to the directory containing the model.	required
load_posterior	bool	bool (default=False) If True, load the model's posterior trace from the saved inference data.	False

Source code in src/spectranorm/snm.py

@classmethod
def load_model(
    cls,
    directory: Path,
    *,
    load_posterior: bool = False,
) -> DirectNormativeModel:
    """
    Load the model and its posterior from a directory.
    The model will be loaded from a subdirectory named 'saved_model'.

    Args:
        directory: Path
            Path to the directory containing the model.
        load_posterior: bool (default=False)
            If True, load the model's posterior trace from the saved inference data.
    """
    # Validate the load directory
    directory = Path(directory)
    saved_model_dir = utils.general.validate_load_directory(
        directory,
        "saved_model",
    )

    # Load the saved model dict
    model_dict = joblib.load(saved_model_dir / "model_dict.joblib")

    # Create an instance of the class
    instance = cls(
        spec=model_dict["spec"],
    )

    # Set the attributes from the loaded model dictionary
    instance.defaults.update(model_dict["defaults"])
    if "model_params" in model_dict:
        instance.model_params = model_dict["model_params"]
        if load_posterior:
            instance.model_inference_data = az.from_netcdf(  # type: ignore[no-untyped-call]
                saved_model_dir / "model_inference_data.nc",
            )

    return instance

predict(test_covariates: pd.DataFrame, *, extended: bool = False, model_params: dict[str, Any] | None = None, predict_without: list[str] | None = None) -> NormativePredictions

Predict normative moments (mean, std) for new data using the fitted model.

Parameters:

Name	Type	Description	Default
test_covariates	DataFrame	pd.DataFrame DataFrame containing the new covariate data to predict. This must include all specified covariates. Note: covariates listed in predict_without will be ignored and are hence not required.	required
extended	bool	bool If True, return additional stats such as log-likelihood, centiles, etc. Note that extended predictions require variable_of_interest to be provided in the test_covariates DataFrame.	False
model_params	dict[str, Any] | None	dict | None Optional dictionary of model parameters to use. If not provided, the stored parameters from model.fit() will be used.	None
predict_without	list[str] | None	list[str] | None Optional list of covariate names to ignore during prediction. This can be used to check the effect of removing certain covariates from the model.	None

Returns:

Name	Type	Description
NormativePredictions	NormativePredictions	Object containing the predicted moments (mean, std) for the variable of interest.

Source code in src/spectranorm/snm.py

def predict(
    self,
    test_covariates: pd.DataFrame,
    *,
    extended: bool = False,
    model_params: dict[str, Any] | None = None,
    predict_without: list[str] | None = None,
) -> NormativePredictions:
    """
    Predict normative moments (mean, std) for new data using the fitted model.

    Args:
        test_covariates: pd.DataFrame
            DataFrame containing the new covariate data to predict.
            This must include all specified covariates.
            Note: covariates listed in predict_without will be ignored and are
            hence not required.
        extended: bool
            If True, return additional stats such as log-likelihood, centiles, etc.
            Note that extended predictions require variable_of_interest to be
            provided in the test_covariates DataFrame.
        model_params: dict | None
            Optional dictionary of model parameters to use. If not provided,
            the stored parameters from model.fit() will be used.
        predict_without: list[str] | None
            Optional list of covariate names to ignore during prediction.
            This can be used to check the effect of removing certain covariates
            from the model.

    Returns:
        NormativePredictions: Object containing the predicted moments (mean, std)
            for the variable of interest.
    """
    # Validate the new data
    validation_columns = [
        cov.name
        for cov in self.spec.covariates
        if cov.name not in (predict_without or [])
    ]
    if extended:
        validation_columns.append(self.spec.variable_of_interest)
    utils.general.validate_dataframe(test_covariates, validation_columns)

    # Parameters
    model_params = model_params or self.model_params
    if model_params is None:
        err = "No model parameters found. Please provide model_params."
        raise ValueError(err)

    # Calculate mean and variance effects and store in the predictions object
    predictions = NormativePredictions(
        {
            "mu_estimate": self._predict_mu(
                test_covariates,
                model_params,
                (predict_without or []),
            ),
            "std_estimate": self._predict_std(
                test_covariates,
                model_params,
                (predict_without or []),
            ),
        },
    )

    # Check if extended predictions are requested
    if extended:
        # Add extended statistics to predictions (e.g. centiles, log loss, etc.)
        predictions.extend_predictions(
            variable_of_interest=(
                cast(
                    "npt.NDArray[Any]",
                    test_covariates[self.spec.variable_of_interest].to_numpy(),
                )
            ),
        )

    return predictions

save_model(directory: Path, *, save_posterior: bool = False) -> None

Save the fitted model and it's posterior to a directory. The model will be saved in a subdirectory named 'saved_model'. If this directory is not empty, an error is raised.

Parameters:

Name	Type	Description	Default
directory	Path	Path Path to a directory to save the model.	required
save_posterior	bool	bool (default=False) If True, save the model's posterior trace inference data.	False

Source code in src/spectranorm/snm.py

def save_model(self, directory: Path, *, save_posterior: bool = False) -> None:
    """
    Save the fitted model and it's posterior to a directory.
    The model will be saved in a subdirectory named 'saved_model'.
    If this directory is not empty, an error is raised.

    Args:
        directory: Path
            Path to a directory to save the model.
        save_posterior: bool (default=False)
            If True, save the model's posterior trace inference data.
    """
    # Prepare the save directory
    directory = Path(directory)
    saved_model_dir = utils.general.prepare_save_directory(directory, "saved_model")

    model_dict = {
        "spec": self.spec,
        "defaults": self.defaults,
    }
    if hasattr(self, "model_params"):
        model_dict["model_params"] = self.model_params
        if hasattr(self, "model_inference_data") and save_posterior:
            self.model_inference_data.to_netcdf(
                saved_model_dir / "model_inference_data.nc",
            )
    joblib.dump(model_dict, saved_model_dir / "model_dict.joblib")

NormativeModelSpec dataclass

General specification of a normative model.

Attributes:

Name	Type	Description
variable_of_interest	str	str Name of the target variable to model (e.g., "thickness").
covariates	list[CovariateSpec]	list[CovariateSpec] Listing all model covariates and specifying how each covariate is modeled.
influencing_mean	list[str]	list[str] List of covariate names that influence the mean of the variable of interest.
influencing_variance	list[str]	list[str] List of covariate names that influence the variance of the variable of interest.

Source code in src/spectranorm/snm.py

@dataclass
class NormativeModelSpec:
    """
    General specification of a normative model.

    Attributes:
        variable_of_interest: str
            Name of the target variable to model (e.g., "thickness").
        covariates: list[CovariateSpec]
            Listing all model covariates and specifying how each covariate is modeled.
        influencing_mean: list[str]
            List of covariate names that influence the mean of the variable of interest.
        influencing_variance: list[str]
            List of covariate names that influence the variance of the variable of
            interest.
    """

    variable_of_interest: str
    covariates: list[CovariateSpec]
    influencing_mean: list[str]
    influencing_variance: list[str]

    def __post_init__(self) -> None:
        if not isinstance(self.variable_of_interest, str):
            err = "variable_of_interest must be a string."
            raise TypeError(err)
        if not isinstance(self.covariates, list):
            err = "covariates must be a list of CovariateSpec instances."
            raise TypeError(err)
        if not all(isinstance(cov, CovariateSpec) for cov in self.covariates):
            err = "All items in covariates must be CovariateSpec instances."
            raise TypeError(err)
        if not isinstance(self.influencing_mean, list):
            err = "influencing_mean must be a list of covariate names."
            raise TypeError(err)
        if not isinstance(self.influencing_variance, list):
            err = "influencing_variance must be a list of covariate names."
            raise TypeError(err)

NormativePredictions dataclass

Container for the results of model.predict() function.

Attributes:

Name	Type	Description
predictions	dict[str, NDArray[floating[Any]]]	dict Dictionary containing the model's predictions, including - Predictions of mean (mu_estimate). - Predictions of standard deviation (std_estimate). - [Optional] The observed variable of interest (the name of which is provided in the function argument). - [Optional] Additional evaluation metrics for the predictions.

Source code in src/spectranorm/snm.py

@dataclass
class NormativePredictions:
    """
    Container for the results of model.predict() function.

    Attributes:
        predictions: dict
            Dictionary containing the model's predictions, including
            - Predictions of mean (mu_estimate).
            - Predictions of standard deviation (std_estimate).
            - [Optional] The observed variable of interest (the name of which is
              provided in the function argument).
            - [Optional] Additional evaluation metrics for the predictions.
    """

    predictions: dict[str, npt.NDArray[np.floating[Any]]]
    evaluations: dict[str, npt.NDArray[np.floating[Any]] | float] = field(
        default_factory=dict,
    )

    def extend_predictions(
        self,
        variable_of_interest: npt.NDArray[np.floating[Any]],
        *,
        likelihood_censoring_quantile: float = 0.01,
    ) -> NormativePredictions:
        """
        Extend the NormativePredictions (predictions dictionary) with additional
        statistics.

        Args:
            variable_of_interest: np.ndarray
                The observed values for the variable(s) of interest.
            likelihood_censoring_quantile: float (default=0.01)
                Quantile below which log-likelihoods are censored for evaluation.
                Note: by default, a censored log-likelihood is computed to avoid
                extreme log-likelihood values for outliers. This affects several
                resulting metrics that are based on log-likelihoods (e.g. MSLL).
                If you want to compute the full (uncensored) log-likelihoods, set
                `likelihood_censoring_quantile` to 0.

        Returns:
            NormativePredictions
                Extended NormativePredictions with additional statistics.
        """
        self.predictions["z-score"] = (
            variable_of_interest - self.predictions["mu_estimate"]
        ) / self.predictions["std_estimate"]
        self.predictions["log-likelihood"] = (
            utils.stats.compute_censored_log_likelihood(
                variable_of_interest,
                self.predictions["mu_estimate"],
                self.predictions["std_estimate"],
                censored_quantile=likelihood_censoring_quantile,
            )
        )
        self.predictions["centiles"] = utils.stats.compute_centiles_from_z_scores(
            self.predictions["z-score"],
        )

        self.predictions["variable_of_interest"] = variable_of_interest

        return self

    def evaluate_predictions(
        self,
        variable_of_interest: npt.NDArray[np.floating[Any]],
        train_mean: npt.NDArray[np.floating[Any]],
        train_std: npt.NDArray[np.floating[Any]],
        n_params: int | None = None,
        msll_censored_quantile: float = 0.01,
    ) -> NormativePredictions:
        """
        Evaluate the predictions against the observed variable of interest.

        This function computes a battery of evaluation metrics implemented
        in `snm.utils.metrics`. Namely the evaluations include:
            - Mean Absolute Error (MAE)
            - Mean Squared Error (MSE)
            - Root Mean Squared Error (RMSE)
            - Mean Absolute Percentage Error (MAPE)
            - R-squared
            - Explained Variance Score
            - Mean Standardized Log Loss (MSLL)

        Args:
            variable_of_interest: np.ndarray
                The observed values for the variable(s) of interest.
            train_mean: np.ndarray
                Mean(s) of the variable(s) of interest from the training data.
            train_std: np.ndarray
                Standard deviation(s) of the variable(s) of interest from the training
                data.
            n_params: int
                Number of free parameters in the model.
            msll_censored_quantile: float (default=0.02)
                Quantile below which log-likelihoods are censored for MSLL.

        Returns:
            NormativePredictions
                Object containing the evaluation results.
        """
        self.extend_predictions(
            variable_of_interest,
            likelihood_censoring_quantile=msll_censored_quantile,
        )
        # Mean Absolute Error (MAE)
        self.evaluations["MAE"] = utils.metrics.compute_mae(
            y=self.predictions["variable_of_interest"],
            y_pred=self.predictions["mu_estimate"],
        )
        # Mean Squared Error (MSE)
        self.evaluations["MSE"] = utils.metrics.compute_mse(
            y=self.predictions["variable_of_interest"],
            y_pred=self.predictions["mu_estimate"],
        )
        # Root Mean Squared Error (RMSE)
        self.evaluations["RMSE"] = utils.metrics.compute_rmse(
            y=self.predictions["variable_of_interest"],
            y_pred=self.predictions["mu_estimate"],
        )
        # Mean Absolute Percentage Error (MAPE)
        self.evaluations["MAPE"] = utils.metrics.compute_mape(
            y=self.predictions["variable_of_interest"],
            y_pred=self.predictions["mu_estimate"],
        )
        # R-squared
        self.evaluations["R-squared"] = utils.metrics.compute_r2(
            y=self.predictions["variable_of_interest"],
            y_pred=self.predictions["mu_estimate"],
        )
        # Explained Variance Score
        self.evaluations["Explained Variance"] = utils.metrics.compute_expv(
            y=self.predictions["variable_of_interest"],
            y_pred=self.predictions["mu_estimate"],
        )
        # Mean Standardized Log Loss (MSLL)
        self.evaluations["MSLL"] = utils.metrics.compute_msll(
            model_log_likelihoods=self.predictions["log-likelihood"],
            baseline_log_likelihoods=utils.stats.compute_censored_log_likelihood(
                self.predictions["variable_of_interest"],
                train_mean,
                train_std,
                censored_quantile=msll_censored_quantile,
            ),
        )
        _ = n_params  # keep for future use (e.g. information criteria calculations)

        return self

    def to_array(self, keys: list[str] | None = None) -> npt.NDArray[np.floating[Any]]:
        """
        Return prediction results as a list of NumPy arrays.

        Args:
            keys: list[str]
                Optional list of keys to return.
                Defaults to ["mu_estimate", "std_estimate"].

        Returns:
            list[np.ndarray]
                NumPy arrays for the requested predictions
        """
        keys = keys or ["mu_estimate", "std_estimate"]
        return np.array([self.predictions[key] for key in keys])

    def to_dataframe(
        self,
        index: pd.Index[Any] | list[Any] | None = None,
    ) -> pd.DataFrame:
        """
        Return prediction results as a DataFrame.

        Args:
            index: pd.Index | list | None
                Optional index for the DataFrame (defaults to None)

        Returns:
            pd.DataFrame
                DataFrame containing the predictions
        """
        predictions = self.predictions.copy()
        # Flatten the predictions dictionary if multiple queries are predicted
        for key in predictions:
            if predictions[key].ndim > 1:
                if predictions[key].shape[1] == 1:
                    predictions[key] = predictions[key].flatten()
                else:
                    for i in range(predictions[key].shape[1]):
                        predictions[f"{key}_{i + 1}"] = predictions[key][:, i]
                    # delete the key
                    del predictions[key]

        # Make a new DataFrame for the predictions dictionary
        return pd.DataFrame(predictions, index=index)

evaluate_predictions(variable_of_interest: npt.NDArray[np.floating[Any]], train_mean: npt.NDArray[np.floating[Any]], train_std: npt.NDArray[np.floating[Any]], n_params: int | None = None, msll_censored_quantile: float = 0.01) -> NormativePredictions

Evaluate the predictions against the observed variable of interest.

This function computes a battery of evaluation metrics implemented in snm.utils.metrics. Namely the evaluations include: - Mean Absolute Error (MAE) - Mean Squared Error (MSE) - Root Mean Squared Error (RMSE) - Mean Absolute Percentage Error (MAPE) - R-squared - Explained Variance Score - Mean Standardized Log Loss (MSLL)

Parameters:

Name	Type	Description	Default
variable_of_interest	NDArray[floating[Any]]	np.ndarray The observed values for the variable(s) of interest.	required
train_mean	NDArray[floating[Any]]	np.ndarray Mean(s) of the variable(s) of interest from the training data.	required
train_std	NDArray[floating[Any]]	np.ndarray Standard deviation(s) of the variable(s) of interest from the training data.	required
n_params	int | None	int Number of free parameters in the model.	None
msll_censored_quantile	float	float (default=0.02) Quantile below which log-likelihoods are censored for MSLL.	0.01

Returns:

Type	Description
NormativePredictions	NormativePredictions Object containing the evaluation results.

Source code in src/spectranorm/snm.py

def evaluate_predictions(
    self,
    variable_of_interest: npt.NDArray[np.floating[Any]],
    train_mean: npt.NDArray[np.floating[Any]],
    train_std: npt.NDArray[np.floating[Any]],
    n_params: int | None = None,
    msll_censored_quantile: float = 0.01,
) -> NormativePredictions:
    """
    Evaluate the predictions against the observed variable of interest.

    This function computes a battery of evaluation metrics implemented
    in `snm.utils.metrics`. Namely the evaluations include:
        - Mean Absolute Error (MAE)
        - Mean Squared Error (MSE)
        - Root Mean Squared Error (RMSE)
        - Mean Absolute Percentage Error (MAPE)
        - R-squared
        - Explained Variance Score
        - Mean Standardized Log Loss (MSLL)

    Args:
        variable_of_interest: np.ndarray
            The observed values for the variable(s) of interest.
        train_mean: np.ndarray
            Mean(s) of the variable(s) of interest from the training data.
        train_std: np.ndarray
            Standard deviation(s) of the variable(s) of interest from the training
            data.
        n_params: int
            Number of free parameters in the model.
        msll_censored_quantile: float (default=0.02)
            Quantile below which log-likelihoods are censored for MSLL.

    Returns:
        NormativePredictions
            Object containing the evaluation results.
    """
    self.extend_predictions(
        variable_of_interest,
        likelihood_censoring_quantile=msll_censored_quantile,
    )
    # Mean Absolute Error (MAE)
    self.evaluations["MAE"] = utils.metrics.compute_mae(
        y=self.predictions["variable_of_interest"],
        y_pred=self.predictions["mu_estimate"],
    )
    # Mean Squared Error (MSE)
    self.evaluations["MSE"] = utils.metrics.compute_mse(
        y=self.predictions["variable_of_interest"],
        y_pred=self.predictions["mu_estimate"],
    )
    # Root Mean Squared Error (RMSE)
    self.evaluations["RMSE"] = utils.metrics.compute_rmse(
        y=self.predictions["variable_of_interest"],
        y_pred=self.predictions["mu_estimate"],
    )
    # Mean Absolute Percentage Error (MAPE)
    self.evaluations["MAPE"] = utils.metrics.compute_mape(
        y=self.predictions["variable_of_interest"],
        y_pred=self.predictions["mu_estimate"],
    )
    # R-squared
    self.evaluations["R-squared"] = utils.metrics.compute_r2(
        y=self.predictions["variable_of_interest"],
        y_pred=self.predictions["mu_estimate"],
    )
    # Explained Variance Score
    self.evaluations["Explained Variance"] = utils.metrics.compute_expv(
        y=self.predictions["variable_of_interest"],
        y_pred=self.predictions["mu_estimate"],
    )
    # Mean Standardized Log Loss (MSLL)
    self.evaluations["MSLL"] = utils.metrics.compute_msll(
        model_log_likelihoods=self.predictions["log-likelihood"],
        baseline_log_likelihoods=utils.stats.compute_censored_log_likelihood(
            self.predictions["variable_of_interest"],
            train_mean,
            train_std,
            censored_quantile=msll_censored_quantile,
        ),
    )
    _ = n_params  # keep for future use (e.g. information criteria calculations)

    return self

extend_predictions(variable_of_interest: npt.NDArray[np.floating[Any]], *, likelihood_censoring_quantile: float = 0.01) -> NormativePredictions

Extend the NormativePredictions (predictions dictionary) with additional statistics.

Parameters:

Name	Type	Description	Default
variable_of_interest	NDArray[floating[Any]]	np.ndarray The observed values for the variable(s) of interest.	required
likelihood_censoring_quantile	float	float (default=0.01) Quantile below which log-likelihoods are censored for evaluation. Note: by default, a censored log-likelihood is computed to avoid extreme log-likelihood values for outliers. This affects several resulting metrics that are based on log-likelihoods (e.g. MSLL). If you want to compute the full (uncensored) log-likelihoods, set likelihood_censoring_quantile to 0.	0.01

Returns:

Type	Description
NormativePredictions	NormativePredictions Extended NormativePredictions with additional statistics.

Source code in src/spectranorm/snm.py

def extend_predictions(
    self,
    variable_of_interest: npt.NDArray[np.floating[Any]],
    *,
    likelihood_censoring_quantile: float = 0.01,
) -> NormativePredictions:
    """
    Extend the NormativePredictions (predictions dictionary) with additional
    statistics.

    Args:
        variable_of_interest: np.ndarray
            The observed values for the variable(s) of interest.
        likelihood_censoring_quantile: float (default=0.01)
            Quantile below which log-likelihoods are censored for evaluation.
            Note: by default, a censored log-likelihood is computed to avoid
            extreme log-likelihood values for outliers. This affects several
            resulting metrics that are based on log-likelihoods (e.g. MSLL).
            If you want to compute the full (uncensored) log-likelihoods, set
            `likelihood_censoring_quantile` to 0.

    Returns:
        NormativePredictions
            Extended NormativePredictions with additional statistics.
    """
    self.predictions["z-score"] = (
        variable_of_interest - self.predictions["mu_estimate"]
    ) / self.predictions["std_estimate"]
    self.predictions["log-likelihood"] = (
        utils.stats.compute_censored_log_likelihood(
            variable_of_interest,
            self.predictions["mu_estimate"],
            self.predictions["std_estimate"],
            censored_quantile=likelihood_censoring_quantile,
        )
    )
    self.predictions["centiles"] = utils.stats.compute_centiles_from_z_scores(
        self.predictions["z-score"],
    )

    self.predictions["variable_of_interest"] = variable_of_interest

    return self

to_array(keys: list[str] | None = None) -> npt.NDArray[np.floating[Any]]

Return prediction results as a list of NumPy arrays.

Parameters:

Name	Type	Description	Default
keys	list[str] | None	list[str] Optional list of keys to return. Defaults to ["mu_estimate", "std_estimate"].	None

Returns:

Type	Description
NDArray[floating[Any]]	list[np.ndarray] NumPy arrays for the requested predictions

Source code in src/spectranorm/snm.py

def to_array(self, keys: list[str] | None = None) -> npt.NDArray[np.floating[Any]]:
    """
    Return prediction results as a list of NumPy arrays.

    Args:
        keys: list[str]
            Optional list of keys to return.
            Defaults to ["mu_estimate", "std_estimate"].

    Returns:
        list[np.ndarray]
            NumPy arrays for the requested predictions
    """
    keys = keys or ["mu_estimate", "std_estimate"]
    return np.array([self.predictions[key] for key in keys])

to_dataframe(index: pd.Index[Any] | list[Any] | None = None) -> pd.DataFrame

Return prediction results as a DataFrame.

Parameters:

Name	Type	Description	Default
index	Index[Any] | list[Any] | None	pd.Index | list | None Optional index for the DataFrame (defaults to None)	None

Returns:

Type	Description
DataFrame	pd.DataFrame DataFrame containing the predictions

Source code in src/spectranorm/snm.py

def to_dataframe(
    self,
    index: pd.Index[Any] | list[Any] | None = None,
) -> pd.DataFrame:
    """
    Return prediction results as a DataFrame.

    Args:
        index: pd.Index | list | None
            Optional index for the DataFrame (defaults to None)

    Returns:
        pd.DataFrame
            DataFrame containing the predictions
    """
    predictions = self.predictions.copy()
    # Flatten the predictions dictionary if multiple queries are predicted
    for key in predictions:
        if predictions[key].ndim > 1:
            if predictions[key].shape[1] == 1:
                predictions[key] = predictions[key].flatten()
            else:
                for i in range(predictions[key].shape[1]):
                    predictions[f"{key}_{i + 1}"] = predictions[key][:, i]
                # delete the key
                del predictions[key]

    # Make a new DataFrame for the predictions dictionary
    return pd.DataFrame(predictions, index=index)

SpectralNormativeModel dataclass

Spectral normative model implementation.

This class implements the spectral normative modeling approach, which utilizes a base direct model to generalize normative modeling to any arbitrary variable of interest reconstructed from a graph spectral embedding. It can be used to fit a normative model to high-dimensional data and predict normative centiles for arbitrary variables of interest.

Attributes:

Name	Type	Description
eigenmode_basis	EigenmodeBasis	utils.gsp.EigenmodeBasis The eigenmode basis used for spectral normative modeling. This should be an instance of utils.gsp.EigenmodeBasis.
base_model	DirectNormativeModel	DirectNormativeModel The base (direct) normative model used for spectral normative modeling.

Source code in src/spectranorm/snm.py

@dataclass
class SpectralNormativeModel:
    """
    Spectral normative model implementation.

    This class implements the spectral normative modeling approach, which
    utilizes a base direct model to generalize normative modeling to any
    arbitrary variable of interest reconstructed from a graph spectral
    embedding. It can be used to fit a normative model to high-dimensional
    data and predict normative centiles for arbitrary variables of interest.

    Attributes:
        eigenmode_basis: utils.gsp.EigenmodeBasis
            The eigenmode basis used for spectral normative modeling. This should be an
            instance of utils.gsp.EigenmodeBasis.
        base_model: DirectNormativeModel
            The base (direct) normative model used for spectral normative modeling.
    """

    eigenmode_basis: utils.gsp.EigenmodeBasis
    base_model: DirectNormativeModel

    @classmethod
    def build_from_dataframe(
        cls,
        eigenmode_basis: utils.gsp.EigenmodeBasis,
        model_type: ModelType,
        covariates_dataframe: pd.DataFrame,
        numerical_covariates: list[str] | None = None,
        categorical_covariates: list[str] | None = None,
        batch_covariates: list[str] | None = None,
        nonlinear_covariates: list[str] | None = None,
        influencing_mean: list[str] | None = None,
        influencing_variance: list[str] | None = None,
        spline_kwargs: dict[str, Any] | None = None,
    ) -> SpectralNormativeModel:
        """
        Initialize SNM with an eigenmode basis and a base direct model built from a
        pandas DataFrame containing all covariates.

        This uses the from_dataframe method of the DirectNormativeModel class
        to populate the direct model specification of SNM. Given that SNM does not
        require a fixed variable of interest, this method assigns a dummy name
        to the variable_of_interest parameter of the DirectNormativeModel. As such,
        the provided dataframe should not contain a column with "dummy_VOI" as name.

        Essentially, the provided dataframe should contain all covariates as columns.

        Args:
            eigenmode_basis: utils.gsp.EigenmodeBasis
                The eigenmode basis to be used for spectral normative modeling.
            model_type: ModelType
                Type of the model to create, either "HBR" (Hierarchical Bayesian
                Regression) or "BLR" (Bayesian Linear Regression).
            covariates_dataframe: pd.DataFrame
                DataFrame containing the data for all covariates and all samples.
            numerical_covariates: list[str] | None
                List of numerical covariate names.
            categorical_covariates: list[str] | None
                List of categorical covariate names.
            batch_covariates: list[str] | None
                List of batch covariate names which should also be included in
                categorical_covariates.
            nonlinear_covariates: list[str] | None
                List of covariate names to be modeled as nonlinear effects.
                These should also be included in numerical_covariates.
            influencing_mean: list[str] | None
                List of covariate names that influence the mean of the variable
                of interest. These should be included in either numerical_covariates
                or categorical_covariates.
            influencing_variance: list[str] | None
                List of covariate names that influence the variance of the variable
                of interest. These should be included in either numerical_covariates
                or categorical_covariates.
            spline_kwargs: dict
                Additional keyword arguments for spline specification, such as
                `df`, `degree`, and `knots`. These are passed to the
                `create_spline_spec` method to create spline specifications for
                nonlinear covariates.

        Returns:
            SpectralNormativeModel
                An instance of SpectralNormativeModel with base model specs initialized
                based on the provided data.
        """
        # Add a dummy variable of interest to the covariates_dataframe
        covariates_dataframe = covariates_dataframe.copy()
        covariates_dataframe["dummy_VOI"] = 0.0  # Dummy variable of interest
        # Specify the base model from the dataframe
        return cls(
            eigenmode_basis=eigenmode_basis,
            base_model=DirectNormativeModel.from_dataframe(
                model_type=model_type,
                dataframe=covariates_dataframe,
                variable_of_interest="dummy_VOI",  # Dummy variable of interest
                numerical_covariates=(numerical_covariates or []),
                categorical_covariates=(categorical_covariates or []),
                batch_covariates=(batch_covariates or []),
                nonlinear_covariates=(nonlinear_covariates or []),
                influencing_mean=(influencing_mean or []),
                influencing_variance=(influencing_variance or []),
                spline_kwargs=(spline_kwargs or {}),
            ),
        )

    def save_model(self, directory: Path) -> None:
        """
        Save the fitted spectral normative model to the specified directory.

        Args:
            directory: Path
                Directory to save the fitted model. A subdirectory named
                "spectral_normative_model" will be created within this directory.
        """
        # Prepare the save directory
        directory = Path(directory)
        saved_model_dir = utils.general.prepare_save_directory(
            directory,
            "spectral_normative_model",
        )

        # Save the eigenmode basis separately
        self.eigenmode_basis.save(str(saved_model_dir / "eigenmode_basis.joblib"))

        # Save the model
        model_dict = {
            "spec": self.base_model.spec,
            "defaults": self.base_model.defaults,
        }
        if hasattr(self, "model_params"):
            model_dict["model_params"] = self.model_params
        joblib.dump(model_dict, saved_model_dir / "spectral_model_dict.joblib")

    @classmethod
    def load_model(
        cls,
        directory: Path,
        mmap_mode: MmapMode | None = "r",
    ) -> SpectralNormativeModel:
        """
        Load a spectral normative model instance from the specified save directory.

        Args:
            directory: Path
                Directory to load the fitted model from. A subdirectory named
                "spectral_normative_model" will be searched within this directory.
            mmap_mode: MmapMode | None
                Memory mapping mode for joblib (default: "r").
                You can set this to None to disable memory-mapping.
        """
        # Validate the load directory
        directory = Path(directory)
        saved_model_dir = utils.general.validate_load_directory(
            directory,
            "spectral_normative_model",
        )

        # Check if the pickled joblib file exists in this directory
        for filename in ["spectral_model_dict.joblib", "eigenmode_basis.joblib"]:
            if not (saved_model_dir / filename).exists():
                err = f"Model Load Error: Required file '{filename}' does not exist."
                raise FileNotFoundError(err)

        # Load the pickled model dictionary
        model_dict = joblib.load(saved_model_dir / "spectral_model_dict.joblib")

        # Load the eigenmode basis
        eigenmode_basis = utils.gsp.EigenmodeBasis.load(
            str(saved_model_dir / "eigenmode_basis.joblib"),
            mmap_mode=mmap_mode,
        )

        # Create an instance of the class
        instance = cls(
            eigenmode_basis=eigenmode_basis,
            base_model=DirectNormativeModel(
                spec=model_dict["spec"],
                defaults=model_dict["defaults"],
            ),
        )

        if "model_params" in model_dict:
            instance.model_params = model_dict["model_params"]

        return instance

    def _validate_fit_input(
        self,
        spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
        n_modes: int,
    ) -> None:
        """
        Internal method to validate input data for fitting the spectral normative model.
        """
        # Validate the input data
        if not isinstance(spectral_coeff_train_data, np.ndarray):
            err = "spectral_coeff_train_data must be a numpy array."
            raise TypeError(err)
        if spectral_coeff_train_data.shape[1] < n_modes:
            err = (
                f"spectral_coeff_train_data must have at least"
                f" {n_modes} columns (n_modes)."
            )
            raise ValueError(err)
        if self.eigenmode_basis.n_modes < n_modes:
            err = (
                f"Eigenmode basis has only {self.eigenmode_basis.n_modes}"
                f" modes, while {n_modes} were requested."
            )
            raise ValueError(err)

    def identify_sparse_covariance_structure(
        self,
        data: npt.NDArray[np.floating[Any]],
        sparsity_threshold: float = 1,
    ) -> npt.NDArray[np.integer[Any]]:
        """
        Identify the sparse cross-basis covariance structure in the phenotype.
        This method analyzes the phenotype's spectral coefficients to determine the
        covariance pairs that need to be modeled.

        Note: if the batches become too small, this estimate can become less stable
        in which case it is recommended to provide the sparse covariance structure
        to the model instead.

        Args:
            data: np.ndarray
                Spectral coefficients of training data representing the phenotype in
                the graph frequency domain
                :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
                as a numpy array (n_samples, n_modes).
            sparsity_threshold: float
                Number of strongest correlations to keep (proportional to the number
                of modes). Defaults to 1, meaning that the number of sparse covariance
                pairs will be equal to the number of modes. If set to a lower value,
                fewer covariance pairs will be retained.

        Returns:
            np.ndarray:
                A (N, 2) array: the rows and columns of the
                identified sparse covariance structure.
        """
        # Start with correlation structure across the whole sample
        correlations = np.corrcoef(data.T)

        # Remove self-correlations
        np.fill_diagonal(correlations, 0)

        # Extract the upper triangle of the correlation matrix
        upper_triangle_indices = np.triu_indices(correlations.shape[0], k=1)

        # Determine the number of correlations to keep
        n_correlations_to_keep = int(
            sparsity_threshold * correlations.shape[0],
        )

        # Find the cutoff value for the top correlations
        if n_correlations_to_keep < len(upper_triangle_indices[0]):
            cutoff_value = np.partition(
                np.abs(correlations[upper_triangle_indices]),
                -n_correlations_to_keep,
            )[-n_correlations_to_keep]
        else:
            cutoff_value = 0
            # Warn the user if they are keeping all correlations
            logger.warning(
                "Sparsity threshold is high, keeping all correlations.",
            )

        # Now compute the sparsity structure based on the resulting matrix
        rows, cols = np.where(np.abs(correlations) > cutoff_value)

        # Remove redundant and duplicate pairs
        rows_lim = rows[rows < cols]
        cols_lim = cols[rows < cols]

        return np.array([rows_lim, cols_lim]).T

    @staticmethod
    def _is_valid_covariance_structure(
        covariance_structure: npt.NDArray[np.integer[Any]] | float,
    ) -> bool:
        """
        Verify the validity of the sparse covariance structure.
        """
        # Check it's a 2D array with two columns
        expected_ndims = 2
        expected_ncols = 2
        if not (
            isinstance(covariance_structure, np.ndarray)
            and covariance_structure.ndim == expected_ndims
            and covariance_structure.shape[1] == expected_ncols
        ):
            return False
        return np.issubdtype(covariance_structure.dtype, np.integer)

    def fit_single_direct(
        self,
        variable_of_interest: npt.NDArray[np.floating[Any]],
        covariates_dataframe: pd.DataFrame,
        *,
        save_directory: Path | None = None,
        return_model_params: bool = True,
        adapt: dict[str, Any] | None = None,
    ) -> dict[str, Any] | None:
        """
        Fit a direct normative model for a single spectral eigenmode.
        This method fits the base direct model to the provided variable of interest
        and covariates dataframe, allowing for the model to be trained on a specific
        eigenmode of the spectral embedding.

        Args:
            variable_of_interest: np.ndarray
                The loading vector capturing the variance within training data that
                corresponds to a single eigenmode.
            covariates_dataframe: pd.DataFrame
                DataFrame containing the covariates for the samples.
            save_directory: Path | None
                Directory to save the fitted model. If None, the model is not saved.
            return_model_params: bool
                If True, return the fitted model parameters.
            adapt: dict[str, Any] | None
                Adaptation parameters from a previously fitted model. If provided,
                the model will be adapted using these parameters during fitting.

        Returns:
            dict:
                If `return_model_params` is True, return the fitted model parameters
                in a dictionary.
        """
        # Prepare the data for fitting
        train_data = covariates_dataframe.copy()
        # Add the mode loading as the variable of interest
        train_data["VOI"] = variable_of_interest

        # Instantiate a direct normative model from the base model
        direct_model = DirectNormativeModel(
            spec=NormativeModelSpec(
                variable_of_interest="VOI",  # Use the added VOI column
                covariates=self.base_model.spec.covariates,
                influencing_mean=self.base_model.spec.influencing_mean,
                influencing_variance=self.base_model.spec.influencing_variance,
            ),
            defaults=self.base_model.defaults,
        )

        # Fit the model silently
        with utils.general.suppress_output():
            direct_model.fit(
                train_data=train_data,
                save_directory=save_directory,
                progress_bar=False,
                adapt=adapt,
            )

        # Return the fitted model parameters if requested
        if return_model_params:
            return direct_model.model_params

        # If not returning model parameters, return None
        return None

    def fit_single_covariance(
        self,
        variable_of_interest_1: npt.NDArray[np.floating[Any]],
        variable_of_interest_2: npt.NDArray[np.floating[Any]],
        direct_model_params_1: dict[str, Any],
        direct_model_params_2: dict[str, Any],
        covariates_dataframe: pd.DataFrame,
        *,
        save_directory: Path | None = None,
        return_model_params: bool = True,
        defaults_overwrite: dict[str, Any] | None = None,
        adapt: dict[str, Any] | None = None,
    ) -> dict[str, Any] | None:
        """
        Fit a covariance normative model between a single pair of eigenmodes.
        This method fits a covariance model to the provided pair of variables
        and covariates dataframe, considering the direct model fits for each
        eigenmode, while allowing for the cross-eigenmode covariance to vary
        normatively.

        Args:
            variable_of_interest_1: np.ndarray
                The loading vector capturing the variance within training data that
                corresponds to a single eigenmode.
            variable_of_interest_2: np.ndarray
                The loading vector capturing the variance within training data that
                corresponds to a second eigenmode.
            direct_model_params_1: dict
                The parameters of the direct model fitted to the first eigenmode.
            direct_model_params_2: dict
                The parameters of the direct model fitted to the second eigenmode.
            covariates_dataframe: pd.DataFrame
                DataFrame containing the covariates for the samples.
            save_directory: Path | None
                Directory to save the fitted model. If None, the model is not saved.
            return_model_params: bool
                If True, return the fitted model parameters.
            defaults_overwrite: dict (default={})
                Dictionary of default values to overwrite in the model fitting process.
            adapt: dict[str, Any] | None = None
                Adaptation parameters from a previously fitted model. If provided,
                the model will be adapted using these parameters during fitting.

        Returns:
            dict:
                If `return_model_params` is True, return the fitted model parameters
                in a dictionary.
        """
        # Prepare the data for fitting
        train_data = covariates_dataframe.copy()
        # Add the respective mode loadings as the variables of interest
        train_data["VOI_1"] = variable_of_interest_1
        train_data["VOI_2"] = variable_of_interest_2
        train_data[["VOI_1_mu_estimate", "VOI_1_std_estimate"]] = (
            self.base_model.predict(
                train_data,
                model_params=direct_model_params_1,
            )
            .to_array()
            .T
        )  # Add the direct model predictions
        train_data[["VOI_2_mu_estimate", "VOI_2_std_estimate"]] = (
            self.base_model.predict(
                train_data,
                model_params=direct_model_params_2,
            )
            .to_array()
            .T
        )  # Add the direct model predictions

        # Instantiate a covariance normative model from the base model
        covariance_model = CovarianceNormativeModel.from_direct_model(
            self.base_model,
            variable_of_interest_1="VOI_1",
            variable_of_interest_2="VOI_2",
            defaults_overwrite=(defaults_overwrite or {}),
        )

        # Fit the model silently
        with utils.general.suppress_output():
            covariance_model.fit(
                train_data=train_data,
                save_directory=save_directory,
                progress_bar=False,
                adapt=adapt,
            )

        # Return the fitted model parameters if requested
        if return_model_params:
            return covariance_model.model_params

        # If not returning model parameters, return None
        return None

    def fit_all_direct(
        self,
        spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
        covariates_dataframe: pd.DataFrame,
        *,
        n_modes: int = -1,
        n_jobs: int = -1,
        save_directory: Path | None = None,
        save_separate: bool = False,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Fit the direct models for all specified eigenmodes.

        Args:
            spectral_coeff_train_data: np.ndarray
                Spectral coefficients of training data
                :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
                as a numpy array (n_samples, n_modes).
            covariates_dataframe: pd.DataFrame
                DataFrame containing the covariates for the samples.
                It must include all specified covariates in the model specification.
            n_modes: int (default=-1)
                Number of eigenmodes to fit the model to. If -1, all modes are
                used. If a positive integer, only the first n_modes are used.
                Note that the spectral_coeff_train_data and the eigenmode basis should
                have at least n_modes columns/eigenvectors.
            n_jobs: int (default=-1)
                Number of parallel jobs to use for fitting the model. If -1, all
                available CPU cores are used. If 1, no parallelization is used.
            save_directory: Path | None
                Directory to save the fitted model. If None, the model is not saved.
                A subdirectory named "spectral_normative_model" will be created
                within the specified save_directory.
            save_separate: bool (default=False)
                Whether to save the fitted direct model parameters separately for each
                eigenmode as individual files. This is only applicable if
                `save_directory` is provided.
            adapt: dict[str, Any] | None
                Adaptation parameters from a previously fitted model. If provided,
                the model will be adapted using these parameters during fitting.
        """
        # Setup the save directory if needed
        if save_directory is not None:
            save_directory = Path(save_directory)

        # Evaluate the number of modes to fit
        if n_modes == -1:
            n_modes = self.eigenmode_basis.n_modes

        # Fit the base direct model for each eigenmode using parallel processing
        tasks = (
            joblib.delayed(self.fit_single_direct)(
                variable_of_interest=spectral_coeff_train_data[:, i],
                covariates_dataframe=covariates_dataframe,
                save_directory=(
                    utils.general.ensure_dir(
                        save_directory
                        / "spectral_normative_model"
                        / "direct_models"
                        / f"mode_{i + 1}",
                    )
                    if save_directory is not None and save_separate
                    else None
                ),
                adapt=(
                    None
                    if adapt is None
                    else {
                        "covariate_to_adapt": adapt["covariate_to_adapt"],
                        "new_category_names": adapt["new_category_names"],
                        "pretrained_model_params": adapt["pretrained_model_params"][
                            "direct_model_params"
                        ][i],
                    }
                ),
            )
            for i in range(n_modes)
        )
        self.direct_model_params = list(
            utils.parallel.ParallelTqdm(
                n_jobs=n_jobs,
                total_tasks=n_modes,
                desc="Fitting direct models",
            )(tasks),  # pyright: ignore[reportCallIssue]
        )

    def identify_covariance_structure(
        self,
        spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
        covariates_dataframe: pd.DataFrame,
        n_modes: int,
        covariance_structure: npt.NDArray[np.floating[Any]] | float = 0.5,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Identify and set the sparse covariance structure for the spectral normative
        model based on the provided training data and covariance structure input.

        Args:
            spectral_coeff_train_data: np.ndarray
                Spectral coefficients of training data
                :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
                as a numpy array (n_samples, n_modes).
            covariates_dataframe: pd.DataFrame
                DataFrame containing the covariates for the samples.
            n_modes: int
                Number of eigenmodes to consider.
            covariance_structure: np.ndarray | float
                Sparse covariance structure to use for the model fitting. If a
                (2, n_pairs) array of row and column indices are provided, the model
                will use this structure. If float, the model will estimate the
                covariance structure based on the training data and the float value
                will be used as the sparsity threshold for the number of covariance
                pairs to keep proportional to the number of modes. Defaults to 0.5,
                meaning that the number of modeled sparse covariance pairs will be
                half the number of modes.
            adapt: dict[str, Any] | None
                Adaptation parameters from a previously fitted model. If provided,
                the sparse covariance structure from the pretrained model parameters
                will be used instead of estimating a new one.
        """
        if adapt is not None:
            covariance_structure = adapt["pretrained_model_params"][
                "sparse_covariance_structure"
            ]

        # Identify sparse covariance structure if a float value is given
        if isinstance(covariance_structure, float):
            # Use trained models to compute z-scores
            spectral_train_z_scores = np.array(
                [
                    self.base_model.predict(
                        test_covariates=covariates_dataframe,
                        model_params=self.direct_model_params[x],
                    )
                    .extend_predictions(
                        variable_of_interest=spectral_coeff_train_data[:, x],
                    )
                    .predictions["z-score"]
                    for x in range(n_modes)
                ],
            ).T

            self.sparse_covariance_structure = (
                self.identify_sparse_covariance_structure(
                    spectral_train_z_scores,
                    covariance_structure,
                )
            )
        else:
            self.sparse_covariance_structure = np.array(covariance_structure)

        # Verify that the covariance structure is valid
        if not self._is_valid_covariance_structure(self.sparse_covariance_structure):
            err = "Invalid sparse covariance structure."
            raise ValueError(err)

    def fit_all_covariance(
        self,
        spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
        covariates_dataframe: pd.DataFrame,
        *,
        n_jobs: int = -1,
        save_directory: Path | None = None,
        save_separate: bool = False,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Fit the direct models for all specified eigenmodes.

        Args:
            spectral_coeff_train_data: np.ndarray
                Spectral coefficients of training data
                :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
                as a numpy array (n_samples, n_modes).
            covariates_dataframe: pd.DataFrame
                DataFrame containing the covariates for the samples.
                It must include all specified covariates in the model specification.
            n_jobs: int (default=-1)
                Number of parallel jobs to use for fitting the model. If -1, all
                available CPU cores are used. If 1, no parallelization is used.
            save_directory: Path | None
                Directory to save the fitted model. If None, the model is not saved.
                A subdirectory named "spectral_normative_model" will be created
                within the specified save_directory.
            save_separate: bool (default=False)
                Whether to save the fitted direct model parameters separately for each
                eigenmode as individual files. This is only applicable if
                `save_directory` is provided.
            adapt: dict[str, Any] | None
                Adaptation parameters from a previously fitted model. If provided,
                the model will be adapted using these parameters during fitting.
        """
        # Setup the save directory if needed
        if save_directory is not None:
            save_directory = Path(save_directory)

        # Fit the base covariance models for selected eigenmode pairs in parallel
        tasks = (
            joblib.delayed(self.fit_single_covariance)(
                variable_of_interest_1=spectral_coeff_train_data[
                    :,
                    self.sparse_covariance_structure[i, 0],
                ],
                variable_of_interest_2=spectral_coeff_train_data[
                    :,
                    self.sparse_covariance_structure[i, 1],
                ],
                direct_model_params_1=self.direct_model_params[
                    self.sparse_covariance_structure[i, 0]
                ],
                direct_model_params_2=self.direct_model_params[
                    self.sparse_covariance_structure[i, 1]
                ],
                covariates_dataframe=covariates_dataframe,
                save_directory=(
                    utils.general.ensure_dir(
                        save_directory
                        / "spectral_normative_model"
                        / "covariance_models"
                        / (
                            f"mode_{self.sparse_covariance_structure[i, 0] + 1},"
                            f"mode_{self.sparse_covariance_structure[i, 1] + 1}"
                        ),
                    )
                    if save_directory is not None and save_separate
                    else None
                ),
                adapt=(
                    None
                    if adapt is None
                    else {
                        "covariate_to_adapt": adapt["covariate_to_adapt"],
                        "new_category_names": adapt["new_category_names"],
                        "pretrained_model_params": adapt["pretrained_model_params"][
                            "covariance_model_params"
                        ][i],
                    }
                ),
            )
            for i in range(self.sparse_covariance_structure.shape[0])
        )
        self.covariance_model_params = utils.parallel.ParallelTqdm(
            n_jobs=n_jobs,
            total_tasks=self.sparse_covariance_structure.shape[0],
            desc="Fitting covariance models",
        )(tasks)  # pyright: ignore[reportCallIssue]

    def fit(
        self,
        spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
        covariates_dataframe: pd.DataFrame,
        *,
        n_modes: int = -1,
        n_jobs: int = -1,
        save_directory: Path | None = None,
        save_separate: bool = False,
        covariance_structure: npt.NDArray[np.floating[Any]] | float = 0.5,
        adapt: dict[str, Any] | None = None,
    ) -> None:
        """
        Fit the spectral normative model to the provided spectral coefficient
        training data.

        Args:
            spectral_coeff_train_data: np.ndarray
                Spectral coefficients of training data
                :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
                as a numpy array (n_samples, n_modes).
            covariates_dataframe: pd.DataFrame
                DataFrame containing the covariates for the samples.
                It must include all specified covariates in the model specification.
            n_modes: int (default=-1)
                Number of eigenmodes to fit the model to. If -1, all modes are
                used. If a positive integer, only the first n_modes are used.
                Note that the spectral_coeff_train_data and the eigenmode basis should
                have at least n_modes columns/eigenvectors.
            n_jobs: int (default=-1)
                Number of parallel jobs to use for fitting the model. If -1, all
                available CPU cores are used. If 1, no parallelization is used.
            save_directory: Path | None
                Directory to save the fitted model. If None, the model is not saved.
                A subdirectory named "spectral_normative_model" will be created
                within the specified save_directory.
            save_separate: bool (default=False)
                Whether to save the fitted direct model parameters separately for each
                eigenmode as individual files. This is only applicable if
                `save_directory` is provided.
            covariance_structure: np.ndarray | float
                Sparse covariance structure to use for the model fitting. If a
                (2, n_pairs) array of row and column indices are provided, the model
                will use this structure. If float, the model will estimate the
                covariance structure based on the training data and the float value
                will be used as the sparsity threshold for the number of covariance
                pairs to keep proportional to the number of modes. Defaults to 0.5,
                meaning that the number of modeled sparse covariance pairs will be
                half the number of modes.
                Note: If using a small number of nodes, it is advisable to increase
                the sparsity threshold to ensure a stable estimation of the covariance
                structure. In contrast, when using a large number of nodes, a lower
                sparsity threshold should be used to ensure sparse modeling of the
                covariance structure.
            adapt: dict[str, Any] | None (default=None)
                If provided, adapt a pre-trained model to a new covariate.
                Note: We recommended using the `adapt_fit` method, and not directly
                changing this argument, unless you know what you are doing.
        """
        logger.info("Starting SNM model fitting:")
        # Evaluate the number of modes to fit
        if n_modes == -1:
            n_modes = self.eigenmode_basis.n_modes
        # Validate the input data
        if not isinstance(spectral_coeff_train_data, np.ndarray):
            err = "spectral_coeff_train_data must be a numpy array."
            raise TypeError(err)
        if spectral_coeff_train_data.shape[1] < n_modes:
            err = (
                f"spectral_coeff_train_data must have at least {n_modes}"
                " columns (n_modes)."
            )
            raise ValueError(err)
        if self.eigenmode_basis.n_modes < n_modes:
            err = (
                f"Eigenmode basis has only {self.eigenmode_basis.n_modes}"
                f" modes, while {n_modes} were requested."
            )
            raise ValueError(err)

        # Setup the save directory if needed
        if save_directory is not None:
            # Prepare the save directory
            save_directory = Path(save_directory)
            utils.general.prepare_save_directory(
                save_directory,
                "spectral_normative_model",
            )

        logger.info("Step 1; direct models for each eigenmode (%s modes)", n_modes)

        self.fit_all_direct(
            spectral_coeff_train_data=spectral_coeff_train_data,
            covariates_dataframe=covariates_dataframe,
            n_modes=n_modes,
            n_jobs=n_jobs,
            save_directory=save_directory,
            save_separate=save_separate,
            adapt=adapt,
        )

        logger.info("Step 2; identify sparse covariance structure")

        self.identify_covariance_structure(
            spectral_coeff_train_data=spectral_coeff_train_data,
            covariates_dataframe=covariates_dataframe,
            n_modes=n_modes,
            covariance_structure=covariance_structure,
            adapt=adapt,
        )

        # Verify that the covariance structure is valid
        if not self._is_valid_covariance_structure(self.sparse_covariance_structure):
            err = "Invalid sparse covariance structure."
            raise ValueError(err)

        # Model cross basis sparse covariance structure
        logger.info(
            "Step 3; cross-eigenmode dependency modeling (%s pairs)",
            self.sparse_covariance_structure.shape[0],
        )

        self.fit_all_covariance(
            spectral_coeff_train_data=spectral_coeff_train_data,
            covariates_dataframe=covariates_dataframe,
            n_jobs=n_jobs,
            save_directory=save_directory,
            save_separate=save_separate,
            adapt=adapt,
        )

        # Save SNM model parameters
        sample_size = spectral_coeff_train_data.shape[0]
        if adapt is not None:
            sample_size += adapt["pretrained_model_params"]["sample_size"]
        self.model_params = {
            "n_modes": n_modes,
            "sample_size": sample_size,
            "direct_model_params": self.direct_model_params,
            "sparse_covariance_structure": self.sparse_covariance_structure,
            "covariance_model_params": self.covariance_model_params,
        }
        if (self.direct_model_params[0] is not None) and (
            "n_params" in self.direct_model_params[0]
        ):
            self.model_params["n_params"] = self.direct_model_params[0]["n_params"]
        else:
            err = "Direct model parameters are not valid."
            raise ValueError(err)

        # Save the model if a save path is provided
        if save_directory is not None:
            self.save_model(save_directory)

    def adapt_fit(
        self,
        covariate_to_adapt: str,
        new_category_names: npt.NDArray[np.str_],
        spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
        covariates_dataframe: pd.DataFrame,
        *,
        pretrained_model_params: dict[str, Any] | None = None,
        n_jobs: int = -1,
        save_directory: Path | None = None,
        save_separate: bool = False,
    ) -> None:
        """
        Using a previously fitted spectral normative model, adapt to a new
        batch.
        This method enables adaptation (fine-tuning) of the model to data
        from a new batch/site by freezing all fitted parameters, and only
        estimating new parameters for the new batch/site category.

        Args:
            covariate_to_adapt: str
                Name of the categorical covariate representing the batch/site
                to which the model should be adapted.
                Note: This covariate must have been specified in the original
                model.
            new_category_names: list[str]
                Names of the new categories in the covariate_to_adapt representing
                the new batch/site labels (e.g. names of the new site).
                Note: These names must not have been present in the original
                fitted model.
            spectral_coeff_train_data: np.ndarray
                Spectral coefficients of training data
                :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
                as a numpy array (n_samples, n_modes).
            covariates_dataframe: pd.DataFrame
                DataFrame containing the covariates for the samples.
                It must include all specified covariates in the model specification.
                Note: The covariate_to_adapt column must only contain the
                new_category_names (no new data from previously trained batches).
            pretrained_model_params: dict[str, Any] | None
                The model parameters from a previously fitted model to adapt.
                If None, the model parameters from the current instance will be used
                (assuming fitting was done).
            n_jobs: int (default=-1)
                Number of parallel jobs to use for fitting the model. If -1, all
                available CPU cores are used. If 1, no parallelization is used.
            save_directory: Path | None
                A path to a directory to save the adapted model. If provided,
                the fitted model will be saved to this path.
            save_separate: bool (default=False)
                Whether to save the fitted direct model parameters separately for each
                eigenmode as individual files. This is only applicable if
                `save_directory` is provided.
        """
        # Locate the covariate to adapt
        cov_to_adapt_index = [
            cov.name for cov in self.base_model.spec.covariates
        ].index(covariate_to_adapt)

        # Extend the covariate categories to include the new categories
        self.base_model.spec.covariates[cov_to_adapt_index].extend_categories(
            new_category_names,
        )

        # Extract the pre-trained model parameters
        if pretrained_model_params is None:
            if not hasattr(self, "model_params") or self.model_params is None:
                err = (
                    "No pretrained model parameters found. "
                    "Please provide pretrained_model_params or fit the model first."
                )
                raise ValueError(err)
            pretrained_model_params = copy.deepcopy(self.model_params)

        # Fit the adapted model
        self.fit(
            spectral_coeff_train_data,
            covariates_dataframe,
            n_modes=pretrained_model_params["n_modes"],
            n_jobs=n_jobs,
            save_directory=save_directory,
            save_separate=save_separate,
            covariance_structure=pretrained_model_params["sparse_covariance_structure"],
            adapt={
                "covariate_to_adapt": covariate_to_adapt,
                "new_category_names": new_category_names,
                "pretrained_model_params": pretrained_model_params,
            },
        )

    def _predict_from_spectral_estimates(
        self,
        encoded_query: npt.NDArray[np.floating[Any]],
        eigenmode_mu_estimates: npt.NDArray[np.floating[Any]],
        eigenmode_std_estimates: npt.NDArray[np.floating[Any]],
        rho_estimates: npt.NDArray[np.floating[Any]],
        model_params: dict[str, Any],
        n_modes: int,
    ) -> NormativePredictions:
        """
        Internal method to predict only the mean and sigma for new data using the fitted
        spectral moments.
        """
        # Constrain mu and std estimates to the number of modes
        eigenmode_mu_estimates = eigenmode_mu_estimates[:, :n_modes]
        eigenmode_std_estimates = eigenmode_std_estimates[:, :n_modes]

        # Prepare the predictions
        predictions_dict = {}
        predictions_dict["mu_estimate"] = eigenmode_mu_estimates @ encoded_query

        # Load sparse covariance structure
        row_indices = model_params["sparse_covariance_structure"][:, 0]
        col_indices = model_params["sparse_covariance_structure"][:, 1]

        # Select indices that are both within n_modes
        corr_index_valid = (row_indices < n_modes) & (col_indices < n_modes)

        # Mask valid values
        valid_rho_estimates = rho_estimates[:, corr_index_valid]
        valid_row_indices = row_indices[corr_index_valid]
        valid_col_indices = col_indices[corr_index_valid]

        # number of dimensions
        n_samples = eigenmode_mu_estimates.shape[0]
        n_queries = encoded_query.shape[1]

        # Pre-compute squared encoded query
        # to have shape of (n_modes, n_queries)
        # Avoids recomputing for every sample
        encoded_query_squared = encoded_query**2

        # Empty matrix to fill in loop
        sample_query_stds = np.empty((n_samples, n_queries), dtype=rho_estimates.dtype)

        # Implement the equivalent of sparse matrix multiplications methodically
        # to avoid large memory consumption and optimize speed
        # sparse multiplication requires (n_samples, n_modes, n_queries) memory
        for sample_idx in range(n_samples):
            # Build weighted mode stds on the fly — small memory (n_modes, n_queries)
            weighted_mode_stds_sample = (
                eigenmode_std_estimates[sample_idx, :, None] * encoded_query
            )

            # diagonal term (within mode variances)
            diagonal_term = (
                eigenmode_std_estimates[sample_idx] ** 2
            ) @ encoded_query_squared
            # off-diagonal term (cross-mode covariances)
            # weighted stds for valid pairs (n_edges, n_queries)
            weighted_mode_stds_sample_row = weighted_mode_stds_sample[valid_row_indices]
            weighted_mode_stds_sample_col = weighted_mode_stds_sample[valid_col_indices]
            off_diagonal_term = 2 * (
                valid_rho_estimates[sample_idx, :, None]
                * (weighted_mode_stds_sample_row * weighted_mode_stds_sample_col)
            ).sum(axis=0)
            # avoid negative values due to numerical issues
            # if it happens, ignore off-diagonal term
            try:
                sample_query_stds[sample_idx] = np.sqrt(
                    diagonal_term + off_diagonal_term,
                )
            except RuntimeWarning:
                sample_query_stds[sample_idx] = np.sqrt(diagonal_term)

        predictions_dict["std_estimate"] = sample_query_stds

        # Create a the predictions object
        return NormativePredictions(predictions=predictions_dict)

    @staticmethod
    def _predict_single_mode_estimates(
        direct_model_spec: NormativeModelSpec,
        direct_model_defaults: dict[str, Any],
        test_covariates: pd.DataFrame,
        model_params: dict[str, Any],
        predict_without: list[str] | None = None,
    ) -> npt.NDArray[np.floating[Any]]:
        """
        Internal method to predict single mode estimates for new data using the fitted
        spectral normative model.
        """
        # Instantiate a direct normative model from the base model
        direct_model = DirectNormativeModel(
            spec=NormativeModelSpec(
                variable_of_interest="VOI",  # Use the added VOI column
                covariates=direct_model_spec.covariates,
                influencing_mean=direct_model_spec.influencing_mean,
                influencing_variance=direct_model_spec.influencing_variance,
            ),
            defaults=direct_model_defaults,
        )
        return (
            direct_model.predict(
                test_covariates,
                model_params=model_params,
                predict_without=predict_without,
            )
            .to_array(["mu_estimate", "std_estimate"])
            .T
        )

    def _predict_all_mode_estimates(
        self,
        test_covariates: pd.DataFrame,
        model_params: dict[str, Any],
        n_modes: int,
        n_jobs: int = -1,
        predict_without: list[str] | None = None,
    ) -> tuple[
        npt.NDArray[np.floating[Any]],
        npt.NDArray[np.floating[Any]],
    ]:
        """
        Internal method to predict all direct estimates for new data using the fitted
        spectral normative model.
        """
        # direct normative predictions for each eigenmode
        tasks = (
            joblib.delayed(self._predict_single_mode_estimates)(
                self.base_model.spec,
                self.base_model.defaults,
                test_covariates,
                model_params=direct_model_params,
                predict_without=predict_without,
            )
            for direct_model_params in model_params["direct_model_params"][:n_modes]
        )

        results = list(
            utils.parallel.ParallelTqdm(
                n_jobs=n_jobs,
                total_tasks=n_modes,
                desc="Computing direct eigenmode estimates",
            )(tasks),  # pyright: ignore[reportCallIssue]
        )

        # Unpack results, estimates have a shape of (n_samples, n_modes)
        eigenmode_mu_estimates, eigenmode_std_estimates = np.array(results).T

        return eigenmode_mu_estimates, eigenmode_std_estimates

    @staticmethod
    def _predict_single_covariance_estimates(
        covariance_model_spec: CovarianceModelSpec,
        covariance_model_defaults: dict[str, Any],
        test_covariates: pd.DataFrame,
        model_params: dict[str, Any],
        predict_without: list[str] | None = None,
    ) -> npt.NDArray[np.floating[Any]]:
        """
        Internal method to predict single covariance estimates for new data using the
        fitted spectral normative model.
        """
        # create a dummy covariance model
        covariance_model = CovarianceNormativeModel(
            spec=CovarianceModelSpec(
                variable_of_interest_1="VOI_1",
                variable_of_interest_2="VOI_2",
                covariates=covariance_model_spec.covariates,
                influencing_covariance=covariance_model_spec.influencing_covariance,
            ),
            defaults=covariance_model_defaults,
        )
        return (
            covariance_model.predict(
                test_covariates,
                model_params=model_params,
                predict_without=predict_without,
            )
            .to_array(["correlation_estimate"])
            .T
        )

    def _predict_all_covariance_estimates(
        self,
        test_covariates: pd.DataFrame,
        model_params: dict[str, Any],
        n_modes: int,
        n_jobs: int = -1,
        predict_without: list[str] | None = None,
    ) -> npt.NDArray[np.floating[Any]]:
        """
        Internal method to predict all covariance estimates for new data using the
        fitted spectral normative model.
        """
        # create a dummy covariance model
        covariance_model = CovarianceNormativeModel.from_direct_model(
            self.base_model,
            variable_of_interest_1="dummy_VOI_1",  # Dummy variable of interest
            variable_of_interest_2="dummy_VOI_2",  # Dummy variable of interest
        )

        # Check sparse covariance structure
        row_indices = model_params["sparse_covariance_structure"][:, 0]
        col_indices = model_params["sparse_covariance_structure"][:, 1]
        # Select indices that are within n_modes
        corr_index_valid = (row_indices < n_modes) & (col_indices < n_modes)

        # cross-mode dependence structure for valid pairs
        tasks = (
            joblib.delayed(self._predict_single_covariance_estimates)(
                covariance_model.spec,
                covariance_model.defaults,
                test_covariates,
                model_params=covariance_model_params,
                predict_without=predict_without,
            )
            for i, covariance_model_params in enumerate(
                model_params["covariance_model_params"],
            )
            if corr_index_valid[i]
        )

        results = list(
            utils.parallel.ParallelTqdm(
                n_jobs=n_jobs,
                total_tasks=np.sum(corr_index_valid),
                desc="Computing cross-mode dependence estimates",
            )(tasks),  # pyright: ignore[reportCallIssue]
        )

        # Unpack results, (n_samples, n_valid_covariance_pairs)
        valid_rho_estimates = np.array(results).T[0]

        # Now fill in the full set of rho estimates with NaNs for the invalid pairs
        rho_estimates = np.full(
            (
                test_covariates.shape[0],
                model_params["sparse_covariance_structure"].shape[0],
            ),
            np.nan,
        )
        rho_estimates[:, corr_index_valid] = valid_rho_estimates
        # final estimates have a shape of (n_samples, n_covariance_pairs)

        return rho_estimates

    def compute_spectral_predictions(
        self,
        test_covariates: pd.DataFrame,
        *,
        model_params: dict[str, Any] | None = None,
        n_modes: int | None = None,
        n_jobs: int = -1,
        predict_without: list[str] | None = None,
    ) -> dict[str, npt.NDArray[np.floating[Any]]]:
        """
        Predict normative moments (mean, std) of the eigenmode basis for new data
        using the fitted spectral normative model.

        This function requires a dataframe of covariates (test_covariates) to compute
        a set of spectral predictions that can subsequently be combined to efficiently
        estimate normative predictions for any query(ies).

        Args:
            test_covariates: pd.DataFrame
                DataFrame containing the new covariate data to predict.
                This must include all specified covariates.
                Note: covariates listed in predict_without will be ignored and are
                hence not required.
            model_params: dict | None
                Optional dictionary of model parameters to use. If not provided,
                the stored parameters from model.fit() will be used.
            n_modes: int | None
                Optional number of modes to use for the prediction. If not provided,
                the number of modes from model_params will be used.
            n_jobs: int (default=-1)
                Number of parallel jobs to utilize. If -1, all available CPU cores are
                used. If 1, no parallelization is used.
            predict_without: list[str] | None
                Optional list of covariate names to ignore during prediction.
                This can be used to check the effect of removing certain covariates
                from the model.

        Returns:
            dict:
                A dictionary containing:
                - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
                - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
                - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
        """
        # Parameters
        if model_params is None:
            model_params = self.model_params

        # Find n_modes
        if n_modes is None:
            n_modes = int(model_params["n_modes"])

        if self.base_model.spec is None:
            err = "The base model is not specified. Cannot predict new data."
            raise ValueError(err)

        # Validate the covariate data
        validation_columns = [
            cov.name
            for cov in self.base_model.spec.covariates
            if cov.name not in (predict_without or [])
        ]
        utils.general.validate_dataframe(test_covariates, validation_columns)

        # direct normative predictions for each eigenmode
        (
            eigenmode_mu_estimates,
            eigenmode_std_estimates,
        ) = self._predict_all_mode_estimates(
            test_covariates,
            model_params,
            n_modes,
            n_jobs=n_jobs,
            predict_without=predict_without,
        )  # estimates have a shape of (n_samples, n_modes)

        # cross-mode dependence structure
        rho_estimates = self._predict_all_covariance_estimates(
            test_covariates,
            model_params,
            n_modes,
            n_jobs=n_jobs,
            predict_without=predict_without,
        )  # estimates have a shape of (n_samples, n_covariance_pairs)

        return {
            "eigenmode_mu_estimates": eigenmode_mu_estimates,
            "eigenmode_std_estimates": eigenmode_std_estimates,
            "rho_estimates": rho_estimates,
        }

    def _validate_spectral_predictions(
        self,
        spectral_predictions: dict[str, npt.NDArray[np.floating[Any]]],
    ) -> None:
        """
        Internal method to validate the spectral predictions dictionary.
        """
        required_keys = [
            "eigenmode_mu_estimates",
            "eigenmode_std_estimates",
            "rho_estimates",
        ]
        if not all(key in spectral_predictions for key in required_keys):
            err = (
                "spectral_predictions must contain 'eigenmode_mu_estimates',"
                " 'eigenmode_std_estimates', and 'rho_estimates'."
            )
            raise ValueError(err)

    def predict(
        self,
        encoded_query: npt.NDArray[np.floating[Any]],
        *,
        spectral_predictions: dict[str, npt.NDArray[np.floating[Any]]] | None = None,
        test_covariates: pd.DataFrame | None = None,
        extended: bool = False,
        model_params: dict[str, Any] | None = None,
        spectral_coeff_test_data: npt.NDArray[np.floating[Any]] | None = None,
        n_modes: int | None = None,
        predict_without: list[str] | None = None,
    ) -> NormativePredictions:
        """
        Predict normative moments (mean, std) for new data using the fitted spectral
        normative model.
        Spectral normative modeling can estimate the normative distribution of any
        variable of interest defined as a spatial query encoded in the latent low-pass
        graph spectral space.

        As such, the predict method requires:
            - The encoded query(ies) defining the variable(s) of interest.

        In addition, the method requires either:
            - A dataframe of covariates (test_covariates) to be used for inference
              of a set of spectral predictions that will subsequently be combined
              to yield the normative predictions for the encoded query(ies).
            OR
            - A dictionary of precomputed spectral predictions (spectral_predictions)
              to be used for efficiently predicting the encoded query(ies).

        The precomputed spectral predictions can be obtained using the
        'compute_spectral_predictions' function. This is particularly useful when
        predicting multiple queries or when the same covariate set is used for
        multiple predictions, as it avoids redundant computations.

        Args:
            encoded_query: np.ndarray
                Encoded query data defining the normative variable of interest.
                Can be provided as:
                - shape = (n_modes) for a single query vector
                - shape = (n_modes, n_queries) for multiple queries predicted at once
            spectral_predictions: dict | None
                Optional dictionary of precomputed spectral predictions to use for
                the prediction. If not provided, test_covariates must be provided
                instead to compute the spectral predictions.
                The dictionary should contain:
                - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
                - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
                - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
                This can be obtained using the 'compute_spectral_predictions' method.
            test_covariates: pd.DataFrame | None
                DataFrame containing the new covariate data to predict.
                This must include all specified covariates.
                Note: covariates listed in predict_without will be ignored and are
                hence not required.
            extended: bool (default: False)
                If True, return additional stats such as log-likelihood, centiles, etc.
                Note that extended predictions require spectral_coeff_test_data to be
                provided in addition to the covariates.
            model_params: dict | None
                Optional dictionary of model parameters to use. If not provided,
                the stored parameters from model.fit() will be used.
            spectral_coeff_test_data: np.ndarray | None
                Optional spectral coefficient of test data for the phenotype being
                modeled :math:`(T_{test} \\Psi_{(k)}) \\in R^{N_{test} \\times k}`
                (only required for extended predictions).
                Expects a numpy array (n_samples, n_modes)
            n_modes: int | None
                Optional number of modes to use for the prediction. If not provided,
                the number of modes from model_params will be used.
            predict_without: list[str] | None
                Optional list of covariate names to ignore during prediction.
                This can be used to check the effect of removing certain covariates
                from the model.

        Returns:
            pd.DataFrame: DataFrame containing the predicted moments (mean, std) for
                the variable of interest defined by the encoded query.
        """
        # Parameters
        if model_params is None:
            model_params = self.model_params

        # Find n_modes
        if n_modes is None:
            n_modes = int(model_params["n_modes"])

        if self.base_model.spec is None:
            err = "The base model is not specified. Cannot predict new data."
            raise ValueError(err)

        if spectral_predictions is None:
            if test_covariates is None:
                err = "Either test_covariates or spectral_predictions must be provided."
                raise ValueError(err)

            # Compute spectral predictions if not provided
            spectral_predictions = self.compute_spectral_predictions(
                test_covariates=test_covariates,
                model_params=model_params,
                n_modes=n_modes,
                predict_without=predict_without,
            )
        elif test_covariates is not None:
            logger.warning(
                "Both test_covariates and spectral_predictions are provided."
                " Ignoring test_covariates and using spectral_predictions.",
            )
            if predict_without is not None:
                logger.warning(
                    "predict_without is ignored when spectral_predictions"
                    " are provided directly.",
                )

        # Unpack spectral predictions
        self._validate_spectral_predictions(spectral_predictions)
        eigenmode_mu_estimates = spectral_predictions["eigenmode_mu_estimates"]
        eigenmode_std_estimates = spectral_predictions["eigenmode_std_estimates"]
        rho_estimates = spectral_predictions["rho_estimates"]

        # Reformat encoded queries (for efficiency)
        encoded_query = np.asarray(encoded_query[:n_modes])
        encoded_query = encoded_query.reshape(n_modes, -1, order="F")

        # Compute the predictions
        predictions = self._predict_from_spectral_estimates(
            encoded_query=encoded_query,
            eigenmode_mu_estimates=eigenmode_mu_estimates,
            eigenmode_std_estimates=eigenmode_std_estimates,
            rho_estimates=rho_estimates,
            model_params=model_params,
            n_modes=n_modes,
        )

        # Check if extended predictions are requested
        if extended:
            if spectral_coeff_test_data is None:
                err = (
                    "Extended predictions require spectral_coeff_test_data"
                    " to be provided."
                )
                raise ValueError(err)
            # Add extended statistics to predictions (e.g. centiles, log-loss, etc.)
            predictions.extend_predictions(
                variable_of_interest=spectral_coeff_test_data @ encoded_query,
            )

        return predictions

    def evaluate(
        self,
        encoded_query: npt.NDArray[np.floating[Any]],
        spectral_coeff_test_data: npt.NDArray[np.floating[Any]],
        *,
        spectral_predictions: dict[str, npt.NDArray[np.floating[Any]]] | None = None,
        test_covariates: pd.DataFrame | None = None,
        query_train_moments: npt.NDArray[np.floating[Any]] | None = None,
        model_params: dict[str, Any] | None = None,
        n_modes: int | None = None,
    ) -> NormativePredictions:
        """
        Evaluate the model on new data and return predictions along with evaluation
        metrics.

        Args:
            encoded_query: np.ndarray
                Encoded query data defining the normative variable of interest.
                Can be provided as:
                - shape = (n_modes) for a single query vector
                - shape = (n_modes, n_queries) for multiple queries predicted at once
            spectral_coeff_test_data: np.ndarray | None
                Spectral coefficient of test data for the phenotype being modeled
                :math:`(T_{test} \\Psi_{(k)}) \\in R^{N_{test} \\times k}`
                (only required for extended predictions).
                Expects a numpy array (n_samples, n_modes)
            spectral_predictions: dict | None
                Optional dictionary of precomputed spectral predictions to use for
                the evaluation. If not provided, test_covariates must be provided
                instead to compute the spectral predictions.
                The dictionary should contain:
                - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
                - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
                - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
                This can be obtained using the 'compute_spectral_predictions' method.
            test_covariates: pd.DataFrame | None
                DataFrame containing the new covariate data to predict.
                This must include all specified covariates.
                Note: This is only required if spectral_predictions was not provided.
            query_train_moments: np.ndarray | None
                A (2, n_queries) array containing the query moments (mean, std) directly
                measured in the training data. While optional, providing these moments
                is strongly recommended for accurate evaluation of the model's MSLL.
                If not provided, the model will use the test data moments as an
                approximation, which may lead to overestimating MSLL. This is made
                optional to allow evaluating MSLL when the training data is not
                accessible (e.g. using a pre-trained model).
            model_params: dict | None
                Optional dictionary of model parameters to use. If not provided,
                the stored parameters from model.fit() will be used.
            n_modes: int | None
                Optional number of modes to use for the prediction. If not provided,
                the stored number of modes from model.fit() will be used.

        Returns:
            NormativePredictions:
                Object containing the predicted moments (mean, std) for
                the variable of interest defined by the encoded query, along with
                evaluation metrics.
        """
        # Find n_modes
        if n_modes is None:
            n_modes = int(self.model_params["n_modes"])

        # Parameters
        if model_params is None:
            model_params = self.model_params

        # Reformat encoded queries (for efficiency)
        encoded_query = np.asarray(encoded_query[:n_modes])
        encoded_query = encoded_query.reshape(n_modes, -1, order="F")

        # Run extended predictions
        predictions = self.predict(
            encoded_query=encoded_query,
            spectral_predictions=spectral_predictions,
            test_covariates=test_covariates,
            extended=True,
            model_params=model_params,
            spectral_coeff_test_data=spectral_coeff_test_data,
            n_modes=n_modes,
        )
        if query_train_moments is None:
            logger.warning(
                "Query moments not provided. Using test data moments as an"
                " approximation, which may lead to overestimating MSLL.",
            )
            query_train_moments = np.array(
                [
                    np.mean(spectral_coeff_test_data @ encoded_query, axis=0),
                    np.std(spectral_coeff_test_data @ encoded_query, axis=0, ddof=1),
                ],
            )
        return predictions.evaluate_predictions(
            variable_of_interest=spectral_coeff_test_data @ encoded_query,
            train_mean=query_train_moments[0],
            train_std=query_train_moments[1],
            n_params=model_params["n_params"],
        )

    def harmonize(
        self,
        encoded_query: npt.NDArray[np.floating[Any]],
        spectral_coeff_data: npt.NDArray[np.floating[Any]],
        *,
        covariates_to_harmonize: list[str] | None = None,
        covariates_dataframe: pd.DataFrame | None = None,
        spectral_predictions_full: dict[str, npt.NDArray[np.floating[Any]]]
        | None = None,
        spectral_predictions_partial: dict[str, npt.NDArray[np.floating[Any]]]
        | None = None,
        model_params: dict[str, Any] | None = None,
        n_modes: int | None = None,
    ) -> npt.NDArray[np.floating[Any]]:
        """
        Harmonize the variables of interest in the data to remove effects of
        certain covariates (e.g. batch). This method uses the spectral normative model
        to harmonize one or several variables of interest defined by the encoded query.

        The harmonization method can be used in two ways:
        - By providing a dataframe of covariates (covariates_dataframe) to compute the
          necessary spectral predictions for both the full model (all covariates)
          and the partial model (excluding covariates to harmonize). In this format,
          you should also provide the covariates_to_harmonize (list of covariate names).
        - By providing precomputed spectral predictions for both the full and partial
          models (spectral_predictions_full and spectral_predictions_partial). In this
          format, the partial spectral predictions should have been computed by
          excluding the covariates to harmonize using the predict_without parameter in
          the compute_spectral_predictions method.
          Note: In the latter, the method will not use the covariates_to_harmonize list.

        Args:
            encoded_query: np.ndarray
                Encoded query data defining the normative variable of interest.
                Can be provided as:
                - shape = (n_modes) for a single query vector
                - shape = (n_modes, n_queries) for multiple queries predicted at once
            spectral_coeff_data: np.ndarray | None
                Spectral coefficient of the the phenotype being modeled
                :math:`(T \\Psi_{(k)}) \\in R^{N_{p} \\times k}`.
                Expects a numpy array (n_samples, n_modes)
            covariates_to_harmonize: list[str] | None
                List of covariate names to harmonize.
                The partial effects of these covariates will be removed from the
                variable of interest, and the harmonized values will be returned.
                Note: This is only required if spectral_predictions_full and
                spectral_predictions_partial were not provided.
            covariates_dataframe: pd.DataFrame | None
                DataFrame containing covariate information for the data to harmonize.
                This must include all specified covariates. The dataframe is expected
                to have all covariates as columns and samples as rows.
                Note: This is only required if spectral_predictions_full and
                spectral_predictions_partial were not provided. Alternatively,
                if any of the aforementioned spectral predictions were previously
                computed, then they could be passed to this method to avoid
                recomputation.
            spectral_predictions_full: dict | None
                Optional dictionary of precomputed spectral predictions to use for
                the harmonization. If not provided, covariates_dataframe must be
                provided instead to compute the spectral predictions.
                These predictions use all set of covariates.
                The dictionary should contain:
                - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
                - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
                - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
                This can be obtained using the 'compute_spectral_predictions' method.
            spectral_predictions_partial: dict | None
                Optional dictionary of precomputed spectral predictions to use for
                the harmonization. If not provided, covariates_dataframe must be
                provided instead to compute the partial spectral predictions.
                These predictions use all set of covariates except those to harmonize.
                The covariates to harmonize need to be partialed out using the
                predict_without parameter.
                The dictionary should contain:
                - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
                - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
                - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
                This can be obtained using the 'compute_spectral_predictions' method.
            model_params: dict | None
                Optional dictionary of model parameters to use. If not provided,
                the stored parameters from model.fit() will be used.
            n_modes: int | None
                Optional number of modes to use for the prediction. If not provided,
                the stored number of modes from model.fit() will be used.

        Returns:
            npt.NDArray[np.floating[Any]]: Array of harmonized values for the
                variable of interest.
        """
        if (spectral_predictions_full is None) or (
            spectral_predictions_partial is None
        ):
            if (covariates_dataframe is None) or (covariates_to_harmonize is None):
                err = (
                    "Either [covariates_dataframe and covariates_to_harmonize] or "
                    "both spectral_predictions_full "
                    "and spectral_predictions_partial must be provided."
                )
                raise ValueError(err)
            # Validate the new data
            validation_columns = [cov.name for cov in self.base_model.spec.covariates]
            utils.general.validate_dataframe(covariates_dataframe, validation_columns)

        # Find n_modes
        if n_modes is None:
            n_modes = int(self.model_params["n_modes"])

        # Parameters
        if model_params is None:
            model_params = self.model_params

        # Reformat encoded queries (for efficiency)
        encoded_query = np.asarray(encoded_query[:n_modes])
        encoded_query = encoded_query.reshape(n_modes, -1, order="F")

        # Predict the mean and std with all covariates
        full_predictions = self.predict(
            encoded_query=encoded_query,
            spectral_predictions=spectral_predictions_full,
            test_covariates=covariates_dataframe,
            model_params=model_params,
            n_modes=n_modes,
            predict_without=[],
        )

        # Predict the mean and std without the covariates to harmonize
        reduced_predictions = self.predict(
            encoded_query=encoded_query,
            spectral_predictions=spectral_predictions_partial,
            test_covariates=covariates_dataframe,
            model_params=model_params,
            n_modes=n_modes,
            predict_without=covariates_to_harmonize,
        )

        # Reconstruct observed phenotype for query from spectral coefficients
        observed_phenotype = spectral_coeff_data @ encoded_query[:n_modes]

        # First standardize the variable of interest based on the full model
        vois_standardized = (
            observed_phenotype - full_predictions.predictions["mu_estimate"]
        ) / full_predictions.predictions["std_estimate"]

        # Then return the harmonized values based on the reduced model
        return np.asarray(
            (
                vois_standardized * reduced_predictions.predictions["std_estimate"]
                + reduced_predictions.predictions["mu_estimate"]
            ),
            dtype=np.float64,
        )

    def reduce_model(
        self,
        n_modes: int,
        *,
        inplace: bool = False,
    ) -> SpectralNormativeModel:
        """
        Create a reduced spectral normative model using only the first n_modes.

        Args:
            n_modes: int
                Number of modes to retain in the reduced model. Must be less than or
                equal to the current number of modes considered by the model.
            inplace: bool (default: False)
                If True, modify the current model instance to reduce its modes. If
                False, return a new SpectralNormativeModel instance with the reduced
                modes.
        Returns:
            SpectralNormativeModel
                A new SpectralNormativeModel instance with reduced number of modes.
        """
        if (n_modes > self.eigenmode_basis.n_modes) or (
            hasattr(self, "model_params") and n_modes > self.model_params["n_modes"]
        ):
            available_modes = self.eigenmode_basis.n_modes
            if hasattr(self, "model_params"):
                available_modes = min(available_modes, self.model_params["n_modes"])
            err = f"Cannot reduce to {n_modes} modes, only {available_modes} available."
            raise ValueError(err)

        # Create a reduced eigenbasis
        reduced_eigenbasis = self.eigenmode_basis.reduce(n_modes)

        if inplace:
            return_model = self
            return_model.eigenmode_basis = reduced_eigenbasis
        else:
            return_model = SpectralNormativeModel(
                base_model=self.base_model,
                eigenmode_basis=reduced_eigenbasis,
            )
            return_model.model_params = self.model_params  # Copy model parameters

        # Update model parameters to reflect reduced modes
        if hasattr(self, "model_params"):
            new_model_params: dict[str, Any] = {}
            new_model_params["n_modes"] = n_modes
            new_model_params["sample_size"] = self.model_params["sample_size"]
            new_model_params["direct_model_params"] = self.model_params[
                "direct_model_params"
            ][:n_modes]
            valid_cov_indices = np.where(
                (
                    return_model.model_params["sparse_covariance_structure"][:, 0]
                    < n_modes
                )
                & (
                    return_model.model_params["sparse_covariance_structure"][:, 1]
                    < n_modes
                ),
            )[0]
            new_model_params["sparse_covariance_structure"] = return_model.model_params[
                "sparse_covariance_structure"
            ][valid_cov_indices]
            new_model_params["covariance_model_params"] = [
                return_model.model_params["covariance_model_params"][i]
                for i in valid_cov_indices
            ]
            new_model_params["n_params"] = self.model_params.get("n_params", None)
            return_model.model_params = new_model_params

        return return_model

adapt_fit(covariate_to_adapt: str, new_category_names: npt.NDArray[np.str_], spectral_coeff_train_data: npt.NDArray[np.floating[Any]], covariates_dataframe: pd.DataFrame, *, pretrained_model_params: dict[str, Any] | None = None, n_jobs: int = -1, save_directory: Path | None = None, save_separate: bool = False) -> None

Using a previously fitted spectral normative model, adapt to a new batch. This method enables adaptation (fine-tuning) of the model to data from a new batch/site by freezing all fitted parameters, and only estimating new parameters for the new batch/site category.

Parameters:

Name	Type	Description	Default
covariate_to_adapt	str	str Name of the categorical covariate representing the batch/site to which the model should be adapted. Note: This covariate must have been specified in the original model.	required
new_category_names	NDArray[str_]	list[str] Names of the new categories in the covariate_to_adapt representing the new batch/site labels (e.g. names of the new site). Note: These names must not have been present in the original fitted model.	required
spectral_coeff_train_data	NDArray[floating[Any]]	np.ndarray Spectral coefficients of training data :math:(T_{train} \Psi_{(k)}) \in R^{N_p \times k} as a numpy array (n_samples, n_modes).	required
covariates_dataframe	DataFrame	pd.DataFrame DataFrame containing the covariates for the samples. It must include all specified covariates in the model specification. Note: The covariate_to_adapt column must only contain the new_category_names (no new data from previously trained batches).	required
pretrained_model_params	dict[str, Any] | None	dict[str, Any] | None The model parameters from a previously fitted model to adapt. If None, the model parameters from the current instance will be used (assuming fitting was done).	None
n_jobs	int	int (default=-1) Number of parallel jobs to use for fitting the model. If -1, all available CPU cores are used. If 1, no parallelization is used.	-1
save_directory	Path | None	Path | None A path to a directory to save the adapted model. If provided, the fitted model will be saved to this path.	None
save_separate	bool	bool (default=False) Whether to save the fitted direct model parameters separately for each eigenmode as individual files. This is only applicable if save_directory is provided.	False

Source code in src/spectranorm/snm.py

def adapt_fit(
    self,
    covariate_to_adapt: str,
    new_category_names: npt.NDArray[np.str_],
    spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
    covariates_dataframe: pd.DataFrame,
    *,
    pretrained_model_params: dict[str, Any] | None = None,
    n_jobs: int = -1,
    save_directory: Path | None = None,
    save_separate: bool = False,
) -> None:
    """
    Using a previously fitted spectral normative model, adapt to a new
    batch.
    This method enables adaptation (fine-tuning) of the model to data
    from a new batch/site by freezing all fitted parameters, and only
    estimating new parameters for the new batch/site category.

    Args:
        covariate_to_adapt: str
            Name of the categorical covariate representing the batch/site
            to which the model should be adapted.
            Note: This covariate must have been specified in the original
            model.
        new_category_names: list[str]
            Names of the new categories in the covariate_to_adapt representing
            the new batch/site labels (e.g. names of the new site).
            Note: These names must not have been present in the original
            fitted model.
        spectral_coeff_train_data: np.ndarray
            Spectral coefficients of training data
            :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
            as a numpy array (n_samples, n_modes).
        covariates_dataframe: pd.DataFrame
            DataFrame containing the covariates for the samples.
            It must include all specified covariates in the model specification.
            Note: The covariate_to_adapt column must only contain the
            new_category_names (no new data from previously trained batches).
        pretrained_model_params: dict[str, Any] | None
            The model parameters from a previously fitted model to adapt.
            If None, the model parameters from the current instance will be used
            (assuming fitting was done).
        n_jobs: int (default=-1)
            Number of parallel jobs to use for fitting the model. If -1, all
            available CPU cores are used. If 1, no parallelization is used.
        save_directory: Path | None
            A path to a directory to save the adapted model. If provided,
            the fitted model will be saved to this path.
        save_separate: bool (default=False)
            Whether to save the fitted direct model parameters separately for each
            eigenmode as individual files. This is only applicable if
            `save_directory` is provided.
    """
    # Locate the covariate to adapt
    cov_to_adapt_index = [
        cov.name for cov in self.base_model.spec.covariates
    ].index(covariate_to_adapt)

    # Extend the covariate categories to include the new categories
    self.base_model.spec.covariates[cov_to_adapt_index].extend_categories(
        new_category_names,
    )

    # Extract the pre-trained model parameters
    if pretrained_model_params is None:
        if not hasattr(self, "model_params") or self.model_params is None:
            err = (
                "No pretrained model parameters found. "
                "Please provide pretrained_model_params or fit the model first."
            )
            raise ValueError(err)
        pretrained_model_params = copy.deepcopy(self.model_params)

    # Fit the adapted model
    self.fit(
        spectral_coeff_train_data,
        covariates_dataframe,
        n_modes=pretrained_model_params["n_modes"],
        n_jobs=n_jobs,
        save_directory=save_directory,
        save_separate=save_separate,
        covariance_structure=pretrained_model_params["sparse_covariance_structure"],
        adapt={
            "covariate_to_adapt": covariate_to_adapt,
            "new_category_names": new_category_names,
            "pretrained_model_params": pretrained_model_params,
        },
    )

build_from_dataframe(eigenmode_basis: utils.gsp.EigenmodeBasis, model_type: ModelType, covariates_dataframe: pd.DataFrame, numerical_covariates: list[str] | None = None, categorical_covariates: list[str] | None = None, batch_covariates: list[str] | None = None, nonlinear_covariates: list[str] | None = None, influencing_mean: list[str] | None = None, influencing_variance: list[str] | None = None, spline_kwargs: dict[str, Any] | None = None) -> SpectralNormativeModel classmethod

Initialize SNM with an eigenmode basis and a base direct model built from a pandas DataFrame containing all covariates.

This uses the from_dataframe method of the DirectNormativeModel class to populate the direct model specification of SNM. Given that SNM does not require a fixed variable of interest, this method assigns a dummy name to the variable_of_interest parameter of the DirectNormativeModel. As such, the provided dataframe should not contain a column with "dummy_VOI" as name.

Essentially, the provided dataframe should contain all covariates as columns.

Parameters:

Name	Type	Description	Default
eigenmode_basis	EigenmodeBasis	utils.gsp.EigenmodeBasis The eigenmode basis to be used for spectral normative modeling.	required
model_type	ModelType	ModelType Type of the model to create, either "HBR" (Hierarchical Bayesian Regression) or "BLR" (Bayesian Linear Regression).	required
covariates_dataframe	DataFrame	pd.DataFrame DataFrame containing the data for all covariates and all samples.	required
numerical_covariates	list[str] | None	list[str] | None List of numerical covariate names.	None
categorical_covariates	list[str] | None	list[str] | None List of categorical covariate names.	None
batch_covariates	list[str] | None	list[str] | None List of batch covariate names which should also be included in categorical_covariates.	None
nonlinear_covariates	list[str] | None	list[str] | None List of covariate names to be modeled as nonlinear effects. These should also be included in numerical_covariates.	None
influencing_mean	list[str] | None	list[str] | None List of covariate names that influence the mean of the variable of interest. These should be included in either numerical_covariates or categorical_covariates.	None
influencing_variance	list[str] | None	list[str] | None List of covariate names that influence the variance of the variable of interest. These should be included in either numerical_covariates or categorical_covariates.	None
spline_kwargs	dict[str, Any] | None	dict Additional keyword arguments for spline specification, such as df, degree, and knots. These are passed to the create_spline_spec method to create spline specifications for nonlinear covariates.	None

Returns:

Type	Description
SpectralNormativeModel	SpectralNormativeModel An instance of SpectralNormativeModel with base model specs initialized based on the provided data.

Source code in src/spectranorm/snm.py

@classmethod
def build_from_dataframe(
    cls,
    eigenmode_basis: utils.gsp.EigenmodeBasis,
    model_type: ModelType,
    covariates_dataframe: pd.DataFrame,
    numerical_covariates: list[str] | None = None,
    categorical_covariates: list[str] | None = None,
    batch_covariates: list[str] | None = None,
    nonlinear_covariates: list[str] | None = None,
    influencing_mean: list[str] | None = None,
    influencing_variance: list[str] | None = None,
    spline_kwargs: dict[str, Any] | None = None,
) -> SpectralNormativeModel:
    """
    Initialize SNM with an eigenmode basis and a base direct model built from a
    pandas DataFrame containing all covariates.

    This uses the from_dataframe method of the DirectNormativeModel class
    to populate the direct model specification of SNM. Given that SNM does not
    require a fixed variable of interest, this method assigns a dummy name
    to the variable_of_interest parameter of the DirectNormativeModel. As such,
    the provided dataframe should not contain a column with "dummy_VOI" as name.

    Essentially, the provided dataframe should contain all covariates as columns.

    Args:
        eigenmode_basis: utils.gsp.EigenmodeBasis
            The eigenmode basis to be used for spectral normative modeling.
        model_type: ModelType
            Type of the model to create, either "HBR" (Hierarchical Bayesian
            Regression) or "BLR" (Bayesian Linear Regression).
        covariates_dataframe: pd.DataFrame
            DataFrame containing the data for all covariates and all samples.
        numerical_covariates: list[str] | None
            List of numerical covariate names.
        categorical_covariates: list[str] | None
            List of categorical covariate names.
        batch_covariates: list[str] | None
            List of batch covariate names which should also be included in
            categorical_covariates.
        nonlinear_covariates: list[str] | None
            List of covariate names to be modeled as nonlinear effects.
            These should also be included in numerical_covariates.
        influencing_mean: list[str] | None
            List of covariate names that influence the mean of the variable
            of interest. These should be included in either numerical_covariates
            or categorical_covariates.
        influencing_variance: list[str] | None
            List of covariate names that influence the variance of the variable
            of interest. These should be included in either numerical_covariates
            or categorical_covariates.
        spline_kwargs: dict
            Additional keyword arguments for spline specification, such as
            `df`, `degree`, and `knots`. These are passed to the
            `create_spline_spec` method to create spline specifications for
            nonlinear covariates.

    Returns:
        SpectralNormativeModel
            An instance of SpectralNormativeModel with base model specs initialized
            based on the provided data.
    """
    # Add a dummy variable of interest to the covariates_dataframe
    covariates_dataframe = covariates_dataframe.copy()
    covariates_dataframe["dummy_VOI"] = 0.0  # Dummy variable of interest
    # Specify the base model from the dataframe
    return cls(
        eigenmode_basis=eigenmode_basis,
        base_model=DirectNormativeModel.from_dataframe(
            model_type=model_type,
            dataframe=covariates_dataframe,
            variable_of_interest="dummy_VOI",  # Dummy variable of interest
            numerical_covariates=(numerical_covariates or []),
            categorical_covariates=(categorical_covariates or []),
            batch_covariates=(batch_covariates or []),
            nonlinear_covariates=(nonlinear_covariates or []),
            influencing_mean=(influencing_mean or []),
            influencing_variance=(influencing_variance or []),
            spline_kwargs=(spline_kwargs or {}),
        ),
    )

compute_spectral_predictions(test_covariates: pd.DataFrame, *, model_params: dict[str, Any] | None = None, n_modes: int | None = None, n_jobs: int = -1, predict_without: list[str] | None = None) -> dict[str, npt.NDArray[np.floating[Any]]]

Predict normative moments (mean, std) of the eigenmode basis for new data using the fitted spectral normative model.

This function requires a dataframe of covariates (test_covariates) to compute a set of spectral predictions that can subsequently be combined to efficiently estimate normative predictions for any query(ies).

Parameters:

Name	Type	Description	Default
test_covariates	DataFrame	pd.DataFrame DataFrame containing the new covariate data to predict. This must include all specified covariates. Note: covariates listed in predict_without will be ignored and are hence not required.	required
model_params	dict[str, Any] | None	dict | None Optional dictionary of model parameters to use. If not provided, the stored parameters from model.fit() will be used.	None
n_modes	int | None	int | None Optional number of modes to use for the prediction. If not provided, the number of modes from model_params will be used.	None
n_jobs	int	int (default=-1) Number of parallel jobs to utilize. If -1, all available CPU cores are used. If 1, no parallelization is used.	-1
predict_without	list[str] | None	list[str] | None Optional list of covariate names to ignore during prediction. This can be used to check the effect of removing certain covariates from the model.	None

Returns:

Name	Type	Description
dict	dict[str, NDArray[floating[Any]]]	A dictionary containing: - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes) - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes) - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)

Source code in src/spectranorm/snm.py

def compute_spectral_predictions(
    self,
    test_covariates: pd.DataFrame,
    *,
    model_params: dict[str, Any] | None = None,
    n_modes: int | None = None,
    n_jobs: int = -1,
    predict_without: list[str] | None = None,
) -> dict[str, npt.NDArray[np.floating[Any]]]:
    """
    Predict normative moments (mean, std) of the eigenmode basis for new data
    using the fitted spectral normative model.

    This function requires a dataframe of covariates (test_covariates) to compute
    a set of spectral predictions that can subsequently be combined to efficiently
    estimate normative predictions for any query(ies).

    Args:
        test_covariates: pd.DataFrame
            DataFrame containing the new covariate data to predict.
            This must include all specified covariates.
            Note: covariates listed in predict_without will be ignored and are
            hence not required.
        model_params: dict | None
            Optional dictionary of model parameters to use. If not provided,
            the stored parameters from model.fit() will be used.
        n_modes: int | None
            Optional number of modes to use for the prediction. If not provided,
            the number of modes from model_params will be used.
        n_jobs: int (default=-1)
            Number of parallel jobs to utilize. If -1, all available CPU cores are
            used. If 1, no parallelization is used.
        predict_without: list[str] | None
            Optional list of covariate names to ignore during prediction.
            This can be used to check the effect of removing certain covariates
            from the model.

    Returns:
        dict:
            A dictionary containing:
            - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
            - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
            - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
    """
    # Parameters
    if model_params is None:
        model_params = self.model_params

    # Find n_modes
    if n_modes is None:
        n_modes = int(model_params["n_modes"])

    if self.base_model.spec is None:
        err = "The base model is not specified. Cannot predict new data."
        raise ValueError(err)

    # Validate the covariate data
    validation_columns = [
        cov.name
        for cov in self.base_model.spec.covariates
        if cov.name not in (predict_without or [])
    ]
    utils.general.validate_dataframe(test_covariates, validation_columns)

    # direct normative predictions for each eigenmode
    (
        eigenmode_mu_estimates,
        eigenmode_std_estimates,
    ) = self._predict_all_mode_estimates(
        test_covariates,
        model_params,
        n_modes,
        n_jobs=n_jobs,
        predict_without=predict_without,
    )  # estimates have a shape of (n_samples, n_modes)

    # cross-mode dependence structure
    rho_estimates = self._predict_all_covariance_estimates(
        test_covariates,
        model_params,
        n_modes,
        n_jobs=n_jobs,
        predict_without=predict_without,
    )  # estimates have a shape of (n_samples, n_covariance_pairs)

    return {
        "eigenmode_mu_estimates": eigenmode_mu_estimates,
        "eigenmode_std_estimates": eigenmode_std_estimates,
        "rho_estimates": rho_estimates,
    }

evaluate(encoded_query: npt.NDArray[np.floating[Any]], spectral_coeff_test_data: npt.NDArray[np.floating[Any]], *, spectral_predictions: dict[str, npt.NDArray[np.floating[Any]]] | None = None, test_covariates: pd.DataFrame | None = None, query_train_moments: npt.NDArray[np.floating[Any]] | None = None, model_params: dict[str, Any] | None = None, n_modes: int | None = None) -> NormativePredictions

Evaluate the model on new data and return predictions along with evaluation metrics.

Parameters:

Name	Type	Description	Default
encoded_query	NDArray[floating[Any]]	np.ndarray Encoded query data defining the normative variable of interest. Can be provided as: - shape = (n_modes) for a single query vector - shape = (n_modes, n_queries) for multiple queries predicted at once	required
spectral_coeff_test_data	NDArray[floating[Any]]	np.ndarray | None Spectral coefficient of test data for the phenotype being modeled :math:(T_{test} \Psi_{(k)}) \in R^{N_{test} \times k} (only required for extended predictions). Expects a numpy array (n_samples, n_modes)	required
spectral_predictions	dict[str, NDArray[floating[Any]]] | None	dict | None Optional dictionary of precomputed spectral predictions to use for the evaluation. If not provided, test_covariates must be provided instead to compute the spectral predictions. The dictionary should contain: - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes) - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes) - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs) This can be obtained using the 'compute_spectral_predictions' method.	None
test_covariates	DataFrame | None	pd.DataFrame | None DataFrame containing the new covariate data to predict. This must include all specified covariates. Note: This is only required if spectral_predictions was not provided.	None
query_train_moments	NDArray[floating[Any]] | None	np.ndarray | None A (2, n_queries) array containing the query moments (mean, std) directly measured in the training data. While optional, providing these moments is strongly recommended for accurate evaluation of the model's MSLL. If not provided, the model will use the test data moments as an approximation, which may lead to overestimating MSLL. This is made optional to allow evaluating MSLL when the training data is not accessible (e.g. using a pre-trained model).	None
model_params	dict[str, Any] | None	dict | None Optional dictionary of model parameters to use. If not provided, the stored parameters from model.fit() will be used.	None
n_modes	int | None	int | None Optional number of modes to use for the prediction. If not provided, the stored number of modes from model.fit() will be used.	None

Returns:

Name	Type	Description
NormativePredictions	NormativePredictions	Object containing the predicted moments (mean, std) for the variable of interest defined by the encoded query, along with evaluation metrics.

Source code in src/spectranorm/snm.py

def evaluate(
    self,
    encoded_query: npt.NDArray[np.floating[Any]],
    spectral_coeff_test_data: npt.NDArray[np.floating[Any]],
    *,
    spectral_predictions: dict[str, npt.NDArray[np.floating[Any]]] | None = None,
    test_covariates: pd.DataFrame | None = None,
    query_train_moments: npt.NDArray[np.floating[Any]] | None = None,
    model_params: dict[str, Any] | None = None,
    n_modes: int | None = None,
) -> NormativePredictions:
    """
    Evaluate the model on new data and return predictions along with evaluation
    metrics.

    Args:
        encoded_query: np.ndarray
            Encoded query data defining the normative variable of interest.
            Can be provided as:
            - shape = (n_modes) for a single query vector
            - shape = (n_modes, n_queries) for multiple queries predicted at once
        spectral_coeff_test_data: np.ndarray | None
            Spectral coefficient of test data for the phenotype being modeled
            :math:`(T_{test} \\Psi_{(k)}) \\in R^{N_{test} \\times k}`
            (only required for extended predictions).
            Expects a numpy array (n_samples, n_modes)
        spectral_predictions: dict | None
            Optional dictionary of precomputed spectral predictions to use for
            the evaluation. If not provided, test_covariates must be provided
            instead to compute the spectral predictions.
            The dictionary should contain:
            - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
            - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
            - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
            This can be obtained using the 'compute_spectral_predictions' method.
        test_covariates: pd.DataFrame | None
            DataFrame containing the new covariate data to predict.
            This must include all specified covariates.
            Note: This is only required if spectral_predictions was not provided.
        query_train_moments: np.ndarray | None
            A (2, n_queries) array containing the query moments (mean, std) directly
            measured in the training data. While optional, providing these moments
            is strongly recommended for accurate evaluation of the model's MSLL.
            If not provided, the model will use the test data moments as an
            approximation, which may lead to overestimating MSLL. This is made
            optional to allow evaluating MSLL when the training data is not
            accessible (e.g. using a pre-trained model).
        model_params: dict | None
            Optional dictionary of model parameters to use. If not provided,
            the stored parameters from model.fit() will be used.
        n_modes: int | None
            Optional number of modes to use for the prediction. If not provided,
            the stored number of modes from model.fit() will be used.

    Returns:
        NormativePredictions:
            Object containing the predicted moments (mean, std) for
            the variable of interest defined by the encoded query, along with
            evaluation metrics.
    """
    # Find n_modes
    if n_modes is None:
        n_modes = int(self.model_params["n_modes"])

    # Parameters
    if model_params is None:
        model_params = self.model_params

    # Reformat encoded queries (for efficiency)
    encoded_query = np.asarray(encoded_query[:n_modes])
    encoded_query = encoded_query.reshape(n_modes, -1, order="F")

    # Run extended predictions
    predictions = self.predict(
        encoded_query=encoded_query,
        spectral_predictions=spectral_predictions,
        test_covariates=test_covariates,
        extended=True,
        model_params=model_params,
        spectral_coeff_test_data=spectral_coeff_test_data,
        n_modes=n_modes,
    )
    if query_train_moments is None:
        logger.warning(
            "Query moments not provided. Using test data moments as an"
            " approximation, which may lead to overestimating MSLL.",
        )
        query_train_moments = np.array(
            [
                np.mean(spectral_coeff_test_data @ encoded_query, axis=0),
                np.std(spectral_coeff_test_data @ encoded_query, axis=0, ddof=1),
            ],
        )
    return predictions.evaluate_predictions(
        variable_of_interest=spectral_coeff_test_data @ encoded_query,
        train_mean=query_train_moments[0],
        train_std=query_train_moments[1],
        n_params=model_params["n_params"],
    )

fit(spectral_coeff_train_data: npt.NDArray[np.floating[Any]], covariates_dataframe: pd.DataFrame, *, n_modes: int = -1, n_jobs: int = -1, save_directory: Path | None = None, save_separate: bool = False, covariance_structure: npt.NDArray[np.floating[Any]] | float = 0.5, adapt: dict[str, Any] | None = None) -> None

Fit the spectral normative model to the provided spectral coefficient training data.

Parameters:

Name	Type	Description	Default
spectral_coeff_train_data	NDArray[floating[Any]]	np.ndarray Spectral coefficients of training data :math:(T_{train} \Psi_{(k)}) \in R^{N_p \times k} as a numpy array (n_samples, n_modes).	required
covariates_dataframe	DataFrame	pd.DataFrame DataFrame containing the covariates for the samples. It must include all specified covariates in the model specification.	required
n_modes	int	int (default=-1) Number of eigenmodes to fit the model to. If -1, all modes are used. If a positive integer, only the first n_modes are used. Note that the spectral_coeff_train_data and the eigenmode basis should have at least n_modes columns/eigenvectors.	-1
n_jobs	int	int (default=-1) Number of parallel jobs to use for fitting the model. If -1, all available CPU cores are used. If 1, no parallelization is used.	-1
save_directory	Path | None	Path | None Directory to save the fitted model. If None, the model is not saved. A subdirectory named "spectral_normative_model" will be created within the specified save_directory.	None
save_separate	bool	bool (default=False) Whether to save the fitted direct model parameters separately for each eigenmode as individual files. This is only applicable if save_directory is provided.	False
covariance_structure	NDArray[floating[Any]] | float	np.ndarray | float Sparse covariance structure to use for the model fitting. If a (2, n_pairs) array of row and column indices are provided, the model will use this structure. If float, the model will estimate the covariance structure based on the training data and the float value will be used as the sparsity threshold for the number of covariance pairs to keep proportional to the number of modes. Defaults to 0.5, meaning that the number of modeled sparse covariance pairs will be half the number of modes. Note: If using a small number of nodes, it is advisable to increase the sparsity threshold to ensure a stable estimation of the covariance structure. In contrast, when using a large number of nodes, a lower sparsity threshold should be used to ensure sparse modeling of the covariance structure.	0.5
adapt	dict[str, Any] | None	dict[str, Any] | None (default=None) If provided, adapt a pre-trained model to a new covariate. Note: We recommended using the adapt_fit method, and not directly changing this argument, unless you know what you are doing.	None

Source code in src/spectranorm/snm.py

def fit(
    self,
    spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
    covariates_dataframe: pd.DataFrame,
    *,
    n_modes: int = -1,
    n_jobs: int = -1,
    save_directory: Path | None = None,
    save_separate: bool = False,
    covariance_structure: npt.NDArray[np.floating[Any]] | float = 0.5,
    adapt: dict[str, Any] | None = None,
) -> None:
    """
    Fit the spectral normative model to the provided spectral coefficient
    training data.

    Args:
        spectral_coeff_train_data: np.ndarray
            Spectral coefficients of training data
            :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
            as a numpy array (n_samples, n_modes).
        covariates_dataframe: pd.DataFrame
            DataFrame containing the covariates for the samples.
            It must include all specified covariates in the model specification.
        n_modes: int (default=-1)
            Number of eigenmodes to fit the model to. If -1, all modes are
            used. If a positive integer, only the first n_modes are used.
            Note that the spectral_coeff_train_data and the eigenmode basis should
            have at least n_modes columns/eigenvectors.
        n_jobs: int (default=-1)
            Number of parallel jobs to use for fitting the model. If -1, all
            available CPU cores are used. If 1, no parallelization is used.
        save_directory: Path | None
            Directory to save the fitted model. If None, the model is not saved.
            A subdirectory named "spectral_normative_model" will be created
            within the specified save_directory.
        save_separate: bool (default=False)
            Whether to save the fitted direct model parameters separately for each
            eigenmode as individual files. This is only applicable if
            `save_directory` is provided.
        covariance_structure: np.ndarray | float
            Sparse covariance structure to use for the model fitting. If a
            (2, n_pairs) array of row and column indices are provided, the model
            will use this structure. If float, the model will estimate the
            covariance structure based on the training data and the float value
            will be used as the sparsity threshold for the number of covariance
            pairs to keep proportional to the number of modes. Defaults to 0.5,
            meaning that the number of modeled sparse covariance pairs will be
            half the number of modes.
            Note: If using a small number of nodes, it is advisable to increase
            the sparsity threshold to ensure a stable estimation of the covariance
            structure. In contrast, when using a large number of nodes, a lower
            sparsity threshold should be used to ensure sparse modeling of the
            covariance structure.
        adapt: dict[str, Any] | None (default=None)
            If provided, adapt a pre-trained model to a new covariate.
            Note: We recommended using the `adapt_fit` method, and not directly
            changing this argument, unless you know what you are doing.
    """
    logger.info("Starting SNM model fitting:")
    # Evaluate the number of modes to fit
    if n_modes == -1:
        n_modes = self.eigenmode_basis.n_modes
    # Validate the input data
    if not isinstance(spectral_coeff_train_data, np.ndarray):
        err = "spectral_coeff_train_data must be a numpy array."
        raise TypeError(err)
    if spectral_coeff_train_data.shape[1] < n_modes:
        err = (
            f"spectral_coeff_train_data must have at least {n_modes}"
            " columns (n_modes)."
        )
        raise ValueError(err)
    if self.eigenmode_basis.n_modes < n_modes:
        err = (
            f"Eigenmode basis has only {self.eigenmode_basis.n_modes}"
            f" modes, while {n_modes} were requested."
        )
        raise ValueError(err)

    # Setup the save directory if needed
    if save_directory is not None:
        # Prepare the save directory
        save_directory = Path(save_directory)
        utils.general.prepare_save_directory(
            save_directory,
            "spectral_normative_model",
        )

    logger.info("Step 1; direct models for each eigenmode (%s modes)", n_modes)

    self.fit_all_direct(
        spectral_coeff_train_data=spectral_coeff_train_data,
        covariates_dataframe=covariates_dataframe,
        n_modes=n_modes,
        n_jobs=n_jobs,
        save_directory=save_directory,
        save_separate=save_separate,
        adapt=adapt,
    )

    logger.info("Step 2; identify sparse covariance structure")

    self.identify_covariance_structure(
        spectral_coeff_train_data=spectral_coeff_train_data,
        covariates_dataframe=covariates_dataframe,
        n_modes=n_modes,
        covariance_structure=covariance_structure,
        adapt=adapt,
    )

    # Verify that the covariance structure is valid
    if not self._is_valid_covariance_structure(self.sparse_covariance_structure):
        err = "Invalid sparse covariance structure."
        raise ValueError(err)

    # Model cross basis sparse covariance structure
    logger.info(
        "Step 3; cross-eigenmode dependency modeling (%s pairs)",
        self.sparse_covariance_structure.shape[0],
    )

    self.fit_all_covariance(
        spectral_coeff_train_data=spectral_coeff_train_data,
        covariates_dataframe=covariates_dataframe,
        n_jobs=n_jobs,
        save_directory=save_directory,
        save_separate=save_separate,
        adapt=adapt,
    )

    # Save SNM model parameters
    sample_size = spectral_coeff_train_data.shape[0]
    if adapt is not None:
        sample_size += adapt["pretrained_model_params"]["sample_size"]
    self.model_params = {
        "n_modes": n_modes,
        "sample_size": sample_size,
        "direct_model_params": self.direct_model_params,
        "sparse_covariance_structure": self.sparse_covariance_structure,
        "covariance_model_params": self.covariance_model_params,
    }
    if (self.direct_model_params[0] is not None) and (
        "n_params" in self.direct_model_params[0]
    ):
        self.model_params["n_params"] = self.direct_model_params[0]["n_params"]
    else:
        err = "Direct model parameters are not valid."
        raise ValueError(err)

    # Save the model if a save path is provided
    if save_directory is not None:
        self.save_model(save_directory)

fit_all_covariance(spectral_coeff_train_data: npt.NDArray[np.floating[Any]], covariates_dataframe: pd.DataFrame, *, n_jobs: int = -1, save_directory: Path | None = None, save_separate: bool = False, adapt: dict[str, Any] | None = None) -> None

Fit the direct models for all specified eigenmodes.

Parameters:

Name	Type	Description	Default
spectral_coeff_train_data	NDArray[floating[Any]]	np.ndarray Spectral coefficients of training data :math:(T_{train} \Psi_{(k)}) \in R^{N_p \times k} as a numpy array (n_samples, n_modes).	required
covariates_dataframe	DataFrame	pd.DataFrame DataFrame containing the covariates for the samples. It must include all specified covariates in the model specification.	required
n_jobs	int	int (default=-1) Number of parallel jobs to use for fitting the model. If -1, all available CPU cores are used. If 1, no parallelization is used.	-1
save_directory	Path | None	Path | None Directory to save the fitted model. If None, the model is not saved. A subdirectory named "spectral_normative_model" will be created within the specified save_directory.	None
save_separate	bool	bool (default=False) Whether to save the fitted direct model parameters separately for each eigenmode as individual files. This is only applicable if save_directory is provided.	False
adapt	dict[str, Any] | None	dict[str, Any] | None Adaptation parameters from a previously fitted model. If provided, the model will be adapted using these parameters during fitting.	None

Source code in src/spectranorm/snm.py

def fit_all_covariance(
    self,
    spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
    covariates_dataframe: pd.DataFrame,
    *,
    n_jobs: int = -1,
    save_directory: Path | None = None,
    save_separate: bool = False,
    adapt: dict[str, Any] | None = None,
) -> None:
    """
    Fit the direct models for all specified eigenmodes.

    Args:
        spectral_coeff_train_data: np.ndarray
            Spectral coefficients of training data
            :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
            as a numpy array (n_samples, n_modes).
        covariates_dataframe: pd.DataFrame
            DataFrame containing the covariates for the samples.
            It must include all specified covariates in the model specification.
        n_jobs: int (default=-1)
            Number of parallel jobs to use for fitting the model. If -1, all
            available CPU cores are used. If 1, no parallelization is used.
        save_directory: Path | None
            Directory to save the fitted model. If None, the model is not saved.
            A subdirectory named "spectral_normative_model" will be created
            within the specified save_directory.
        save_separate: bool (default=False)
            Whether to save the fitted direct model parameters separately for each
            eigenmode as individual files. This is only applicable if
            `save_directory` is provided.
        adapt: dict[str, Any] | None
            Adaptation parameters from a previously fitted model. If provided,
            the model will be adapted using these parameters during fitting.
    """
    # Setup the save directory if needed
    if save_directory is not None:
        save_directory = Path(save_directory)

    # Fit the base covariance models for selected eigenmode pairs in parallel
    tasks = (
        joblib.delayed(self.fit_single_covariance)(
            variable_of_interest_1=spectral_coeff_train_data[
                :,
                self.sparse_covariance_structure[i, 0],
            ],
            variable_of_interest_2=spectral_coeff_train_data[
                :,
                self.sparse_covariance_structure[i, 1],
            ],
            direct_model_params_1=self.direct_model_params[
                self.sparse_covariance_structure[i, 0]
            ],
            direct_model_params_2=self.direct_model_params[
                self.sparse_covariance_structure[i, 1]
            ],
            covariates_dataframe=covariates_dataframe,
            save_directory=(
                utils.general.ensure_dir(
                    save_directory
                    / "spectral_normative_model"
                    / "covariance_models"
                    / (
                        f"mode_{self.sparse_covariance_structure[i, 0] + 1},"
                        f"mode_{self.sparse_covariance_structure[i, 1] + 1}"
                    ),
                )
                if save_directory is not None and save_separate
                else None
            ),
            adapt=(
                None
                if adapt is None
                else {
                    "covariate_to_adapt": adapt["covariate_to_adapt"],
                    "new_category_names": adapt["new_category_names"],
                    "pretrained_model_params": adapt["pretrained_model_params"][
                        "covariance_model_params"
                    ][i],
                }
            ),
        )
        for i in range(self.sparse_covariance_structure.shape[0])
    )
    self.covariance_model_params = utils.parallel.ParallelTqdm(
        n_jobs=n_jobs,
        total_tasks=self.sparse_covariance_structure.shape[0],
        desc="Fitting covariance models",
    )(tasks)  # pyright: ignore[reportCallIssue]

fit_all_direct(spectral_coeff_train_data: npt.NDArray[np.floating[Any]], covariates_dataframe: pd.DataFrame, *, n_modes: int = -1, n_jobs: int = -1, save_directory: Path | None = None, save_separate: bool = False, adapt: dict[str, Any] | None = None) -> None

Fit the direct models for all specified eigenmodes.

Parameters:

Name	Type	Description	Default
spectral_coeff_train_data	NDArray[floating[Any]]	np.ndarray Spectral coefficients of training data :math:(T_{train} \Psi_{(k)}) \in R^{N_p \times k} as a numpy array (n_samples, n_modes).	required
covariates_dataframe	DataFrame	pd.DataFrame DataFrame containing the covariates for the samples. It must include all specified covariates in the model specification.	required
n_modes	int	int (default=-1) Number of eigenmodes to fit the model to. If -1, all modes are used. If a positive integer, only the first n_modes are used. Note that the spectral_coeff_train_data and the eigenmode basis should have at least n_modes columns/eigenvectors.	-1
n_jobs	int	int (default=-1) Number of parallel jobs to use for fitting the model. If -1, all available CPU cores are used. If 1, no parallelization is used.	-1
save_directory	Path | None	Path | None Directory to save the fitted model. If None, the model is not saved. A subdirectory named "spectral_normative_model" will be created within the specified save_directory.	None
save_separate	bool	bool (default=False) Whether to save the fitted direct model parameters separately for each eigenmode as individual files. This is only applicable if save_directory is provided.	False
adapt	dict[str, Any] | None	dict[str, Any] | None Adaptation parameters from a previously fitted model. If provided, the model will be adapted using these parameters during fitting.	None

Source code in src/spectranorm/snm.py

def fit_all_direct(
    self,
    spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
    covariates_dataframe: pd.DataFrame,
    *,
    n_modes: int = -1,
    n_jobs: int = -1,
    save_directory: Path | None = None,
    save_separate: bool = False,
    adapt: dict[str, Any] | None = None,
) -> None:
    """
    Fit the direct models for all specified eigenmodes.

    Args:
        spectral_coeff_train_data: np.ndarray
            Spectral coefficients of training data
            :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
            as a numpy array (n_samples, n_modes).
        covariates_dataframe: pd.DataFrame
            DataFrame containing the covariates for the samples.
            It must include all specified covariates in the model specification.
        n_modes: int (default=-1)
            Number of eigenmodes to fit the model to. If -1, all modes are
            used. If a positive integer, only the first n_modes are used.
            Note that the spectral_coeff_train_data and the eigenmode basis should
            have at least n_modes columns/eigenvectors.
        n_jobs: int (default=-1)
            Number of parallel jobs to use for fitting the model. If -1, all
            available CPU cores are used. If 1, no parallelization is used.
        save_directory: Path | None
            Directory to save the fitted model. If None, the model is not saved.
            A subdirectory named "spectral_normative_model" will be created
            within the specified save_directory.
        save_separate: bool (default=False)
            Whether to save the fitted direct model parameters separately for each
            eigenmode as individual files. This is only applicable if
            `save_directory` is provided.
        adapt: dict[str, Any] | None
            Adaptation parameters from a previously fitted model. If provided,
            the model will be adapted using these parameters during fitting.
    """
    # Setup the save directory if needed
    if save_directory is not None:
        save_directory = Path(save_directory)

    # Evaluate the number of modes to fit
    if n_modes == -1:
        n_modes = self.eigenmode_basis.n_modes

    # Fit the base direct model for each eigenmode using parallel processing
    tasks = (
        joblib.delayed(self.fit_single_direct)(
            variable_of_interest=spectral_coeff_train_data[:, i],
            covariates_dataframe=covariates_dataframe,
            save_directory=(
                utils.general.ensure_dir(
                    save_directory
                    / "spectral_normative_model"
                    / "direct_models"
                    / f"mode_{i + 1}",
                )
                if save_directory is not None and save_separate
                else None
            ),
            adapt=(
                None
                if adapt is None
                else {
                    "covariate_to_adapt": adapt["covariate_to_adapt"],
                    "new_category_names": adapt["new_category_names"],
                    "pretrained_model_params": adapt["pretrained_model_params"][
                        "direct_model_params"
                    ][i],
                }
            ),
        )
        for i in range(n_modes)
    )
    self.direct_model_params = list(
        utils.parallel.ParallelTqdm(
            n_jobs=n_jobs,
            total_tasks=n_modes,
            desc="Fitting direct models",
        )(tasks),  # pyright: ignore[reportCallIssue]
    )

fit_single_covariance(variable_of_interest_1: npt.NDArray[np.floating[Any]], variable_of_interest_2: npt.NDArray[np.floating[Any]], direct_model_params_1: dict[str, Any], direct_model_params_2: dict[str, Any], covariates_dataframe: pd.DataFrame, *, save_directory: Path | None = None, return_model_params: bool = True, defaults_overwrite: dict[str, Any] | None = None, adapt: dict[str, Any] | None = None) -> dict[str, Any] | None

Fit a covariance normative model between a single pair of eigenmodes. This method fits a covariance model to the provided pair of variables and covariates dataframe, considering the direct model fits for each eigenmode, while allowing for the cross-eigenmode covariance to vary normatively.

Parameters:

Name	Type	Description	Default
variable_of_interest_1	NDArray[floating[Any]]	np.ndarray The loading vector capturing the variance within training data that corresponds to a single eigenmode.	required
variable_of_interest_2	NDArray[floating[Any]]	np.ndarray The loading vector capturing the variance within training data that corresponds to a second eigenmode.	required
direct_model_params_1	dict[str, Any]	dict The parameters of the direct model fitted to the first eigenmode.	required
direct_model_params_2	dict[str, Any]	dict The parameters of the direct model fitted to the second eigenmode.	required
covariates_dataframe	DataFrame	pd.DataFrame DataFrame containing the covariates for the samples.	required
save_directory	Path | None	Path | None Directory to save the fitted model. If None, the model is not saved.	None
return_model_params	bool	bool If True, return the fitted model parameters.	True
defaults_overwrite	dict[str, Any] | None	dict (default={}) Dictionary of default values to overwrite in the model fitting process.	None
adapt	dict[str, Any] | None	dict[str, Any] | None = None Adaptation parameters from a previously fitted model. If provided, the model will be adapted using these parameters during fitting.	None

Returns:

Name	Type	Description
dict	dict[str, Any] | None	If return_model_params is True, return the fitted model parameters in a dictionary.

Source code in src/spectranorm/snm.py

def fit_single_covariance(
    self,
    variable_of_interest_1: npt.NDArray[np.floating[Any]],
    variable_of_interest_2: npt.NDArray[np.floating[Any]],
    direct_model_params_1: dict[str, Any],
    direct_model_params_2: dict[str, Any],
    covariates_dataframe: pd.DataFrame,
    *,
    save_directory: Path | None = None,
    return_model_params: bool = True,
    defaults_overwrite: dict[str, Any] | None = None,
    adapt: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
    """
    Fit a covariance normative model between a single pair of eigenmodes.
    This method fits a covariance model to the provided pair of variables
    and covariates dataframe, considering the direct model fits for each
    eigenmode, while allowing for the cross-eigenmode covariance to vary
    normatively.

    Args:
        variable_of_interest_1: np.ndarray
            The loading vector capturing the variance within training data that
            corresponds to a single eigenmode.
        variable_of_interest_2: np.ndarray
            The loading vector capturing the variance within training data that
            corresponds to a second eigenmode.
        direct_model_params_1: dict
            The parameters of the direct model fitted to the first eigenmode.
        direct_model_params_2: dict
            The parameters of the direct model fitted to the second eigenmode.
        covariates_dataframe: pd.DataFrame
            DataFrame containing the covariates for the samples.
        save_directory: Path | None
            Directory to save the fitted model. If None, the model is not saved.
        return_model_params: bool
            If True, return the fitted model parameters.
        defaults_overwrite: dict (default={})
            Dictionary of default values to overwrite in the model fitting process.
        adapt: dict[str, Any] | None = None
            Adaptation parameters from a previously fitted model. If provided,
            the model will be adapted using these parameters during fitting.

    Returns:
        dict:
            If `return_model_params` is True, return the fitted model parameters
            in a dictionary.
    """
    # Prepare the data for fitting
    train_data = covariates_dataframe.copy()
    # Add the respective mode loadings as the variables of interest
    train_data["VOI_1"] = variable_of_interest_1
    train_data["VOI_2"] = variable_of_interest_2
    train_data[["VOI_1_mu_estimate", "VOI_1_std_estimate"]] = (
        self.base_model.predict(
            train_data,
            model_params=direct_model_params_1,
        )
        .to_array()
        .T
    )  # Add the direct model predictions
    train_data[["VOI_2_mu_estimate", "VOI_2_std_estimate"]] = (
        self.base_model.predict(
            train_data,
            model_params=direct_model_params_2,
        )
        .to_array()
        .T
    )  # Add the direct model predictions

    # Instantiate a covariance normative model from the base model
    covariance_model = CovarianceNormativeModel.from_direct_model(
        self.base_model,
        variable_of_interest_1="VOI_1",
        variable_of_interest_2="VOI_2",
        defaults_overwrite=(defaults_overwrite or {}),
    )

    # Fit the model silently
    with utils.general.suppress_output():
        covariance_model.fit(
            train_data=train_data,
            save_directory=save_directory,
            progress_bar=False,
            adapt=adapt,
        )

    # Return the fitted model parameters if requested
    if return_model_params:
        return covariance_model.model_params

    # If not returning model parameters, return None
    return None

fit_single_direct(variable_of_interest: npt.NDArray[np.floating[Any]], covariates_dataframe: pd.DataFrame, *, save_directory: Path | None = None, return_model_params: bool = True, adapt: dict[str, Any] | None = None) -> dict[str, Any] | None

Fit a direct normative model for a single spectral eigenmode. This method fits the base direct model to the provided variable of interest and covariates dataframe, allowing for the model to be trained on a specific eigenmode of the spectral embedding.

Parameters:

Name	Type	Description	Default
variable_of_interest	NDArray[floating[Any]]	np.ndarray The loading vector capturing the variance within training data that corresponds to a single eigenmode.	required
covariates_dataframe	DataFrame	pd.DataFrame DataFrame containing the covariates for the samples.	required
save_directory	Path | None	Path | None Directory to save the fitted model. If None, the model is not saved.	None
return_model_params	bool	bool If True, return the fitted model parameters.	True
adapt	dict[str, Any] | None	dict[str, Any] | None Adaptation parameters from a previously fitted model. If provided, the model will be adapted using these parameters during fitting.	None

Returns:

Name	Type	Description
dict	dict[str, Any] | None	If return_model_params is True, return the fitted model parameters in a dictionary.

Source code in src/spectranorm/snm.py

def fit_single_direct(
    self,
    variable_of_interest: npt.NDArray[np.floating[Any]],
    covariates_dataframe: pd.DataFrame,
    *,
    save_directory: Path | None = None,
    return_model_params: bool = True,
    adapt: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
    """
    Fit a direct normative model for a single spectral eigenmode.
    This method fits the base direct model to the provided variable of interest
    and covariates dataframe, allowing for the model to be trained on a specific
    eigenmode of the spectral embedding.

    Args:
        variable_of_interest: np.ndarray
            The loading vector capturing the variance within training data that
            corresponds to a single eigenmode.
        covariates_dataframe: pd.DataFrame
            DataFrame containing the covariates for the samples.
        save_directory: Path | None
            Directory to save the fitted model. If None, the model is not saved.
        return_model_params: bool
            If True, return the fitted model parameters.
        adapt: dict[str, Any] | None
            Adaptation parameters from a previously fitted model. If provided,
            the model will be adapted using these parameters during fitting.

    Returns:
        dict:
            If `return_model_params` is True, return the fitted model parameters
            in a dictionary.
    """
    # Prepare the data for fitting
    train_data = covariates_dataframe.copy()
    # Add the mode loading as the variable of interest
    train_data["VOI"] = variable_of_interest

    # Instantiate a direct normative model from the base model
    direct_model = DirectNormativeModel(
        spec=NormativeModelSpec(
            variable_of_interest="VOI",  # Use the added VOI column
            covariates=self.base_model.spec.covariates,
            influencing_mean=self.base_model.spec.influencing_mean,
            influencing_variance=self.base_model.spec.influencing_variance,
        ),
        defaults=self.base_model.defaults,
    )

    # Fit the model silently
    with utils.general.suppress_output():
        direct_model.fit(
            train_data=train_data,
            save_directory=save_directory,
            progress_bar=False,
            adapt=adapt,
        )

    # Return the fitted model parameters if requested
    if return_model_params:
        return direct_model.model_params

    # If not returning model parameters, return None
    return None

harmonize(encoded_query: npt.NDArray[np.floating[Any]], spectral_coeff_data: npt.NDArray[np.floating[Any]], *, covariates_to_harmonize: list[str] | None = None, covariates_dataframe: pd.DataFrame | None = None, spectral_predictions_full: dict[str, npt.NDArray[np.floating[Any]]] | None = None, spectral_predictions_partial: dict[str, npt.NDArray[np.floating[Any]]] | None = None, model_params: dict[str, Any] | None = None, n_modes: int | None = None) -> npt.NDArray[np.floating[Any]]

Harmonize the variables of interest in the data to remove effects of certain covariates (e.g. batch). This method uses the spectral normative model to harmonize one or several variables of interest defined by the encoded query.

The harmonization method can be used in two ways: - By providing a dataframe of covariates (covariates_dataframe) to compute the necessary spectral predictions for both the full model (all covariates) and the partial model (excluding covariates to harmonize). In this format, you should also provide the covariates_to_harmonize (list of covariate names). - By providing precomputed spectral predictions for both the full and partial models (spectral_predictions_full and spectral_predictions_partial). In this format, the partial spectral predictions should have been computed by excluding the covariates to harmonize using the predict_without parameter in the compute_spectral_predictions method. Note: In the latter, the method will not use the covariates_to_harmonize list.

Parameters:

Name	Type	Description	Default
encoded_query	NDArray[floating[Any]]	np.ndarray Encoded query data defining the normative variable of interest. Can be provided as: - shape = (n_modes) for a single query vector - shape = (n_modes, n_queries) for multiple queries predicted at once	required
spectral_coeff_data	NDArray[floating[Any]]	np.ndarray | None Spectral coefficient of the the phenotype being modeled :math:(T \Psi_{(k)}) \in R^{N_{p} \times k}. Expects a numpy array (n_samples, n_modes)	required
covariates_to_harmonize	list[str] | None	list[str] | None List of covariate names to harmonize. The partial effects of these covariates will be removed from the variable of interest, and the harmonized values will be returned. Note: This is only required if spectral_predictions_full and spectral_predictions_partial were not provided.	None
covariates_dataframe	DataFrame | None	pd.DataFrame | None DataFrame containing covariate information for the data to harmonize. This must include all specified covariates. The dataframe is expected to have all covariates as columns and samples as rows. Note: This is only required if spectral_predictions_full and spectral_predictions_partial were not provided. Alternatively, if any of the aforementioned spectral predictions were previously computed, then they could be passed to this method to avoid recomputation.	None
spectral_predictions_full	dict[str, NDArray[floating[Any]]] | None	dict | None Optional dictionary of precomputed spectral predictions to use for the harmonization. If not provided, covariates_dataframe must be provided instead to compute the spectral predictions. These predictions use all set of covariates. The dictionary should contain: - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes) - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes) - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs) This can be obtained using the 'compute_spectral_predictions' method.	None
spectral_predictions_partial	dict[str, NDArray[floating[Any]]] | None	dict | None Optional dictionary of precomputed spectral predictions to use for the harmonization. If not provided, covariates_dataframe must be provided instead to compute the partial spectral predictions. These predictions use all set of covariates except those to harmonize. The covariates to harmonize need to be partialed out using the predict_without parameter. The dictionary should contain: - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes) - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes) - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs) This can be obtained using the 'compute_spectral_predictions' method.	None
model_params	dict[str, Any] | None	dict | None Optional dictionary of model parameters to use. If not provided, the stored parameters from model.fit() will be used.	None
n_modes	int | None	int | None Optional number of modes to use for the prediction. If not provided, the stored number of modes from model.fit() will be used.	None

Returns:

Type	Description
NDArray[floating[Any]]	npt.NDArray[np.floating[Any]]: Array of harmonized values for the variable of interest.

Source code in src/spectranorm/snm.py

def harmonize(
    self,
    encoded_query: npt.NDArray[np.floating[Any]],
    spectral_coeff_data: npt.NDArray[np.floating[Any]],
    *,
    covariates_to_harmonize: list[str] | None = None,
    covariates_dataframe: pd.DataFrame | None = None,
    spectral_predictions_full: dict[str, npt.NDArray[np.floating[Any]]]
    | None = None,
    spectral_predictions_partial: dict[str, npt.NDArray[np.floating[Any]]]
    | None = None,
    model_params: dict[str, Any] | None = None,
    n_modes: int | None = None,
) -> npt.NDArray[np.floating[Any]]:
    """
    Harmonize the variables of interest in the data to remove effects of
    certain covariates (e.g. batch). This method uses the spectral normative model
    to harmonize one or several variables of interest defined by the encoded query.

    The harmonization method can be used in two ways:
    - By providing a dataframe of covariates (covariates_dataframe) to compute the
      necessary spectral predictions for both the full model (all covariates)
      and the partial model (excluding covariates to harmonize). In this format,
      you should also provide the covariates_to_harmonize (list of covariate names).
    - By providing precomputed spectral predictions for both the full and partial
      models (spectral_predictions_full and spectral_predictions_partial). In this
      format, the partial spectral predictions should have been computed by
      excluding the covariates to harmonize using the predict_without parameter in
      the compute_spectral_predictions method.
      Note: In the latter, the method will not use the covariates_to_harmonize list.

    Args:
        encoded_query: np.ndarray
            Encoded query data defining the normative variable of interest.
            Can be provided as:
            - shape = (n_modes) for a single query vector
            - shape = (n_modes, n_queries) for multiple queries predicted at once
        spectral_coeff_data: np.ndarray | None
            Spectral coefficient of the the phenotype being modeled
            :math:`(T \\Psi_{(k)}) \\in R^{N_{p} \\times k}`.
            Expects a numpy array (n_samples, n_modes)
        covariates_to_harmonize: list[str] | None
            List of covariate names to harmonize.
            The partial effects of these covariates will be removed from the
            variable of interest, and the harmonized values will be returned.
            Note: This is only required if spectral_predictions_full and
            spectral_predictions_partial were not provided.
        covariates_dataframe: pd.DataFrame | None
            DataFrame containing covariate information for the data to harmonize.
            This must include all specified covariates. The dataframe is expected
            to have all covariates as columns and samples as rows.
            Note: This is only required if spectral_predictions_full and
            spectral_predictions_partial were not provided. Alternatively,
            if any of the aforementioned spectral predictions were previously
            computed, then they could be passed to this method to avoid
            recomputation.
        spectral_predictions_full: dict | None
            Optional dictionary of precomputed spectral predictions to use for
            the harmonization. If not provided, covariates_dataframe must be
            provided instead to compute the spectral predictions.
            These predictions use all set of covariates.
            The dictionary should contain:
            - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
            - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
            - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
            This can be obtained using the 'compute_spectral_predictions' method.
        spectral_predictions_partial: dict | None
            Optional dictionary of precomputed spectral predictions to use for
            the harmonization. If not provided, covariates_dataframe must be
            provided instead to compute the partial spectral predictions.
            These predictions use all set of covariates except those to harmonize.
            The covariates to harmonize need to be partialed out using the
            predict_without parameter.
            The dictionary should contain:
            - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
            - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
            - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
            This can be obtained using the 'compute_spectral_predictions' method.
        model_params: dict | None
            Optional dictionary of model parameters to use. If not provided,
            the stored parameters from model.fit() will be used.
        n_modes: int | None
            Optional number of modes to use for the prediction. If not provided,
            the stored number of modes from model.fit() will be used.

    Returns:
        npt.NDArray[np.floating[Any]]: Array of harmonized values for the
            variable of interest.
    """
    if (spectral_predictions_full is None) or (
        spectral_predictions_partial is None
    ):
        if (covariates_dataframe is None) or (covariates_to_harmonize is None):
            err = (
                "Either [covariates_dataframe and covariates_to_harmonize] or "
                "both spectral_predictions_full "
                "and spectral_predictions_partial must be provided."
            )
            raise ValueError(err)
        # Validate the new data
        validation_columns = [cov.name for cov in self.base_model.spec.covariates]
        utils.general.validate_dataframe(covariates_dataframe, validation_columns)

    # Find n_modes
    if n_modes is None:
        n_modes = int(self.model_params["n_modes"])

    # Parameters
    if model_params is None:
        model_params = self.model_params

    # Reformat encoded queries (for efficiency)
    encoded_query = np.asarray(encoded_query[:n_modes])
    encoded_query = encoded_query.reshape(n_modes, -1, order="F")

    # Predict the mean and std with all covariates
    full_predictions = self.predict(
        encoded_query=encoded_query,
        spectral_predictions=spectral_predictions_full,
        test_covariates=covariates_dataframe,
        model_params=model_params,
        n_modes=n_modes,
        predict_without=[],
    )

    # Predict the mean and std without the covariates to harmonize
    reduced_predictions = self.predict(
        encoded_query=encoded_query,
        spectral_predictions=spectral_predictions_partial,
        test_covariates=covariates_dataframe,
        model_params=model_params,
        n_modes=n_modes,
        predict_without=covariates_to_harmonize,
    )

    # Reconstruct observed phenotype for query from spectral coefficients
    observed_phenotype = spectral_coeff_data @ encoded_query[:n_modes]

    # First standardize the variable of interest based on the full model
    vois_standardized = (
        observed_phenotype - full_predictions.predictions["mu_estimate"]
    ) / full_predictions.predictions["std_estimate"]

    # Then return the harmonized values based on the reduced model
    return np.asarray(
        (
            vois_standardized * reduced_predictions.predictions["std_estimate"]
            + reduced_predictions.predictions["mu_estimate"]
        ),
        dtype=np.float64,
    )

identify_covariance_structure(spectral_coeff_train_data: npt.NDArray[np.floating[Any]], covariates_dataframe: pd.DataFrame, n_modes: int, covariance_structure: npt.NDArray[np.floating[Any]] | float = 0.5, adapt: dict[str, Any] | None = None) -> None

Identify and set the sparse covariance structure for the spectral normative model based on the provided training data and covariance structure input.

Parameters:

Name	Type	Description	Default
spectral_coeff_train_data	NDArray[floating[Any]]	np.ndarray Spectral coefficients of training data :math:(T_{train} \Psi_{(k)}) \in R^{N_p \times k} as a numpy array (n_samples, n_modes).	required
covariates_dataframe	DataFrame	pd.DataFrame DataFrame containing the covariates for the samples.	required
n_modes	int	int Number of eigenmodes to consider.	required
covariance_structure	NDArray[floating[Any]] | float	np.ndarray | float Sparse covariance structure to use for the model fitting. If a (2, n_pairs) array of row and column indices are provided, the model will use this structure. If float, the model will estimate the covariance structure based on the training data and the float value will be used as the sparsity threshold for the number of covariance pairs to keep proportional to the number of modes. Defaults to 0.5, meaning that the number of modeled sparse covariance pairs will be half the number of modes.	0.5
adapt	dict[str, Any] | None	dict[str, Any] | None Adaptation parameters from a previously fitted model. If provided, the sparse covariance structure from the pretrained model parameters will be used instead of estimating a new one.	None

Source code in src/spectranorm/snm.py

def identify_covariance_structure(
    self,
    spectral_coeff_train_data: npt.NDArray[np.floating[Any]],
    covariates_dataframe: pd.DataFrame,
    n_modes: int,
    covariance_structure: npt.NDArray[np.floating[Any]] | float = 0.5,
    adapt: dict[str, Any] | None = None,
) -> None:
    """
    Identify and set the sparse covariance structure for the spectral normative
    model based on the provided training data and covariance structure input.

    Args:
        spectral_coeff_train_data: np.ndarray
            Spectral coefficients of training data
            :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
            as a numpy array (n_samples, n_modes).
        covariates_dataframe: pd.DataFrame
            DataFrame containing the covariates for the samples.
        n_modes: int
            Number of eigenmodes to consider.
        covariance_structure: np.ndarray | float
            Sparse covariance structure to use for the model fitting. If a
            (2, n_pairs) array of row and column indices are provided, the model
            will use this structure. If float, the model will estimate the
            covariance structure based on the training data and the float value
            will be used as the sparsity threshold for the number of covariance
            pairs to keep proportional to the number of modes. Defaults to 0.5,
            meaning that the number of modeled sparse covariance pairs will be
            half the number of modes.
        adapt: dict[str, Any] | None
            Adaptation parameters from a previously fitted model. If provided,
            the sparse covariance structure from the pretrained model parameters
            will be used instead of estimating a new one.
    """
    if adapt is not None:
        covariance_structure = adapt["pretrained_model_params"][
            "sparse_covariance_structure"
        ]

    # Identify sparse covariance structure if a float value is given
    if isinstance(covariance_structure, float):
        # Use trained models to compute z-scores
        spectral_train_z_scores = np.array(
            [
                self.base_model.predict(
                    test_covariates=covariates_dataframe,
                    model_params=self.direct_model_params[x],
                )
                .extend_predictions(
                    variable_of_interest=spectral_coeff_train_data[:, x],
                )
                .predictions["z-score"]
                for x in range(n_modes)
            ],
        ).T

        self.sparse_covariance_structure = (
            self.identify_sparse_covariance_structure(
                spectral_train_z_scores,
                covariance_structure,
            )
        )
    else:
        self.sparse_covariance_structure = np.array(covariance_structure)

    # Verify that the covariance structure is valid
    if not self._is_valid_covariance_structure(self.sparse_covariance_structure):
        err = "Invalid sparse covariance structure."
        raise ValueError(err)

identify_sparse_covariance_structure(data: npt.NDArray[np.floating[Any]], sparsity_threshold: float = 1) -> npt.NDArray[np.integer[Any]]

Identify the sparse cross-basis covariance structure in the phenotype. This method analyzes the phenotype's spectral coefficients to determine the covariance pairs that need to be modeled.

Note: if the batches become too small, this estimate can become less stable in which case it is recommended to provide the sparse covariance structure to the model instead.

Parameters:

Name	Type	Description	Default
data	NDArray[floating[Any]]	np.ndarray Spectral coefficients of training data representing the phenotype in the graph frequency domain :math:(T_{train} \Psi_{(k)}) \in R^{N_p \times k} as a numpy array (n_samples, n_modes).	required
sparsity_threshold	float	float Number of strongest correlations to keep (proportional to the number of modes). Defaults to 1, meaning that the number of sparse covariance pairs will be equal to the number of modes. If set to a lower value, fewer covariance pairs will be retained.	1

Returns:

Type	Description
NDArray[integer[Any]]	np.ndarray: A (N, 2) array: the rows and columns of the identified sparse covariance structure.

Source code in src/spectranorm/snm.py

def identify_sparse_covariance_structure(
    self,
    data: npt.NDArray[np.floating[Any]],
    sparsity_threshold: float = 1,
) -> npt.NDArray[np.integer[Any]]:
    """
    Identify the sparse cross-basis covariance structure in the phenotype.
    This method analyzes the phenotype's spectral coefficients to determine the
    covariance pairs that need to be modeled.

    Note: if the batches become too small, this estimate can become less stable
    in which case it is recommended to provide the sparse covariance structure
    to the model instead.

    Args:
        data: np.ndarray
            Spectral coefficients of training data representing the phenotype in
            the graph frequency domain
            :math:`(T_{train} \\Psi_{(k)}) \\in R^{N_p \\times k}`
            as a numpy array (n_samples, n_modes).
        sparsity_threshold: float
            Number of strongest correlations to keep (proportional to the number
            of modes). Defaults to 1, meaning that the number of sparse covariance
            pairs will be equal to the number of modes. If set to a lower value,
            fewer covariance pairs will be retained.

    Returns:
        np.ndarray:
            A (N, 2) array: the rows and columns of the
            identified sparse covariance structure.
    """
    # Start with correlation structure across the whole sample
    correlations = np.corrcoef(data.T)

    # Remove self-correlations
    np.fill_diagonal(correlations, 0)

    # Extract the upper triangle of the correlation matrix
    upper_triangle_indices = np.triu_indices(correlations.shape[0], k=1)

    # Determine the number of correlations to keep
    n_correlations_to_keep = int(
        sparsity_threshold * correlations.shape[0],
    )

    # Find the cutoff value for the top correlations
    if n_correlations_to_keep < len(upper_triangle_indices[0]):
        cutoff_value = np.partition(
            np.abs(correlations[upper_triangle_indices]),
            -n_correlations_to_keep,
        )[-n_correlations_to_keep]
    else:
        cutoff_value = 0
        # Warn the user if they are keeping all correlations
        logger.warning(
            "Sparsity threshold is high, keeping all correlations.",
        )

    # Now compute the sparsity structure based on the resulting matrix
    rows, cols = np.where(np.abs(correlations) > cutoff_value)

    # Remove redundant and duplicate pairs
    rows_lim = rows[rows < cols]
    cols_lim = cols[rows < cols]

    return np.array([rows_lim, cols_lim]).T

load_model(directory: Path, mmap_mode: MmapMode | None = 'r') -> SpectralNormativeModel classmethod

Load a spectral normative model instance from the specified save directory.

Parameters:

Name	Type	Description	Default
directory	Path	Path Directory to load the fitted model from. A subdirectory named "spectral_normative_model" will be searched within this directory.	required
mmap_mode	MmapMode | None	MmapMode | None Memory mapping mode for joblib (default: "r"). You can set this to None to disable memory-mapping.	'r'

Source code in src/spectranorm/snm.py

@classmethod
def load_model(
    cls,
    directory: Path,
    mmap_mode: MmapMode | None = "r",
) -> SpectralNormativeModel:
    """
    Load a spectral normative model instance from the specified save directory.

    Args:
        directory: Path
            Directory to load the fitted model from. A subdirectory named
            "spectral_normative_model" will be searched within this directory.
        mmap_mode: MmapMode | None
            Memory mapping mode for joblib (default: "r").
            You can set this to None to disable memory-mapping.
    """
    # Validate the load directory
    directory = Path(directory)
    saved_model_dir = utils.general.validate_load_directory(
        directory,
        "spectral_normative_model",
    )

    # Check if the pickled joblib file exists in this directory
    for filename in ["spectral_model_dict.joblib", "eigenmode_basis.joblib"]:
        if not (saved_model_dir / filename).exists():
            err = f"Model Load Error: Required file '{filename}' does not exist."
            raise FileNotFoundError(err)

    # Load the pickled model dictionary
    model_dict = joblib.load(saved_model_dir / "spectral_model_dict.joblib")

    # Load the eigenmode basis
    eigenmode_basis = utils.gsp.EigenmodeBasis.load(
        str(saved_model_dir / "eigenmode_basis.joblib"),
        mmap_mode=mmap_mode,
    )

    # Create an instance of the class
    instance = cls(
        eigenmode_basis=eigenmode_basis,
        base_model=DirectNormativeModel(
            spec=model_dict["spec"],
            defaults=model_dict["defaults"],
        ),
    )

    if "model_params" in model_dict:
        instance.model_params = model_dict["model_params"]

    return instance

predict(encoded_query: npt.NDArray[np.floating[Any]], *, spectral_predictions: dict[str, npt.NDArray[np.floating[Any]]] | None = None, test_covariates: pd.DataFrame | None = None, extended: bool = False, model_params: dict[str, Any] | None = None, spectral_coeff_test_data: npt.NDArray[np.floating[Any]] | None = None, n_modes: int | None = None, predict_without: list[str] | None = None) -> NormativePredictions

Predict normative moments (mean, std) for new data using the fitted spectral normative model. Spectral normative modeling can estimate the normative distribution of any variable of interest defined as a spatial query encoded in the latent low-pass graph spectral space.

As such, the predict method requires: - The encoded query(ies) defining the variable(s) of interest.

In addition, the method requires either: - A dataframe of covariates (test_covariates) to be used for inference of a set of spectral predictions that will subsequently be combined to yield the normative predictions for the encoded query(ies). OR - A dictionary of precomputed spectral predictions (spectral_predictions) to be used for efficiently predicting the encoded query(ies).

The precomputed spectral predictions can be obtained using the 'compute_spectral_predictions' function. This is particularly useful when predicting multiple queries or when the same covariate set is used for multiple predictions, as it avoids redundant computations.

Parameters:

Name	Type	Description	Default
encoded_query	NDArray[floating[Any]]	np.ndarray Encoded query data defining the normative variable of interest. Can be provided as: - shape = (n_modes) for a single query vector - shape = (n_modes, n_queries) for multiple queries predicted at once	required
spectral_predictions	dict[str, NDArray[floating[Any]]] | None	dict | None Optional dictionary of precomputed spectral predictions to use for the prediction. If not provided, test_covariates must be provided instead to compute the spectral predictions. The dictionary should contain: - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes) - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes) - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs) This can be obtained using the 'compute_spectral_predictions' method.	None
test_covariates	DataFrame | None	pd.DataFrame | None DataFrame containing the new covariate data to predict. This must include all specified covariates. Note: covariates listed in predict_without will be ignored and are hence not required.	None
extended	bool	bool (default: False) If True, return additional stats such as log-likelihood, centiles, etc. Note that extended predictions require spectral_coeff_test_data to be provided in addition to the covariates.	False
model_params	dict[str, Any] | None	dict | None Optional dictionary of model parameters to use. If not provided, the stored parameters from model.fit() will be used.	None
spectral_coeff_test_data	NDArray[floating[Any]] | None	np.ndarray | None Optional spectral coefficient of test data for the phenotype being modeled :math:(T_{test} \Psi_{(k)}) \in R^{N_{test} \times k} (only required for extended predictions). Expects a numpy array (n_samples, n_modes)	None
n_modes	int | None	int | None Optional number of modes to use for the prediction. If not provided, the number of modes from model_params will be used.	None
predict_without	list[str] | None	list[str] | None Optional list of covariate names to ignore during prediction. This can be used to check the effect of removing certain covariates from the model.	None

Returns:

Type	Description
NormativePredictions	pd.DataFrame: DataFrame containing the predicted moments (mean, std) for the variable of interest defined by the encoded query.

Source code in src/spectranorm/snm.py

def predict(
    self,
    encoded_query: npt.NDArray[np.floating[Any]],
    *,
    spectral_predictions: dict[str, npt.NDArray[np.floating[Any]]] | None = None,
    test_covariates: pd.DataFrame | None = None,
    extended: bool = False,
    model_params: dict[str, Any] | None = None,
    spectral_coeff_test_data: npt.NDArray[np.floating[Any]] | None = None,
    n_modes: int | None = None,
    predict_without: list[str] | None = None,
) -> NormativePredictions:
    """
    Predict normative moments (mean, std) for new data using the fitted spectral
    normative model.
    Spectral normative modeling can estimate the normative distribution of any
    variable of interest defined as a spatial query encoded in the latent low-pass
    graph spectral space.

    As such, the predict method requires:
        - The encoded query(ies) defining the variable(s) of interest.

    In addition, the method requires either:
        - A dataframe of covariates (test_covariates) to be used for inference
          of a set of spectral predictions that will subsequently be combined
          to yield the normative predictions for the encoded query(ies).
        OR
        - A dictionary of precomputed spectral predictions (spectral_predictions)
          to be used for efficiently predicting the encoded query(ies).

    The precomputed spectral predictions can be obtained using the
    'compute_spectral_predictions' function. This is particularly useful when
    predicting multiple queries or when the same covariate set is used for
    multiple predictions, as it avoids redundant computations.

    Args:
        encoded_query: np.ndarray
            Encoded query data defining the normative variable of interest.
            Can be provided as:
            - shape = (n_modes) for a single query vector
            - shape = (n_modes, n_queries) for multiple queries predicted at once
        spectral_predictions: dict | None
            Optional dictionary of precomputed spectral predictions to use for
            the prediction. If not provided, test_covariates must be provided
            instead to compute the spectral predictions.
            The dictionary should contain:
            - 'eigenmode_mu_estimates': np.ndarray (n_samples, n_modes)
            - 'eigenmode_std_estimates': np.ndarray (n_samples, n_modes)
            - 'rho_estimates': np.ndarray (n_samples, n_covariance_pairs)
            This can be obtained using the 'compute_spectral_predictions' method.
        test_covariates: pd.DataFrame | None
            DataFrame containing the new covariate data to predict.
            This must include all specified covariates.
            Note: covariates listed in predict_without will be ignored and are
            hence not required.
        extended: bool (default: False)
            If True, return additional stats such as log-likelihood, centiles, etc.
            Note that extended predictions require spectral_coeff_test_data to be
            provided in addition to the covariates.
        model_params: dict | None
            Optional dictionary of model parameters to use. If not provided,
            the stored parameters from model.fit() will be used.
        spectral_coeff_test_data: np.ndarray | None
            Optional spectral coefficient of test data for the phenotype being
            modeled :math:`(T_{test} \\Psi_{(k)}) \\in R^{N_{test} \\times k}`
            (only required for extended predictions).
            Expects a numpy array (n_samples, n_modes)
        n_modes: int | None
            Optional number of modes to use for the prediction. If not provided,
            the number of modes from model_params will be used.
        predict_without: list[str] | None
            Optional list of covariate names to ignore during prediction.
            This can be used to check the effect of removing certain covariates
            from the model.

    Returns:
        pd.DataFrame: DataFrame containing the predicted moments (mean, std) for
            the variable of interest defined by the encoded query.
    """
    # Parameters
    if model_params is None:
        model_params = self.model_params

    # Find n_modes
    if n_modes is None:
        n_modes = int(model_params["n_modes"])

    if self.base_model.spec is None:
        err = "The base model is not specified. Cannot predict new data."
        raise ValueError(err)

    if spectral_predictions is None:
        if test_covariates is None:
            err = "Either test_covariates or spectral_predictions must be provided."
            raise ValueError(err)

        # Compute spectral predictions if not provided
        spectral_predictions = self.compute_spectral_predictions(
            test_covariates=test_covariates,
            model_params=model_params,
            n_modes=n_modes,
            predict_without=predict_without,
        )
    elif test_covariates is not None:
        logger.warning(
            "Both test_covariates and spectral_predictions are provided."
            " Ignoring test_covariates and using spectral_predictions.",
        )
        if predict_without is not None:
            logger.warning(
                "predict_without is ignored when spectral_predictions"
                " are provided directly.",
            )

    # Unpack spectral predictions
    self._validate_spectral_predictions(spectral_predictions)
    eigenmode_mu_estimates = spectral_predictions["eigenmode_mu_estimates"]
    eigenmode_std_estimates = spectral_predictions["eigenmode_std_estimates"]
    rho_estimates = spectral_predictions["rho_estimates"]

    # Reformat encoded queries (for efficiency)
    encoded_query = np.asarray(encoded_query[:n_modes])
    encoded_query = encoded_query.reshape(n_modes, -1, order="F")

    # Compute the predictions
    predictions = self._predict_from_spectral_estimates(
        encoded_query=encoded_query,
        eigenmode_mu_estimates=eigenmode_mu_estimates,
        eigenmode_std_estimates=eigenmode_std_estimates,
        rho_estimates=rho_estimates,
        model_params=model_params,
        n_modes=n_modes,
    )

    # Check if extended predictions are requested
    if extended:
        if spectral_coeff_test_data is None:
            err = (
                "Extended predictions require spectral_coeff_test_data"
                " to be provided."
            )
            raise ValueError(err)
        # Add extended statistics to predictions (e.g. centiles, log-loss, etc.)
        predictions.extend_predictions(
            variable_of_interest=spectral_coeff_test_data @ encoded_query,
        )

    return predictions

reduce_model(n_modes: int, *, inplace: bool = False) -> SpectralNormativeModel

Create a reduced spectral normative model using only the first n_modes.

Parameters:

Name	Type	Description	Default
n_modes	int	int Number of modes to retain in the reduced model. Must be less than or equal to the current number of modes considered by the model.	required
inplace	bool	bool (default: False) If True, modify the current model instance to reduce its modes. If False, return a new SpectralNormativeModel instance with the reduced modes.	False

Returns: SpectralNormativeModel A new SpectralNormativeModel instance with reduced number of modes.

Source code in src/spectranorm/snm.py

def reduce_model(
    self,
    n_modes: int,
    *,
    inplace: bool = False,
) -> SpectralNormativeModel:
    """
    Create a reduced spectral normative model using only the first n_modes.

    Args:
        n_modes: int
            Number of modes to retain in the reduced model. Must be less than or
            equal to the current number of modes considered by the model.
        inplace: bool (default: False)
            If True, modify the current model instance to reduce its modes. If
            False, return a new SpectralNormativeModel instance with the reduced
            modes.
    Returns:
        SpectralNormativeModel
            A new SpectralNormativeModel instance with reduced number of modes.
    """
    if (n_modes > self.eigenmode_basis.n_modes) or (
        hasattr(self, "model_params") and n_modes > self.model_params["n_modes"]
    ):
        available_modes = self.eigenmode_basis.n_modes
        if hasattr(self, "model_params"):
            available_modes = min(available_modes, self.model_params["n_modes"])
        err = f"Cannot reduce to {n_modes} modes, only {available_modes} available."
        raise ValueError(err)

    # Create a reduced eigenbasis
    reduced_eigenbasis = self.eigenmode_basis.reduce(n_modes)

    if inplace:
        return_model = self
        return_model.eigenmode_basis = reduced_eigenbasis
    else:
        return_model = SpectralNormativeModel(
            base_model=self.base_model,
            eigenmode_basis=reduced_eigenbasis,
        )
        return_model.model_params = self.model_params  # Copy model parameters

    # Update model parameters to reflect reduced modes
    if hasattr(self, "model_params"):
        new_model_params: dict[str, Any] = {}
        new_model_params["n_modes"] = n_modes
        new_model_params["sample_size"] = self.model_params["sample_size"]
        new_model_params["direct_model_params"] = self.model_params[
            "direct_model_params"
        ][:n_modes]
        valid_cov_indices = np.where(
            (
                return_model.model_params["sparse_covariance_structure"][:, 0]
                < n_modes
            )
            & (
                return_model.model_params["sparse_covariance_structure"][:, 1]
                < n_modes
            ),
        )[0]
        new_model_params["sparse_covariance_structure"] = return_model.model_params[
            "sparse_covariance_structure"
        ][valid_cov_indices]
        new_model_params["covariance_model_params"] = [
            return_model.model_params["covariance_model_params"][i]
            for i in valid_cov_indices
        ]
        new_model_params["n_params"] = self.model_params.get("n_params", None)
        return_model.model_params = new_model_params

    return return_model

save_model(directory: Path) -> None

Save the fitted spectral normative model to the specified directory.

Parameters:

Name	Type	Description	Default
directory	Path	Path Directory to save the fitted model. A subdirectory named "spectral_normative_model" will be created within this directory.	required

Source code in src/spectranorm/snm.py

def save_model(self, directory: Path) -> None:
    """
    Save the fitted spectral normative model to the specified directory.

    Args:
        directory: Path
            Directory to save the fitted model. A subdirectory named
            "spectral_normative_model" will be created within this directory.
    """
    # Prepare the save directory
    directory = Path(directory)
    saved_model_dir = utils.general.prepare_save_directory(
        directory,
        "spectral_normative_model",
    )

    # Save the eigenmode basis separately
    self.eigenmode_basis.save(str(saved_model_dir / "eigenmode_basis.joblib"))

    # Save the model
    model_dict = {
        "spec": self.base_model.spec,
        "defaults": self.base_model.defaults,
    }
    if hasattr(self, "model_params"):
        model_dict["model_params"] = self.model_params
    joblib.dump(model_dict, saved_model_dir / "spectral_model_dict.joblib")

SplineSpec dataclass

Specification for spline basis construction.

Attributes:

Name	Type	Description
df	int	int Degrees of freedom (number of basis functions).
degree	int	int Degree of the spline (e.g., 3 for cubic splines).
lower_bound	float	float Lower boundary for the spline domain.
upper_bound	float	float Upper boundary for the spline domain.
knots	list[float] | None	Optional[List[float]] Optional list of internal knot locations within the spline domain (excluding the boundary knots). Must be strictly increasing and contain exactly df - degree - 1 values. If unspecified, then equally spaced quantiles of the input data are used.

Source code in src/spectranorm/snm.py

@dataclass
class SplineSpec:
    """
    Specification for spline basis construction.

    Attributes:
        df: int
            Degrees of freedom (number of basis functions).
        degree: int
            Degree of the spline (e.g., 3 for cubic splines).
        lower_bound: float
            Lower boundary for the spline domain.
        upper_bound: float
            Upper boundary for the spline domain.
        knots: Optional[List[float]]
            Optional list of internal knot locations within the spline domain
            (excluding the boundary knots). Must be strictly increasing and
            contain exactly `df - degree - 1` values. If unspecified, then
            equally spaced quantiles of the input data are used.
    """

    lower_bound: float
    upper_bound: float
    df: int = DEFAULT_SPLINE_DF
    degree: int = DEFAULT_SPLINE_DEGREE
    knots: list[float] | None = None

    # Validation checks for the spline specification.
    def __post_init__(self) -> None:
        # Check that df (degrees of freedom) is greater than degree
        if self.df <= self.degree:
            err = "df (degrees of freedom) must be greater than degree."
            raise ValueError(err)
        # Check that degree is at least 1
        if self.degree < 1:
            err = "degree must be at least 1."
            raise ValueError(err)
        # Check that lower_bound and upper_bound are numeric
        if self.lower_bound >= self.upper_bound:
            err = "lower_bound must be less than upper_bound."
            raise ValueError(err)
        if self.knots is not None:
            if not all(isinstance(k, (int, float)) for k in self.knots):
                err = "All knots must be numeric (int or float)."
                raise TypeError(err)
            # Check if knots are strictly increasing
            if not all(x < y for x, y in zip(self.knots, self.knots[1:])):
                err = "Knots must be strictly increasing."
                raise ValueError(err)
            # Check if knots are within bounds
            if any(k < self.lower_bound or k > self.upper_bound for k in self.knots):
                err = (
                    "All knots must be within the bounds defined by "
                    "lower_bound and upper_bound."
                )
                raise ValueError(err)
            # Check if the number of knots is correct
            if len(self.knots) != (self.df - self.degree - 1):
                err = (
                    f"knots must contain exactly {self.df - self.degree - 1} "
                    f"values, got {len(self.knots)}."
                )
                raise ValueError(err)

    @classmethod
    def create_spline_spec(
        cls,
        values: pd.Series[float],
        df: int = DEFAULT_SPLINE_DF,
        degree: int = DEFAULT_SPLINE_DEGREE,
        knots: list[float] | None = None,
        extrapolation_factor: float = DEFAULT_SPLINE_EXTRAPOLATION_FACTOR,
        lower_bound: float | None = None,
        upper_bound: float | None = None,
    ) -> SplineSpec:
        """
        Create a spline specification from a pandas Series.

        Args:
            values: pd.Series
                The list of input values to make the spline.
            df: int
                Degrees of freedom for the spline (default is 5).
            degree: int
                Degree of the spline (default is 3).
            knots: list[float] | None
                [Optional] List of internal knot locations within the spline domain.
                If None, equally spaced quantiles of the input data are used.
            extrapolation_factor: float, positive, default is 0.1
                [Optional] Factor to extend the lower and upper bounds of the spline
                domain.
            lower_bound: float | None
                [Optional] Lower boundary for the spline domain. If None, it is set to
                `values.min() - extrapolation_factor * (values.max() - values.min())`.
            upper_bound: float | None
                [Optional] Upper boundary for the spline domain. If None, it is set to
                `values.max() + extrapolation_factor * (values.max() - values.min())`.

        Returns:
            SplineSpec
                The created spline specification.
        """
        extrapolation = extrapolation_factor * (values.max() - values.min())
        if lower_bound is None:
            lower_bound = values.min() - extrapolation
        if upper_bound is None:
            upper_bound = values.max() + extrapolation
        if knots is None:
            # Use equally spaced quantiles as knots
            knots = np.linspace(lower_bound, upper_bound, df - degree + 1)[
                1:-1
            ].tolist()
        return cls(
            lower_bound=lower_bound,
            upper_bound=upper_bound,
            df=df,
            degree=degree,
            knots=knots,
        )

create_spline_spec(values: pd.Series[float], df: int = DEFAULT_SPLINE_DF, degree: int = DEFAULT_SPLINE_DEGREE, knots: list[float] | None = None, extrapolation_factor: float = DEFAULT_SPLINE_EXTRAPOLATION_FACTOR, lower_bound: float | None = None, upper_bound: float | None = None) -> SplineSpec classmethod

Create a spline specification from a pandas Series.

Parameters:

Name	Type	Description	Default
values	Series[float]	pd.Series The list of input values to make the spline.	required
df	int	int Degrees of freedom for the spline (default is 5).	DEFAULT_SPLINE_DF
degree	int	int Degree of the spline (default is 3).	DEFAULT_SPLINE_DEGREE
knots	list[float] | None	list[float] | None [Optional] List of internal knot locations within the spline domain. If None, equally spaced quantiles of the input data are used.	None
extrapolation_factor	float	float, positive, default is 0.1 [Optional] Factor to extend the lower and upper bounds of the spline domain.	DEFAULT_SPLINE_EXTRAPOLATION_FACTOR
lower_bound	float | None	float | None [Optional] Lower boundary for the spline domain. If None, it is set to values.min() - extrapolation_factor * (values.max() - values.min()).	None
upper_bound	float | None	float | None [Optional] Upper boundary for the spline domain. If None, it is set to values.max() + extrapolation_factor * (values.max() - values.min()).	None

Returns:

Type	Description
SplineSpec	SplineSpec The created spline specification.

Source code in src/spectranorm/snm.py

@classmethod
def create_spline_spec(
    cls,
    values: pd.Series[float],
    df: int = DEFAULT_SPLINE_DF,
    degree: int = DEFAULT_SPLINE_DEGREE,
    knots: list[float] | None = None,
    extrapolation_factor: float = DEFAULT_SPLINE_EXTRAPOLATION_FACTOR,
    lower_bound: float | None = None,
    upper_bound: float | None = None,
) -> SplineSpec:
    """
    Create a spline specification from a pandas Series.

    Args:
        values: pd.Series
            The list of input values to make the spline.
        df: int
            Degrees of freedom for the spline (default is 5).
        degree: int
            Degree of the spline (default is 3).
        knots: list[float] | None
            [Optional] List of internal knot locations within the spline domain.
            If None, equally spaced quantiles of the input data are used.
        extrapolation_factor: float, positive, default is 0.1
            [Optional] Factor to extend the lower and upper bounds of the spline
            domain.
        lower_bound: float | None
            [Optional] Lower boundary for the spline domain. If None, it is set to
            `values.min() - extrapolation_factor * (values.max() - values.min())`.
        upper_bound: float | None
            [Optional] Upper boundary for the spline domain. If None, it is set to
            `values.max() + extrapolation_factor * (values.max() - values.min())`.

    Returns:
        SplineSpec
            The created spline specification.
    """
    extrapolation = extrapolation_factor * (values.max() - values.min())
    if lower_bound is None:
        lower_bound = values.min() - extrapolation
    if upper_bound is None:
        upper_bound = values.max() + extrapolation
    if knots is None:
        # Use equally spaced quantiles as knots
        knots = np.linspace(lower_bound, upper_bound, df - degree + 1)[
            1:-1
        ].tolist()
    return cls(
        lower_bound=lower_bound,
        upper_bound=upper_bound,
        df=df,
        degree=degree,
        knots=knots,
    )