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, i8>;
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>, i8> {
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 stochastic mass-action system.
145    pub fn build_stochastic_system(
146        &self,
147        model: &ModalDblModel,
148        data: MassActionProblemData,
149    ) -> StochasticMassActionAnalysis {
150        let ob_generators: Vec<_> = model.ob_generators_with_type(&self.place_ob_type).collect();
151
152        let initial: Vec<_> = ob_generators
153            .iter()
154            .map(|id| data.initial_values.get(id).copied().unwrap_or_default() as isize)
155            .collect();
156        let mut problem = gillespie::Gillespie::new(initial, false);
157
158        for mor in model.mor_generators_with_type(&self.transition_mor_type) {
159            let inputs = model
160                .get_dom(&mor)
161                .and_then(|ob| ob.clone().collect_product(None))
162                .unwrap_or_default();
163            let outputs = model
164                .get_cod(&mor)
165                .and_then(|ob| ob.clone().collect_product(None))
166                .unwrap_or_default();
167
168            // 1. convert the inputs/outputs to sequences of counts
169            let input_vec = ob_generators.iter().map(|id| {
170                inputs
171                    .iter()
172                    .filter(|&ob| matches!(ob, ModalOb::Generator(id2) if id2 == id))
173                    .count() as u32
174            });
175            let output_vec = ob_generators.iter().map(|id| {
176                outputs
177                    .iter()
178                    .filter(|&ob| matches!(ob, ModalOb::Generator(id2) if id2 == id))
179                    .count() as isize
180            });
181
182            // 2. output := output - input
183            let input_vec: Vec<_> = input_vec.collect();
184            let output_vec: Vec<_> = output_vec
185                .zip(input_vec.iter().copied())
186                .map(|(o, i)| o - (i as isize))
187                .collect();
188            if let Some(rate) = data.rates.get(&mor) {
189                problem.add_reaction(gillespie::Rate::lma(*rate as f64, input_vec), output_vec)
190            }
191        }
192
193        let variable_index: IndexMap<_, _> =
194            ob_generators.into_iter().enumerate().map(|(i, x)| (x, i)).collect();
195
196        StochasticMassActionAnalysis {
197            problem,
198            variable_index,
199            initial_values: data.initial_values,
200            duration: data.duration,
201        }
202    }
203}
204
205/// Mass-action ODE analysis for stock-flow models.
206pub struct StockFlowMassActionAnalysis {
207    /// Object type for stocks.
208    pub stock_ob_type: TabObType,
209    /// Morphism type for flows between stocks.
210    pub flow_mor_type: TabMorType,
211    /// Morphism type for positive links from stocks to flows.
212    pub pos_link_mor_type: TabMorType,
213    /// Morphism type for negative links from stocks to flows.
214    pub neg_link_mor_type: TabMorType,
215}
216
217impl Default for StockFlowMassActionAnalysis {
218    fn default() -> Self {
219        let stock_ob_type = TabObType::Basic(name("Object"));
220        let flow_mor_type = TabMorType::Hom(Box::new(stock_ob_type.clone()));
221        Self {
222            stock_ob_type,
223            flow_mor_type,
224            pos_link_mor_type: TabMorType::Basic(name("Link")),
225            neg_link_mor_type: TabMorType::Basic(name("NegativeLink")),
226        }
227    }
228}
229
230impl StockFlowMassActionAnalysis {
231    /// Creates a mass-action system with symbolic rate coefficients.
232    pub fn build_system(
233        &self,
234        model: &DiscreteTabModel,
235    ) -> PolynomialSystem<QualifiedName, Parameter<QualifiedName>, i8> {
236        let mut terms: HashMap<QualifiedName, Monomial<QualifiedName, i8>> = model
237            .mor_generators_with_type(&self.flow_mor_type)
238            .map(|flow| {
239                let dom = model.mor_generator_dom(&flow).unwrap_basic();
240                (flow, Monomial::generator(dom))
241            })
242            .collect();
243
244        let mut multiply_for_link = |link: QualifiedName, exponent: i8| {
245            let dom = model.mor_generator_dom(&link).unwrap_basic();
246            let path = model.mor_generator_cod(&link).unwrap_tabulated();
247            let Some(TabEdge::Basic(cod)) = path.only() else {
248                panic!("Codomain of link should be basic morphism");
249            };
250            if let Some(term) = terms.get_mut(&cod) {
251                let mon: Monomial<_, i8> = [(dom, exponent)].into_iter().collect();
252                *term = std::mem::take(term) * mon;
253            } else {
254                panic!("Codomain of link does not belong to model");
255            };
256        };
257
258        for link in model.mor_generators_with_type(&self.pos_link_mor_type) {
259            multiply_for_link(link, 1);
260        }
261        for link in model.mor_generators_with_type(&self.neg_link_mor_type) {
262            multiply_for_link(link, -1);
263        }
264
265        let terms: Vec<_> = terms
266            .into_iter()
267            .map(|(flow, term)| {
268                let param = Parameter::generator(flow.clone());
269                let poly: Polynomial<_, _, _> = [(param, term)].into_iter().collect();
270                (flow, poly)
271            })
272            .collect();
273
274        let mut sys = PolynomialSystem::new();
275        for ob in model.ob_generators_with_type(&self.stock_ob_type) {
276            sys.add_term(ob, Polynomial::zero());
277        }
278        for (flow, term) in terms.iter() {
279            let dom = model.mor_generator_dom(flow).unwrap_basic();
280            sys.add_term(dom, -term.clone());
281        }
282        for (flow, term) in terms {
283            let cod = model.mor_generator_cod(&flow).unwrap_basic();
284            sys.add_term(cod, term);
285        }
286        sys
287    }
288}
289
290/// Substitutes numerical rate coefficients into a symbolic mass-action system.
291pub fn extend_mass_action_scalars(
292    sys: PolynomialSystem<QualifiedName, Parameter<QualifiedName>, i8>,
293    data: &MassActionProblemData,
294) -> PolynomialSystem<QualifiedName, f32, i8> {
295    sys.extend_scalars(|poly| poly.eval(|flow| data.rates.get(flow).copied().unwrap_or_default()))
296}
297
298/// Builds the numerical ODE analysis for a mass-action system whose scalars have been substituted.
299pub fn into_mass_action_analysis(
300    sys: PolynomialSystem<QualifiedName, f32, i8>,
301    data: MassActionProblemData,
302) -> ODEAnalysis<NumericalPolynomialSystem<i8>> {
303    let ob_index: IndexMap<_, _> =
304        sys.components.keys().cloned().enumerate().map(|(i, x)| (x, i)).collect();
305    let n = ob_index.len();
306
307    let initial_values = ob_index
308        .keys()
309        .map(|ob| data.initial_values.get(ob).copied().unwrap_or_default());
310    let x0 = DVector::from_iterator(n, initial_values);
311
312    let num_sys = sys.to_numerical();
313    let problem = ODEProblem::new(num_sys, x0).end_time(data.duration);
314
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::simulate::ode::LatexEquation;
325    use crate::stdlib::{models::*, theories::*};
326
327    #[test]
328    fn backward_link_dynamics() {
329        let th = Rc::new(th_category_links());
330        let model = backward_link(th);
331        let sys = StockFlowMassActionAnalysis::default().build_system(&model);
332        let expected = expect!([r#"
333            dx = (-f) x y
334            dy = f x y
335        "#]);
336        expected.assert_eq(&sys.to_string());
337    }
338
339    #[test]
340    fn positive_backward_link_dynamics() {
341        let th = Rc::new(th_category_signed_links());
342        let model = positive_backward_link(th);
343        let sys = StockFlowMassActionAnalysis::default().build_system(&model);
344        let expected = expect!([r#"
345            dx = (-f) x y
346            dy = f x y
347        "#]);
348        expected.assert_eq(&sys.to_string());
349    }
350
351    #[test]
352    fn negative_backward_link_dynamics() {
353        let th = Rc::new(th_category_signed_links());
354        let model = negative_backward_link(th);
355        let sys = StockFlowMassActionAnalysis::default().build_system(&model);
356        let expected = expect!([r#"
357            dx = (-f) x y^{-1}
358            dy = f x y^{-1}
359        "#]);
360        expected.assert_eq(&sys.to_string());
361    }
362
363    #[test]
364    fn catalysis_dynamics() {
365        let th = Rc::new(th_sym_monoidal_category());
366        let model = catalyzed_reaction(th);
367        let sys = PetriNetMassActionAnalysis::default().build_system(&model);
368        let expected = expect!([r#"
369            dx = (-f) c x
370            dy = f c x
371            dc = 0
372        "#]);
373        expected.assert_eq(&sys.to_string());
374    }
375
376    #[test]
377    fn sir_petri_stochastic_dynamics() {
378        let th = Rc::new(th_sym_monoidal_category());
379        let model = sir_petri(th);
380        let data = MassActionProblemData {
381            rates: HashMap::from_iter([(name("infect"), 1e-5f32), (name("recover"), 1e-2f32)]),
382            initial_values: HashMap::from_iter([
383                (name("S"), 1e5f32),
384                (name("I"), 1f32),
385                (name("R"), 0f32),
386            ]),
387            duration: 10f32,
388        };
389        let sys = PetriNetMassActionAnalysis::default().build_stochastic_system(&model, data);
390        assert_eq!(2, sys.problem.nb_reactions());
391        assert_eq!(3, sys.problem.nb_species());
392    }
393
394    #[test]
395    fn to_latex() {
396        let th = Rc::new(th_category_links());
397        let model = backward_link(th);
398        let sys = StockFlowMassActionAnalysis::default().build_system(&model);
399        let expected = vec![
400            LatexEquation {
401                lhs: "\\dot{x}".to_string(),
402                rhs: "(-f) x y".to_string(),
403            },
404            LatexEquation {
405                lhs: "\\dot{y}".to_string(),
406                rhs: "f x y".to_string(),
407            },
408        ];
409        assert_eq!(expected, sys.to_latex_equations());
410    }
411}
catlog/stdlib/analyses/ode/mass_action.rs

catlog/stdlib/analyses/ode/
mass_action.rs
