Multi start example

In [ ]:

# SPDX-License-Identifier: MIT
"""
Multi-start optimization example
=================================

Demonstrates how to use ``MultiStart`` to escape local minima when fitting
a non-convex objective.

Problem
-------
Damped spring-mass system.  ISE objective in c (damping coefficient).
The true optimum is c ≈ 1.65.  Starting far from it (e.g. c = 3.0) causes
gradient-based methods to converge slowly — multi-start with random restarts
reliably finds the global minimum.

Usage::

    python docs/examples/optimization/multi_start_example.py
"""

# SPDX-License-Identifier: MIT """ Multi-start optimization example ================================= Demonstrates how to use ``MultiStart`` to escape local minima when fitting a non-convex objective. Problem ------- Damped spring-mass system. ISE objective in c (damping coefficient). The true optimum is c ≈ 1.65. Starting far from it (e.g. c = 3.0) causes gradient-based methods to converge slowly — multi-start with random restarts reliably finds the global minimum. Usage:: python docs/examples/optimization/multi_start_example.py """

In [ ]:

import numpy as np

import numpy as np

In [ ]:

from jaxonomy import DiagramBuilder, Parameter, SimulatorOptions
from jaxonomy.library import Adder, Gain, Integrator, Power
from jaxonomy.optimization import Optimizable, Scipy, MultiStart

from jaxonomy import DiagramBuilder, Parameter, SimulatorOptions from jaxonomy.library import Adder, Gain, Integrator, Power from jaxonomy.optimization import Optimizable, Scipy, MultiStart

── build diagram ─────────────────────────────────────────────────────────────

In [ ]:

params = {"c": Parameter(np.array(1.0))}
b = DiagramBuilder()
k_x = b.add(Gain(1.0,         name="k_x"))
c_v = b.add(Gain(params["c"], name="c_v"))
add = b.add(Adder(2, operators="--", name="adder"))
inv = b.add(Gain(1.0,         name="inv_m"))
v   = b.add(Integrator(0.1,   name="v"))
x   = b.add(Integrator(1.0,   name="x"))
b.connect(k_x.output_ports[0], add.input_ports[0])
b.connect(c_v.output_ports[0], add.input_ports[1])
b.connect(add.output_ports[0], inv.input_ports[0])
b.connect(inv.output_ports[0], v.input_ports[0])
b.connect(v.output_ports[0],   x.input_ports[0])
b.connect(v.output_ports[0],   c_v.input_ports[0])
b.connect(x.output_ports[0],   k_x.input_ports[0])
sq_v = b.add(Power(2.0, name="sq_v"))
sq_x = b.add(Power(2.0, name="sq_x"))
cv   = b.add(Integrator(0.0, name="cv"))
cx   = b.add(Integrator(0.0, name="cx"))
obj  = b.add(Adder(2, operators="++", name="obj"))
b.connect(v.output_ports[0],   sq_v.input_ports[0])
b.connect(x.output_ports[0],   sq_x.input_ports[0])
b.connect(sq_v.output_ports[0], cv.input_ports[0])
b.connect(sq_x.output_ports[0], cx.input_ports[0])
b.connect(cv.output_ports[0],  obj.input_ports[0])
b.connect(cx.output_ports[0],  obj.input_ports[1])
diagram = b.build(parameters=params)

params = {"c": Parameter(np.array(1.0))} b = DiagramBuilder() k_x = b.add(Gain(1.0, name="k_x")) c_v = b.add(Gain(params["c"], name="c_v")) add = b.add(Adder(2, operators="--", name="adder")) inv = b.add(Gain(1.0, name="inv_m")) v = b.add(Integrator(0.1, name="v")) x = b.add(Integrator(1.0, name="x")) b.connect(k_x.output_ports[0], add.input_ports[0]) b.connect(c_v.output_ports[0], add.input_ports[1]) b.connect(add.output_ports[0], inv.input_ports[0]) b.connect(inv.output_ports[0], v.input_ports[0]) b.connect(v.output_ports[0], x.input_ports[0]) b.connect(v.output_ports[0], c_v.input_ports[0]) b.connect(x.output_ports[0], k_x.input_ports[0]) sq_v = b.add(Power(2.0, name="sq_v")) sq_x = b.add(Power(2.0, name="sq_x")) cv = b.add(Integrator(0.0, name="cv")) cx = b.add(Integrator(0.0, name="cx")) obj = b.add(Adder(2, operators="++", name="obj")) b.connect(v.output_ports[0], sq_v.input_ports[0]) b.connect(x.output_ports[0], sq_x.input_ports[0]) b.connect(sq_v.output_ports[0], cv.input_ports[0]) b.connect(sq_x.output_ports[0], cx.input_ports[0]) b.connect(cv.output_ports[0], obj.input_ports[0]) b.connect(cx.output_ports[0], obj.input_ports[1]) diagram = b.build(parameters=params)

── define optimizable ────────────────────────────────────────────────────────

In [ ]:

class SpringOptimizable(Optimizable):
    def __init__(self, *a, **kw):
        super().__init__(*a, **kw)
        self._obj = diagram["obj"].output_ports[0]

    def optimizable_params(self, ctx):
        return {"c": ctx.parameters["c"]}

    def objective_from_context(self, ctx):
        return self._obj.eval(ctx)

    def prepare_context(self, ctx, p):
        return ctx.with_parameters(p)

class SpringOptimizable(Optimizable): def __init__(self, *a, **kw): super().__init__(*a, **kw) self._obj = diagram["obj"].output_ports[0] def optimizable_params(self, ctx): return {"c": ctx.parameters["c"]} def objective_from_context(self, ctx): return self._obj.eval(ctx) def prepare_context(self, ctx, p): return ctx.with_parameters(p)

In [ ]:

model = SpringOptimizable(
    diagram, diagram.create_context(),
    params_0={"c": 0.5},
    sim_t_span=(0.0, 2.0),
    sim_options=SimulatorOptions(max_major_steps=1),
)

model = SpringOptimizable( diagram, diagram.create_context(), params_0={"c": 0.5}, sim_t_span=(0.0, 2.0), sim_options=SimulatorOptions(max_major_steps=1), )

── single start (for comparison) ────────────────────────────────────────────

In [ ]:

single_optim = Scipy(model, "L-BFGS-B",
                     opt_method_config={"maxiter": 40},
                     use_autodiff_grad=True)
single_result = single_optim.optimize()
print("=== Single-start result ===")
print(f"  c = {float(single_result['c']):.4f}  (expected ≈ 1.65)")
print(f"  success={single_result.success}, nit={single_result.nit}, "
      f"final_loss={single_result.final_loss:.6g}")

single_optim = Scipy(model, "L-BFGS-B", opt_method_config={"maxiter": 40}, use_autodiff_grad=True) single_result = single_optim.optimize() print("=== Single-start result ===") print(f" c = {float(single_result['c']):.4f} (expected ≈ 1.65)") print(f" success={single_result.success}, nit={single_result.nit}, " f"final_loss={single_result.final_loss:.6g}")

── multi-start ───────────────────────────────────────────────────────────────

In [ ]:

def factory(opt):
    return Scipy(opt, "L-BFGS-B",
                 opt_method_config={"maxiter": 40},
                 use_autodiff_grad=True)

def factory(opt): return Scipy(opt, "L-BFGS-B", opt_method_config={"maxiter": 40}, use_autodiff_grad=True)

In [ ]:

ms = MultiStart(
    model,
    factory,
    n_starts=6,
    sample_scale=1.5,   # search window: ±1.5×|c0|
    seed=42,
    include_initial=True,
)

ms = MultiStart( model, factory, n_starts=6, sample_scale=1.5, # search window: ±1.5×|c0| seed=42, include_initial=True, )

In [ ]:

ms_result = ms.run()
print("\n=== Multi-start result ===")
print(ms_result.summary())
print(f"\nBest c = {float(ms_result.best_result['c']):.4f}  (expected ≈ 1.65)")

ms_result = ms.run() print("\n=== Multi-start result ===") print(ms_result.summary()) print(f"\nBest c = {float(ms_result.best_result['c']):.4f} (expected ≈ 1.65)")