Skip to main content

runmat_runtime/builtins/math/optim/
fsolve.rs

1//! MATLAB-compatible `fsolve` builtin for nonlinear systems.
2
3use 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}