causal_hub/models/

mod.rs

mod bayesian_network;
use std::{
    borrow::Cow,
    ops::{DivAssign, MulAssign},
};

use approx::{AbsDiffEq, RelativeEq};
pub use bayesian_network::*;

mod continuous_time_bayesian_network;
pub use continuous_time_bayesian_network::*;
use itertools::Either;
use rand::Rng;

mod graphs;
use std::fmt::Debug;

pub use graphs::*;

use crate::types::{Error, Labels, Result, Set};

/// A trait for models with labelled variables.
pub trait HasLabels {
    /// Returns the labels of the variables.
    ///
    /// # Returns
    ///
    /// A reference to the labels.
    ///
    fn labels(&self) -> &Labels;

    /// Return the variable index for a given label.
    ///
    /// # Arguments
    ///
    /// * `x` - The label of the variable.
    ///
    /// # Errors
    ///
    /// * If the label is not found.
    ///
    /// # Returns
    ///
    /// The index of the variable.
    ///
    #[inline]
    fn label_to_index(&self, x: &str) -> Result<usize> {
        self.labels()
            .get_index_of(x)
            .ok_or_else(|| Error::MissingLabel(x))
    }

    /// Return the label for a given variable index.
    ///
    /// # Arguments
    ///
    /// * `x` - The index of the variable.
    ///
    /// # Errors
    ///
    /// * If the index is out of bounds.
    ///
    /// # Returns
    ///
    /// The label of the variable.
    ///
    #[inline]
    fn index_to_label(&self, x: usize) -> Result<&str> {
        self.labels()
            .get_index(x)
            .map(|x| x.as_str())
            .ok_or_else(|| Error::IndexOutOfBounds(x))
    }

    /// Maps an index from this model to another model with the same label.
    ///
    /// # Arguments
    ///
    /// * `x` - The index in this model.
    /// * `other` - The labels of the other model.
    ///
    /// # Errors
    ///
    /// * If the index is out of bounds.
    /// * If the label does not exist in the other model.
    ///
    /// # Returns
    ///
    /// The index in the other model.
    ///
    #[inline]
    fn index_to(&self, x: usize, other: &Labels) -> Result<usize> {
        // Get the label of the variable in this model.
        let label = self.index_to_label(x)?;
        // Get the index of the variable in the other model.
        other
            .get_index_of(label)
            .ok_or_else(|| Error::MissingLabel(label))
    }

    /// Maps a set of indices from this model to another model with the same labels.
    ///
    /// # Arguments
    ///
    /// * `x` - The set of indices in this model.
    /// * `other` - The labels of the other model.
    ///
    /// # Errors
    ///
    /// * If any index is out of bounds.
    /// * If any label does not exist in the other model.
    ///
    /// # Returns
    ///
    /// The set of indices in the other model.
    ///
    #[inline]
    fn indices_to(&self, x: &Set<usize>, other: &Labels) -> Result<Set<usize>> {
        x.iter().map(|&x| self.index_to(x, other)).collect()
    }

    /// Maps an index from another model to this model with the same label.
    ///
    /// # Arguments
    ///
    /// * `x` - The index in the other model.
    /// * `other` - The labels of the other model.
    ///
    /// # Errors
    ///
    /// * If the index is out of bounds.
    /// * If the label does not exist in this model.
    ///
    /// # Returns
    ///
    /// The index in this model.
    ///
    #[inline]
    fn index_from(&self, x: usize, other: &Labels) -> Result<usize> {
        // Get the label of the variable in the other model.
        let label = other
            .get_index(x)
            .ok_or_else(|| Error::IndexOutOfBounds(x))?;
        // Get the index of the variable in this model.
        self.labels()
            .get_index_of(label)
            .ok_or_else(|| Error::MissingLabel(label))
    }

    /// Maps a set of indices from another model to this model with the same labels.
    ///
    /// # Arguments
    ///
    /// * `x` - The set of indices in the other model.
    /// * `other` - The labels of the other model.
    ///
    /// # Errors
    ///
    /// * If any index is out of bounds.
    /// * If any label does not exist in this model.
    ///
    /// # Returns
    ///
    /// The set of indices in this model.
    ///
    #[inline]
    fn indices_from(&self, x: &Set<usize>, other: &Labels) -> Result<Set<usize>> {
        x.iter().map(|&x| self.index_from(x, other)).collect()
    }
}

impl<L, R> HasLabels for Either<L, R>
where
    L: HasLabels,
    R: HasLabels,
{
    fn labels(&self) -> &Labels {
        match self {
            Either::Left(l) => l.labels(),
            Either::Right(r) => r.labels(),
        }
    }
}

/// A trait for conditional probability distributions.
pub trait CPD: Clone + Debug + HasLabels + PartialEq + AbsDiffEq + RelativeEq {
    /// The type of the samples.
    type Sample;
    /// The type of the support metadata (variable ranges / state sets).
    type Support: Clone;
    /// The type of the parameters.
    type Parameters;
    /// The type of the sufficient statistics.
    type Statistics: Clone;

    /// Returns the labels of the conditioned variables.
    ///
    /// # Returns
    ///
    /// A reference to the conditioning labels.
    ///
    fn conditioning_labels(&self) -> &Labels;

    /// Returns the support metadata of the CPD.
    ///
    /// For categorical CPDs this returns the discrete support (possible states per variable).
    /// For Gaussian CPDs this returns the range of each variable (defaults to (-inf, +inf)).
    ///
    /// # Returns
    ///
    /// A reference to the support metadata.
    ///
    fn support(&self) -> Cow<'_, Self::Support>;

    /// Returns the conditioning support metadata of the CPD.
    ///
    /// For categorical CPDs this returns the discrete support of the conditioning variables.
    /// For Gaussian CPDs this returns the range of each conditioning variable.
    ///
    /// # Returns
    ///
    /// A reference to the conditioning support metadata.
    ///
    fn conditioning_support(&self) -> Cow<'_, Self::Support>;

    /// Returns the parameters.
    ///
    /// # Returns
    ///
    /// A reference to the parameters.
    ///
    fn parameters(&self) -> &Self::Parameters;

    /// Returns the parameters size.
    ///
    /// # Returns
    ///
    /// The parameters size.
    ///
    fn parameters_size(&self) -> usize;

    /// Returns the sufficient statistics, if any.
    ///
    /// # Returns
    ///
    /// An option containing the sufficient statistics, either borrowed or owned.
    ///
    fn fitted_statistics(&self) -> Option<Cow<'_, Self::Statistics>>;

    /// Returns the log-likelihood of the fitted dataset, if any.
    ///
    /// # Returns
    ///
    /// An option containing the log-likelihood.
    ///
    fn fitted_log_likelihood(&self) -> Option<f64>;

    /// Returns the value of probability (mass or density) function for P(X = x | Z = z).
    ///
    /// # Arguments
    ///
    /// * `x` - The value of the conditioned variables.
    /// * `z` - The value of the conditioning variables.
    ///
    /// # Errors
    ///
    /// * If the value of the conditioned variables is out of bounds.
    /// * If the value of the conditioning variables is out of bounds.
    ///
    /// # Returns
    ///
    /// The probability P(X = x | Z = z).
    ///
    fn pf(&self, x: &Self::Sample, z: &Self::Sample) -> Result<f64>;

    /// Samples from the conditional distribution P(X | Z = z).
    ///
    /// # Arguments
    ///
    /// * `rng` - A mutable reference to a random number generator.
    /// * `z` - The value of the conditioning variables.
    ///
    /// # Errors
    ///
    /// * If the value of the conditioning variables is out of bounds.
    ///
    /// # Returns
    ///
    /// A sample from P(X | Z = z).
    ///
    fn sample<R: Rng>(&self, rng: &mut R, z: &Self::Sample) -> Result<Self::Sample>;
}

/// A trait for conditional intensity matrices.
pub trait CIM: Clone + Debug + HasLabels + PartialEq + AbsDiffEq + RelativeEq {
    /// The type of the samples.
    type Sample;
    /// The type of the support metadata (variable ranges / state sets).
    type Support: Clone;
    /// The type of the parameters.
    type Parameters;
    /// The type of the sufficient statistics.
    type Statistics: Clone;

    /// Returns the labels of the conditioned variables.
    ///
    /// # Returns
    ///
    /// A reference to the conditioning labels.
    ///
    fn conditioning_labels(&self) -> &Labels;

    /// Returns the support metadata of the CIM.
    ///
    /// # Returns
    ///
    /// A reference to the support metadata.
    ///
    fn support(&self) -> Cow<'_, Self::Support>;

    /// Returns the conditioning support metadata of the CIM.
    ///
    /// # Returns
    ///
    /// A reference to the conditioning support metadata.
    ///
    fn conditioning_support(&self) -> Cow<'_, Self::Support>;

    /// Returns the parameters.
    ///
    /// # Returns
    ///
    /// A reference to the parameters.
    ///
    fn parameters(&self) -> &Self::Parameters;

    /// Returns the parameters size.
    ///
    /// # Returns
    ///
    /// The parameters size.
    ///
    fn parameters_size(&self) -> usize;

    /// Returns the sufficient statistics, if any.
    ///
    /// # Returns
    ///
    /// An option containing the sufficient statistics, either borrowed or owned.
    ///
    fn fitted_statistics(&self) -> Option<Cow<'_, Self::Statistics>>;

    /// Returns the log-likelihood of the fitted dataset, if any.
    ///
    /// # Returns
    ///
    /// An option containing the log-likelihood.
    ///
    fn fitted_log_likelihood(&self) -> Option<f64>;
}

/// A trait for potential functions.
pub trait Phi:
    Clone
    + Debug
    + HasLabels
    + PartialEq
    + AbsDiffEq
    + RelativeEq
    + for<'a> MulAssign<&'a Self>
    + for<'a> DivAssign<&'a Self>
{
    /// The type of the CPD.
    type CPD;
    /// The type of the support metadata (variable ranges / state sets).
    type Support: Clone;
    /// The type of the parameters.
    type Parameters;
    /// The type of the evidence.
    type Evidence;

    /// Returns the support metadata of the potential.
    ///
    /// # Returns
    ///
    /// A reference to the support metadata.
    ///
    fn support(&self) -> Cow<'_, Self::Support>;

    /// Returns the parameters.
    ///
    /// # Returns
    ///
    /// A reference to the parameters.
    ///
    fn parameters(&self) -> &Self::Parameters;

    /// Returns the parameters size.
    ///
    /// # Returns
    ///
    /// The parameters size.
    ///
    fn parameters_size(&self) -> usize;

    /// Conditions the potential on a set of variables.
    ///
    /// # Arguments
    ///
    /// * `e` - A map from variable indices to their observed support.
    ///
    /// # Returns
    ///
    /// A new potential instance.
    ///
    fn condition(&self, evidence: &Self::Evidence) -> Result<Self>;

    /// Marginalizes the potential over a set of variables.
    ///
    /// # Arguments
    ///
    /// * `x` - A set of variable indices to marginalize over.
    ///
    /// # Returns
    ///
    /// A new potential instance.
    ///
    fn marginalize(&self, x: &Set<usize>) -> Result<Self>;

    /// Normalizes the potential.
    ///
    /// # Returns
    ///
    /// The normalized potential.
    ///
    fn normalize(&self) -> Result<Self>;

    /// Converts a CPD P(X | Z) to a potential \phi(X \cup Z).
    ///
    /// # Arguments
    ///
    /// * `cpd` - The CPD to convert.
    ///
    /// # Returns
    ///
    /// The corresponding potential.
    ///
    fn from_cpd(distribution: Self::CPD) -> Result<Self>;

    /// Converts a potential \phi(X \cup Z) to a CPD P(X | Z).
    ///
    /// # Arguments
    ///
    /// * `x` - The set of variables.
    /// * `z` - The set of conditioning variables.
    ///
    /// # Returns
    ///
    /// The corresponding CPD.
    ///
    fn into_cpd(self, x: &Set<usize>, z: &Set<usize>) -> Result<Self::CPD>;
}