Example-01: Optimizaion (import/export API)

# Loading API facilitates (suboptimal) interface with different optimization libraries
# In this example, quadrupole gradients are used to fit beta functions

from pathlib import Path

from model.command.external import load_sdds
from model.command.external import load_lattice
from model.command.external import text_lattice

import numpy
from numpy import ndarray as Array

from scipy.optimize import minimize

# Set quadrupole gradient and compute and return twiss parameters

def evaluate(knobs:Array) -> Array:
    kf, kd = knobs
    path:Path = Path('optimize.lte')
    lattice:dict[str, dict[str, str | int | float | dict]] = load_lattice(path)
    lattice['QF']['K1'] = float(kf)
    lattice['QD']['K1'] = float(kd)
    with path.open('w') as stream:
        stream.write(text_lattice('LTE', lattice))
    !elegant 'optimize.ele' > /dev/null
    !sddsconvert -ascii 'binary.twiss' 'optimize.twiss'
    path:Path = Path('optimize.twiss')
    _, columns = load_sdds(path)
    return numpy.asarray([[data['betax'], data['betay']] for location, data in columns.items()]).T

# Set target beta functions

target:Array = numpy.asarray([+0.21, -0.19])
result:Array = evaluate(target)

# Set objetive function to minimize

def objective(knobs:Array) -> Array:
    return numpy.sum((evaluate(knobs) - result)**2)

objective(target)

0.0

# Optimize

knobs:Array = numpy.asarray([+0.20, -0.20])

minimize(objective, knobs, method='Nelder-Mead')

       message: Optimization terminated successfully.
       success: True
        status: 0
           fun: 1.3489521479279302e-28
             x: [ 2.100e-01 -1.900e-01]
           nit: 20
          nfev: 40
 final_simplex: (array([[ 2.100e-01, -1.900e-01],
                       [ 2.100e-01, -1.900e-01],
                       [ 2.100e-01, -1.900e-01]]), array([ 1.349e-28,  2.891e-05,  6.010e-05]))