cellular_raza-concepts 0.5.0

use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};

use crate::errors::{BoundaryError, DecomposeError};

/// Identifier for voxels used internally to get rid of user-defined ones.
#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
#[derive(Clone, Copy, Debug, Deserialize, Hash, PartialEq, Eq, Ord, PartialOrd, Serialize)]
pub struct VoxelPlainIndex(pub usize);

/// This is mainly used by the simulation_flow::cocmmunicator for testing purposes
#[allow(unused)]
#[doc(hidden)]
impl VoxelPlainIndex {
    pub fn new(value: usize) -> Self {
        Self(value)
    }
}

#[cfg(feature = "pyo3")]
#[pyo3::pymethods]
impl VoxelPlainIndex {
    #[new]
    fn __new__(index: usize) -> Self {
        Self::new(index)
    }
}

/// Identifier or subdomains
#[derive(Clone, Copy, Debug, Deserialize, Hash, PartialEq, Eq, Ord, PartialOrd, Serialize)]
#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
pub struct SubDomainPlainIndex(pub usize);

/// Provides an abstraction of the physical total simulation domain.
///
/// [cellular_raza](https://github.com/jonaspleyer/cellular_raza) uses domain-decomposition
/// algorithms to split up the computational workload over multiple physical regions.
/// That's why the domain itself is mostly responsible for being deconstructed
/// into smaller [SubDomains](SubDomain) which can then be used to numerically solve our system.
///
/// This trait can be automatically implemented when the [SortCells], [DomainRngSeed],
/// and [DomainCreateSubDomains] are satisfied together with a small number of trait bounds to hash
/// and compare indices.
pub trait Domain<C, S, Ci = Vec<C>> {
    /// Subdomains can be identified by their unique [SubDomainIndex](Domain::SubDomainIndex).
    /// The backend uses this property to construct a mapping (graph) between subdomains.
    type SubDomainIndex;

    /// Similarly to the [SubDomainIndex](Domain::SubDomainIndex), voxels can be accessed by
    /// their unique index. The backend will use this information to construct a mapping
    /// (graph) between voxels inside their respective subdomains.
    type VoxelIndex;

    /// Deconstructs the [Domain] into its respective subdomains.
    ///
    /// When using the blanket implementation of this function, the following steps are carried
    /// out:
    /// Its functionality consists of the following steps:
    /// 1. Decompose the Domain into [Subdomains](SubDomain)
    /// 2. Build a neighbor map between [SubDomains](SubDomain)
    /// 3. Sort cells to their respective [SubDomain]
    ///
    /// However, to increase performance or avoid trait bounds, one can also opt to implement this
    /// trait directly.
    fn decompose(
        self,
        n_subdomains: core::num::NonZeroUsize,
        cells: Ci,
    ) -> Result<DecomposedDomain<Self::SubDomainIndex, S, C>, DecomposeError>;
}

/// Manage the current rng seed of a [Domain]
pub trait DomainRngSeed {
    // fn set_rng_seed(&mut self, seed: u64);

    /// Obtains the current rng seed
    fn get_rng_seed(&self) -> u64;
}

/// Generate [SubDomains](SubDomain) from an existing [Domain]
pub trait DomainCreateSubDomains<S> {
    /// This should always be identical to [Domain::SubDomainIndex].
    type SubDomainIndex;
    /// This should always be identical to [Domain::VoxelIndex].
    type VoxelIndex;

    /// Generates at most `n_subdomains`. This function can also return a lower amount of
    /// subdomains but never less than 1.
    fn create_subdomains(
        &self,
        n_subdomains: core::num::NonZeroUsize,
    ) -> Result<
        impl IntoIterator<Item = (Self::SubDomainIndex, S, Vec<Self::VoxelIndex>)>,
        DecomposeError,
    >;
}

/// Generated by the [decompose](Domain::decompose) method. The backend will know how to
/// deal with this type and crate a working simulation from it.
pub struct DecomposedDomain<I, S, C> {
    /// Number of spawned [SubDomains](SubDomain). This number is guaranteed to be
    /// smaller or equal to the number may be different to the one given to the
    /// [Domain::decompose] method.
    /// Such behaviour can result from not being able to construct as many subdomains as desired.
    /// Note that this function will attempt to construct more [SubDomains](SubDomain)
    /// than available CPUs if given a larger number.
    pub n_subdomains: core::num::NonZeroUsize,
    /// Vector containing properties of individual [SubDomains](SubDomain).
    /// Entries are [Domain::SubDomainIndex], [SubDomain], and a vector of cells.
    // TODO can we use another iterator than Vec<(I, S, Vec<C>)>?
    pub index_subdomain_cells: Vec<(I, S, Vec<C>)>,
    /// Encapsulates how the subdomains are linked to each other.
    /// Eg. two subdomains without any boundary will never appear in each others collection
    /// of neighbors.
    /// For the future, we might opt to change to an undirected graph rather than a [BTreeMap].
    pub neighbor_map: BTreeMap<I, BTreeSet<I>>,
    /// Initial seed of the simulation for random number generation.
    pub rng_seed: u64,
}

/// Subdomains are produced by decomposing a [Domain] into multiple physical regions.
///
/// # Derivation
/// ```
/// # use cellular_raza_concepts::*;
/// struct MySubDomain {
///     x_min: f32,
///     x_max: f32,
///     n: usize,
/// }
///
/// impl SubDomain for MySubDomain {
///     type VoxelIndex = usize;
///
///     fn get_neighbor_voxel_indices(
///         &self,
///         voxel_index: &Self::VoxelIndex
///     ) -> Vec<Self::VoxelIndex> {
///         (voxel_index.saturating_sub(1)..voxel_index.saturating_add(1).min(self.n)+1)
///             .filter(|k| k!=voxel_index)
///             .collect()
///     }
///
///     fn get_all_indices(&self) -> Vec<Self::VoxelIndex> {
///         (0..self.n).collect()
///     }
/// }
///
/// #[derive(SubDomain)]
/// struct MyNewSubDomain {
///     #[Base]
///     base: MySubDomain,
/// }
/// # let _my_sdm = MyNewSubDomain {
/// #     base: MySubDomain {
/// #         x_min: -20.0,
/// #         x_max: -11.0,
/// #         n: 20,
/// #     }
/// # };
/// # assert_eq!(_my_sdm.get_all_indices(), (0..20).collect::<Vec<_>>());
/// # assert_eq!(_my_sdm.get_neighbor_voxel_indices(&0), vec![1]);
/// # assert_eq!(_my_sdm.get_neighbor_voxel_indices(&3), vec![2,4]);
/// # assert_eq!(_my_sdm.get_neighbor_voxel_indices(&7), vec![6,8]);
/// ```
pub trait SubDomain {
    /// Individual Voxels inside each subdomain can be accessed by this index.
    type VoxelIndex;

    /// Obtains the neighbor voxels of the specified voxel index. This function behaves similarly
    /// to [SortCells::get_voxel_index_of] in that it also has to return
    /// indices which are in other [SubDomains](SubDomain).
    fn get_neighbor_voxel_indices(&self, voxel_index: &Self::VoxelIndex) -> Vec<Self::VoxelIndex>;

    // fn apply_boundary(&self, cell: &mut C) -> Result<(), BoundaryError>;

    /// Get all voxel indices of this [SubDomain].
    fn get_all_indices(&self) -> Vec<Self::VoxelIndex>;
}

/// Assign an [VoxelIndex](SortCells::VoxelIndex) to a given cell.
///
/// This trait is used by the [Domain] and [SubDomain] trait to assign a [Domain::SubDomainIndex]
/// and [SubDomain::VoxelIndex] respectively.
///
/// # [SubDomain]
/// This trait is supposed to return the correct voxel index of the cell
/// even if this index is inside another [SubDomain].
/// This restriction might be lifted in the future but is still
/// required now.
pub trait SortCells<C> {
    /// An index which determines to which next smaller unit the cell should be assigned.
    type VoxelIndex;

    /// If given a cell, we can sort this cell into the corresponding sub unit.
    fn get_voxel_index_of(&self, cell: &C) -> Result<Self::VoxelIndex, BoundaryError>;
}

/// Apply boundary conditions to a cells position and velocity.
///
/// # Derivation
/// ```
/// # use cellular_raza_concepts::*;
/// # use cellular_raza_concepts::BoundaryError;
/// struct MyMechanics {
///     x_min: f64,
///     x_max: f64,
/// }
///
/// impl SubDomainMechanics<f64, f64> for MyMechanics {
///     fn apply_boundary(&self, pos: &mut f64, vel: &mut f64) -> Result<(), BoundaryError> {
///         if *pos < self.x_min {
///             *vel = vel.abs();
///         }
///         if *pos > self.x_max {
///             *vel = -vel.abs();
///         }
///         *pos = pos.clamp(self.x_min, self.x_max);
///         Ok(())
///     }
/// }
///
/// #[derive(SubDomain)]
/// struct MySubDomain {
///     #[Mechanics]
///     mechanics: MyMechanics,
/// }
/// # let _my_sdm = MySubDomain {
/// #     mechanics: MyMechanics {
/// #         x_min: 1.0,
/// #         x_max: 33.0,
/// #     }
/// # };
/// # let mut pos = 0.0;
/// # let mut vel = - 0.1;
/// # _my_sdm.apply_boundary(&mut pos, &mut vel).unwrap();
/// # assert_eq!(pos, 1.0);
/// # assert_eq!(vel, 0.1);
/// ```
pub trait SubDomainMechanics<Pos, Vel> {
    /// If the subdomain has boundary conditions, this function will enforce them onto the cells.
    /// For the future, we plan to replace this function to additionally obtain information
    /// about the previous and current location of the cell.
    fn apply_boundary(&self, pos: &mut Pos, vel: &mut Vel) -> Result<(), BoundaryError>;
}

/// Apply a force on a cell depending on its position and velocity.
///
/// # Derivation
/// ```
/// # use cellular_raza_concepts::*;
/// struct MyForce {
///     damping: f64,
/// }
///
/// impl SubDomainForce<f64, f64, f64, f64> for MyForce {
///     fn calculate_custom_force(&self, _: &f64, vel: &f64, _: &f64) -> Result<f64, CalcError> {
///         Ok(- self.damping * vel)
///     }
/// }
///
/// #[derive(SubDomain)]
/// struct MySubDomain {
///     #[Force]
///     force: MyForce,
/// }
/// # let _my_sdm = MySubDomain {
/// #     force: MyForce {
/// #         damping: 0.1,
/// #     }
/// # };
/// # let calculated_force = _my_sdm.calculate_custom_force(&0.0, &1.0, &0.0).unwrap();
/// # assert_eq!(calculated_force, -0.1);
/// ```
pub trait SubDomainForce<Pos, Vel, For, Inf> {
    /// Calculates the force which acts on the cell agent
    fn calculate_custom_force(
        &self,
        pos: &Pos,
        vel: &Vel,
        inf: &Inf,
    ) -> Result<For, crate::CalcError>;
}

/// Describes extracellular reactions and fluid dynamics
///
/// # Derivation
/// ```
/// # use cellular_raza_concepts::*;
///
/// #[derive(Clone, Debug)]
/// struct MyReactions<const N: usize> {
///     values: Vec<f32>,
///     pos: [f32; N],
/// }
///
/// impl<const N: usize> SubDomainReactions<[f32; N], Vec<f32>, f32> for MyReactions<N> {
///     type NeighborValue = Vec<f32>;
///     type BorderInfo = Self;
///
///     fn treat_increments<I, J>(
///         &mut self,
///         neighbors: I,
///         sources: J,
///     ) -> Result<(), CalcError>
///     where
///         I: IntoIterator<Item = Self::NeighborValue>,
///         J: IntoIterator<Item = ([f32; N], Vec<f32>)>,
///     {
///         Ok(())
///     }
///
///     fn update_fluid_dynamics(&mut self, dt: f32) -> Result<(), CalcError> {
///         Ok(())
///     }
///
///     fn get_extracellular_at_pos(&self, pos: &[f32; N]) -> Result<Vec<f32>, CalcError> {
///         Ok(self.values.clone())
///     }
///
///     fn get_neighbor_value(&self, border_info: Self::BorderInfo) -> Self::NeighborValue {
///         self.values.clone()
///     }
///
///     fn get_border_info(&self) -> Self::BorderInfo {
///         self.clone()
///     }
/// }
///
/// #[derive(SubDomain)]
/// struct DerivedSubDomain<const N: usize> {
///     #[Reactions]
///     reactions: MyReactions<N>,
/// }
/// ```
pub trait SubDomainReactions<Pos, Re, Float> {
    /// Extracellular value of neighbor
    type NeighborValue;
    /// Exchanged information to locate neighboring subdomains.
    type BorderInfo;

    /// Combines increments which have been obtained by neighbors and cell-sources
    fn treat_increments<I, J>(&mut self, neighbors: I, sources: J) -> Result<(), crate::CalcError>
    where
        I: IntoIterator<Item = Self::NeighborValue>,
        J: IntoIterator<Item = (Pos, Re)>;

    /// Main update function to calculate new values of extracellular concentrations.
    fn update_fluid_dynamics(&mut self, dt: Float) -> Result<(), crate::CalcError>;

    /// Obtain extracellular concentrations at given point.
    fn get_extracellular_at_pos(&self, pos: &Pos) -> Result<Re, crate::CalcError>;
    /// Obtains the [SubDomainReactions::NeighborValue] which should be sent to the neighbor which
    /// has exposed the given [SubDomainReactions::BorderInfo].
    fn get_neighbor_value(&self, border_info: Self::BorderInfo) -> Self::NeighborValue;
    /// Obtains the [SubDomainReactions::BorderInfo] used to retrieve the
    /// [SubDomainReactions::NeighborValue].
    fn get_border_info(&self) -> Self::BorderInfo;
}

/// This trait derives the different aspects of a [SubDomain].
///
/// It serves similarly as the [cellular_raza_concepts_derive::CellAgent] trait to quickly
/// build new structures from already existing functionality.
///
/// | Attribute | Trait | Implemented |
/// | ---  | --- |:---:|
/// | `Base` | [SubDomain] | ✅ |
/// | `SortCells` | [SortCells] | ✅ |
/// | `Mechanics` | [SubDomainMechanics] | ✅ |
/// | `Force` | [SubDomainForce] | ✅  |
/// | `Reactions` | [SubDomainReactions] | ❌ |
///
/// # Example Usage
/// ```
/// # use cellular_raza_concepts::*;
/// # struct MySubDomain;
/// # impl SubDomain for MySubDomain {
/// #     type VoxelIndex = usize;
/// #     fn get_neighbor_voxel_indices(&self, voxel_index: &Self::VoxelIndex) -> Vec<usize> {
/// #         Vec::new()
/// #     }
/// #     fn get_all_indices(&self) -> Vec<Self::VoxelIndex> {
/// #         Vec::new()
/// #     }
/// # }
/// #[derive(SubDomain)]
/// struct MyDerivedSubDomain {
///     #[Base]
///     s: MySubDomain,
/// }
/// # let derived_subdomain = MyDerivedSubDomain {
/// #     s: MySubDomain,
/// # };
/// # let all_indices = derived_subdomain.get_all_indices();
/// # assert_eq!(all_indices.len(), 0);
/// ```
#[doc(inline)]
pub use cellular_raza_concepts_derive::SubDomain;

/// Derives aspects related to the simulation [Domain]
///
/// | Attribute | Trait | Implemented |
/// | -- | --- |:---:|
/// | `Base` | [Domain] | ✅ |
/// | `DomainPartialDerive` | [Domain] | ✅ |
/// | `DomainRngSeed` | [DomainRngSeed] | ✅ |
/// | `DomainCreateSubDomains` | [DomainCreateSubDomains] | ✅ |
/// | `SortCells` | [SortCells] | ✅ |
#[doc(inline)]
pub use cellular_raza_concepts_derive::Domain;