1use runmat_builtins::{
4 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinExtensionDescriptor,
5 BuiltinExtensionMode, BuiltinIntegerBackendRule, BuiltinIntegerCapabilityDescriptor,
6 BuiltinIntegerComputationDomain, BuiltinIntegerInputAvailability,
7 BuiltinIntegerInputCapability, BuiltinIntegerOutputClassRule, BuiltinIntegerOverflowRule,
8 BuiltinIntegerOverloadKind, BuiltinIntegerScalarDoubleRule, BuiltinOutputMode,
9 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
10};
11use runmat_macros::runtime_builtin;
12use runmat_value::{StructValue, Value};
13
14use crate::builtins::common::spec::{
15 BroadcastSemantics, BuiltinFusionSpec, BuiltinGpuSpec, ConstantStrategy, GpuOpKind,
16 ReductionNaN, ResidencyPolicy, ShapeRequirements,
17};
18use crate::builtins::math::optim::common::{
19 call_function, option_f64, option_string, option_usize, value_to_real_vector, vector_to_value,
20};
21use crate::builtins::math::optim::least_squares::{
22 solve_least_squares, LeastSquaresBounds, LeastSquaresEvaluator, LeastSquaresOptions,
23 LeastSquaresResult, ResidualFuture,
24};
25use crate::builtins::math::optim::type_resolvers::nonlinear_solve_type;
26use crate::{build_runtime_error, BuiltinResult, RuntimeError};
27
28const NAME: &str = "fsolve";
29const DEFAULT_TOL_X: f64 = 1.0e-6;
30const DEFAULT_TOL_FUN: f64 = 1.0e-6;
31const DEFAULT_MAX_ITER: usize = 400;
32
33const FSOLVE_OUTPUT_X: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
34 name: "x",
35 ty: BuiltinParamType::NumericArray,
36 arity: BuiltinParamArity::Required,
37 default: None,
38 description: "Approximate solution vector/scalar.",
39}];
40
41macro_rules! fsolve_output {
42 ($name:literal, $ty:expr, $description:literal) => {
43 BuiltinParamDescriptor {
44 name: $name,
45 ty: $ty,
46 arity: BuiltinParamArity::Required,
47 default: None,
48 description: $description,
49 }
50 };
51}
52const FSOLVE_OUTPUT_X_FVAL: [BuiltinParamDescriptor; 2] = [
53 fsolve_output!("x", BuiltinParamType::NumericArray, "Approximate solution."),
54 fsolve_output!(
55 "fval",
56 BuiltinParamType::NumericArray,
57 "Function value at x."
58 ),
59];
60const FSOLVE_OUTPUT_X_FVAL_EXITFLAG: [BuiltinParamDescriptor; 3] = [
61 fsolve_output!("x", BuiltinParamType::NumericArray, "Approximate solution."),
62 fsolve_output!(
63 "fval",
64 BuiltinParamType::NumericArray,
65 "Function value at x."
66 ),
67 fsolve_output!(
68 "exitflag",
69 BuiltinParamType::NumericScalar,
70 "Convergence status code."
71 ),
72];
73const FSOLVE_OUTPUT_X_FVAL_EXITFLAG_OUTPUT: [BuiltinParamDescriptor; 4] = [
74 fsolve_output!("x", BuiltinParamType::NumericArray, "Approximate solution."),
75 fsolve_output!(
76 "fval",
77 BuiltinParamType::NumericArray,
78 "Function value at x."
79 ),
80 fsolve_output!(
81 "exitflag",
82 BuiltinParamType::NumericScalar,
83 "Convergence status code."
84 ),
85 fsolve_output!("output", BuiltinParamType::Any, "Solver diagnostics."),
86];
87const FSOLVE_OUTPUT_ALL: [BuiltinParamDescriptor; 5] = [
88 fsolve_output!("x", BuiltinParamType::NumericArray, "Approximate solution."),
89 fsolve_output!(
90 "fval",
91 BuiltinParamType::NumericArray,
92 "Function value at x."
93 ),
94 fsolve_output!(
95 "exitflag",
96 BuiltinParamType::NumericScalar,
97 "Convergence status code."
98 ),
99 fsolve_output!("output", BuiltinParamType::Any, "Solver diagnostics."),
100 fsolve_output!("jacobian", BuiltinParamType::NumericArray, "Jacobian at x."),
101];
102
103const FSOLVE_INPUTS_CORE: [BuiltinParamDescriptor; 2] = [
104 BuiltinParamDescriptor {
105 name: "fun",
106 ty: BuiltinParamType::Any,
107 arity: BuiltinParamArity::Required,
108 default: None,
109 description: "System residual callback.",
110 },
111 BuiltinParamDescriptor {
112 name: "x0",
113 ty: BuiltinParamType::Any,
114 arity: BuiltinParamArity::Required,
115 default: None,
116 description: "Initial guess scalar/vector.",
117 },
118];
119
120const FSOLVE_INPUTS_WITH_OPTIONS: [BuiltinParamDescriptor; 3] = [
121 BuiltinParamDescriptor {
122 name: "fun",
123 ty: BuiltinParamType::Any,
124 arity: BuiltinParamArity::Required,
125 default: None,
126 description: "System residual callback.",
127 },
128 BuiltinParamDescriptor {
129 name: "x0",
130 ty: BuiltinParamType::Any,
131 arity: BuiltinParamArity::Required,
132 default: None,
133 description: "Initial guess scalar/vector.",
134 },
135 BuiltinParamDescriptor {
136 name: "options",
137 ty: BuiltinParamType::Any,
138 arity: BuiltinParamArity::Optional,
139 default: None,
140 description: "Options struct from optimset.",
141 },
142];
143
144const FSOLVE_SIGNATURES: [BuiltinSignatureDescriptor; 10] = [
145 BuiltinSignatureDescriptor {
146 label: "x = fsolve(fun, x0)",
147 inputs: &FSOLVE_INPUTS_CORE,
148 outputs: &FSOLVE_OUTPUT_X,
149 },
150 BuiltinSignatureDescriptor {
151 label: "x = fsolve(fun, x0, options)",
152 inputs: &FSOLVE_INPUTS_WITH_OPTIONS,
153 outputs: &FSOLVE_OUTPUT_X,
154 },
155 BuiltinSignatureDescriptor {
156 label: "[x,fval] = fsolve(fun, x0)",
157 inputs: &FSOLVE_INPUTS_CORE,
158 outputs: &FSOLVE_OUTPUT_X_FVAL,
159 },
160 BuiltinSignatureDescriptor {
161 label: "[x,fval] = fsolve(fun, x0, options)",
162 inputs: &FSOLVE_INPUTS_WITH_OPTIONS,
163 outputs: &FSOLVE_OUTPUT_X_FVAL,
164 },
165 BuiltinSignatureDescriptor {
166 label: "[x,fval,exitflag] = fsolve(fun, x0)",
167 inputs: &FSOLVE_INPUTS_CORE,
168 outputs: &FSOLVE_OUTPUT_X_FVAL_EXITFLAG,
169 },
170 BuiltinSignatureDescriptor {
171 label: "[x,fval,exitflag] = fsolve(fun, x0, options)",
172 inputs: &FSOLVE_INPUTS_WITH_OPTIONS,
173 outputs: &FSOLVE_OUTPUT_X_FVAL_EXITFLAG,
174 },
175 BuiltinSignatureDescriptor {
176 label: "[x,fval,exitflag,output] = fsolve(fun, x0)",
177 inputs: &FSOLVE_INPUTS_CORE,
178 outputs: &FSOLVE_OUTPUT_X_FVAL_EXITFLAG_OUTPUT,
179 },
180 BuiltinSignatureDescriptor {
181 label: "[x,fval,exitflag,output] = fsolve(fun, x0, options)",
182 inputs: &FSOLVE_INPUTS_WITH_OPTIONS,
183 outputs: &FSOLVE_OUTPUT_X_FVAL_EXITFLAG_OUTPUT,
184 },
185 BuiltinSignatureDescriptor {
186 label: "[x,fval,exitflag,output,jacobian] = fsolve(fun, x0)",
187 inputs: &FSOLVE_INPUTS_CORE,
188 outputs: &FSOLVE_OUTPUT_ALL,
189 },
190 BuiltinSignatureDescriptor {
191 label: "[x,fval,exitflag,output,jacobian] = fsolve(fun, x0, options)",
192 inputs: &FSOLVE_INPUTS_WITH_OPTIONS,
193 outputs: &FSOLVE_OUTPUT_ALL,
194 },
195];
196
197const FSOLVE_ERROR_INVALID_ARGUMENT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
198 code: "RM.FSOLVE.INVALID_ARGUMENT",
199 identifier: Some("RunMat:fsolve:InvalidArgument"),
200 when: "Argument grammar/options configuration is invalid.",
201 message: "fsolve: invalid argument",
202};
203
204const FSOLVE_ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
205 code: "RM.FSOLVE.INVALID_INPUT",
206 identifier: Some("RunMat:fsolve:InvalidInput"),
207 when: "Initial guess/callback/iteration semantics are invalid.",
208 message: "fsolve: invalid input",
209};
210
211const FSOLVE_ERRORS: [BuiltinErrorDescriptor; 2] =
212 [FSOLVE_ERROR_INVALID_ARGUMENT, FSOLVE_ERROR_INVALID_INPUT];
213
214const FSOLVE_INTEGER_INITIAL_INPUTS: [BuiltinIntegerInputCapability; 1] =
215 [BuiltinIntegerInputCapability {
216 name: "x0",
217 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
218 availability: BuiltinIntegerInputAvailability::RunMatOnly,
219 scalar_double: BuiltinIntegerScalarDoubleRule::NotApplicable,
220 notes: "Typed-integer initial guesses are independently gated and every element must convert exactly to binary64.",
221 }];
222const FSOLVE_INTEGER_RESIDUAL_INPUTS: [BuiltinIntegerInputCapability; 1] =
223 [BuiltinIntegerInputCapability {
224 name: "fun(x) residual",
225 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
226 availability: BuiltinIntegerInputAvailability::RunMatOnly,
227 scalar_double: BuiltinIntegerScalarDoubleRule::NotApplicable,
228 notes: "Typed-integer residuals are gated before resident gather and every element must convert exactly to binary64.",
229 }];
230const FSOLVE_INTEGER_TOLERANCE_INPUTS: [BuiltinIntegerInputCapability; 2] = [
231 BuiltinIntegerInputCapability {
232 name: "TolX",
233 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
234 availability: BuiltinIntegerInputAvailability::RunMatOnly,
235 scalar_double: BuiltinIntegerScalarDoubleRule::Allowed,
236 notes: "Typed-integer TolX is independently gated and must convert exactly to a positive binary64 scalar.",
237 },
238 BuiltinIntegerInputCapability {
239 name: "TolFun",
240 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
241 availability: BuiltinIntegerInputAvailability::RunMatOnly,
242 scalar_double: BuiltinIntegerScalarDoubleRule::Allowed,
243 notes: "Typed-integer TolFun is independently gated and must convert exactly to a positive binary64 scalar.",
244 },
245];
246const FSOLVE_INTEGER_COUNT_INPUTS: [BuiltinIntegerInputCapability; 2] = [
247 BuiltinIntegerInputCapability {
248 name: "MaxIter",
249 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
250 availability: BuiltinIntegerInputAvailability::RunMatOnly,
251 scalar_double: BuiltinIntegerScalarDoubleRule::Allowed,
252 notes: "Typed-integer iteration counts are independently gated and decoded exactly through platform bounds.",
253 },
254 BuiltinIntegerInputCapability {
255 name: "MaxFunEvals",
256 classes: &crate::builtins::common::integer_capability::ALL_INTEGER_CLASSES,
257 availability: BuiltinIntegerInputAvailability::RunMatOnly,
258 scalar_double: BuiltinIntegerScalarDoubleRule::Allowed,
259 notes: "Typed-integer evaluation counts are independently gated and decoded exactly through platform bounds.",
260 },
261];
262pub const FSOLVE_INTEGER_CAPABILITIES: [BuiltinIntegerCapabilityDescriptor; 4] = [
263 BuiltinIntegerCapabilityDescriptor {
264 form: "x = fsolve(fun, integer_x0, options)",
265 inputs: &FSOLVE_INTEGER_INITIAL_INPUTS,
266 computation_domain: BuiltinIntegerComputationDomain::FloatingPoint,
267 output_class: BuiltinIntegerOutputClassRule::FunctionSpecific,
268 overflow: BuiltinIntegerOverflowRule::Error,
269 backend: BuiltinIntegerBackendRule::GatherFallback,
270 overload: BuiltinIntegerOverloadKind::Multiple,
271 notes: "Documented x0 is double. Strict compatibility rejects typed integers; RunMat mode admits exact binary64 values and returns double x, fval, exitflag, diagnostics, and Jacobian outputs.",
272 },
273 BuiltinIntegerCapabilityDescriptor {
274 form: "fsolve callback returns an integer residual",
275 inputs: &FSOLVE_INTEGER_RESIDUAL_INPUTS,
276 computation_domain: BuiltinIntegerComputationDomain::FloatingPoint,
277 output_class: BuiltinIntegerOutputClassRule::FunctionSpecific,
278 overflow: BuiltinIntegerOverflowRule::Error,
279 backend: BuiltinIntegerBackendRule::GatherFallback,
280 overload: BuiltinIntegerOverloadKind::Multiple,
281 notes: "Strict compatibility rejects typed residuals before provider access. RunMat mode converts exact values to binary64; the one-to-five output contract remains function-specific and double-valued for numeric results.",
282 },
283 BuiltinIntegerCapabilityDescriptor {
284 form: "fsolve(..., options.TolX=integer, options.TolFun=integer)",
285 inputs: &FSOLVE_INTEGER_TOLERANCE_INPUTS,
286 computation_domain: BuiltinIntegerComputationDomain::FloatingPoint,
287 output_class: BuiltinIntegerOutputClassRule::FunctionSpecific,
288 overflow: BuiltinIntegerOverflowRule::Error,
289 backend: BuiltinIntegerBackendRule::HostOnly,
290 overload: BuiltinIntegerOverloadKind::StructuralParameter,
291 notes: "Typed-integer tolerance controls are RunMat-only and do not change output classes.",
292 },
293 BuiltinIntegerCapabilityDescriptor {
294 form: "fsolve(..., options.MaxIter=integer, options.MaxFunEvals=integer)",
295 inputs: &FSOLVE_INTEGER_COUNT_INPUTS,
296 computation_domain: BuiltinIntegerComputationDomain::Structural,
297 output_class: BuiltinIntegerOutputClassRule::FunctionSpecific,
298 overflow: BuiltinIntegerOverflowRule::Error,
299 backend: BuiltinIntegerBackendRule::HostOnly,
300 overload: BuiltinIntegerOverloadKind::StructuralParameter,
301 notes: "The documented controls are integer-valued numeric scalars, not documented typed-integer classes. RunMat's typed forms preserve exact counts without a binary64 round trip.",
302 },
303];
304
305pub(crate) const FSOLVE_INPUT_NUMERIC_EXTENSION: BuiltinExtensionDescriptor =
306 BuiltinExtensionDescriptor {
307 id: "fsolve-nonfloating-initial-point",
308 mode: BuiltinExtensionMode::RunMatOnly,
309 description: "fsolve with a typed-integer or logical initial point is a RunMat extension",
310 error_identifier: Some("RunMat:compatibility:FsolveNumericInputExtension"),
311 };
312pub(crate) const FSOLVE_CALLBACK_NUMERIC_EXTENSION: BuiltinExtensionDescriptor =
313 BuiltinExtensionDescriptor {
314 id: "fsolve-nonfloating-callback-output",
315 mode: BuiltinExtensionMode::RunMatOnly,
316 description: "fsolve with typed-integer or logical residual output is a RunMat extension",
317 error_identifier: Some("RunMat:compatibility:FsolveCallbackExtension"),
318 };
319pub(crate) const FSOLVE_OPTION_EXTENSION: BuiltinExtensionDescriptor = BuiltinExtensionDescriptor {
320 id: "fsolve-typed-option-controls",
321 mode: BuiltinExtensionMode::RunMatOnly,
322 description: "fsolve with typed-integer option controls is a RunMat extension",
323 error_identifier: Some("RunMat:compatibility:FsolveOptionExtension"),
324};
325pub(crate) const FSOLVE_RESIDENT_EXTENSION: BuiltinExtensionDescriptor =
326 BuiltinExtensionDescriptor {
327 id: "fsolve-resident-fallback",
328 mode: BuiltinExtensionMode::RunMatOnly,
329 description:
330 "fsolve with provider-resident numeric input or callback output is a RunMat extension",
331 error_identifier: Some("RunMat:compatibility:FsolveResidentExtension"),
332 };
333pub const FSOLVE_EXTENSIONS: [BuiltinExtensionDescriptor; 4] = [
334 FSOLVE_INPUT_NUMERIC_EXTENSION,
335 FSOLVE_CALLBACK_NUMERIC_EXTENSION,
336 FSOLVE_OPTION_EXTENSION,
337 FSOLVE_RESIDENT_EXTENSION,
338];
339
340pub const FSOLVE_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
341 signatures: &FSOLVE_SIGNATURES,
342 output_mode: BuiltinOutputMode::ByRequestedOutputCount,
343 completion_policy: BuiltinCompletionPolicy::Public,
344 errors: &FSOLVE_ERRORS,
345};
346
347fn fsolve_error_with_detail(
348 error: &'static BuiltinErrorDescriptor,
349 detail: impl AsRef<str>,
350) -> RuntimeError {
351 let detail = detail.as_ref();
352 let message = if detail.starts_with("fsolve:") {
353 detail.to_string()
354 } else {
355 format!("{}: {detail}", error.message)
356 };
357 let mut builder = build_runtime_error(message).with_builtin(NAME);
358 if let Some(identifier) = error.identifier {
359 builder = builder.with_identifier(identifier);
360 }
361 builder.build()
362}
363
364fn fsolve_map_error(err: RuntimeError, fallback: &'static BuiltinErrorDescriptor) -> RuntimeError {
365 if err.identifier().is_some() {
366 err
367 } else {
368 fsolve_error_with_detail(fallback, err.message())
369 }
370}
371
372#[runmat_macros::register_gpu_spec(builtin_path = "crate::builtins::math::optim::fsolve")]
373pub const GPU_SPEC: BuiltinGpuSpec = BuiltinGpuSpec {
374 name: "fsolve",
375 op_kind: GpuOpKind::Custom("nonlinear-solve"),
376 supported_precisions: &[],
377 broadcast: BroadcastSemantics::None,
378 provider_hooks: &[],
379 constant_strategy: ConstantStrategy::InlineLiteral,
380 residency: ResidencyPolicy::GatherImmediately,
381 nan_mode: ReductionNaN::Include,
382 two_pass_threshold: None,
383 workgroup_size: None,
384 accepts_nan_mode: false,
385 notes: "Host finite-difference Levenberg-Marquardt solver. Callback computations may use GPU-aware builtins.",
386};
387
388#[runmat_macros::register_fusion_spec(builtin_path = "crate::builtins::math::optim::fsolve")]
389pub const FUSION_SPEC: BuiltinFusionSpec = BuiltinFusionSpec {
390 name: "fsolve",
391 shape: ShapeRequirements::Any,
392 constant_strategy: ConstantStrategy::InlineLiteral,
393 elementwise: None,
394 reduction: None,
395 emits_nan: false,
396 notes: "Nonlinear solving repeatedly invokes user code and terminates fusion planning.",
397};
398
399#[runtime_builtin(
400 name = "fsolve",
401 category = "math/optim",
402 summary = "Solve nonlinear equation systems.",
403 keywords = "fsolve,nonlinear solve,root finding,levenberg-marquardt,jacobian",
404 accel = "sink",
405 type_resolver(nonlinear_solve_type),
406 descriptor(crate::builtins::math::optim::fsolve::FSOLVE_DESCRIPTOR),
407 extensions(crate::builtins::math::optim::fsolve::FSOLVE_EXTENSIONS),
408 integer_capabilities(crate::builtins::math::optim::fsolve::FSOLVE_INTEGER_CAPABILITIES),
409 builtin_path = "crate::builtins::math::optim::fsolve"
410)]
411async fn fsolve_builtin(function: Value, x0: Value, rest: Vec<Value>) -> BuiltinResult<Value> {
412 if rest.len() > 1 {
413 return Err(fsolve_error_with_detail(
414 &FSOLVE_ERROR_INVALID_ARGUMENT,
415 "too many input arguments",
416 ));
417 }
418 let options = parse_options(rest.first())
419 .map_err(|err| fsolve_map_error(err, &FSOLVE_ERROR_INVALID_ARGUMENT))?;
420 let opts = FsolveOptions::from_struct(options.as_ref())
421 .map_err(|err| fsolve_map_error(err, &FSOLVE_ERROR_INVALID_ARGUMENT))?;
422 let guess = crate::builtins::math::optim::common::initial_guess_with_extensions(
423 NAME,
424 x0,
425 &FSOLVE_INPUT_NUMERIC_EXTENSION,
426 &FSOLVE_RESIDENT_EXTENSION,
427 )
428 .await
429 .map_err(|err| fsolve_map_error(err, &FSOLVE_ERROR_INVALID_INPUT))?;
430 let outcome = solve(&function, guess.values, &guess.shape, guess.scalar, &opts)
431 .await
432 .map_err(|err| fsolve_map_error(err, &FSOLVE_ERROR_INVALID_INPUT))?;
433 finalize(outcome, &guess.shape, guess.scalar)
434 .map_err(|err| fsolve_map_error(err, &FSOLVE_ERROR_INVALID_INPUT))
435}
436
437fn parse_options(value: Option<&Value>) -> BuiltinResult<Option<StructValue>> {
438 match value {
439 None => Ok(None),
440 Some(Value::Struct(options)) => Ok(Some(options.clone())),
441 Some(other) => Err(fsolve_error_with_detail(
442 &FSOLVE_ERROR_INVALID_ARGUMENT,
443 format!("options must be a struct, got {other:?}"),
444 )),
445 }
446}
447
448#[derive(Clone, Copy)]
449struct FsolveOptions {
450 tol_x: f64,
451 tol_fun: f64,
452 max_iter: usize,
453 max_fun_evals: usize,
454}
455
456impl FsolveOptions {
457 fn from_struct(options: Option<&StructValue>) -> BuiltinResult<Self> {
458 crate::builtins::math::optim::common::ensure_option_extensions(
459 NAME,
460 options,
461 &FSOLVE_OPTION_EXTENSION,
462 &FSOLVE_RESIDENT_EXTENSION,
463 )?;
464 let display = option_string(options, "Display", "off")?;
465 if !matches!(display.as_str(), "off" | "none" | "final" | "iter") {
466 return Err(fsolve_error_with_detail(
467 &FSOLVE_ERROR_INVALID_ARGUMENT,
468 "option Display must be 'off', 'none', 'final', or 'iter'",
469 ));
470 }
471 let tol_x = option_f64(NAME, options, "TolX", DEFAULT_TOL_X)?;
472 let tol_fun = option_f64(NAME, options, "TolFun", DEFAULT_TOL_FUN)?;
473 if tol_x <= 0.0 || tol_fun <= 0.0 {
474 return Err(fsolve_error_with_detail(
475 &FSOLVE_ERROR_INVALID_ARGUMENT,
476 "options TolX and TolFun must be positive",
477 ));
478 }
479 let max_iter = option_usize(NAME, options, "MaxIter", DEFAULT_MAX_ITER)?.max(1);
480 let max_fun_evals = option_usize(NAME, options, "MaxFunEvals", 100 * max_iter)?.max(1);
481 Ok(Self {
482 tol_x,
483 tol_fun,
484 max_iter,
485 max_fun_evals,
486 })
487 }
488}
489
490async fn solve(
491 function: &Value,
492 x: Vec<f64>,
493 shape: &[usize],
494 scalar: bool,
495 options: &FsolveOptions,
496) -> BuiltinResult<FsolveOutcome> {
497 let mut evaluator = FsolveEvaluator {
498 function,
499 shape: shape.to_vec(),
500 scalar,
501 residual_shape: None,
502 residual_scalar: false,
503 };
504 let variable_len = x.len();
505 let mut result = solve_least_squares(
506 NAME,
507 &mut evaluator,
508 x,
509 &LeastSquaresBounds::unbounded(variable_len),
510 &LeastSquaresOptions {
511 tol_x: options.tol_x,
512 tol_fun: options.tol_fun,
513 max_iter: options.max_iter,
514 max_fun_evals: options.max_fun_evals,
515 final_jacobian: true,
516 },
517 )
518 .await?;
519 if result.exitflag > 0
520 && result
521 .residual
522 .iter()
523 .fold(0.0_f64, |norm, value| norm.max(value.abs()))
524 > options.tol_fun
525 {
526 result.exitflag = -2;
527 result.message =
528 "Equation not solved. The solver converged to a point with a nonzero residual."
529 .to_string();
530 }
531 Ok(FsolveOutcome {
532 result,
533 residual_shape: evaluator.residual_shape.unwrap_or_else(|| vec![1, 1]),
534 residual_scalar: evaluator.residual_scalar,
535 })
536}
537
538struct FsolveOutcome {
539 result: LeastSquaresResult,
540 residual_shape: Vec<usize>,
541 residual_scalar: bool,
542}
543
544struct FsolveEvaluator<'a> {
545 function: &'a Value,
546 shape: Vec<usize>,
547 scalar: bool,
548 residual_shape: Option<Vec<usize>>,
549 residual_scalar: bool,
550}
551
552impl LeastSquaresEvaluator for FsolveEvaluator<'_> {
553 fn residual<'a>(&'a mut self, x: &'a [f64]) -> ResidualFuture<'a> {
554 Box::pin(async move {
555 let arg = if self.scalar {
556 Value::Num(x[0])
557 } else {
558 Value::Tensor(
559 runmat_value::Tensor::new(x.to_vec(), self.shape.clone())
560 .map_err(|e| fsolve_error_with_detail(&FSOLVE_ERROR_INVALID_INPUT, e))?,
561 )
562 };
563 let value = call_function(self.function, vec![arg]).await?;
564 let value = crate::builtins::math::optim::common::prepare_floating_value(
565 NAME,
566 value,
567 &FSOLVE_CALLBACK_NUMERIC_EXTENSION,
568 &FSOLVE_RESIDENT_EXTENSION,
569 "function residual",
570 )
571 .await?;
572 let (shape, scalar) = match &value {
573 Value::Num(_) | Value::Int(_) | Value::Bool(_) => (vec![1, 1], true),
574 Value::Tensor(tensor) => (tensor.shape.clone(), false),
575 Value::LogicalArray(array) => (array.shape.clone(), false),
576 _ => (vec![1, 1], false),
577 };
578 self.residual_shape = Some(shape);
579 self.residual_scalar = scalar;
580 let residual = value_to_real_vector(NAME, value).await?;
581 if residual.is_empty() {
582 Err(fsolve_error_with_detail(
583 &FSOLVE_ERROR_INVALID_INPUT,
584 "function value must not be empty",
585 ))
586 } else {
587 Ok(residual)
588 }
589 })
590 }
591}
592
593fn finalize(outcome: FsolveOutcome, x_shape: &[usize], x_scalar: bool) -> BuiltinResult<Value> {
594 let result = outcome.result;
595 let x = vector_to_value(NAME, result.x.clone(), x_shape, x_scalar)?;
596 let fval = vector_to_value(
597 NAME,
598 result.residual.clone(),
599 &outcome.residual_shape,
600 outcome.residual_scalar,
601 )?;
602 let exitflag = Value::Num(result.exitflag as f64);
603 let mut fields = StructValue::new();
604 fields.insert("iterations", Value::Num(result.iterations as f64));
605 fields.insert("funcCount", Value::Num(result.func_count as f64));
606 fields.insert("algorithm", Value::from("levenberg-marquardt"));
607 fields.insert("firstorderopt", Value::Num(result.first_order_optimality));
608 fields.insert("stepsize", Value::Num(result.step_size));
609 fields.insert("message", Value::from(result.message.clone()));
610 let output = Value::Struct(fields);
611 let mut jacobian_data = Vec::with_capacity(result.jacobian.len());
612 for column in 0..result.variable_len {
613 for row in 0..result.residual_len {
614 jacobian_data.push(result.jacobian[row * result.variable_len + column]);
615 }
616 }
617 let jacobian = Value::Tensor(
618 runmat_value::Tensor::new(
619 jacobian_data,
620 vec![result.residual_len, result.variable_len],
621 )
622 .map_err(|error| fsolve_error_with_detail(&FSOLVE_ERROR_INVALID_INPUT, error))?,
623 );
624 match crate::output_count::current_output_count() {
625 None => Ok(x),
626 Some(0) => Ok(Value::OutputList(Vec::new())),
627 Some(1) => Ok(crate::output_count::output_list_with_padding(1, vec![x])),
628 Some(2) => Ok(crate::output_count::output_list_with_padding(
629 2,
630 vec![x, fval],
631 )),
632 Some(3) => Ok(crate::output_count::output_list_with_padding(
633 3,
634 vec![x, fval, exitflag],
635 )),
636 Some(4) => Ok(crate::output_count::output_list_with_padding(
637 4,
638 vec![x, fval, exitflag, output],
639 )),
640 Some(5) => Ok(crate::output_count::output_list_with_padding(
641 5,
642 vec![x, fval, exitflag, output, jacobian],
643 )),
644 Some(_) => Err(fsolve_error_with_detail(
645 &FSOLVE_ERROR_INVALID_ARGUMENT,
646 "too many output arguments; maximum is 5",
647 )),
648 }
649}
650
651#[cfg(test)]
652mod tests {
653 use super::*;
654 use futures::executor::block_on;
655 use runmat_value::Tensor;
656 use std::sync::{Arc, Mutex};
657
658 #[test]
659 fn fsolve_scalar_builtin_handle() {
660 let root = block_on(fsolve_builtin(
661 Value::FunctionHandle("sin".into()),
662 Value::Num(3.0),
663 Vec::new(),
664 ))
665 .unwrap();
666 match root {
667 Value::Num(n) => assert!((n - std::f64::consts::PI).abs() < 1.0e-5),
668 other => panic!("unexpected value {other:?}"),
669 }
670 }
671
672 #[test]
673 fn fsolve_five_output_form_returns_residual_status_diagnostics_and_jacobian() {
674 let _outputs = crate::output_count::push_output_count(Some(5));
675 let result = block_on(fsolve_builtin(
676 Value::FunctionHandle("sin".into()),
677 Value::Num(3.0),
678 Vec::new(),
679 ))
680 .unwrap();
681 let Value::OutputList(outputs) = result else {
682 panic!("expected five outputs");
683 };
684 assert_eq!(outputs.len(), 5);
685 assert!(matches!(outputs[0], Value::Num(_)));
686 assert!(matches!(outputs[1], Value::Num(_)));
687 assert!(matches!(outputs[2], Value::Num(_)));
688 assert!(matches!(outputs[3], Value::Struct(_)));
689 assert!(matches!(&outputs[4], Value::Tensor(tensor) if tensor.shape == vec![1, 1]));
690 }
691
692 #[test]
693 fn fsolve_returns_stationary_non_root_with_solver_status() {
694 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(Arc::new(
695 |_function, args, _requested_outputs| {
696 let x = match &args[0] {
697 Value::Num(value) => *value,
698 other => panic!("expected scalar numeric argument, got {other:?}"),
699 };
700 Box::pin(async move { Ok(Value::Num(x * x + 1.0)) })
701 },
702 )));
703 let _outputs = crate::output_count::push_output_count(Some(3));
704 let result = block_on(fsolve_builtin(
705 Value::BoundFunctionHandle {
706 name: "no_real_root".to_string(),
707 function: 44,
708 },
709 Value::Num(0.0),
710 Vec::new(),
711 ))
712 .unwrap();
713 let Value::OutputList(outputs) = result else {
714 panic!("expected outputs")
715 };
716 assert!(matches!(outputs[2], Value::Num(flag) if flag <= 0.0));
717 }
718
719 #[test]
720 fn fsolve_vector_system_via_semantic_resolver() {
721 let _resolver =
722 crate::user_functions::install_semantic_function_resolver(Some(Arc::new(|_name| {
723 Some(0)
724 })));
725 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(
726 std::sync::Arc::new(|_function, args, _requested_outputs| {
727 let x = match &args[0] {
728 Value::Tensor(t) => t.materialize_f64().clone(),
729 _ => panic!("expected tensor input"),
730 };
731 Box::pin(async move {
732 Ok(Value::Tensor(
733 Tensor::new(
734 vec![x[0] * x[0] + x[1] * x[1] - 4.0, x[0] * x[1] - 1.0],
735 vec![2, 1],
736 )
737 .unwrap(),
738 ))
739 })
740 }),
741 ));
742 let x0 = Tensor::new(vec![1.0, 1.0], vec![2, 1]).unwrap();
743 let root = block_on(fsolve_builtin(
744 Value::FunctionHandle("system".into()),
745 Value::Tensor(x0),
746 Vec::new(),
747 ))
748 .unwrap();
749 match root {
750 Value::Tensor(t) => {
751 assert!(
752 (t.materialize_f64()[0] * t.materialize_f64()[0]
753 + t.materialize_f64()[1] * t.materialize_f64()[1]
754 - 4.0)
755 .abs()
756 < 1.0e-5
757 );
758 assert!((t.materialize_f64()[0] * t.materialize_f64()[1] - 1.0).abs() < 1.0e-5);
759 }
760 other => panic!("unexpected value {other:?}"),
761 }
762 }
763
764 #[test]
765 fn fsolve_preserves_row_vector_shape_for_callback() {
766 let seen_shapes = Arc::new(Mutex::new(Vec::new()));
767 let seen_shapes_for_invoker = Arc::clone(&seen_shapes);
768 let _resolver =
769 crate::user_functions::install_semantic_function_resolver(Some(Arc::new(|_name| {
770 Some(0)
771 })));
772 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(Arc::new(
773 move |_function, args, _requested_outputs| {
774 let (x, shape) = match &args[0] {
775 Value::Tensor(t) => (t.materialize_f64().clone(), t.shape.clone()),
776 other => panic!("expected tensor input, got {other:?}"),
777 };
778 assert_eq!(shape, vec![1, 2]);
779 seen_shapes_for_invoker.lock().unwrap().push(shape.clone());
780 Box::pin(async move {
781 Ok(Value::Tensor(
782 Tensor::new(vec![x[0] - 3.0, x[1] - 4.0], shape).unwrap(),
783 ))
784 })
785 },
786 )));
787 let x0 = Tensor::new(vec![0.0, 0.0], vec![1, 2]).unwrap();
788 let root = block_on(fsolve_builtin(
789 Value::FunctionHandle("row_system".into()),
790 Value::Tensor(x0),
791 Vec::new(),
792 ))
793 .unwrap();
794 match root {
795 Value::Tensor(t) => {
796 assert_eq!(t.shape, vec![1, 2]);
797 assert!((t.materialize_f64()[0] - 3.0).abs() < 1.0e-5);
798 assert!((t.materialize_f64()[1] - 4.0).abs() < 1.0e-5);
799 }
800 other => panic!("unexpected value {other:?}"),
801 }
802 assert!(!seen_shapes.lock().unwrap().is_empty());
803 }
804
805 #[test]
806 fn fsolve_preserves_matrix_shape_for_callback() {
807 let seen_shapes = Arc::new(Mutex::new(Vec::new()));
808 let seen_shapes_for_invoker = Arc::clone(&seen_shapes);
809 let _resolver =
810 crate::user_functions::install_semantic_function_resolver(Some(Arc::new(|_name| {
811 Some(0)
812 })));
813 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(Arc::new(
814 move |_function, args, _requested_outputs| {
815 let (x, shape) = match &args[0] {
816 Value::Tensor(t) => (t.materialize_f64().clone(), t.shape.clone()),
817 other => panic!("expected tensor input, got {other:?}"),
818 };
819 assert_eq!(shape, vec![2, 2]);
820 seen_shapes_for_invoker.lock().unwrap().push(shape.clone());
821 Box::pin(async move {
822 Ok(Value::Tensor(
823 Tensor::new(vec![x[0] - 1.0, x[1] - 2.0, x[2] - 3.0, x[3] - 4.0], shape)
824 .unwrap(),
825 ))
826 })
827 },
828 )));
829 let x0 = Tensor::new(vec![0.0, 0.0, 0.0, 0.0], vec![2, 2]).unwrap();
830 let root = block_on(fsolve_builtin(
831 Value::FunctionHandle("matrix_system".into()),
832 Value::Tensor(x0),
833 Vec::new(),
834 ))
835 .unwrap();
836 match root {
837 Value::Tensor(t) => {
838 assert_eq!(t.shape, vec![2, 2]);
839 assert!((t.materialize_f64()[0] - 1.0).abs() < 1.0e-5);
840 assert!((t.materialize_f64()[1] - 2.0).abs() < 1.0e-5);
841 assert!((t.materialize_f64()[2] - 3.0).abs() < 1.0e-5);
842 assert!((t.materialize_f64()[3] - 4.0).abs() < 1.0e-5);
843 }
844 other => panic!("unexpected value {other:?}"),
845 }
846 assert!(!seen_shapes.lock().unwrap().is_empty());
847 }
848
849 #[test]
850 fn fsolve_accepts_semantic_function_handle_callback() {
851 let _invoker = crate::user_functions::install_semantic_function_invoker(Some(Arc::new(
852 |function, args, requested_outputs| {
853 assert_eq!(function, 43);
854 assert_eq!(requested_outputs, 1);
855 let x = match &args[0] {
856 Value::Num(value) => *value,
857 other => panic!("expected scalar numeric argument, got {other:?}"),
858 };
859 Box::pin(async move { Ok(Value::Num(x - 3.0)) })
860 },
861 )));
862 let root = block_on(fsolve_builtin(
863 Value::BoundFunctionHandle {
864 name: "system_function".to_string(),
865 function: 43,
866 },
867 Value::Num(1.0),
868 Vec::new(),
869 ))
870 .unwrap();
871 match root {
872 Value::Num(n) => assert!((n - 3.0).abs() < 1.0e-5),
873 other => panic!("unexpected value {other:?}"),
874 }
875 }
876
877 #[test]
878 fn fsolve_descriptor_signatures_cover_core_forms() {
879 let labels: Vec<&str> = FSOLVE_DESCRIPTOR
880 .signatures
881 .iter()
882 .map(|signature| signature.label)
883 .collect();
884 assert_eq!(
885 labels,
886 vec![
887 "x = fsolve(fun, x0)",
888 "x = fsolve(fun, x0, options)",
889 "[x,fval] = fsolve(fun, x0)",
890 "[x,fval] = fsolve(fun, x0, options)",
891 "[x,fval,exitflag] = fsolve(fun, x0)",
892 "[x,fval,exitflag] = fsolve(fun, x0, options)",
893 "[x,fval,exitflag,output] = fsolve(fun, x0)",
894 "[x,fval,exitflag,output] = fsolve(fun, x0, options)",
895 "[x,fval,exitflag,output,jacobian] = fsolve(fun, x0)",
896 "[x,fval,exitflag,output,jacobian] = fsolve(fun, x0, options)",
897 ]
898 );
899
900 let codes: Vec<&str> = FSOLVE_DESCRIPTOR
901 .errors
902 .iter()
903 .map(|error| error.code)
904 .collect();
905 assert_eq!(
906 codes,
907 vec!["RM.FSOLVE.INVALID_ARGUMENT", "RM.FSOLVE.INVALID_INPUT"]
908 );
909 }
910
911 #[test]
912 fn fsolve_too_many_args_uses_stable_identifier() {
913 let err = block_on(fsolve_builtin(
914 Value::FunctionHandle("sin".into()),
915 Value::Num(1.0),
916 vec![
917 Value::Struct(StructValue::new()),
918 Value::Struct(StructValue::new()),
919 ],
920 ))
921 .unwrap_err();
922 assert_eq!(err.identifier(), Some("RunMat:fsolve:InvalidArgument"));
923 }
924}