From 0fcce03c3f2533e6df6cff9c9835e0110582cbdb Mon Sep 17 00:00:00 2001 From: Johanna Haffner <38662446+johannahaffner@users.noreply.github.com> Date: Sun, 9 Feb 2025 12:05:30 +0100 Subject: [PATCH] Adding `__init__` arguments to documentation for descents (#109) * fix typo * fix typo * fix typo * fix typo * document fix for rewrite functions, such that closure can be checked in implicit_jvp * add __init__ method to documentation for descents * correct a typo and add a sentence * fix typos * adding a wrapper such that converted fields now inherit the correct type --------- Co-authored-by: Johanna Haffner --- docs/api/searches/descents.md | 12 ++++++------ optimistix/_fixed_point.py | 2 +- optimistix/_least_squares.py | 2 +- optimistix/_minimise.py | 5 +++-- optimistix/_root_find.py | 2 +- optimistix/_solver/learning_rate.py | 6 +++++- optimistix/_solver/nonlinear_cg.py | 11 ++++++----- 7 files changed, 23 insertions(+), 17 deletions(-) diff --git a/docs/api/searches/descents.md b/docs/api/searches/descents.md index 60afe69..c00bbd2 100644 --- a/docs/api/searches/descents.md +++ b/docs/api/searches/descents.md @@ -10,39 +10,39 @@ ::: optimistix.SteepestDescent selection: members: - false + - __init__ --- ::: optimistix.NonlinearCGDescent selection: members: - false + - __init__ --- ::: optimistix.NewtonDescent selection: members: - false + - __init__ --- ::: optimistix.DampedNewtonDescent selection: members: - false + - __init__ --- ::: optimistix.IndirectDampedNewtonDescent selection: members: - false + - __init__ --- ::: optimistix.DoglegDescent selection: members: - false + - __init__ diff --git a/optimistix/_fixed_point.py b/optimistix/_fixed_point.py index 9daa657..6916273 100644 --- a/optimistix/_fixed_point.py +++ b/optimistix/_fixed_point.py @@ -101,7 +101,7 @@ def fixed_point( an error. If `False` then the returned solution object will have a `result` field indicating whether any failures occured. (See [`optimistix.Solution`][].) Keyword only argument. - - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing the + - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing any structure of the Jacobian of `y -> fn(y, args) - y` with respect to y. (That is, the structure of the matrix `dfn/dy - I`.) Used with [`optimistix.ImplicitAdjoint`][] to implement the implicit function theorem as diff --git a/optimistix/_least_squares.py b/optimistix/_least_squares.py index bdad3e4..3477706 100644 --- a/optimistix/_least_squares.py +++ b/optimistix/_least_squares.py @@ -87,7 +87,7 @@ def least_squares( an error. If `False` then the returned solution object will have a `result` field indicating whether any failures occured. (See [`optimistix.Solution`][].) Keyword only argument. - - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing the + - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing any structure of the Hessian of `y -> sum(fn(y, args)**2)` with respect to y. Used with [`optimistix.ImplicitAdjoint`][] to implement the implicit function theorem as efficiently as possible. Keyword only argument. diff --git a/optimistix/_minimise.py b/optimistix/_minimise.py index 31b67e5..8eb6780 100644 --- a/optimistix/_minimise.py +++ b/optimistix/_minimise.py @@ -30,7 +30,8 @@ def min_no_aux(x): return jax.grad(min_no_aux)(minimum) -# Keep `optx.implicit_jvp` is happy. +# Keep `optx.implicit_jvp` happy. +# https://github.com/patrick-kidger/optimistix/issues/102#event-15786001854 if _rewrite_fn.__globals__["__name__"].startswith("jaxtyping"): _rewrite_fn = _rewrite_fn.__wrapped__ # pyright: ignore[reportFunctionMemberAccess] @@ -75,7 +76,7 @@ def minimise( an error. If `False` then the returned solution object will have a `result` field indicating whether any failures occured. (See [`optimistix.Solution`][].) Keyword only argument. - - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing the + - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing any structure of the Hessian of `fn` with respect to `y`. Used with [`optimistix.ImplicitAdjoint`][] to implement the implicit function theorem as efficiently as possible. Keyword only argument. diff --git a/optimistix/_root_find.py b/optimistix/_root_find.py index de9d591..c63e5e4 100644 --- a/optimistix/_root_find.py +++ b/optimistix/_root_find.py @@ -163,7 +163,7 @@ def root_find( an error. If `False` then the returned solution object will have a `result` field indicating whether any failures occured. (See [`optimistix.Solution`][].) Keyword only argument. - - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing the + - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing any structure of the Jacobian of `fn` with respect to `y`. Used with some solvers (e.g. [`optimistix.Newton`][]), and with some adjoint methods (e.g. [`optimistix.ImplicitAdjoint`][]) to improve the efficiency of linear solves. diff --git a/optimistix/_solver/learning_rate.py b/optimistix/_solver/learning_rate.py index 76ab21b..029227e 100644 --- a/optimistix/_solver/learning_rate.py +++ b/optimistix/_solver/learning_rate.py @@ -9,10 +9,14 @@ from .._solution import RESULTS +def _typed_asarray(x: ScalarLike) -> Array: + return jnp.asarray(x) + + class LearningRate(AbstractSearch[Y, FunctionInfo, FunctionInfo, None], strict=True): """Move downhill by taking a step of the fixed size `learning_rate`.""" - learning_rate: ScalarLike = eqx.field(converter=jnp.asarray) + learning_rate: ScalarLike = eqx.field(converter=_typed_asarray) def init(self, y: Y, f_info_struct: FunctionInfo) -> None: return None diff --git a/optimistix/_solver/nonlinear_cg.py b/optimistix/_solver/nonlinear_cg.py index 2479b15..af0e155 100644 --- a/optimistix/_solver/nonlinear_cg.py +++ b/optimistix/_solver/nonlinear_cg.py @@ -164,11 +164,12 @@ def step( NonlinearCGDescent.__init__.__doc__ = """**Arguments:** - `method`: A callable `method(vector, vector_prev, diff_prev)` describing how to - calculate the beta parameter of nonlinear CG. Each of these inputs has the meaning - described above. The "beta parameter" is the sake as can be described as e.g. the - β_n value - [on Wikipedia](https://en.wikipedia.org/wiki/Nonlinear_conjugate_gradient_method). - In practice Optimistix includes four built-in methods: + calculate the beta parameter of nonlinear CG. Nonlinear CG uses the previous search + direction, scaled by beta, and subtracts the gradient to find the next search + direction. This parameter, in the nonlinear case, is the same as the parameter β_n + described e.g. [on Wikipedia](https://en.wikipedia.org/wiki/Nonlinear_conjugate_gradient_method) + for the linear case. + Defaults to `polak_ribiere`. Optimistix includes four built-in methods: [`optimistix.polak_ribiere`][], [`optimistix.fletcher_reeves`][], [`optimistix.hestenes_stiefel`][], and [`optimistix.dai_yuan`][]. """