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 mass-action system with numerical rate coefficients.
145    pub fn build_numerical_system(
146        &self,
147        model: &ModalDblModel,
148        data: MassActionProblemData,
149    ) -> ODEAnalysis<NumericalPolynomialSystem<i8>> {
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 type for flows between stocks.
219    pub flow_mor_type: TabMorType,
220    /// Morphism type for positive links from stocks to flows.
221    pub pos_link_mor_type: TabMorType,
222    /// Morphism type for negative links from stocks to flows.
223    pub neg_link_mor_type: TabMorType,
224}
225
226impl Default for StockFlowMassActionAnalysis {
227    fn default() -> Self {
228        let stock_ob_type = TabObType::Basic(name("Object"));
229        let flow_mor_type = TabMorType::Hom(Box::new(stock_ob_type.clone()));
230        Self {
231            stock_ob_type,
232            flow_mor_type,
233            pos_link_mor_type: TabMorType::Basic(name("Link")),
234            neg_link_mor_type: TabMorType::Basic(name("NegativeLink")),
235        }
236    }
237}
238
239impl StockFlowMassActionAnalysis {
240    /// Creates a mass-action system with symbolic rate coefficients.
241    pub fn build_system(
242        &self,
243        model: &DiscreteTabModel,
244    ) -> PolynomialSystem<QualifiedName, Parameter<QualifiedName>, i8> {
245        let mut terms: HashMap<QualifiedName, Monomial<QualifiedName, i8>> = model
246            .mor_generators_with_type(&self.flow_mor_type)
247            .map(|flow| {
248                let dom = model.mor_generator_dom(&flow).unwrap_basic();
249                (flow, Monomial::generator(dom))
250            })
251            .collect();
252
253        let mut multiply_for_link = |link: QualifiedName, exponent: i8| {
254            let dom = model.mor_generator_dom(&link).unwrap_basic();
255            let path = model.mor_generator_cod(&link).unwrap_tabulated();
256            let Some(TabEdge::Basic(cod)) = path.only() else {
257                panic!("Codomain of link should be basic morphism");
258            };
259            if let Some(term) = terms.get_mut(&cod) {
260                let mon: Monomial<_, i8> = [(dom, exponent)].into_iter().collect();
261                *term = std::mem::take(term) * mon;
262            } else {
263                panic!("Codomain of link does not belong to model");
264            };
265        };
266
267        for link in model.mor_generators_with_type(&self.pos_link_mor_type) {
268            multiply_for_link(link, 1);
269        }
270        for link in model.mor_generators_with_type(&self.neg_link_mor_type) {
271            multiply_for_link(link, -1);
272        }
273
274        let terms: Vec<_> = terms
275            .into_iter()
276            .map(|(flow, term)| {
277                let param = Parameter::generator(flow.clone());
278                let poly: Polynomial<_, _, _> = [(param, term)].into_iter().collect();
279                (flow, poly)
280            })
281            .collect();
282
283        let mut sys = PolynomialSystem::new();
284        for ob in model.ob_generators_with_type(&self.stock_ob_type) {
285            sys.add_term(ob, Polynomial::zero());
286        }
287        for (flow, term) in terms.iter() {
288            let dom = model.mor_generator_dom(flow).unwrap_basic();
289            sys.add_term(dom, -term.clone());
290        }
291        for (flow, term) in terms {
292            let cod = model.mor_generator_cod(&flow).unwrap_basic();
293            sys.add_term(cod, term);
294        }
295        sys
296    }
297
298    /// Creates a mass-action system with numerical rate coefficients.
299    pub fn build_numerical_system(
300        &self,
301        model: &DiscreteTabModel,
302        data: MassActionProblemData,
303    ) -> ODEAnalysis<NumericalPolynomialSystem<i8>> {
304        into_numerical_system(self.build_system(model), data)
305    }
306}
307
308fn into_numerical_system(
309    sys: PolynomialSystem<QualifiedName, Parameter<QualifiedName>, i8>,
310    data: MassActionProblemData,
311) -> ODEAnalysis<NumericalPolynomialSystem<i8>> {
312    let ob_index: IndexMap<_, _> =
313        sys.components.keys().cloned().enumerate().map(|(i, x)| (x, i)).collect();
314    let n = ob_index.len();
315
316    let initial_values = ob_index
317        .keys()
318        .map(|ob| data.initial_values.get(ob).copied().unwrap_or_default());
319    let x0 = DVector::from_iterator(n, initial_values);
320
321    let sys = sys
322        .extend_scalars(|poly| poly.eval(|flow| data.rates.get(flow).copied().unwrap_or_default()))
323        .to_numerical();
324
325    let problem = ODEProblem::new(sys, x0).end_time(data.duration);
326    ODEAnalysis::new(problem, ob_index)
327}
328
329#[cfg(test)]
330mod tests {
331    use expect_test::expect;
332    use std::rc::Rc;
333
334    use super::*;
335    use crate::stdlib::{models::*, theories::*};
336
337    #[test]
338    fn backward_link_dynamics() {
339        let th = Rc::new(th_category_links());
340        let model = backward_link(th);
341        let sys = StockFlowMassActionAnalysis::default().build_system(&model);
342        let expected = expect!([r#"
343            dx = ((-1) f) x y
344            dy = f x y
345        "#]);
346        expected.assert_eq(&sys.to_string());
347    }
348
349    #[test]
350    fn positive_backward_link_dynamics() {
351        let th = Rc::new(th_category_signed_links());
352        let model = positive_backward_link(th);
353        let sys = StockFlowMassActionAnalysis::default().build_system(&model);
354        let expected = expect!([r#"
355            dx = ((-1) f) x y
356            dy = f x y
357        "#]);
358        expected.assert_eq(&sys.to_string());
359    }
360
361    #[test]
362    fn negative_backward_link_dynamics() {
363        let th = Rc::new(th_category_signed_links());
364        let model = negative_backward_link(th);
365        let sys = StockFlowMassActionAnalysis::default().build_system(&model);
366        let expected = expect!([r#"
367            dx = ((-1) f) x y^{-1}
368            dy = f x y^{-1}
369        "#]);
370        expected.assert_eq(&sys.to_string());
371    }
372
373    #[test]
374    fn catalysis_dynamics() {
375        let th = Rc::new(th_sym_monoidal_category());
376        let model = catalyzed_reaction(th);
377        let sys = PetriNetMassActionAnalysis::default().build_system(&model);
378        let expected = expect!([r#"
379            dc = 0
380            dx = ((-1) f) c x
381            dy = f c x
382        "#]);
383        expected.assert_eq(&sys.to_string());
384    }
385
386    #[test]
387    fn sir_petri_stochastic_dynamics() {
388        let th = Rc::new(th_sym_monoidal_category());
389        let model = sir_petri(th);
390        let data = MassActionProblemData {
391            rates: HashMap::from_iter([(name("infection"), 1e-5f32), (name("recovery"), 1e-2f32)]),
392            initial_values: HashMap::from_iter([
393                (name("S"), 1e5f32),
394                (name("I"), 1f32),
395                (name("R"), 0f32),
396            ]),
397            duration: 10f32,
398        };
399        let sys = PetriNetMassActionAnalysis::default().build_stochastic_system(&model, data);
400        assert_eq!(2, sys.problem.nb_reactions());
401        assert_eq!(3, sys.problem.nb_species());
402    }
403}
catlog/stdlib/analyses/ode/mass_action.rs

catlog/stdlib/analyses/ode/
mass_action.rs
