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

catlog/stdlib/analyses/ode/
mass_action.rs
