assembly-theory 0.6.1

//! Compute assembly indices of molecules.
//!
//! # Example
//! ```
//! # use std::{fs, path::PathBuf};
//! use assembly_theory::assembly::index;
//! use assembly_theory::loader::parse_molfile_str;
//!
//! # fn main() -> Result<(), std::io::Error> {
//! // Load a molecule from a .mol file.
//! let path = PathBuf::from(format!("./data/checks/anthracene.mol"));
//! let molfile = fs::read_to_string(path)?;
//! let anthracene = parse_molfile_str(&molfile).expect("Parsing failure.");
//!
//! // Compute the molecule's assembly index.
//! assert_eq!(index(&anthracene), 6);
//! # Ok(())
//! # }
//! ```

use std::{
    sync::{
        atomic::{AtomicUsize, Ordering::Relaxed},
        Arc,
    },
    time::Duration,
};

use bit_set::BitSet;
use clap::ValueEnum;
use rayon::iter::{ParallelBridge, ParallelIterator};
use tokio::{runtime::Runtime, sync::oneshot, time::timeout as tktimeout};

use crate::{
    bounds::{state_bounds, Bound},
    canonize::CanonizeMode,
    kernels::KernelMode,
    matches::Matches,
    memoize::{Cache, MemoizeMode},
    molecule::Molecule,
    state::State,
    utils::connected_components_under_edges,
};

/// Parallelization strategy for the recursive search phase.
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
pub enum ParallelMode {
    /// No parallelism.
    None,
    /// Create a task pool from the recursion's first level only.
    DepthOne,
    /// Spawn a new thread at every recursive call.
    Always,
}

/// Compute assembly depth; see
/// [Pagel et al. (2024)](https://arxiv.org/abs/2409.05993).
///
/// Note: This function is currently very (unusably) slow. It primarily exists
/// in this crate as a proof of concept.
///
/// # Example
/// ```
/// # use std::{fs, path::PathBuf};
/// use assembly_theory::assembly::depth;
/// use assembly_theory::loader::parse_molfile_str;
///
/// # fn main() -> Result<(), std::io::Error> {
/// // Load a molecule from a .mol file.
/// let path = PathBuf::from(format!("./data/checks/benzene.mol"));
/// let molfile = fs::read_to_string(path)?;
/// let benzene = parse_molfile_str(&molfile).expect("Parsing failure.");
///
/// // Compute the molecule's assembly depth.
/// assert_eq!(depth(&benzene), 3);
/// # Ok(())
/// # }
/// ```
pub fn depth(mol: &Molecule) -> u32 {
    let mut ix = u32::MAX;
    for (left, right) in mol.partitions().unwrap() {
        let l = if left.is_basic_unit() {
            0
        } else {
            depth(&left)
        };

        let r = if right.is_basic_unit() {
            0
        } else {
            depth(&right)
        };

        ix = ix.min(l.max(r) + 1)
    }
    ix
}

/// Determine the fragments produced from the given assembly state by removing
/// the given pair of edge-disjoint, isomorphic subgraphs and then adding one
/// back; return `None` if not possible.
fn fragments(mol: &Molecule, state: &[BitSet], h1: &BitSet, h2: &BitSet) -> Option<Vec<BitSet>> {
    // Attempt to find fragments f1 and f2 containing h1 and h2, respectively;
    // if either do not exist, exit without further fragmentation.
    let f1 = state.iter().enumerate().find(|(_, c)| h1.is_subset(c));
    let f2 = state.iter().enumerate().find(|(_, c)| h2.is_subset(c));
    let (Some((i1, f1)), Some((i2, f2))) = (f1, f2) else {
        return None;
    };

    let mut fragments = state.to_owned();

    // If the same fragment f1 (== f2) contains both h1 and h2, replace this
    // one fragment f1 with all connected components comprising f1 - (h1 U h2).
    // Otherwise, replace fragments f1 and f2 with all connected components
    // comprising f1 - h1 and f2 - h2, respectively.
    if i1 == i2 {
        let mut union = h1.clone();
        union.union_with(h2);
        let mut difference = f1.clone();
        difference.difference_with(&union);
        let c = connected_components_under_edges(mol.graph(), &difference);
        fragments.extend(c);
        fragments.swap_remove(i1);
    } else {
        let mut diff1 = f1.clone();
        diff1.difference_with(h1);
        let c1 = connected_components_under_edges(mol.graph(), &diff1);
        fragments.extend(c1);

        let mut diff2 = f2.clone();
        diff2.difference_with(h2);
        let c2 = connected_components_under_edges(mol.graph(), &diff2);
        fragments.extend(c2);

        fragments.swap_remove(i1.max(i2));
        fragments.swap_remove(i1.min(i2));
    }

    // Drop any singleton fragments, add h1 as a fragment, and return.
    fragments.retain(|i| i.len() > 1);
    fragments.push(h1.clone());
    Some(fragments)
}

/// Recursive helper for [`index_search`], only public for benchmarking.
///
/// Inputs:
/// - `mol`: The molecule whose assembly index is being calculated.
/// - `matches`: Structural information about the molecule's matched fragments.
/// - `state`: The current assembly state.
/// - `best_index`: The smallest assembly index for all assembly states so far.
/// - `bounds`: The list of bounding strategies to apply.
/// - `cache`: Memoization cache storing previously searched assembly states.
/// - `parallel_mode`: The parallelism mode for this state's match iteration.
///
/// Returns, from this assembly state and any of its descendents:
/// - `usize`: An updated upper bound on the assembly index. (Note: If this
///   state is pruned by bounds or deemed redundant by memoization, then the
///   upper bound returned is unchanged.)
/// - `usize`: The number of assembly states searched.
pub fn recurse_index_search(
    mol: &Molecule,
    matches: &Matches,
    state: &State,
    best_index: Arc<AtomicUsize>,
    bounds: &[Bound],
    cache: &mut Cache,
    parallel_mode: ParallelMode,
) -> (usize, usize) {
    // If any bounds would prune this assembly state or if memoization is
    // enabled and this assembly state is preempted by the cached state, halt.
    if state_bounds(mol, state, best_index.load(Relaxed), bounds) || cache.memoize_state(mol, state)
    {
        return (state.index(), 1);
    }

    // Generate a list of matches (i.e., pairs of edge-disjoint, isomorphic
    // fragments) to remove from this state.
    let (intermediate_frags, matches_to_remove): (Vec<BitSet>, Vec<usize>) =
        matches.matches_to_remove(mol, state, best_index.load(Relaxed), bounds);

    // Keep track of the best assembly index found in any of this assembly
    // state's children and the number of states searched, including this one.
    let best_child_index = AtomicUsize::from(state.index());
    let states_searched = AtomicUsize::from(1);

    // Define a closure that handles recursing to a new assembly state based on
    // the given match.
    let recurse_on_match = |i: usize, match_ix: usize| {
        let (h1, h2) = matches.match_fragments(match_ix);

        if let Some(fragments) = fragments(mol, &intermediate_frags, h1, h2) {
            // If using depth-one parallelism, all descendant states should be
            // computed serially.
            let new_parallel = if parallel_mode == ParallelMode::DepthOne {
                ParallelMode::None
            } else {
                parallel_mode
            };

            // Recurse using the remaining matches and updated fragments.
            let (child_index, child_states_searched) = recurse_index_search(
                mol,
                matches,
                &state.update(fragments, i, match_ix, h1.len()),
                best_index.clone(),
                bounds,
                &mut cache.clone(),
                new_parallel,
            );

            // Update the best assembly indices (across children states and the
            // entire search) and the number of descendant states searched.
            best_child_index.fetch_min(child_index, Relaxed);
            best_index.fetch_min(best_child_index.load(Relaxed), Relaxed);
            states_searched.fetch_add(child_states_searched, Relaxed);
        }
    };

    // Use the iterator type corresponding to the specified parallelism mode.
    if parallel_mode == ParallelMode::None {
        matches_to_remove
            .iter()
            .enumerate()
            .for_each(|(i, match_ix)| recurse_on_match(i, *match_ix));
    } else {
        matches_to_remove
            .iter()
            .enumerate()
            .par_bridge()
            .for_each(|(i, match_ix)| recurse_on_match(i, *match_ix));
    }

    (
        best_child_index.load(Relaxed),
        states_searched.load(Relaxed),
    )
}

/// Compute a molecule's assembly index and related information using a
/// top-down recursive algorithm, parameterized by the specified options.
///
/// If `timeout` is `None`, run until the assembly index is found. Otherwise,
/// stop after `timeout` milliseconds and return the best upper bound on the
/// assembly index found so far.
///
/// See [`CanonizeMode`], [`ParallelMode`], [`KernelMode`], and [`Bound`] for
/// details on how to customize the algorithm. Notably, bounds are applied in
/// the order they appear in the `bounds` slice. It is generally better to
/// provide bounds that are quick to compute first.
///
/// The results returned are:
/// - The molecule's `u32` assembly index (or an upper bound if timed out).
/// - The molecule's `u32` number of edge-disjoint isomorphic subgraph pairs.
/// - The `usize` total number of assembly [`State`]s searched if search
///   completes, and `None` otherwise (i.e., if search timed out).
///
/// # Example
/// ```
/// # use std::{fs, path::PathBuf};
/// use assembly_theory::{
///     assembly::{index_search, ParallelMode},
///     bounds::Bound,
///     canonize::CanonizeMode,
///     kernels::KernelMode,
///     loader::parse_molfile_str,
///     memoize::MemoizeMode,
/// };
///
/// # fn main() -> Result<(), std::io::Error> {
/// // Load a molecule from a .mol file.
/// let path = PathBuf::from(format!("./data/checks/anthracene.mol"));
/// let molfile = fs::read_to_string(path)?;
/// let anthracene = parse_molfile_str(&molfile).expect("Parsing failure.");
///
/// // Compute the molecule's assembly index without parallelism, memoization,
/// // kernelization, or bounds.
/// let (slow_index, _, _) = index_search(
///     &anthracene,
///     None,
///     CanonizeMode::TreeNauty,
///     ParallelMode::None,
///     MemoizeMode::None,
///     KernelMode::None,
///     &[],
/// );
///
/// // Compute the molecule's assembly index with parallelism, memoization, and
/// // some bounds.
/// let (fast_index, _, _) = index_search(
///     &anthracene,
///     None,
///     CanonizeMode::TreeNauty,
///     ParallelMode::DepthOne,
///     MemoizeMode::CanonIndex,
///     KernelMode::None,
///     &[Bound::Log, Bound::Int],
/// );
///
/// // Limit search to 1 ms, which should time out.
/// let (index_bound, _, states_searched) = index_search(
///     &anthracene,
///     Some(1),
///     CanonizeMode::TreeNauty,
///     ParallelMode::None,
///     MemoizeMode::None,
///     KernelMode::None,
///     &[],
/// );
///
/// assert_eq!(slow_index, 6);
/// assert_eq!(fast_index, 6);
/// assert!(index_bound >= fast_index && states_searched == None);
/// # Ok(())
/// # }
/// ```
pub fn index_search(
    mol: &Molecule,
    timeout: Option<u64>,
    canonize_mode: CanonizeMode,
    parallel_mode: ParallelMode,
    memoize_mode: MemoizeMode,
    kernel_mode: KernelMode,
    bounds: &[Bound],
) -> (u32, u32, Option<usize>) {
    // Catch not-yet-implemented modes.
    if kernel_mode != KernelMode::None {
        panic!("The chosen --kernel mode is not implemented yet!")
    }

    // Create the initial assembly state and memoization cache.
    let state = State::new(mol);
    let mut cache = Cache::new(memoize_mode, canonize_mode);

    // Enumerate matches (i.e., pairs of edge-disjoint isomorphic fragments).
    let matches = Matches::new(mol, canonize_mode);

    // Use an `Arc` to track the best assembly index across parallel threads.
    let best_index = Arc::new(AtomicUsize::from(mol.graph().edge_count() - 1));

    // Search for the shortest assembly pathway recursively.
    if let Some(timeout) = timeout {
        // If a timeout is provided, we will search within an asynchronous task
        // that can be interrupted after the specified duration (see below). To
        // avoid subsequent scope issues, make copies of various variables.
        let best_index_copy = best_index.clone();
        let mol = mol.clone();
        let bounds = bounds.to_vec();
        let num_matches = matches.len();

        // Search within a dedicated asynchronous runtime.
        let rt = Runtime::new().unwrap();
        let result = rt.block_on(async {
            let (send, recv) = oneshot::channel();
            rayon::spawn(move || {
                let _ = send.send(recurse_index_search(
                    &mol,
                    &matches,
                    &state,
                    best_index_copy,
                    &bounds,
                    &mut cache,
                    parallel_mode,
                ));
            });
            tktimeout(Duration::from_millis(timeout), recv).await
        });

        // If the search completes before the timeout, return the true assembly
        // index. Otherwise, return the best upper bound on the assembly index
        // found before timing out.
        let (index, states_searched) = match result {
            Ok(Ok((index, states_searched))) => (index, Some(states_searched)),
            Err(_) => (best_index.load(Relaxed), None),
            _ => panic!("An unexpected error occurred in async index_search"),
        };
        (index as u32, num_matches as u32, states_searched)
    } else {
        // Otherwise, if no timeout is provided, run the search normally.
        let (index, states_searched) = recurse_index_search(
            mol,
            &matches,
            &state,
            best_index,
            bounds,
            &mut cache,
            parallel_mode,
        );
        (index as u32, matches.len() as u32, Some(states_searched))
    }
}

/// Compute a molecule's assembly index using an efficient default strategy.
///
/// To customize assembly index calculation beyond the default strategy, see
/// [`index_search`].
///
/// # Example
/// ```
/// # use std::{fs, path::PathBuf};
/// use assembly_theory::assembly::index;
/// use assembly_theory::loader::parse_molfile_str;
///
/// # fn main() -> Result<(), std::io::Error> {
/// // Load a molecule from a .mol file.
/// let path = PathBuf::from(format!("./data/checks/anthracene.mol"));
/// let molfile = fs::read_to_string(path)?;
/// let anthracene = parse_molfile_str(&molfile).expect("Parsing failure.");
///
/// // Compute the molecule's assembly index.
/// assert_eq!(index(&anthracene), 6);
/// # Ok(())
/// # }
/// ```
pub fn index(mol: &Molecule) -> u32 {
    index_search(
        mol,
        None,
        CanonizeMode::TreeNauty,
        ParallelMode::DepthOne,
        MemoizeMode::CanonIndex,
        KernelMode::None,
        &[Bound::Int, Bound::MatchableEdges],
    )
    .0
}