Skip to main content

ocas_py/
ode.rs

1//! Python `ode` module — ordinary differential equation solvers.
2
3use ocas_atom::Symbol;
4use ocas_calc::ode::{
5    ODE, ODESolution, ODEType, classify_ode as rs_classify, dsolve as rs_dsolve,
6    dsolve_ivp as rs_dsolve_ivp,
7};
8use pyo3::exceptions::PyValueError;
9use pyo3::prelude::*;
10
11use crate::expression::Expression;
12
13/// Classify an ODE and return the applicable method names.
14///
15/// `equation` is the ODE written as an expression equal to zero (e.g.
16/// `"Derivative(y(x), x) - y(x)"`), `func` the unknown function name
17/// (e.g. `"y"`), and `var` the independent variable (e.g. `"x"`).
18///
19/// Returns a list of method names such as `"LinearFirst"`, `"Separable"`,
20/// `"LinearConstantCoeff"`, etc.
21#[pyfunction]
22#[pyo3(name = "classify_ode")]
23pub fn py_classify_ode(equation: &Expression, func: &str, var: &str) -> PyResult<Vec<String>> {
24    let (ctx, ode) = build_ode(equation, func, var)?;
25    Ok(rs_classify(ctx, ode)
26        .iter()
27        .map(ode_type_name)
28        .map(str::to_owned)
29        .collect())
30}
31
32/// Solve an ODE symbolically.
33///
34/// - `equation`: expression equal to zero, e.g. `"Derivative(y(x), x) - y(x)"`.
35/// - `func`: unknown function name, e.g. `"y"`.
36/// - `var`: independent variable, e.g. `"x"`.
37/// - `hint`: optional method name (one of the strings returned by
38///   `classify_ode`) to force a specific solver.
39///
40/// Returns a string describing the solution: an explicit solution
41/// `y = ...`, an implicit form, a truncated series, or the unevaluated ODE.
42#[pyfunction]
43#[pyo3(name = "dsolve", signature = (equation, func, var, hint=None))]
44pub fn py_dsolve(
45    equation: &Expression,
46    func: &str,
47    var: &str,
48    hint: Option<&str>,
49) -> PyResult<String> {
50    let hint_type = hint.map(parse_ode_type).transpose()?;
51    let (ctx, ode) = build_ode(equation, func, var)?;
52    let sol = rs_dsolve(ctx, ode, hint_type);
53    Ok(format_solution(&sol))
54}
55
56/// Solve a first- or second-order linear constant-coefficient IVP via the
57/// Laplace transform.
58///
59/// - `y0`: value `y(0)` as a string expression (e.g. `"1"`).
60/// - `y1`: value `y'(0)` (required for second-order problems).
61///
62/// Returns an explicit solution string with no free constants.
63#[pyfunction]
64#[pyo3(name = "dsolve_ivp", signature = (equation, func, var, y0, y1=None))]
65pub fn py_dsolve_ivp(
66    equation: &Expression,
67    func: &str,
68    var: &str,
69    y0: &str,
70    y1: Option<&str>,
71) -> PyResult<String> {
72    let (ctx, ode) = build_ode(equation, func, var)?;
73    let y0_atom = parse_in(ctx, y0)?;
74    let y1_atom = y1.map(|s| parse_in(ctx, s)).transpose()?;
75    let sol = rs_dsolve_ivp(ctx, ode, y0_atom, y1_atom);
76    Ok(format_solution(&sol))
77}
78
79// ---------------------------------------------------------------------------
80// Helpers
81// ---------------------------------------------------------------------------
82
83/// Build an `ODE` inside the expression's arena, returning the arena
84/// reference and the constructed ODE.
85fn build_ode(
86    equation: &Expression,
87    func: &str,
88    var: &str,
89) -> PyResult<(&'static ocas_atom::AtomArena<'static>, ODE<'static>)> {
90    let ctx = equation.ctx_ref();
91    let x = ctx.var(var);
92    let func_atom = ctx.fun(func, &[x]);
93    let ode = ODE {
94        equation: equation.atom(),
95        func: func_atom,
96        var: Symbol::new(var),
97    };
98    Ok((ctx, ode))
99}
100
101fn parse_in<'a>(ctx: &'a ocas_atom::AtomArena<'a>, input: &str) -> PyResult<ocas_atom::Atom<'a>> {
102    ocas_parse::parse(ctx, input).map_err(|e| PyValueError::new_err(format!("parse error: {e}")))
103}
104
105fn ode_type_name(t: &ODEType) -> &'static str {
106    match t {
107        ODEType::Separable => "Separable",
108        ODEType::LinearFirst => "LinearFirst",
109        ODEType::Bernoulli => "Bernoulli",
110        ODEType::Exact => "Exact",
111        ODEType::Homogeneous => "Homogeneous",
112        ODEType::LinearConstantCoeff => "LinearConstantCoeff",
113        ODEType::CauchyEuler => "CauchyEuler",
114        ODEType::ReductionOfOrder => "ReductionOfOrder",
115        ODEType::PowerSeries => "PowerSeries",
116    }
117}
118
119fn parse_ode_type(name: &str) -> PyResult<ODEType> {
120    match name {
121        "Separable" => Ok(ODEType::Separable),
122        "LinearFirst" => Ok(ODEType::LinearFirst),
123        "Bernoulli" => Ok(ODEType::Bernoulli),
124        "Exact" => Ok(ODEType::Exact),
125        "Homogeneous" => Ok(ODEType::Homogeneous),
126        "LinearConstantCoeff" => Ok(ODEType::LinearConstantCoeff),
127        "CauchyEuler" => Ok(ODEType::CauchyEuler),
128        "ReductionOfOrder" => Ok(ODEType::ReductionOfOrder),
129        "PowerSeries" => Ok(ODEType::PowerSeries),
130        other => Err(PyValueError::new_err(format!(
131            "unknown ODE type hint: {other}"
132        ))),
133    }
134}
135
136fn format_solution(sol: &ODESolution<'_>) -> String {
137    match sol {
138        ODESolution::Explicit(e) => format!("y = {e}"),
139        ODESolution::Implicit(e) => format!("{e} = C"),
140        ODESolution::Parametric(a, b) => format!("x = {a}, y = {b}"),
141        ODESolution::Series(e, n) => format!("series({n} terms): y = {e}"),
142        ODESolution::System(comps) => {
143            let joined: Vec<String> = comps
144                .iter()
145                .enumerate()
146                .map(|(i, c)| format!("y{} = {c}", i + 1))
147                .collect();
148            joined.join(", ")
149        }
150        ODESolution::Unsolved(_) => "unsolved".to_string(),
151    }
152}