1use runmat_builtins::{
4 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
5 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
6};
7use runmat_macros::runtime_builtin;
8use runmat_value::Value;
9
10use crate::builtins::common::spec::{
11 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
12 ReductionNaN, ResidencyPolicy, ShapeRequirements,
13};
14use crate::builtins::math::ode::common::{
15 build_ode_output, define_ode_integer_contract, ode_options_from_struct, parse_ode_input,
16 parse_options, prepare_ode_options, solve_ode, OdeMethod,
17};
18use crate::builtins::math::ode::type_resolvers::ode_solution_type;
19use crate::{build_runtime_error, BuiltinResult, RuntimeError};
20
21const NAME: &str = "ode15s";
22
23define_ode_integer_contract!("ode15s", "Ode15s");
24
25const ODE15S_OUTPUT_Y: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
26 name: "y",
27 ty: BuiltinParamType::NumericArray,
28 arity: BuiltinParamArity::Required,
29 default: None,
30 description: "Solution states evaluated over tspan.",
31}];
32
33const ODE15S_OUTPUT_TY: [BuiltinParamDescriptor; 2] = [
34 BuiltinParamDescriptor {
35 name: "t",
36 ty: BuiltinParamType::NumericArray,
37 arity: BuiltinParamArity::Required,
38 default: None,
39 description: "Time points selected by solver.",
40 },
41 BuiltinParamDescriptor {
42 name: "y",
43 ty: BuiltinParamType::NumericArray,
44 arity: BuiltinParamArity::Required,
45 default: None,
46 description: "Solution states at each returned time point.",
47 },
48];
49
50const ODE15S_INPUTS_CORE: [BuiltinParamDescriptor; 3] = [
51 BuiltinParamDescriptor {
52 name: "odefun",
53 ty: BuiltinParamType::Any,
54 arity: BuiltinParamArity::Required,
55 default: None,
56 description: "ODE right-hand-side callback f(t,y).",
57 },
58 BuiltinParamDescriptor {
59 name: "tspan",
60 ty: BuiltinParamType::Any,
61 arity: BuiltinParamArity::Required,
62 default: None,
63 description: "Time interval or monotonic time vector.",
64 },
65 BuiltinParamDescriptor {
66 name: "y0",
67 ty: BuiltinParamType::Any,
68 arity: BuiltinParamArity::Required,
69 default: None,
70 description: "Initial state vector/value.",
71 },
72];
73
74const ODE15S_INPUTS_WITH_OPTIONS: [BuiltinParamDescriptor; 4] = [
75 BuiltinParamDescriptor {
76 name: "odefun",
77 ty: BuiltinParamType::Any,
78 arity: BuiltinParamArity::Required,
79 default: None,
80 description: "ODE right-hand-side callback f(t,y).",
81 },
82 BuiltinParamDescriptor {
83 name: "tspan",
84 ty: BuiltinParamType::Any,
85 arity: BuiltinParamArity::Required,
86 default: None,
87 description: "Time interval or monotonic time vector.",
88 },
89 BuiltinParamDescriptor {
90 name: "y0",
91 ty: BuiltinParamType::Any,
92 arity: BuiltinParamArity::Required,
93 default: None,
94 description: "Initial state vector/value.",
95 },
96 BuiltinParamDescriptor {
97 name: "options",
98 ty: BuiltinParamType::Any,
99 arity: BuiltinParamArity::Optional,
100 default: None,
101 description: "Optional struct with tolerances and step controls.",
102 },
103];
104
105const ODE15S_SIGNATURES: [BuiltinSignatureDescriptor; 4] = [
106 BuiltinSignatureDescriptor {
107 label: "y = ode15s(odefun, tspan, y0)",
108 inputs: &ODE15S_INPUTS_CORE,
109 outputs: &ODE15S_OUTPUT_Y,
110 },
111 BuiltinSignatureDescriptor {
112 label: "y = ode15s(odefun, tspan, y0, options)",
113 inputs: &ODE15S_INPUTS_WITH_OPTIONS,
114 outputs: &ODE15S_OUTPUT_Y,
115 },
116 BuiltinSignatureDescriptor {
117 label: "[t, y] = ode15s(odefun, tspan, y0)",
118 inputs: &ODE15S_INPUTS_CORE,
119 outputs: &ODE15S_OUTPUT_TY,
120 },
121 BuiltinSignatureDescriptor {
122 label: "[t, y] = ode15s(odefun, tspan, y0, options)",
123 inputs: &ODE15S_INPUTS_WITH_OPTIONS,
124 outputs: &ODE15S_OUTPUT_TY,
125 },
126];
127
128const ODE15S_ERROR_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
129 code: "RM.ODE15S.INVALID_ARGUMENT",
130 identifier: Some("RunMat:ode15s:InvalidArgument"),
131 when: "Input argument count/options struct grammar is invalid.",
132 message: "ode15s: invalid argument",
133};
134
135const ODE15S_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
136 code: "RM.ODE15S.INVALID_INPUT",
137 identifier: Some("RunMat:ode15s:InvalidInput"),
138 when: "ODE input/state/callback semantics are invalid for integration.",
139 message: "ode15s: invalid input",
140};
141
142const ODE15S_ERROR_INTERNAL: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
143 code: "RM.ODE15S.INTERNAL",
144 identifier: Some("RunMat:ode15s:Internal"),
145 when: "Internal output materialization fails.",
146 message: "ode15s: internal runtime failure",
147};
148
149const ODE15S_ERRORS: [BuiltinErrorDescriptor; 3] = [
150 ODE15S_ERROR_INVALID_ARGUMENT,
151 ODE15S_ERROR_INVALID_INPUT,
152 ODE15S_ERROR_INTERNAL,
153];
154
155pub const ODE15S_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
156 signatures: &ODE15S_SIGNATURES,
157 output_mode: BuiltinOutputMode::ByRequestedOutputCount,
158 completion_policy: BuiltinCompletionPolicy::Public,
159 errors: &ODE15S_ERRORS,
160};
161
162fn ode15s_error_with_detail(
163 error: &'static BuiltinErrorDescriptor,
164 detail: impl AsRef<str>,
165) -> RuntimeError {
166 let detail = detail.as_ref();
167 let message = if detail.starts_with("ode15s:") {
168 detail.to_string()
169 } else {
170 format!("{}: {}", error.message, detail)
171 };
172 let mut builder = build_runtime_error(message).with_builtin(NAME);
173 if let Some(identifier) = error.identifier {
174 builder = builder.with_identifier(identifier);
175 }
176 builder.build()
177}
178
179fn ode15s_map_error(err: RuntimeError, fallback: &'static BuiltinErrorDescriptor) -> RuntimeError {
180 if err.identifier().is_some() {
181 err
182 } else {
183 ode15s_error_with_detail(fallback, err.message())
184 }
185}
186
187#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::math::ode::ode15s")]
188pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
189 name: "ode15s",
190 op_kind: GpuOpKind::Custom("ode-solve-stiff"),
191 supported_precisions: &[],
192 broadcast: BroadcastSemantics::None,
193 provider_hooks: &[],
194 constant_strategy: ConstantStrategy::InlineLiteral,
195 residency: ResidencyPolicy::GatherImmediately,
196 nan_mode: ReductionNaN::Include,
197 two_pass_threshold: None,
198 workgroup_size: None,
199 accepts_nan_mode: false,
200 notes: "Stiff ODE integration runs on the host. RHS callbacks may call GPU-aware builtins.",
201};
202
203#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::math::ode::ode15s")]
204pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
205 name: "ode15s",
206 shape: ShapeRequirements::Any,
207 constant_strategy: ConstantStrategy::InlineLiteral,
208 elementwise: None,
209 reduction: None,
210 emits_nan: false,
211 notes: "ODE integration repeatedly invokes user callbacks and terminates fusion planning.",
212};
213
214#[runtime_builtin(
215 name = "ode15s",
216 category = "math/ode",
217 summary = "Solve stiff ODE systems with adaptive implicit integration.",
218 keywords = "ode15s,ode,stiff,implicit,adaptive step",
219 accel = "sink",
220 type_resolver(ode_solution_type),
221 descriptor(crate::builtins::math::ode::ode15s::ODE15S_DESCRIPTOR),
222 extensions(crate::builtins::math::ode::ode15s::EXTENSIONS),
223 integer_capabilities(crate::builtins::math::ode::ode15s::INTEGER_CAPABILITIES),
224 builtin_path = "crate::builtins::math::ode::ode15s"
225)]
226async fn ode15s_builtin(
227 function: Value,
228 tspan: Value,
229 y0: Value,
230 rest: Vec<Value>,
231) -> BuiltinResult<Value> {
232 if rest.len() > 1 {
233 return Err(ode15s_error_with_detail(
234 &ODE15S_ERROR_INVALID_ARGUMENT,
235 "too many input arguments",
236 ));
237 }
238 let options = parse_options(NAME, rest.first())
239 .map_err(|err| ode15s_map_error(err, &ODE15S_ERROR_INVALID_ARGUMENT))?;
240 let options = prepare_ode_options(NAME, options, ODE_COMPATIBILITY_EXTENSIONS)
241 .await
242 .map_err(|err| ode15s_map_error(err, &ODE15S_ERROR_INVALID_ARGUMENT))?;
243 let opts = ode_options_from_struct(NAME, options.as_ref())
244 .map_err(|err| ode15s_map_error(err, &ODE15S_ERROR_INVALID_ARGUMENT))?;
245 let input = parse_ode_input(NAME, tspan, y0, ODE_COMPATIBILITY_EXTENSIONS)
246 .await
247 .map_err(|err| ode15s_map_error(err, &ODE15S_ERROR_INVALID_INPUT))?;
248 let result = solve_ode(NAME, OdeMethod::Ode15s, &function, &input, &opts)
249 .await
250 .map_err(|err| ode15s_map_error(err, &ODE15S_ERROR_INVALID_INPUT))?;
251 build_ode_output(NAME, result).map_err(|err| ode15s_map_error(err, &ODE15S_ERROR_INTERNAL))
252}
253
254#[cfg(test)]
255mod tests {
256 use super::*;
257 use futures::executor::block_on;
258 use runmat_value::{StructValue, Tensor};
259 use std::sync::Arc;
260
261 #[test]
262 fn ode15s_handles_linear_stiff_decay() {
263 let _resolver =
264 crate::user_functions::install_semantic_function_resolver(Some(Arc::new(|_name| {
265 Some(0)
266 })));
267 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(Arc::new(
268 move |_function, args, _requested_outputs| {
269 let y = match &args[1] {
270 Value::Num(n) => *n,
271 other => panic!("expected scalar state, got {other:?}"),
272 };
273 Box::pin(async move { Ok(Value::Num(-15.0 * y)) })
274 },
275 )));
276
277 let out = block_on(ode15s_builtin(
278 Value::FunctionHandle("stiff_decay".into()),
279 Value::Tensor(Tensor::new(vec![0.0, 1.0], vec![1, 2]).unwrap()),
280 Value::Num(1.0),
281 Vec::new(),
282 ))
283 .unwrap();
284
285 match out {
286 Value::Tensor(t) => {
287 assert_eq!(t.cols(), 1);
288 let last = t.materialize_f64()[t.rows() - 1];
289 assert!(last.is_finite());
290 assert!(last > 0.0);
291 assert!(last < 0.1);
292 }
293 other => panic!("unexpected output {other:?}"),
294 }
295 }
296
297 #[test]
298 fn ode15s_accepts_picard_unstable_stiff_step_with_newton() {
299 let _resolver =
300 crate::user_functions::install_semantic_function_resolver(Some(Arc::new(|_name| {
301 Some(0)
302 })));
303 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(Arc::new(
304 move |_function, args, _requested_outputs| {
305 let y = match &args[1] {
306 Value::Num(n) => *n,
307 other => panic!("expected scalar state, got {other:?}"),
308 };
309 Box::pin(async move { Ok(Value::Num(-1000.0 * y)) })
310 },
311 )));
312 let mut options = StructValue::new();
313 options.insert("RelTol", Value::Num(1.0e6));
314 options.insert("AbsTol", Value::Num(1.0e6));
315 options.insert("InitialStep", Value::Num(0.1));
316 options.insert("MaxStep", Value::Num(0.1));
317 options.insert("MaxSteps", Value::Num(2.0));
318
319 let out = block_on(ode15s_builtin(
320 Value::FunctionHandle("very_stiff_decay".into()),
321 Value::Tensor(Tensor::new(vec![0.0, 0.1], vec![1, 2]).unwrap()),
322 Value::Num(1.0),
323 vec![Value::Struct(options)],
324 ))
325 .unwrap();
326
327 match out {
328 Value::Tensor(t) => {
329 assert_eq!(t.cols(), 1);
330 let last = t.materialize_f64()[t.rows() - 1];
331 assert!(last.is_finite());
332 assert!(last > 0.0);
333 assert!(last < 0.02);
334 }
335 other => panic!("unexpected output {other:?}"),
336 }
337 }
338
339 #[test]
340 fn ode15s_accepts_semantic_function_handle_rhs() {
341 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(Arc::new(
342 move |function, args, _requested_outputs| {
343 assert_eq!(function, 57);
344 let y = match &args[1] {
345 Value::Num(n) => *n,
346 other => panic!("expected scalar state, got {other:?}"),
347 };
348 Box::pin(async move { Ok(Value::Num(-15.0 * y)) })
349 },
350 )));
351
352 let out = block_on(ode15s_builtin(
353 Value::BoundFunctionHandle {
354 name: "ode_stiff_decay".to_string(),
355 function: 57,
356 },
357 Value::Tensor(Tensor::new(vec![0.0, 1.0], vec![1, 2]).unwrap()),
358 Value::Num(1.0),
359 Vec::new(),
360 ))
361 .unwrap();
362
363 match out {
364 Value::Tensor(t) => {
365 assert_eq!(t.cols(), 1);
366 let last = t.materialize_f64()[t.rows() - 1];
367 assert!(last.is_finite());
368 assert!(last > 0.0);
369 assert!(last < 0.1);
370 }
371 other => panic!("unexpected output {other:?}"),
372 }
373 }
374
375 #[test]
376 fn ode15s_too_many_inputs_uses_stable_identifier() {
377 let err = block_on(ode15s_builtin(
378 Value::FunctionHandle("stiff_decay".into()),
379 Value::Tensor(Tensor::new(vec![0.0, 1.0], vec![1, 2]).unwrap()),
380 Value::Num(1.0),
381 vec![Value::Num(1.0), Value::Num(2.0)],
382 ))
383 .expect_err("expected too many inputs error");
384 assert_eq!(err.identifier(), ODE15S_ERROR_INVALID_ARGUMENT.identifier);
385 }
386
387 #[test]
388 fn ode15s_descriptor_signatures_cover_surface() {
389 let labels: Vec<&str> = ODE15S_DESCRIPTOR
390 .signatures
391 .iter()
392 .map(|signature| signature.label)
393 .collect();
394 assert_eq!(
395 labels,
396 vec![
397 "y = ode15s(odefun, tspan, y0)",
398 "y = ode15s(odefun, tspan, y0, options)",
399 "[t, y] = ode15s(odefun, tspan, y0)",
400 "[t, y] = ode15s(odefun, tspan, y0, options)",
401 ]
402 );
403 }
404
405 #[test]
406 fn ode15s_descriptor_errors_have_stable_codes() {
407 let codes: Vec<&str> = ODE15S_DESCRIPTOR
408 .errors
409 .iter()
410 .map(|error| error.code)
411 .collect();
412 assert_eq!(
413 codes,
414 vec![
415 "RM.ODE15S.INVALID_ARGUMENT",
416 "RM.ODE15S.INVALID_INPUT",
417 "RM.ODE15S.INTERNAL",
418 ]
419 );
420 }
421}