Add nsolve and nsolve_system to Python API

- Fix digit counting for Decimal

src/api/python.rs

@@ -34,7 +34,7 @@ use crate::{
         algebraic_number::AlgebraicExtension,
         atom::AtomField,
         finite_field::{ToFiniteField, Zp, Z2},
-        float::{Complex, Float},
+        float::{Complex, Float, RealNumberLike, F64},
         integer::{FromFiniteField, Integer, IntegerRing, Z},
         rational::{Rational, RationalField, Q},
         rational_polynomial::{
@@ -1677,6 +1677,7 @@ impl<'a> FromPyObject<'a> for PythonMultiPrecisionFloat {
             let digits = a
                 .chars()
                 .skip_while(|x| *x == '.' || *x == '0')
+                .filter(|x| *x != '.')
                 .take_while(|x| x.is_ascii_digit())
                 .count();
 
@@ -3957,6 +3958,130 @@ impl PythonExpression {
         Ok(res.into_iter().map(|x| x.into()).collect())
     }
 
+    /// Find the root of an expression in `x` numerically over the reals using Newton's method.
+    /// Use `init` as the initial guess for the root.
+    ///
+    /// Examples
+    /// --------
+    /// >>> from symbolica import Expression
+    /// >>> x, y, c = Expression.symbol('x', 'y', 'c')
+    /// >>> f = Expression.symbol('f')
+    /// >>> x_r, y_r = Expression.solve_linear_system([f(c)*x + y/c - 1, y-c/2], [x, y])
+    /// >>> print('x =', x_r, ', y =', y_r)
+    #[pyo3(signature =
+        (variable,
+        init,
+        prec = 1e-4,
+        max_iterations = 1000),
+        )]
+    pub fn nsolve(
+        &self,
+        variable: PythonExpression,
+        init: PythonMultiPrecisionFloat,
+        prec: f64,
+        max_iterations: usize,
+        py: Python,
+    ) -> PyResult<PyObject> {
+        let id = if let AtomView::Var(x) = variable.expr.as_view() {
+            x.get_symbol()
+        } else {
+            return Err(exceptions::PyValueError::new_err(
+                "Expected variable instead of expression",
+            ));
+        };
+
+        if init.0.prec() == 53 {
+            let r = self
+                .expr
+                .nsolve::<F64>(id, init.0.to_f64().into(), prec.into(), max_iterations)
+                .map_err(|e| {
+                    exceptions::PyValueError::new_err(format!("Could not solve system: {}", e))
+                })?;
+            Ok(r.into_inner().into_py(py))
+        } else {
+            Ok(PythonMultiPrecisionFloat(
+                self.expr
+                    .nsolve(id, init.0, prec.into(), max_iterations)
+                    .map_err(|e| {
+                        exceptions::PyValueError::new_err(format!("Could not solve system: {}", e))
+                    })?,
+            )
+            .to_object(py))
+        }
+    }
+
+    /// Find a common root of multiple expressions in `variables` numerically over the reals using Newton's method.
+    /// Use `init` as the initial guess for the root.
+    ///
+    /// Examples
+    /// --------
+    /// >>> from symbolica import Expression
+    /// >>> x, y, c = Expression.symbol('x', 'y', 'c')
+    /// >>> f = Expression.symbol('f')
+    /// >>> x_r, y_r = Expression.solve_linear_system([f(c)*x + y/c - 1, y-c/2], [x, y])
+    /// >>> print('x =', x_r, ', y =', y_r)
+    #[pyo3(signature =
+        (system,
+        variables,
+        init,
+        prec = 1e-4,
+        max_iterations = 1000),
+        )]
+    #[classmethod]
+    pub fn nsolve_system(
+        _cls: &PyType,
+        system: Vec<ConvertibleToExpression>,
+        variables: Vec<PythonExpression>,
+        init: Vec<PythonMultiPrecisionFloat>,
+        prec: f64,
+        max_iterations: usize,
+        py: Python,
+    ) -> PyResult<Vec<PyObject>> {
+        let system: Vec<_> = system.into_iter().map(|x| x.to_expression()).collect();
+        let system_b: Vec<_> = system.iter().map(|x| x.expr.as_view()).collect();
+
+        let mut vars = vec![];
+        for v in variables {
+            match v.expr.as_view() {
+                AtomView::Var(v) => vars.push(v.get_symbol().into()),
+                e => {
+                    Err(exceptions::PyValueError::new_err(format!(
+                        "Expected variable instead of {}",
+                        e
+                    )))?;
+                }
+            }
+        }
+
+        if init[0].0.prec() == 53 {
+            let init: Vec<_> = init.into_iter().map(|x| x.0.to_f64().into()).collect();
+
+            let res: Vec<F64> =
+                AtomView::nsolve_system(&system_b, &vars, &init, prec.into(), max_iterations)
+                    .map_err(|e| {
+                        exceptions::PyValueError::new_err(format!("Could not solve system: {}", e))
+                    })?;
+
+            Ok(res
+                .into_iter()
+                .map(|x| x.into_inner().into_py(py))
+                .collect())
+        } else {
+            let init: Vec<_> = init.into_iter().map(|x| x.0).collect();
+
+            let res: Vec<Float> =
+                AtomView::nsolve_system(&system_b, &vars, &init, prec.into(), max_iterations)
+                    .map_err(|e| {
+                        exceptions::PyValueError::new_err(format!("Could not solve system: {}", e))
+                    })?;
+
+            Ok(res
+                .into_iter()
+                .map(|x| PythonMultiPrecisionFloat(x).to_object(py))
+                .collect())
+        }
+    }
+
     /// Evaluate the expression, using a map of all the constants and
     /// user functions to a float.
     ///

src/domains/float.rs

@@ -530,6 +530,12 @@ impl From<Rational> for f64 {
 #[derive(Debug, Copy, Clone)]
 pub struct F64(f64);
 
+impl F64 {
+    pub fn into_inner(self) -> f64 {
+        self.0
+    }
+}
+
 impl NumericalFloatLike for F64 {
     #[inline(always)]
     fn mul_add(&self, a: &Self, b: &Self) -> Self {

src/solve.rs

@@ -14,6 +14,41 @@ use crate::{
     tensors::matrix::Matrix,
 };
 
+impl Atom {
+    /// Find the root of a function in `x` numerically over the reals using Newton's method.
+    pub fn nsolve<N: SingleFloat + Real + PartialOrd>(
+        &self,
+        x: Symbol,
+        init: N,
+        prec: N,
+        max_iterations: usize,
+    ) -> Result<N, String> {
+        self.as_view().nsolve(x, init, prec, max_iterations)
+    }
+
+    /// Solve a non-linear system numerically over the reals using Newton's method.
+    pub fn nsolve_system<
+        N: SingleFloat + Real + PartialOrd + InternalOrdering + Eq + std::hash::Hash,
+    >(
+        system: &[AtomView],
+        vars: &[Symbol],
+        init: &[N],
+        prec: N,
+        max_iterations: usize,
+    ) -> Result<Vec<N>, String> {
+        AtomView::nsolve_system(system, vars, init, prec, max_iterations)
+    }
+
+    /// Solve a system that is linear in `vars`, if possible.
+    /// Each expression in `system` is understood to yield 0.
+    pub fn solve_linear_system<E: Exponent>(
+        system: &[AtomView],
+        vars: &[Symbol],
+    ) -> Result<Vec<Atom>, String> {
+        AtomView::solve_linear_system::<E>(system, vars)
+    }
+}
+
 impl<'a> AtomView<'a> {
     /// Find the root of a function in `x` numerically over the reals using Newton's method.
     pub fn nsolve<N: SingleFloat + Real + PartialOrd>(

symbolica.pyi

@@ -1144,6 +1144,81 @@ class Expression:
         >>> print('x =', x_r, ', y =', y_r)
         """
 
+    @overload
+    def nsolve(
+        self,
+        variable: Expression,
+        init: float,
+        prec: float = 1e-4,
+        max_iter: int = 10000,
+    ) -> float:
+        """Find the root of an expression in `x` numerically over the reals using Newton's method.
+        Use `init` as the initial guess for the root. This method uses the same precision as `init`.
+
+        Examples
+        --------
+        >>> from symbolica import *
+        >>> a = E("x^2-2").nsolve(E("x"), 1., 0.0001, 1000000)
+        """
+
+    @overload
+    def nsolve(
+        self,
+        variable: Expression,
+        init: Decimal,
+        prec: float = 1e-4,
+        max_iter: int = 10000,
+    ) -> Decimal:
+        """Find the root of an expression in `x` numerically over the reals using Newton's method.
+        Use `init` as the initial guess for the root. This method uses the same precision as `init`.
+
+        Examples
+        --------
+        >>> from symbolica import *
+        >>> a = E("x^2-2").nsolve(E("x"),
+                      Decimal("1.000000000000000000000000000000000000000000000000000000000000000000000000"), 1e-74, 1000000)
+        """
+
+    @overload
+    @classmethod
+    def nsolve_system(
+        _cls,
+        system: Sequence[Expression],
+        variables: Sequence[Expression],
+        init: Sequence[float],
+        prec: float = 1e-4,
+        max_iter: int = 10000,
+    ) -> Sequence[float]:
+        """Find a common root of multiple expressions in `variables` numerically over the reals using Newton's method.
+        Use `init` as the initial guess for the root. This method uses the same precision as `init`.
+
+        Examples
+        --------
+        >>> from symbolica import *
+        >>> a = Expression.nsolve_system([E("5x^2+x*y^2+sin(2y)^2 - 2"), E("exp(2x-y)+4y - 3")], [S("x"), S("y")],
+                             [Decimal("1.00000000000000000"), Decimal("1.00000000000000000")], 1e-20, 1000000)
+        """
+
+    @overload
+    @classmethod
+    def nsolve_system(
+        _cls,
+        system: Sequence[Expression],
+        variables: Sequence[Expression],
+        init: Sequence[Decimal],
+        prec: float = 1e-4,
+        max_iter: int = 10000,
+    ) -> Sequence[Decimal]:
+        """Find a common root of multiple expressions in `variables` numerically over the reals using Newton's method.
+        Use `init` as the initial guess for the root. This method uses the same precision as `init`.
+
+        Examples
+        --------
+        >>> from symbolica import *
+        >>> a = Expression.nsolve_system([E("5x^2+x*y^2+sin(2y)^2 - 2"), E("exp(2x-y)+4y - 3")], [S("x"), S("y")],
+                             [1., 1.], 1e-20, 1000000)
+        """
+
     def evaluate(
         self, constants: dict[Expression, float], funs: dict[Expression, Callable[[Sequence[float]], float]]
     ) -> float: