catlog/stdlib/analyses/ode/
mass_action.rs

1//! Mass-action ODE analysis of models.
2//!
3//! Such ODEs are based on the *law of mass action* familiar from chemistry and
4//! mathematical epidemiology.
5
6use std::collections::HashMap;
7
8use indexmap::IndexMap;
9use nalgebra::DVector;
10use num_traits::Zero;
11use rebop::gillespie;
12
13#[cfg(feature = "serde")]
14use serde::{Deserialize, Serialize};
15#[cfg(feature = "serde-wasm")]
16use tsify::Tsify;
17
18use super::{ODEAnalysis, ODESolution};
19use crate::dbl::{
20    model::{DiscreteTabModel, FgDblModel, ModalDblModel, ModalOb, MutDblModel, TabEdge},
21    theory::{ModalMorType, ModalObType, TabMorType, TabObType},
22};
23use crate::one::FgCategory;
24use crate::simulate::ode::{NumericalPolynomialSystem, ODEProblem, PolynomialSystem};
25use crate::zero::{QualifiedName, alg::Polynomial, name, rig::Monomial};
26
27/// Data defining a mass-action ODE problem for a model.
28#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
29#[cfg_attr(feature = "serde-wasm", derive(Tsify))]
30#[cfg_attr(
31    feature = "serde-wasm",
32    tsify(into_wasm_abi, from_wasm_abi, hashmap_as_object)
33)]
34pub struct MassActionProblemData {
35    /// Map from morphism IDs to rate coefficients (nonnegative reals).
36    rates: HashMap<QualifiedName, f32>,
37
38    /// Map from object IDs to initial values (nonnegative reals).
39    #[cfg_attr(feature = "serde", serde(rename = "initialValues"))]
40    pub initial_values: HashMap<QualifiedName, f32>,
41
42    /// Duration of simulation.
43    pub duration: f32,
44}
45
46/// Stochastic mass-action analysis of a model.
47pub struct StochasticMassActionAnalysis {
48    /// Reaction network for the analysis.
49    pub problem: rebop::gillespie::Gillespie,
50
51    /// Map from object IDs to variable indices.
52    pub variable_index: IndexMap<QualifiedName, usize>,
53
54    /// Map from object IDs to initial values.
55    pub initial_values: HashMap<QualifiedName, f32>,
56
57    /// Duration of simulation.
58    pub duration: f32,
59}
60
61impl StochasticMassActionAnalysis {
62    /// Simulates the stochastic mass-action system and collects the results.
63    pub fn simulate(&mut self) -> ODESolution {
64        let mut time = vec![0.0];
65        let mut states: HashMap<_, _> = self
66            .variable_index
67            .keys()
68            .map(|id| {
69                let initial = self.initial_values.get(id).copied().unwrap_or_default();
70                (id.clone(), vec![initial])
71            })
72            .collect();
73        for t in 0..(self.duration as usize) {
74            self.problem.advance_until(t as f64);
75            time.push(self.problem.get_time() as f32);
76            for (id, idx) in self.variable_index.iter() {
77                states.get_mut(id).unwrap().push(self.problem.get_species(*idx) as f32)
78            }
79        }
80        ODESolution { time, states }
81    }
82}
83
84/// Symbolic parameter in mass-action polynomial system.
85type Parameter<Id> = Polynomial<Id, f32, u8>;
86
87/// Mass-action ODE analysis for Petri nets.
88///
89/// This struct implements the object part of the functorial semantics for reaction
90/// networks (aka, Petri nets) due to [Baez & Pollard](crate::refs::ReactionNets).
91pub struct PetriNetMassActionAnalysis {
92    /// Object type for places.
93    pub place_ob_type: ModalObType,
94    /// Morphism type for transitions.
95    pub transition_mor_type: ModalMorType,
96}
97
98impl Default for PetriNetMassActionAnalysis {
99    fn default() -> Self {
100        let ob_type = ModalObType::new(name("Object"));
101        Self {
102            place_ob_type: ob_type.clone(),
103            transition_mor_type: ModalMorType::Zero(ob_type),
104        }
105    }
106}
107
108impl PetriNetMassActionAnalysis {
109    /// Creates a mass-action system with symbolic rate coefficients.
110    pub fn build_system(
111        &self,
112        model: &ModalDblModel,
113    ) -> PolynomialSystem<QualifiedName, Parameter<QualifiedName>, u8> {
114        let mut sys = PolynomialSystem::new();
115        for ob in model.ob_generators_with_type(&self.place_ob_type) {
116            sys.add_term(ob, Polynomial::zero());
117        }
118        for mor in model.mor_generators_with_type(&self.transition_mor_type) {
119            let inputs = model
120                .get_dom(&mor)
121                .and_then(|ob| ob.clone().collect_product(None))
122                .unwrap_or_default();
123            let outputs = model
124                .get_cod(&mor)
125                .and_then(|ob| ob.clone().collect_product(None))
126                .unwrap_or_default();
127
128            let term: Monomial<_, _> =
129                inputs.iter().map(|ob| (ob.clone().unwrap_generator(), 1)).collect();
130            let term: Polynomial<_, _, _> =
131                [(Parameter::generator(mor), term)].into_iter().collect();
132            for input in inputs {
133                sys.add_term(input.unwrap_generator(), -term.clone());
134            }
135            for output in outputs {
136                sys.add_term(output.unwrap_generator(), term.clone());
137            }
138        }
139
140        // Normalize since terms commonly cancel out in mass-action dynamics.
141        sys.normalize()
142    }
143
144    /// Creates a mass-action system with numerical rate coefficients.
145    pub fn build_numerical_system(
146        &self,
147        model: &ModalDblModel,
148        data: MassActionProblemData,
149    ) -> ODEAnalysis<NumericalPolynomialSystem<u8>> {
150        into_numerical_system(self.build_system(model), data)
151    }
152
153    /// Creates a stochastic mass-action system.
154    pub fn build_stochastic_system(
155        &self,
156        model: &ModalDblModel,
157        data: MassActionProblemData,
158    ) -> StochasticMassActionAnalysis {
159        let ob_generators: Vec<_> = model.ob_generators_with_type(&self.place_ob_type).collect();
160
161        let initial: Vec<_> = ob_generators
162            .iter()
163            .map(|id| data.initial_values.get(id).copied().unwrap_or_default() as isize)
164            .collect();
165        let mut problem = gillespie::Gillespie::new(initial, false);
166
167        for mor in model.mor_generators_with_type(&self.transition_mor_type) {
168            let inputs = model
169                .get_dom(&mor)
170                .and_then(|ob| ob.clone().collect_product(None))
171                .unwrap_or_default();
172            let outputs = model
173                .get_cod(&mor)
174                .and_then(|ob| ob.clone().collect_product(None))
175                .unwrap_or_default();
176
177            // 1. convert the inputs/outputs to sequences of counts
178            let input_vec = ob_generators.iter().map(|id| {
179                inputs
180                    .iter()
181                    .filter(|&ob| matches!(ob, ModalOb::Generator(id2) if id2 == id))
182                    .count() as u32
183            });
184            let output_vec = ob_generators.iter().map(|id| {
185                outputs
186                    .iter()
187                    .filter(|&ob| matches!(ob, ModalOb::Generator(id2) if id2 == id))
188                    .count() as isize
189            });
190
191            // 2. output := output - input
192            let input_vec: Vec<_> = input_vec.collect();
193            let output_vec: Vec<_> = output_vec
194                .zip(input_vec.iter().copied())
195                .map(|(o, i)| o - (i as isize))
196                .collect();
197            if let Some(rate) = data.rates.get(&mor) {
198                problem.add_reaction(gillespie::Rate::lma(*rate as f64, input_vec), output_vec)
199            }
200        }
201
202        let variable_index: IndexMap<_, _> =
203            ob_generators.into_iter().enumerate().map(|(i, x)| (x, i)).collect();
204
205        StochasticMassActionAnalysis {
206            problem,
207            variable_index,
208            initial_values: data.initial_values,
209            duration: data.duration,
210        }
211    }
212}
213
214/// Mass-action ODE analysis for stock-flow models.
215pub struct StockFlowMassActionAnalysis {
216    /// Object type for stocks.
217    pub stock_ob_type: TabObType,
218    /// Morphism types for flows between stocks.
219    pub flow_mor_type: TabMorType,
220    /// Morphism types for links for stocks to flows.
221    pub link_mor_type: TabMorType,
222}
223
224impl Default for StockFlowMassActionAnalysis {
225    fn default() -> Self {
226        let stock_ob_type = TabObType::Basic(name("Object"));
227        let flow_mor_type = TabMorType::Hom(Box::new(stock_ob_type.clone()));
228        Self {
229            stock_ob_type,
230            flow_mor_type,
231            link_mor_type: TabMorType::Basic(name("Link")),
232        }
233    }
234}
235
236impl StockFlowMassActionAnalysis {
237    /// Creates a mass-action system with symbolic rate coefficients.
238    pub fn build_system(
239        &self,
240        model: &DiscreteTabModel,
241    ) -> PolynomialSystem<QualifiedName, Parameter<QualifiedName>, u8> {
242        let mut terms: HashMap<QualifiedName, Monomial<QualifiedName, u8>> = model
243            .mor_generators_with_type(&self.flow_mor_type)
244            .map(|flow| {
245                let dom = model.mor_generator_dom(&flow).unwrap_basic();
246                (flow, Monomial::generator(dom))
247            })
248            .collect();
249
250        for link in model.mor_generators_with_type(&self.link_mor_type) {
251            let dom = model.mor_generator_dom(&link).unwrap_basic();
252            let path = model.mor_generator_cod(&link).unwrap_tabulated();
253            let Some(TabEdge::Basic(cod)) = path.only() else {
254                panic!("Codomain of link should be basic morphism");
255            };
256            if let Some(term) = terms.get_mut(&cod) {
257                *term = std::mem::take(term) * Monomial::generator(dom);
258            } else {
259                panic!("Codomain of link does not belong to model");
260            };
261        }
262
263        let terms: Vec<_> = terms
264            .into_iter()
265            .map(|(flow, term)| {
266                let param = Parameter::generator(flow.clone());
267                let poly: Polynomial<_, _, _> = [(param, term)].into_iter().collect();
268                (flow, poly)
269            })
270            .collect();
271
272        let mut sys = PolynomialSystem::new();
273        for ob in model.ob_generators_with_type(&self.stock_ob_type) {
274            sys.add_term(ob, Polynomial::zero());
275        }
276        for (flow, term) in terms.iter() {
277            let dom = model.mor_generator_dom(flow).unwrap_basic();
278            sys.add_term(dom, -term.clone());
279        }
280        for (flow, term) in terms {
281            let cod = model.mor_generator_cod(&flow).unwrap_basic();
282            sys.add_term(cod, term);
283        }
284        sys
285    }
286
287    /// Creates a mass-action system with numerical rate coefficients.
288    pub fn build_numerical_system(
289        &self,
290        model: &DiscreteTabModel,
291        data: MassActionProblemData,
292    ) -> ODEAnalysis<NumericalPolynomialSystem<u8>> {
293        into_numerical_system(self.build_system(model), data)
294    }
295}
296
297fn into_numerical_system(
298    sys: PolynomialSystem<QualifiedName, Parameter<QualifiedName>, u8>,
299    data: MassActionProblemData,
300) -> ODEAnalysis<NumericalPolynomialSystem<u8>> {
301    let ob_index: IndexMap<_, _> =
302        sys.components.keys().cloned().enumerate().map(|(i, x)| (x, i)).collect();
303    let n = ob_index.len();
304
305    let initial_values = ob_index
306        .keys()
307        .map(|ob| data.initial_values.get(ob).copied().unwrap_or_default());
308    let x0 = DVector::from_iterator(n, initial_values);
309
310    let sys = sys
311        .extend_scalars(|poly| poly.eval(|flow| data.rates.get(flow).copied().unwrap_or_default()))
312        .to_numerical();
313
314    let problem = ODEProblem::new(sys, x0).end_time(data.duration);
315    ODEAnalysis::new(problem, ob_index)
316}
317
318#[cfg(test)]
319mod tests {
320    use expect_test::expect;
321    use std::rc::Rc;
322
323    use super::*;
324    use crate::stdlib::{models::*, theories::*};
325
326    #[test]
327    fn backward_link_dynamics() {
328        let th = Rc::new(th_category_links());
329        let model = backward_link(th);
330        let sys = StockFlowMassActionAnalysis::default().build_system(&model);
331        let expected = expect!([r#"
332            dx = ((-1) f) x y
333            dy = f x y
334        "#]);
335        expected.assert_eq(&sys.to_string());
336    }
337
338    #[test]
339    fn catalysis_dynamics() {
340        let th = Rc::new(th_sym_monoidal_category());
341        let model = catalyzed_reaction(th);
342        let sys = PetriNetMassActionAnalysis::default().build_system(&model);
343        let expected = expect!([r#"
344            dc = 0
345            dx = ((-1) f) c x
346            dy = f c x
347        "#]);
348        expected.assert_eq(&sys.to_string());
349    }
350
351    #[test]
352    fn sir_petri_stochastic_dynamics() {
353        let th = Rc::new(th_sym_monoidal_category());
354        let model = sir_petri(th);
355        let data = MassActionProblemData {
356            rates: HashMap::from_iter([(name("infection"), 1e-5f32), (name("recovery"), 1e-2f32)]),
357            initial_values: HashMap::from_iter([
358                (name("S"), 1e5f32),
359                (name("I"), 1f32),
360                (name("R"), 0f32),
361            ]),
362            duration: 10f32,
363        };
364        let sys = PetriNetMassActionAnalysis::default().build_stochastic_system(&model, data);
365        assert_eq!(2, sys.problem.nb_reactions());
366        assert_eq!(3, sys.problem.nb_species());
367    }
368}
catlog/stdlib/analyses/ode/mass_action.rs

catlog/stdlib/analyses/ode/
mass_action.rs
