From defbeb5258a4ac353cb3b2f9fe9537ce0af91590 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Wed, 1 Oct 2025 18:02:08 -0300
Subject: [PATCH 01/40] OPT.BASE.FIX: missing `self` on `_initialization`

---
 apsuite/optimization/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/apsuite/optimization/base.py b/apsuite/optimization/base.py
index f2d9cc59..3802d492 100644
--- a/apsuite/optimization/base.py
+++ b/apsuite/optimization/base.py
@@ -1,4 +1,4 @@
-"""Simulated Annealing Algorithm for Minimization."""
+"""."""
 import logging as _log
 
 import matplotlib.pyplot as _mplt
@@ -283,7 +283,7 @@ def _optimize(self):
         """Implement here optimization algorithm."""
         raise NotImplementedError()
 
-    def _initialization():
+    def _initialization(self):
         """To be called before optimization starts.
 
         If the return value is False, optimization will not run.

From fd099f4c43fb1d3a5943c9d5488f01b6f6577f8d Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Wed, 1 Oct 2025 18:08:29 -0300
Subject: [PATCH 02/40] OPT.LEASTSQUARES.WIP: initial version of least-squares
 optimizer

- implements  Levenberg-Maquardt/Gauss-Newton style iterative least-squares solver
---
 apsuite/optimization/least_squares.py | 195 ++++++++++++++++++++++++++
 1 file changed, 195 insertions(+)
 create mode 100644 apsuite/optimization/least_squares.py

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
new file mode 100644
index 00000000..742e121e
--- /dev/null
+++ b/apsuite/optimization/least_squares.py
@@ -0,0 +1,195 @@
+"""."""
+
+import numpy as _np
+
+from apsuite.optimization.base import Optimize, OptimizeParams
+
+
+class LeastSquaresParams(OptimizeParams):
+    """."""
+
+    def __init__(self):
+        """."""
+        super().__init__()
+        self._boundary_policy = self.BoundaryPolicy.ToBoundary
+        self.max_number_iters = 10
+        self.max_number_evals = _np.inf
+        self.abs_tol_convergence = 1e-6
+        self.rel_tol_convergence = 1e-3
+        self.damping_constant = 1.0
+        self.damping_factor = 10.0
+        self.ridge_constant = 0.0
+        self.jacobian = None
+        self.jacobian_update_rate = 0
+        self.verbose = True
+        self.learning_rate = 1.0
+        self.min_learning_rate = 1e-3
+        self.backtracking_factor = 2
+        self.patience = 5
+
+
+class LeastSquaresOptimize(Optimize):
+    """Least-squares optimization via Levenberg–Marquardt style loop."""
+
+    def __init__(
+        self,
+        params=None,
+        merit_figure_goal=None,
+        jacobian=None,
+        use_thread=True,
+        isonline=False,
+    ):
+        """."""
+        super().__init__(
+            params=LeastSquaresParams() if params is None else params,
+            use_thread=use_thread,
+            isonline=isonline,
+        )
+        self.merit_figure_goal = merit_figure_goal
+        self.jacobian = jacobian
+
+        self.history_chi2 = []
+
+    def objective_function(self, pos):
+        """."""
+        merit_figure = self.calc_merit_figure(pos)
+        residual = self.calc_residual(merit_figure)
+        return residual
+
+    def calc_chi2(self, residual):
+        """."""
+        chi2 = _np.mean(residual**2)
+        self.history_chi2.append(chi2)
+        return chi2
+
+    def calc_residual(self, merit_figure_meas, merit_figure_goal=None):
+        """."""
+        if merit_figure_goal is None:
+            merit_figure_goal = self.merit_figure_goal
+
+        return merit_figure_meas - merit_figure_goal
+
+    def calc_merit_figure(self, pos):
+        """."""
+        raise NotImplementedError(
+            'Problem-specific figure of merit needs to be implemented'
+        )
+
+    def calc_jacobian(self, pos=None, step=1e-4):
+        """."""
+        jacobian_t = list()
+        pos0 = self.params.initial_position
+        pos = pos0.copy() if pos is None else pos
+
+        for i in range(len(pos)):
+            pos[i] += step / 2
+            figm_pos = self.calc_merit_figure(pos)
+            pos[i] = pos0[i]
+            pos[i] -= step / 2
+            figm_neg = self.calc_merit_figure(pos)
+            pos[i] = pos0[i]
+
+            jac_col = (figm_pos - figm_neg) / step
+            jacobian_t.append(jac_col)
+        return _np.array(jacobian_t).T
+
+    def _optimize(self):
+        """LM-like optimization loop."""
+        niter = self.params.max_number_iters
+        atol = self.params.abs_tol_convergence
+        rtol = self.params.rel_tol_convergence
+        jacobian = self.jacobian
+        damping_constant = self.params.damping_constant
+        damping_factor = self.params.damping_factor
+        ridge_reg = self.params.ridge_constant
+        jacobian_update_rate = self.params.jacobian_update_rate
+        lr = self.params.learning_rate
+        lr_min = self.params.min_learning_rate
+        lr_factor = self.params.backtracking_factor
+        patience = self.params.patience
+        verbose = self.params.verbose
+
+        pos0 = self.params.initial_position
+        pos = pos0.copy()
+
+        res = self._objective_func(pos)
+        chi2 = self.calc_chi2(res)
+        M = self.calc_jacobian(pos) if jacobian is None else jacobian
+
+        # if c:
+        #     chi2 += w * _np.std(res0) * _np.linalg.norm(pos0)
+
+        if verbose:
+            print(f'initial chi²: {chi2:.6g}')
+
+        for it in range(niter):
+            if verbose:
+                print(f'iteration {it:03d}')
+
+            pos_init = pos.copy()
+
+            update_jac = jacobian_update_rate and it
+            update_jac &= not (it % jacobian_update_rate)
+
+            if update_jac:
+                M = self.calc_jacobian(pos)
+
+            MTM = M.T @ M
+            reg = _np.diag(
+                damping_constant * _np.diag(MTM)
+            )  # LM regularization
+            reg += ridge_reg * _np.eye(MTM.shape[0])  # Ridge regularization
+            matrix = MTM + reg
+
+            res = self.objfuncs_evaluated[-1]  # last evaluated residual
+            delta = (
+                _np.linalg.pinv(matrix) @ M.T @ res
+            )  # LM/GN Normal equation
+            # TODO: include processing of pseudo-inverse for singvals selection
+            pos -= lr * delta  # position update
+
+            if _np.any(_np.isnan(pos)):
+                if verbose:
+                    print('\tNaNs detected in pos. Aborting.')
+                break
+
+            fails = 0
+            while fails <= patience:
+                try:
+                    res = self._objective_func(pos)
+                    break
+                except Exception:
+                    if verbose:
+                        print('Merit figure evaluation failed. Back-tracking.')
+                    fails += 1
+                    lr /= lr_factor  # reduce step size
+                    lr = min(lr, lr_min)
+                    pos = pos_init  # restore pos
+                    pos -= lr * delta  # apply smaller step
+
+            chi2_old = self.history_chi2[-1]
+            chi2_new = self.calc_chi2(res)
+            # if c:
+            #     chi2 += w * _np.std(res) * _np.linalg.norm(pos)
+            if verbose:
+                print(f'\tchi²: {chi2_new:.6g}')
+
+            delta_chi2 = chi2_old - chi2_new
+            if abs(delta_chi2) < atol or abs(delta_chi2) / chi2_old < rtol:
+                if verbose:
+                    print('\tConvergence tolerance reached. Exiting.')
+                break
+
+            if damping_constant:
+                if delta_chi2 > 0:
+                    damping_constant /= damping_factor
+                    if verbose:
+                        print(
+                            '\tImproved fit. Decreasing LM damping_constant.'
+                        )
+                else:
+                    damping_constant *= damping_factor
+                    if verbose:
+                        print(
+                            '\tWorsened fit. Increasing LM damping_constant.'
+                        )

From 061954d3eafd93230c262b0dceb8a208fe3acb7e Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Wed, 1 Oct 2025 18:17:12 -0300
Subject: [PATCH 03/40] OPTICSANLY.NOECO.WIP: initial version of NOECO class

- for calibration/fitting  of sextupoles strengths via the Nonlinear Optics from Off-energy Orbits (NOECO) method
---
 apsuite/optics_analysis/noeco.py | 253 +++++++++++++++++++++++++++++++
 1 file changed, 253 insertions(+)
 create mode 100644 apsuite/optics_analysis/noeco.py

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
new file mode 100644
index 00000000..613065cb
--- /dev/null
+++ b/apsuite/optics_analysis/noeco.py
@@ -0,0 +1,253 @@
+"""."""
+
+import matplotlib.pyplot as _mplt
+import numpy as _np
+import pyaccel as _pa
+from pyaccel.optics.miscellaneous import (
+    get_chromaticities,
+    get_mcf,
+    get_rf_frequency
+)
+from pymodels import si as _si
+
+from apsuite.optics_analysis import ChromCorr, TuneCorr
+from apsuite.optimization.least_squares import (
+    LeastSquaresOptimize,
+    LeastSquaresParams
+)
+from apsuite.orbcorr import OrbRespmat
+
+
+class NOECOParams(LeastSquaresParams):
+    """."""
+
+    SEXT_FAMS = (
+        'SDA0',
+        'SDB0',
+        'SDP0',
+        'SFA0',
+        'SFB0',
+        'SFP0',
+        'SDA1',
+        'SDB1',
+        'SDP1',
+        'SDA2',
+        'SDB2',
+        'SDP2',
+        'SDA3',
+        'SDB3',
+        'SDP3',
+        'SFA1',
+        'SFB1',
+        'SFP1',
+        'SFA2',
+        'SFB2',
+        'SFP2',
+    )
+    SEXT_FAMS_ACHROM = SEXT_FAMS[:6]
+    SEXT_FAMS_CHROM = SEXT_FAMS[6:]
+
+    def __init__(self):
+        """."""
+        super().__init__()
+        self.denergy_oeorm_calc = 1e-2  # 1 %
+        self.tunex = 49.16
+        self.tuney = 14.22
+        self.chromx = 3.11
+        self.chromy = 2.5
+        self.sextfams2fit = self.SEXT_FAMS_CHROM
+
+
+class NOECOFit(LeastSquaresOptimize):
+    """."""
+
+    def __init__(
+        self,
+        merit_figure_goal=None,
+        jacobian=None,
+        use_thread=True,
+        isonline=False,
+    ):
+        """."""
+        params = NOECOParams()
+        super().__init__(
+            params, merit_figure_goal, jacobian, use_thread, isonline
+        )
+        self._model = None
+        self.famdata = None
+
+        self._orbmat = None
+        self._tunecorr = None
+        self._chromcorr = None
+
+    @property
+    def model(self):
+        """."""
+        return self._model
+
+    @model.setter
+    def model(self, model):
+        """."""
+        if not isinstance(model, _pa.accelerator.Accelerator):
+            raise TypeError('model must be pyaccel.accelerator.Accelerator')
+        self._model = model
+        self.famdata = _si.get_family_data(self._model)
+        self.create_corr_objects()
+
+    def create_model(self):
+        """."""
+        self.model = _si.create_accelerator()
+
+    def create_corr_objects(self):
+        """."""
+        if self.model is None:
+            raise RuntimeError(
+                'self.model is None. Call create_model() or set model first'
+            )
+
+        if self._orbmat is None:
+            self._orbmat = OrbRespmat(
+                model=self.model, acc='SI', use6dtrack=True
+            )
+        if self._tunecorr is None:
+            self._tunecorr = TuneCorr(model=self.model, acc='SI')
+        if self._chromcorr is None:
+            self._chromcorr = ChromCorr(model=self.model, acc='SI')
+
+    def adjust_model(self, tunes=None, chroms=None):
+        """."""
+        if tunes is None:
+            tunes = self.params.tunex, self.params.tuney
+        if chroms is None:
+            chroms = self.params.chromx, self.params.chromy
+
+        self._tunecorr.correct_parameters(tunes, model=self.model)
+        self._chromcorr.correct_parameters(chroms, model=self.model)
+
+    def _set_energy_offset(self, energy_offset, model=None):
+        """."""
+        if model is None:
+            model = self.model
+        freq0 = get_rf_frequency(model)
+        alpha = get_mcf(model)
+        df = -alpha * freq0 * energy_offset
+        self._set_delta_rf_frequency(df, model)
+
+    def _set_delta_rf_frequency(self, delta_frequency, model=None):
+        if model is None:
+            model = self.model
+        freq0 = get_rf_frequency(model)
+        print(f'Changing RF freq. by {delta_frequency:.4e}')
+        self._set_rf_frequency(freq0 + delta_frequency, model)
+
+    def _set_rf_frequency(self, frequency, model=None):
+        """."""
+        if model is None:
+            model = self.model
+        cav_idx = _si.get_family_data(model)['SRFCav']['index'][0][0]
+        print(f'Setting RF freq. to {frequency:.8e}')
+        model[cav_idx].frequency = frequency
+
+    def get_strengths_(self):
+        """."""
+        strengths = list()
+        for fam in self.params.sextfams2fit:
+            idc = _pa.lattice.flatten(self.famdata[fam]['index'])[0]
+            strengths.append(self.model[idc].SL)
+        return _np.array(strengths)
+
+    def set_strengths_(self, strengths):
+        """."""
+        for stren, fam in zip(strengths, self.params.sextfams2fit):
+            idcs = _pa.lattice.flatten(self.famdata[fam]['index'])
+            _pa.lattice.set_attribute(self.model, 'SL', idcs, stren)
+
+    def get_oeorm(
+        self,
+        model=None,
+        strengths=None,
+        idempotent=True,
+        delta=1e-2,
+        normalize_rf=True,
+        ravel=False,
+    ):
+        """."""
+        if model is None:
+            model = self.model
+
+        if strengths is not None:
+            pos0 = self.get_strengths_()
+            self.set_strengths_(strengths)
+
+        rf_freq0 = get_rf_frequency(self.model)
+        self._set_energy_offset(energy_offset=delta, model=self.model)
+        mat_pos = self._orbmat.get_respm(add_rfline=True)
+
+        self._set_rf_frequency(frequency=rf_freq0, model=self.model)
+
+        self._set_energy_offset(energy_offset=-delta, model=self.model)
+        mat_neg = self._orbmat.get_respm(add_rfline=True)
+
+        self._set_rf_frequency(frequency=rf_freq0, model=self.model)
+
+        oeorm = (mat_pos - mat_neg) / 2 / delta
+
+        if normalize_rf:
+            oeorm[:, -1] *= 1e6  # [um/Hz]
+
+        if ravel:
+            oeorm = _np.ravel(oeorm)
+
+        if strengths is not None and idempotent:
+            self.set_strengths_(pos0)
+
+        return oeorm
+
+    def get_chroms(self, strengths=None, idempotent=True):
+        """."""
+        if strengths is not None:
+            pos0 = self.get_strengths_()
+            self.set_strengths_(strengths)
+
+        chroms = get_chromaticities(self.model)
+
+        if strengths is not None and idempotent:
+            self.set_strengths_(pos0)
+
+        return _np.array(chroms)
+
+    def calc_merit_figure(self, pos):
+        """."""
+        oeorm = self.get_oeorm(
+            strengths=pos,
+            idempotent=True,
+            delta=self.params.denergy_oeorm_calc,
+            normalize_rf=True,
+            ravel=True,
+        )
+        return oeorm
+
+    def get_optimization_pos(self):
+        """."""
+        pos = self.get_strengths_()
+        # other params will go here as well (gains, rolls, etc)
+        return pos
+
+    def plot_strengths(self, strengths=None, fig=None, ax=None, label=None):
+        """."""
+        if fig is None or ax is None:
+            fig, ax = _mplt.subplots(figsize=(12, 4))
+        if strengths is None:
+            strengths = self.get_strengths_()
+
+        ax.plot(strengths, 'o-', mfc='none', label=label)
+        tick_label = self.params.sextfams2fit
+        tick_pos = list(range(len(tick_label)))
+        ax.set_xticks(tick_pos, labels=tick_label, rotation=45)
+        ax.set_ylabel(r'SL [m$^{-2}$]')
+        ax.set_xlabel('chromatic sextupole families')
+        if label:
+            ax.legend()
+        fig.tight_layout()
+        fig.show()
+        return fig, ax

From 83eb61673908cd964e13567754661ec6e7a3e5fd Mon Sep 17 00:00:00 2001
From: Matheus Velloso <vellosok75@gmail.com>
Date: Thu, 2 Oct 2025 09:18:39 -0300
Subject: [PATCH 04/40] OPT.LEASTSQUARES.BUG: avoid division by zero

---
 apsuite/optimization/least_squares.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 742e121e..bfe0e462 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -128,8 +128,8 @@ def _optimize(self):
 
             pos_init = pos.copy()
 
-            update_jac = jacobian_update_rate and it
-            update_jac &= not (it % jacobian_update_rate)
+            if jacobian_update_rate and it:
+                update_jac = not it % jacobian_update_rate
 
             if update_jac:
                 M = self.calc_jacobian(pos)

From 3833d6aea6f95f3a99ba99ccd8e7080a34d7438d Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 2 Oct 2025 09:25:04 -0300
Subject: [PATCH 05/40] OPT.LEASTSQUARES.BUG: var used before assignment

---
 apsuite/optimization/least_squares.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index bfe0e462..5c3cb98c 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -129,10 +129,8 @@ def _optimize(self):
             pos_init = pos.copy()
 
             if jacobian_update_rate and it:
-                update_jac = not it % jacobian_update_rate
-
-            if update_jac:
-                M = self.calc_jacobian(pos)
+                if not it % jacobian_update_rate:
+                    M = self.calc_jacobian(pos)
 
             MTM = M.T @ M
             reg = _np.diag(

From 368c1e830689a420108e04a8da45c968959dc921 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 2 Oct 2025 11:01:00 -0300
Subject: [PATCH 06/40] OPT.LEASTSQUARES.ENH: include error sigma as weights

---
 apsuite/optimization/least_squares.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 5c3cb98c..a1754e1b 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -26,6 +26,7 @@ def __init__(self):
         self.min_learning_rate = 1e-3
         self.backtracking_factor = 2
         self.patience = 5
+        self.errorbars = None
 
 
 class LeastSquaresOptimize(Optimize):
@@ -67,7 +68,10 @@ def calc_residual(self, merit_figure_meas, merit_figure_goal=None):
         if merit_figure_goal is None:
             merit_figure_goal = self.merit_figure_goal
 
-        return merit_figure_meas - merit_figure_goal
+        res = merit_figure_meas - merit_figure_goal
+        if self.params.errorbars is not None:
+            res = res / self.params.errorbars
+        return res
 
     def calc_merit_figure(self, pos):
         """."""
@@ -90,6 +94,8 @@ def calc_jacobian(self, pos=None, step=1e-4):
             pos[i] = pos0[i]
 
             jac_col = (figm_pos - figm_neg) / step
+            if self.params.errorbars is not None:
+                jac_col = jac_col / self.params.errorbars
             jacobian_t.append(jac_col)
         return _np.array(jacobian_t).T
 

From 64aae9cd1578b7a67a0838dac83d3af53aaa261f Mon Sep 17 00:00:00 2001
From: Matheus Velloso <74263753+vellosok75@users.noreply.github.com>
Date: Thu, 2 Oct 2025 14:02:32 -0300
Subject: [PATCH 07/40] Apply suggestions from code review

Co-authored-by: fernandohds564 <fernandohds564@gmail.com>
---
 apsuite/optimization/least_squares.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index a1754e1b..d839ad04 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -101,6 +101,11 @@ def calc_jacobian(self, pos=None, step=1e-4):
 
     def _optimize(self):
         """LM-like optimization loop."""
+        def print_(*args, **kwargs):
+            if not verbose:
+                return
+            print(*args, **kwargs)
+            
         niter = self.params.max_number_iters
         atol = self.params.abs_tol_convergence
         rtol = self.params.rel_tol_convergence
@@ -125,8 +130,7 @@ def _optimize(self):
         # if c:
         #     chi2 += w * _np.std(res0) * _np.linalg.norm(pos0)
 
-        if verbose:
-            print(f'initial chi²: {chi2:.6g}')
+        print_(f'initial chi²: {chi2:.6g}')
 
         for it in range(niter):
             if verbose:

From 367a6aaf3828f558269d4b4823fdd0ad5e4c0538 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 2 Oct 2025 14:20:14 -0300
Subject: [PATCH 08/40] OPT.LEASTSQUARES.ENH: print_ function automatically
 handling verbose param

---
 apsuite/optimization/least_squares.py | 25 ++++++++-----------------
 1 file changed, 8 insertions(+), 17 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index d839ad04..a321d8ce 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -101,8 +101,9 @@ def calc_jacobian(self, pos=None, step=1e-4):
 
     def _optimize(self):
         """LM-like optimization loop."""
+
         def print_(*args, **kwargs):
-            if not verbose:
+            if not self.params.verbose:
                 return
             print(*args, **kwargs)
             
@@ -157,8 +158,7 @@ def print_(*args, **kwargs):
             pos -= lr * delta  # position update
 
             if _np.any(_np.isnan(pos)):
-                if verbose:
-                    print('\tNaNs detected in pos. Aborting.')
+                print_('\tNaNs detected in pos. Aborting.')
                 break
 
             fails = 0
@@ -167,8 +167,7 @@ def print_(*args, **kwargs):
                     res = self._objective_func(pos)
                     break
                 except Exception:
-                    if verbose:
-                        print('Merit figure evaluation failed. Back-tracking.')
+                    print_('Merit figure evaluation failed. Back-tracking.')
                     fails += 1
                     lr /= lr_factor  # reduce step size
                     lr = min(lr, lr_min)
@@ -179,25 +178,17 @@ def print_(*args, **kwargs):
             chi2_new = self.calc_chi2(res)
             # if c:
             #     chi2 += w * _np.std(res) * _np.linalg.norm(pos)
-            if verbose:
-                print(f'\tchi²: {chi2_new:.6g}')
+            print_(f'\tchi²: {chi2_new:.6g}')
 
             delta_chi2 = chi2_old - chi2_new
             if abs(delta_chi2) < atol or abs(delta_chi2) / chi2_old < rtol:
-                if verbose:
-                    print('\tConvergence tolerance reached. Exiting.')
+                print_('\tConvergence tolerance reached. Exiting.')
                 break
 
             if damping_constant:
                 if delta_chi2 > 0:
                     damping_constant /= damping_factor
-                    if verbose:
-                        print(
-                            '\tImproved fit. Decreasing LM damping_constant.'
-                        )
+                    print_('\tImproved fit. Decreasing LM damping_constant.')
                 else:
                     damping_constant *= damping_factor
-                    if verbose:
-                        print(
-                            '\tWorsened fit. Increasing LM damping_constant.'
-                        )
+                    print_('\tWorsened fit. Increasing LM damping_constant.')

From 3d398a66d2cf4062ca651d50ad4b50756e213ccc Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 2 Oct 2025 14:25:51 -0300
Subject: [PATCH 09/40] OPT.LEASTSQUARES.ENH: avoid unnecessary recalculations
 w/ jacobian

- do not calculate M.T M every iter, only when jacobian is update
- same for ridge regularization
---
 apsuite/optimization/least_squares.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index a321d8ce..8cea86b2 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -106,20 +106,19 @@ def print_(*args, **kwargs):
             if not self.params.verbose:
                 return
             print(*args, **kwargs)
-            
+
         niter = self.params.max_number_iters
         atol = self.params.abs_tol_convergence
         rtol = self.params.rel_tol_convergence
         jacobian = self.jacobian
         damping_constant = self.params.damping_constant
         damping_factor = self.params.damping_factor
-        ridge_reg = self.params.ridge_constant
+        ridge_constant = self.params.ridge_constant
         jacobian_update_rate = self.params.jacobian_update_rate
         lr = self.params.learning_rate
         lr_min = self.params.min_learning_rate
         lr_factor = self.params.backtracking_factor
         patience = self.params.patience
-        verbose = self.params.verbose
 
         pos0 = self.params.initial_position
         pos = pos0.copy()
@@ -133,22 +132,26 @@ def print_(*args, **kwargs):
 
         print_(f'initial chi²: {chi2:.6g}')
 
+        MTM = M.T @ M
+        ridge_reg = ridge_constant * _np.eye(
+            MTM.shape[0]
+        )  # Ridge regularization
+
         for it in range(niter):
-            if verbose:
-                print(f'iteration {it:03d}')
+            print_(f'iteration {it:03d}')
 
             pos_init = pos.copy()
 
             if jacobian_update_rate and it:
                 if not it % jacobian_update_rate:
                     M = self.calc_jacobian(pos)
+                    MTM = M.T @ M
 
-            MTM = M.T @ M
-            reg = _np.diag(
+            lm_reg = _np.diag(
                 damping_constant * _np.diag(MTM)
-            )  # LM regularization
-            reg += ridge_reg * _np.eye(MTM.shape[0])  # Ridge regularization
-            matrix = MTM + reg
+            )  # Levenberg-Macquardt regularization
+
+            matrix = MTM + ridge_reg + lm_reg
 
             res = self.objfuncs_evaluated[-1]  # last evaluated residual
             delta = (

From c90ef7b3ad300e0f6dd554c966877ebfc29b6c90 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 2 Oct 2025 16:19:20 -0300
Subject: [PATCH 10/40] OPTICSANLY.NOECO.STY: remove underscore from methods
 names

---
 apsuite/optics_analysis/noeco.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 613065cb..aa44d50e 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -148,7 +148,7 @@ def _set_rf_frequency(self, frequency, model=None):
         print(f'Setting RF freq. to {frequency:.8e}')
         model[cav_idx].frequency = frequency
 
-    def get_strengths_(self):
+    def get_strengths(self):
         """."""
         strengths = list()
         for fam in self.params.sextfams2fit:
@@ -156,7 +156,7 @@ def get_strengths_(self):
             strengths.append(self.model[idc].SL)
         return _np.array(strengths)
 
-    def set_strengths_(self, strengths):
+    def set_strengths(self, strengths):
         """."""
         for stren, fam in zip(strengths, self.params.sextfams2fit):
             idcs = _pa.lattice.flatten(self.famdata[fam]['index'])
@@ -176,8 +176,8 @@ def get_oeorm(
             model = self.model
 
         if strengths is not None:
-            pos0 = self.get_strengths_()
-            self.set_strengths_(strengths)
+            pos0 = self.get_strengths()
+            self.set_strengths(strengths)
 
         rf_freq0 = get_rf_frequency(self.model)
         self._set_energy_offset(energy_offset=delta, model=self.model)
@@ -199,20 +199,20 @@ def get_oeorm(
             oeorm = _np.ravel(oeorm)
 
         if strengths is not None and idempotent:
-            self.set_strengths_(pos0)
+            self.set_strengths(pos0)
 
         return oeorm
 
     def get_chroms(self, strengths=None, idempotent=True):
         """."""
         if strengths is not None:
-            pos0 = self.get_strengths_()
-            self.set_strengths_(strengths)
+            pos0 = self.get_strengths()
+            self.set_strengths(strengths)
 
         chroms = get_chromaticities(self.model)
 
         if strengths is not None and idempotent:
-            self.set_strengths_(pos0)
+            self.set_strengths(pos0)
 
         return _np.array(chroms)
 
@@ -229,7 +229,7 @@ def calc_merit_figure(self, pos):
 
     def get_optimization_pos(self):
         """."""
-        pos = self.get_strengths_()
+        pos = self.get_strengths()
         # other params will go here as well (gains, rolls, etc)
         return pos
 
@@ -238,7 +238,7 @@ def plot_strengths(self, strengths=None, fig=None, ax=None, label=None):
         if fig is None or ax is None:
             fig, ax = _mplt.subplots(figsize=(12, 4))
         if strengths is None:
-            strengths = self.get_strengths_()
+            strengths = self.get_strengths()
 
         ax.plot(strengths, 'o-', mfc='none', label=label)
         tick_label = self.params.sextfams2fit

From 699c81c948712e17b58284c219ec07e6c1b670b3 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Wed, 8 Oct 2025 14:47:33 -0300
Subject: [PATCH 11/40] NOECO.ENH: raveling of OEORM, new properties &
 inclusion of gains transformations

additions:
- automatically raveling of OEORM for `merit_figure_goal`
- setter for goal oeorm
- properties for nr of hor and ver bpms and corrs
- methods to apply gains, parse gains and other params from optimization pos and vice-versa
---
 apsuite/optics_analysis/noeco.py | 103 ++++++++++++++++++++++++++++---
 1 file changed, 93 insertions(+), 10 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index aa44d50e..a37cfb64 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -6,14 +6,14 @@
 from pyaccel.optics.miscellaneous import (
     get_chromaticities,
     get_mcf,
-    get_rf_frequency
+    get_rf_frequency,
 )
 from pymodels import si as _si
 
 from apsuite.optics_analysis import ChromCorr, TuneCorr
 from apsuite.optimization.least_squares import (
     LeastSquaresOptimize,
-    LeastSquaresParams
+    LeastSquaresParams,
 )
 from apsuite.orbcorr import OrbRespmat
 
@@ -62,24 +62,44 @@ class NOECOFit(LeastSquaresOptimize):
     """."""
 
     def __init__(
-        self,
-        merit_figure_goal=None,
-        jacobian=None,
-        use_thread=True,
-        isonline=False,
+        self, oeorm_goal=None, jacobian=None, use_thread=True, isonline=False
     ):
         """."""
         params = NOECOParams()
+
+        merit_figure_goal = None
+        if oeorm_goal is not None and oeorm_goal.squeeze().ndim > 1:
+            merit_figure_goal = oeorm_goal.ravel()
+
         super().__init__(
-            params, merit_figure_goal, jacobian, use_thread, isonline
+            params=params,
+            merit_figure_goal=merit_figure_goal,
+            jacobian=jacobian,
+            use_thread=use_thread,
+            isonline=isonline,
         )
+
+        self._oeorm_goal = None
         self._model = None
         self.famdata = None
 
+        self.nr_sexts = len(self.params.sextfams2fit)
+        self.nr_bpms_total = 320
+        self.nr_bpms_x = 160
+        self.nr_bpms_y = 160
+        self.nr_corrs_total = 281
+        self.nr_corrs_x = 120
+        self.nr_corrs_y = 160
+
         self._orbmat = None
         self._tunecorr = None
         self._chromcorr = None
 
+        self.sext_strengths = None
+        self.gain_corrs = _np.ones(self.nr_corrs_total)
+        self.gain_bpms = _np.zeros(self.nr_bpms_total)  # Gx, Gy
+        self.coup_bpms = _np.zeros(self.nr_bpms_total)  # Cc, Cy
+
     @property
     def model(self):
         """."""
@@ -94,6 +114,19 @@ def model(self, model):
         self.famdata = _si.get_family_data(self._model)
         self.create_corr_objects()
 
+    @property
+    def oeorm_goal(self):
+        """."""
+        return self._oeorm_goal
+
+    @oeorm_goal.setter
+    def oeorm_goal(self, oeorm):
+        """."""
+        self._oeorm_goal = oeorm
+        self.merit_figure_goal = oeorm.ravel()
+        self.nr_bpms_total = oeorm.shape[0] / 2
+        self.nr_corrs_total = oeorm[:-1].shape[-1]
+
     def create_model(self):
         """."""
         self.model = _si.create_accelerator()
@@ -227,10 +260,60 @@ def calc_merit_figure(self, pos):
         )
         return oeorm
 
+    def calc_residual(self, merit_figure_meas, merit_figure_goal=None):
+        """."""
+        if merit_figure_goal is None:
+            merit_figure_goal = self.oeorm_goal
+
+        merit_figure_goal = self.get_real_matrix_from_meas(
+            merit_figure_goal
+        ).ravel()
+
+        return super().calc_residual(merit_figure_meas, merit_figure_goal)
+
+    def get_real_matrix_from_meas(self, oeorm):
+        """."""
+        _, *ret = self.parse_params_from_pos()
+        gains_corrs, gains_bpms, coup_bpms = ret
+        Gb = self.bpms_gains_matrix(gains_bpms, coup_bpms)
+        oeorm_ = Gb @ (oeorm * gains_corrs)
+        return oeorm_
+
+    def bpms_gains_matrix(self, gains, coups):
+        """."""
+        gx = _np.diag(1 + gains[: self.nr_bpms_x])
+        gy = _np.diag(1 + gains[self.nr_bpms_x :])
+        cx = _np.diag(coups[: self.nr_bpms_x])
+        cy = _np.diag(coups[self.nr_bpms_x :])
+        return _np.block([[gx, cx], [cy, gy]])
+
+    def parse_params_from_pos(self, pos=None):
+        """."""
+        if pos is None:
+            pos = (
+                self.positions_evaluated[-1]
+                if len(self.positions_evaluated)
+                else self.get_optimization_pos()
+            )
+        n = self.nr_sexts
+        sexts = pos[:n]
+        gains_corrs = pos[n : n + self.nr_corrs_total]
+        n += self.nr_corrs_total
+        gains_bpms = pos[n : n + self.nr_bpms_total]
+        n += self.nr_bpms_total
+        coup_bpms = pos[n : n + self.nr_bpms_total]
+        return sexts, gains_corrs, gains_bpms, coup_bpms
+
     def get_optimization_pos(self):
         """."""
-        pos = self.get_strengths()
-        # other params will go here as well (gains, rolls, etc)
+        if self.sext_strengths is None:
+            sexts = self.get_strengths()
+        else:
+            sexts = self.sext_strengths
+        gain_bpms = self.gain_bpms
+        gain_corrs = self.gain_corrs
+        coup_bpms = self.coup_bpms
+        pos = _np.r_[sexts, gain_corrs, gain_bpms, coup_bpms]
         return pos
 
     def plot_strengths(self, strengths=None, fig=None, ax=None, label=None):

From 342c807f154c1af0db602691eb4052bef28c6de0 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 16 Oct 2025 13:41:00 -0300
Subject: [PATCH 12/40] OPT.LEASTSQUARES.ENH: add `rcond` param for singval
 selection

-`rcond` added to params for cutting-off singular values less than `rcond * largest_singular_value`
- initialized to `numpy.linalg.pinv` default of `1e-15`
---
 apsuite/optimization/least_squares.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 8cea86b2..2c2d3c72 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -16,6 +16,7 @@ def __init__(self):
         self.max_number_evals = _np.inf
         self.abs_tol_convergence = 1e-6
         self.rel_tol_convergence = 1e-3
+        self.rcond = 1e-15
         self.damping_constant = 1.0
         self.damping_factor = 10.0
         self.ridge_constant = 0.0
@@ -155,7 +156,7 @@ def print_(*args, **kwargs):
 
             res = self.objfuncs_evaluated[-1]  # last evaluated residual
             delta = (
-                _np.linalg.pinv(matrix) @ M.T @ res
+                _np.linalg.pinv(matrix, rcond=self.params.rcond) @ M.T @ res
             )  # LM/GN Normal equation
             # TODO: include processing of pseudo-inverse for singvals selection
             pos -= lr * delta  # position update

From ac6adcc3d6a9cd0b7a193c176b2eebd04010672a Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 16 Oct 2025 13:42:09 -0300
Subject: [PATCH 13/40] OPT.LEASTSQUARES.ENH: require position for calculating
 jacobian

- `ValueError` raised if no pos is passed.
---
 apsuite/optimization/least_squares.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 2c2d3c72..8758c107 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -74,7 +74,7 @@ def calc_residual(self, merit_figure_meas, merit_figure_goal=None):
             res = res / self.params.errorbars
         return res
 
-    def calc_merit_figure(self, pos):
+    def calc_merit_figure(self):
         """."""
         raise NotImplementedError(
             'Problem-specific figure of merit needs to be implemented'
@@ -83,9 +83,9 @@ def calc_merit_figure(self, pos):
     def calc_jacobian(self, pos=None, step=1e-4):
         """."""
         jacobian_t = list()
-        pos0 = self.params.initial_position
-        pos = pos0.copy() if pos is None else pos
-
+        if pos is None:
+            raise ValueError('Specify position to calculate Jacobian.')
+        pos0 = pos.copy()
         for i in range(len(pos)):
             pos[i] += step / 2
             figm_pos = self.calc_merit_figure(pos)

From e3697bc05b97b4a6d63e08a43837e1b17a45aa24 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 16 Oct 2025 14:02:09 -0300
Subject: [PATCH 14/40] OPTICSANLY.NOECO.ENH: major refactor (BPMs & corrs.
 gains, jacobians, error bars)

- refactor module to include fitting of corrs. gains and BPMs gains and couplings, support for bpms noise as error bars
- organization and stylistic changes (func args names)
---
 apsuite/optics_analysis/noeco.py | 357 ++++++++++++++++++++++++-------
 1 file changed, 277 insertions(+), 80 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index a37cfb64..8e49a020 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -9,6 +9,7 @@
     get_rf_frequency,
 )
 from pymodels import si as _si
+from scipy.linalg import block_diag
 
 from apsuite.optics_analysis import ChromCorr, TuneCorr
 from apsuite.optimization.least_squares import (
@@ -55,50 +56,58 @@ def __init__(self):
         self.tuney = 14.22
         self.chromx = 3.11
         self.chromy = 2.5
+
         self.sextfams2fit = self.SEXT_FAMS_CHROM
+        self.fit_gain_bpms = True
+        self.fit_coup_bpms = True
+        self.fit_gain_corr = True
+
+        self.update_jacobian_sexts = False
+        self.update_jacobian_gains = True
+
+        self.nr_sexts = len(self.sextfams2fit)
+        self.nr_bpms = 160
+        self.nr_chs = 120
+        self.nr_cvs = 160
+        self.nr_corrs = self.nr_chs + self.nr_cvs
+
+        self.init_gain_corr = _np.ones(self.nr_corrs)
+        self.init_gain_bpms = _np.ones(2 * self.nr_bpms)  # Gx, Gy
+        self.init_coup_bpms = _np.zeros(2 * self.nr_bpms)  # Cc, Cy
 
 
 class NOECOFit(LeastSquaresOptimize):
     """."""
 
     def __init__(
-        self, oeorm_goal=None, jacobian=None, use_thread=True, isonline=False
+        self, oeorm_goal=None, bpms_noise=None, use_thread=True, isonline=False
     ):
         """."""
         params = NOECOParams()
 
-        merit_figure_goal = None
-        if oeorm_goal is not None and oeorm_goal.squeeze().ndim > 1:
-            merit_figure_goal = oeorm_goal.ravel()
+        self._oeorm_goal = None
+        self._bpms_noise = None
+        self._model = None
+        self.famdata = None
+        self._orbmat = None
+        self._tunecorr = None
+        self._chromcorr = None
+        self.jacobian_sexts = None
+        self.jacobian_gains = None
+        self.jacobian = None
 
         super().__init__(
             params=params,
-            merit_figure_goal=merit_figure_goal,
-            jacobian=jacobian,
+            merit_figure_goal=None,
             use_thread=use_thread,
             isonline=isonline,
         )
 
-        self._oeorm_goal = None
-        self._model = None
-        self.famdata = None
+        if oeorm_goal is not None:
+            self.oeorm_goal = oeorm_goal
 
-        self.nr_sexts = len(self.params.sextfams2fit)
-        self.nr_bpms_total = 320
-        self.nr_bpms_x = 160
-        self.nr_bpms_y = 160
-        self.nr_corrs_total = 281
-        self.nr_corrs_x = 120
-        self.nr_corrs_y = 160
-
-        self._orbmat = None
-        self._tunecorr = None
-        self._chromcorr = None
-
-        self.sext_strengths = None
-        self.gain_corrs = _np.ones(self.nr_corrs_total)
-        self.gain_bpms = _np.zeros(self.nr_bpms_total)  # Gx, Gy
-        self.coup_bpms = _np.zeros(self.nr_bpms_total)  # Cc, Cy
+        if bpms_noise is not None:
+            self.bpms_noise = bpms_noise
 
     @property
     def model(self):
@@ -122,10 +131,20 @@ def oeorm_goal(self):
     @oeorm_goal.setter
     def oeorm_goal(self, oeorm):
         """."""
-        self._oeorm_goal = oeorm
+        n, m = self.params.nr_bpms, self.params.nr_corrs + 1
+        self._oeorm_goal = oeorm.reshape((2 * n, m))
         self.merit_figure_goal = oeorm.ravel()
-        self.nr_bpms_total = oeorm.shape[0] / 2
-        self.nr_corrs_total = oeorm[:-1].shape[-1]
+
+    @property
+    def bpms_noise(self):
+        """."""
+        return self._bpms_noise
+
+    @bpms_noise.setter
+    def bpms_noise(self, bpms_noise):
+        """."""
+        self._bpms_noise = bpms_noise
+        self.params.errorbars = self._get_errorbars(bpms_noise)
 
     def create_model(self):
         """."""
@@ -138,14 +157,11 @@ def create_corr_objects(self):
                 'self.model is None. Call create_model() or set model first'
             )
 
-        if self._orbmat is None:
-            self._orbmat = OrbRespmat(
-                model=self.model, acc='SI', use6dtrack=True
-            )
-        if self._tunecorr is None:
-            self._tunecorr = TuneCorr(model=self.model, acc='SI')
-        if self._chromcorr is None:
-            self._chromcorr = ChromCorr(model=self.model, acc='SI')
+        self._orbmat = OrbRespmat(
+            model=self.model, acc='SI', use6dtrack=True
+        )
+        self._tunecorr = TuneCorr(model=self.model, acc='SI')
+        self._chromcorr = ChromCorr(model=self.model, acc='SI')
 
     def adjust_model(self, tunes=None, chroms=None):
         """."""
@@ -170,7 +186,7 @@ def _set_delta_rf_frequency(self, delta_frequency, model=None):
         if model is None:
             model = self.model
         freq0 = get_rf_frequency(model)
-        print(f'Changing RF freq. by {delta_frequency:.4e}')
+        # print(f'Changing RF freq. by {delta_frequency:.4e}')
         self._set_rf_frequency(freq0 + delta_frequency, model)
 
     def _set_rf_frequency(self, frequency, model=None):
@@ -178,9 +194,44 @@ def _set_rf_frequency(self, frequency, model=None):
         if model is None:
             model = self.model
         cav_idx = _si.get_family_data(model)['SRFCav']['index'][0][0]
-        print(f'Setting RF freq. to {frequency:.8e}')
+        # print(f'Setting RF freq. to {frequency:.8e}')
         model[cav_idx].frequency = frequency
 
+    def _initialization(self):
+        if self.model is None:
+            print('Cannot start optimization without machine model.')
+            return False
+
+        if self.oeorm_goal is None:
+            print('Cannot start optimization without oeorm_goal')
+            return False
+
+        if not len(self.params.initial_position):
+            print('No initial position specified. Using default pos.')
+            pos = self.get_strengths()
+            if self.params.fit_gain_corr:
+                pos = _np.r_[pos, self.params.init_gain_corr]
+            if self.params.fit_gain_bpms:
+                pos = _np.r_[pos, self.params.init_gain_bpms]
+            if self.params.fit_coup_bpms:
+                pos = _np.r_[pos, self.params.init_coup_bpms]
+            self.params.initial_position = pos
+
+        self.params.limit_lower = _np.full(
+            self.params.initial_position.shape, -_np.inf
+        )
+        self.params.limit_upper = _np.full(
+            self.params.initial_position.shape, _np.inf
+        )
+        return True
+
+    def _get_errorbars(self, bpms_noise):
+        n, m = self.params.nr_bpms, self.params.nr_corrs + 1
+        errorbar = _np.ones(2 * n * m)
+        for i, sigma in enumerate(bpms_noise, start=0):
+            errorbar[i * m : (i + 1) * m] = _np.repeat(sigma, m)
+        return errorbar
+
     def get_strengths(self):
         """."""
         strengths = list()
@@ -209,7 +260,7 @@ def get_oeorm(
             model = self.model
 
         if strengths is not None:
-            pos0 = self.get_strengths()
+            strengths0 = self.get_strengths()
             self.set_strengths(strengths)
 
         rf_freq0 = get_rf_frequency(self.model)
@@ -232,27 +283,32 @@ def get_oeorm(
             oeorm = _np.ravel(oeorm)
 
         if strengths is not None and idempotent:
-            self.set_strengths(pos0)
+            self.set_strengths(strengths0)
 
         return oeorm
 
     def get_chroms(self, strengths=None, idempotent=True):
         """."""
         if strengths is not None:
-            pos0 = self.get_strengths()
+            strengths0 = self.get_strengths()
             self.set_strengths(strengths)
 
         chroms = get_chromaticities(self.model)
 
         if strengths is not None and idempotent:
-            self.set_strengths(pos0)
+            self.set_strengths(strengths0)
 
         return _np.array(chroms)
 
-    def calc_merit_figure(self, pos):
+    def calc_merit_figure(self, strengths=None):
         """."""
+        if strengths is None:
+            pos = self.get_default_pos()
+            pos = self.params.initial_position
+            strengths, *_ = self.parse_params_from_pos(pos)
+
         oeorm = self.get_oeorm(
-            strengths=pos,
+            strengths=strengths,
             idempotent=True,
             delta=self.params.denergy_oeorm_calc,
             normalize_rf=True,
@@ -265,56 +321,197 @@ def calc_residual(self, merit_figure_meas, merit_figure_goal=None):
         if merit_figure_goal is None:
             merit_figure_goal = self.oeorm_goal
 
-        merit_figure_goal = self.get_real_matrix_from_meas(
-            merit_figure_goal
-        ).ravel()
+        if merit_figure_goal.ndim == 1:
+            merit_figure_goal = merit_figure_goal.reshape((
+                2 * self.params.nr_bpms,
+                self.params.nr_corrs + 1,
+            ))
+
+        merit_figure_goal = self.apply_gains(merit_figure_goal).ravel()
 
         return super().calc_residual(merit_figure_meas, merit_figure_goal)
 
-    def get_real_matrix_from_meas(self, oeorm):
+    def calc_jacobian_sexts(self, strengths, step):
         """."""
-        _, *ret = self.parse_params_from_pos()
+        jacobian = super().calc_jacobian(strengths, step)
+        self.jacobian_sexts = jacobian
+        return jacobian
+
+    def calc_jacobian_gains(self, pos=None):
+        """."""
+        if pos is None:
+            pos = self.get_default_pos()
+
+        oeorm_meas = self.oeorm_goal
+        jacobian = None
+
+        # corr. gains Jacobian jth col = measured OEORM jth col & zeros
+        # J_gains = [[oeorm_bpm_gains[:, 0].T, zeros, ...      zeros        ],
+        #            [zeros, oeorm_bpm_gains[:, 1].T, ...      zeros        ],
+        #            [zeros, zeros, oeorm_bpm_gains[:, 2].T    zeros        ],
+        #            [  .      .            .                     .         ],
+        #            [  .      .            .                     .         ],
+        #            [zeros,  zeros,       ...    oeorm_bpm_gains[:, 280].T]]
+        # this is implemented below
+        if self.params.fit_gain_corr:
+            # apply BPMs gains and calculate corrs gains jacobian
+            oeorm_bpms_gains = self.apply_gains(
+                oeorm_meas, pos, bpms=True, corrs=False
+            )
+            jacobian_gains_corrs = block_diag(*oeorm_bpms_gains.T).T
+            if self.params.errorbars is not None:
+                jacobian_gains_corrs /= self.params.errorbars[:, None]
+
+            jacobian = jacobian_gains_corrs[:, : self.params.nr_corrs]
+
+        oeorm_corr_gains = self.apply_gains(
+            oeorm_meas, pos, bpms=False, corrs=True
+        )
+        # BPM gains Jacobian jth col = measured OEORM jth line, padded with
+        # zeros
+        # J_gains = [[oeorm_corr_gains[0, :].T, zeros, ...      zeros        ],
+        #            [zeros, oeorm_corr_gains[1, :].T, ...      zeros        ],
+        #            [zeros, zeros, oeorm_corr_gains[2, :].T    zeros        ],
+        #            [  .      .            .                     .          ],
+        #            [  .      .            .                     .          ],
+        #            [zeros,  zeros,       ...    oeorm_corr_gains[320, :].T]]
+        # this is implemented below
+        if self.params.fit_gain_bpms:
+            jacobian_gains_bpms = block_diag(*oeorm_corr_gains).T
+            if self.params.errorbars is not None:
+                jacobian_gains_bpms /= self.params.errorbars[:, None]
+
+            if jacobian is None:
+                jacobian = jacobian_gains_bpms
+            else:
+                jacobian = _np.hstack((jacobian, jacobian_gains_bpms))
+
+        # couplings Jacobian jth col = jth line of OEORM with swapped upper and
+        # lower blocks. Eg:
+        #       OEORM =      [[Mxx, Mxy], [Myx, Myy]],
+        #       OEORM_swap = [[Myx, Myy], [Mxx, Mxy]]
+        # implementeded below
+        if self.params.fit_coup_bpms:
+            jacobian_coup_bpms = block_diag(
+                *_np.vstack((
+                    oeorm_corr_gains[self.params.nr_bpms :],
+                    oeorm_corr_gains[: self.params.nr_bpms],
+                ))
+            ).T
+            if self.params.errorbars is not None:
+                jacobian_coup_bpms /= self.params.errorbars[:, None]
+
+            if jacobian is None:
+                jacobian = jacobian_coup_bpms
+            else:
+                jacobian = _np.hstack((jacobian, jacobian_coup_bpms))
+
+        self.jacobian_gains = jacobian
+
+        return jacobian
+
+    def calc_jacobian(self, pos=None, step=0.0001):
+        """."""
+        if pos is None:
+            pos = self.get_default_pos()
+
+        if self.jacobian_sexts is not None:
+            jacobian_sexts = self.jacobian_sexts
+
+        if self.jacobian_gains is not None:
+            jacobian_gains = self.jacobian_gains
+
+        if self.jacobian_sexts is None or self.params.update_jacobian_sexts:
+            strengths, _, _, _ = self.parse_params_from_pos(pos)
+            jacobian_sexts = self.calc_jacobian_sexts(strengths, step)
+
+        if self.jacobian_gains is None or self.params.update_jacobian_gains:
+            jacobian_gains = self.calc_jacobian_gains(pos)
+
+        gains = self.params.fit_gain_bpms or self.params.fit_coup_bpms
+        gains = gains or self.params.fit_gain_corr
+
+        if gains:
+            jacobian = _np.hstack((jacobian_sexts, jacobian_gains))
+        else:
+            jacobian = jacobian_sexts
+
+        self.jacobian = jacobian
+
+        return jacobian
+
+    def get_oeorm_blocks(self, oeorm):
+        """."""
+        nr_bpms = self.params.nr_bpms
+        nr_ch = self.params.nr_chs
+        nr_cv = self.params.nr_cvs
+
+        xx = oeorm[:nr_bpms, :nr_ch]
+        xy = oeorm[:nr_bpms, nr_ch : nr_ch + nr_cv]
+        dispx = oeorm[:nr_bpms, -1]
+        yx = oeorm[nr_bpms:, :nr_ch]
+        yy = oeorm[nr_bpms:, nr_ch : nr_ch + nr_cv]
+        dispy = oeorm[nr_bpms:, -1]
+
+        return xx, xy, dispx, yx, yy, dispy
+
+    def apply_gains(self, oeorm, pos=None, bpms=True, corrs=True):
+        """."""
+        if pos is None:
+            pos = self.get_default_pos()
+        _, *ret = self.parse_params_from_pos(pos)
         gains_corrs, gains_bpms, coup_bpms = ret
-        Gb = self.bpms_gains_matrix(gains_bpms, coup_bpms)
-        oeorm_ = Gb @ (oeorm * gains_corrs)
+
+        oeorm_ = oeorm.copy()
+
+        if corrs and gains_corrs is not None:
+            oeorm_ = oeorm_ * _np.r_[gains_corrs, [1]]
+
+        if bpms and (gains_bpms is not None or coup_bpms is not None):
+            Gb = self.bpms_gains_matrix(gains_bpms, coup_bpms)
+            oeorm_ = Gb @ oeorm_
+
         return oeorm_
 
     def bpms_gains_matrix(self, gains, coups):
         """."""
-        gx = _np.diag(1 + gains[: self.nr_bpms_x])
-        gy = _np.diag(1 + gains[self.nr_bpms_x :])
-        cx = _np.diag(coups[: self.nr_bpms_x])
-        cy = _np.diag(coups[self.nr_bpms_x :])
+        if gains is None:
+            gains = _np.ones(2 * self.params.nr_bpms)
+        if coups is None:
+            coups = _np.zeros(2 * self.params.nr_bpms)
+
+        gx = _np.diag(gains[: self.params.nr_bpms])
+        gy = _np.diag(gains[self.params.nr_bpms :])
+        cx = _np.diag(coups[: self.params.nr_bpms])
+        cy = _np.diag(coups[self.params.nr_bpms :])
         return _np.block([[gx, cx], [cy, gy]])
 
-    def parse_params_from_pos(self, pos=None):
+    def get_default_pos(self):
         """."""
-        if pos is None:
-            pos = (
-                self.positions_evaluated[-1]
-                if len(self.positions_evaluated)
-                else self.get_optimization_pos()
-            )
-        n = self.nr_sexts
-        sexts = pos[:n]
-        gains_corrs = pos[n : n + self.nr_corrs_total]
-        n += self.nr_corrs_total
-        gains_bpms = pos[n : n + self.nr_bpms_total]
-        n += self.nr_bpms_total
-        coup_bpms = pos[n : n + self.nr_bpms_total]
-        return sexts, gains_corrs, gains_bpms, coup_bpms
+        pos = (
+            self.positions_evaluated[-1]
+            if len(self.positions_evaluated)
+            else self.params.initial_position
+        )
+        return pos
 
-    def get_optimization_pos(self):
+    def parse_params_from_pos(self, pos):
         """."""
-        if self.sext_strengths is None:
-            sexts = self.get_strengths()
-        else:
-            sexts = self.sext_strengths
-        gain_bpms = self.gain_bpms
-        gain_corrs = self.gain_corrs
-        coup_bpms = self.coup_bpms
-        pos = _np.r_[sexts, gain_corrs, gain_bpms, coup_bpms]
-        return pos
+        gains_corrs = None
+        gains_bpms = None
+        coup_bpms = None
+
+        n = self.params.nr_sexts
+        sexts = pos[:n]
+        if self.params.fit_gain_corr:
+            gains_corrs = pos[n : n + self.params.nr_corrs]
+            n += self.params.nr_corrs
+        if self.params.fit_gain_bpms:
+            gains_bpms = pos[n : n + 2 * self.params.nr_bpms]
+            n += 2 * self.params.nr_bpms
+        if self.params.fit_coup_bpms:
+            coup_bpms = pos[n : n + 2 * self.params.nr_bpms]
+        return sexts, gains_corrs, gains_bpms, coup_bpms
 
     def plot_strengths(self, strengths=None, fig=None, ax=None, label=None):
         """."""

From 2eaa15e27afd28d65882ab5f3d3be5f2357763b5 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 23 Oct 2025 08:22:00 -0300
Subject: [PATCH 15/40] OPTANLY.NOECO.FIX: 4D tracking for tune and chrom
 correction

---
 apsuite/optics_analysis/noeco.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 8e49a020..75c04a8c 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -170,8 +170,15 @@ def adjust_model(self, tunes=None, chroms=None):
         if chroms is None:
             chroms = self.params.chromx, self.params.chromy
 
+        cav = self.model.cavity_on
+        rad = self.model.radiation_on
+        self.model.cavity_on = False
+        self.model.radiation_on = False
         self._tunecorr.correct_parameters(tunes, model=self.model)
         self._chromcorr.correct_parameters(chroms, model=self.model)
+        self.model.cavity_on = cav
+        self.model.radiation_on = rad
+
 
     def _set_energy_offset(self, energy_offset, model=None):
         """."""

From 776ab0753d753bd5065fc5cfd5e0af238a6ff74a Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 23 Oct 2025 08:30:58 -0300
Subject: [PATCH 16/40] OPTANLY.ENH: support OEORM block selection

- allows choosing to use or not the diagonal & off-diagonal OEORM blocks as well as the hor. & ver. dispersion col.
---
 apsuite/optics_analysis/noeco.py | 28 +++++++++++++++++++++++++---
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 75c04a8c..0663fb10 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -62,6 +62,11 @@ def __init__(self):
         self.fit_coup_bpms = True
         self.fit_gain_corr = True
 
+        self.use_diag_blocks = True
+        self.use_hor_disp = True
+        self.use_ver_disp = False
+        self.use_offdiag_blocks = False
+
         self.update_jacobian_sexts = False
         self.update_jacobian_gains = True
 
@@ -131,9 +136,9 @@ def oeorm_goal(self):
     @oeorm_goal.setter
     def oeorm_goal(self, oeorm):
         """."""
-        n, m = self.params.nr_bpms, self.params.nr_corrs + 1
-        self._oeorm_goal = oeorm.reshape((2 * n, m))
-        self.merit_figure_goal = oeorm.ravel()
+        oeorm_ = self._select_matrix_blocks(oeorm)
+        self._oeorm_goal = oeorm_
+        self.merit_figure_goal = oeorm_.ravel()
 
     @property
     def bpms_noise(self):
@@ -179,6 +184,21 @@ def adjust_model(self, tunes=None, chroms=None):
         self.model.cavity_on = cav
         self.model.radiation_on = rad
 
+    def _select_matrix_blocks(self, oeorm):
+        n, m = self.params.nr_bpms, self.params.nr_corrs + 1
+        xx, xy, dispx, yx, yy, dispy = self.get_oeorm_blocks(oeorm)
+        oeorm_ = _np.zeros((2 * n, m))
+        if self.params.use_diag_blocks:
+            oeorm_[:n, : self.params.nr_chs] = xx
+            oeorm_[n:, self.params.nr_chs : self.params.nr_corrs] = yy
+        if self.params.use_offdiag_blocks:
+            oeorm_[:n, self.params.nr_chs : self.params.nr_corrs] = xy
+            oeorm_[n:, : self.params.nr_chs] = yx
+        if self.params.use_hor_disp:
+            oeorm_[:n, -1] = dispx
+        if self.params.use_ver_disp:
+            oeorm_[n:, -1] = dispy
+        return oeorm_.reshape((2 * n, m))
 
     def _set_energy_offset(self, energy_offset, model=None):
         """."""
@@ -286,6 +306,8 @@ def get_oeorm(
         if normalize_rf:
             oeorm[:, -1] *= 1e6  # [um/Hz]
 
+        oeorm = self._select_matrix_blocks(oeorm)
+
         if ravel:
             oeorm = _np.ravel(oeorm)
 

From 29dc354db66fbaf359d4c3aeaf8dbc6d56160f75 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 23 Oct 2025 08:41:32 -0300
Subject: [PATCH 17/40] OPTANLY.NOECO.ENH: get initial optimization position 
 from params method and new behavior for default pos

- dedicated method to get positions from init guesses and machine model state
- default pos: if no last position evaluated, use params initial position, if empty, calculate it.
---
 apsuite/optics_analysis/noeco.py | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 0663fb10..8762e3db 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -235,13 +235,10 @@ def _initialization(self):
 
         if not len(self.params.initial_position):
             print('No initial position specified. Using default pos.')
-            pos = self.get_strengths()
-            if self.params.fit_gain_corr:
-                pos = _np.r_[pos, self.params.init_gain_corr]
-            if self.params.fit_gain_bpms:
-                pos = _np.r_[pos, self.params.init_gain_bpms]
-            if self.params.fit_coup_bpms:
-                pos = _np.r_[pos, self.params.init_coup_bpms]
+            pos = self.get_initial_pos()
+            if not len(pos):
+                print('No parameters to fit!')
+                return False
             self.params.initial_position = pos
 
         self.params.limit_lower = _np.full(
@@ -520,10 +517,27 @@ def get_default_pos(self):
         pos = (
             self.positions_evaluated[-1]
             if len(self.positions_evaluated)
-            else self.params.initial_position
+            else (
+                self.params.initial_position
+                if len(self.params.initial_position)
+                else self.get_initial_pos()
+            )
         )
         return pos
 
+    def get_initial_pos(self):
+        """."""
+        pos = []
+        if self.params.fit_sexts:
+            pos = _np.r_[pos, self.get_strengths()]
+        if self.params.fit_gain_corr:
+            pos = _np.r_[pos, self.params.init_gain_corr]
+        if self.params.fit_gain_bpms:
+            pos = _np.r_[pos, self.params.init_gain_bpms]
+        if self.params.fit_coup_bpms:
+            pos = _np.r_[pos, self.params.init_coup_bpms]
+        return pos
+
     def parse_params_from_pos(self, pos):
         """."""
         gains_corrs = None

From 6312e59e6d8d67e1665f2b193811b1de25e8d7ef Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 23 Oct 2025 08:43:29 -0300
Subject: [PATCH 18/40] OPTANLY.NOECO.FEAT: allows choosing whether or not to
 fit sextupoles

- now supports fitting only sextupoles, only corr or bpms gains, or only corr or only bpms gains
---
 apsuite/optics_analysis/noeco.py | 57 ++++++++++++++++++++------------
 1 file changed, 36 insertions(+), 21 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 8762e3db..da23c9d7 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -58,6 +58,7 @@ def __init__(self):
         self.chromy = 2.5
 
         self.sextfams2fit = self.SEXT_FAMS_CHROM
+        self.fit_sexts = True
         self.fit_gain_bpms = True
         self.fit_coup_bpms = True
         self.fit_gain_corr = True
@@ -162,9 +163,7 @@ def create_corr_objects(self):
                 'self.model is None. Call create_model() or set model first'
             )
 
-        self._orbmat = OrbRespmat(
-            model=self.model, acc='SI', use6dtrack=True
-        )
+        self._orbmat = OrbRespmat(model=self.model, acc='SI', use6dtrack=True)
         self._tunecorr = TuneCorr(model=self.model, acc='SI')
         self._chromcorr = ChromCorr(model=self.model, acc='SI')
 
@@ -330,7 +329,6 @@ def calc_merit_figure(self, strengths=None):
         """."""
         if strengths is None:
             pos = self.get_default_pos()
-            pos = self.params.initial_position
             strengths, *_ = self.parse_params_from_pos(pos)
 
         oeorm = self.get_oeorm(
@@ -441,26 +439,41 @@ def calc_jacobian(self, pos=None, step=0.0001):
         if pos is None:
             pos = self.get_default_pos()
 
-        if self.jacobian_sexts is not None:
-            jacobian_sexts = self.jacobian_sexts
+        jacobians = []
+
+        # --- sexts ---
+        if self.params.fit_sexts:
+            if self.jacobian_sexts is not None:
+                jacobian_sexts = self.jacobian_sexts
+
+            update_sexts = self.params.update_jacobian_sexts
+            if self.jacobian_sexts is None or update_sexts:
+                strengths, _, _, _ = self.parse_params_from_pos(pos)
+                jacobian_sexts = self.calc_jacobian_sexts(strengths, step)
 
-        if self.jacobian_gains is not None:
-            jacobian_gains = self.jacobian_gains
+            jacobians.append(jacobian_sexts)
 
-        if self.jacobian_sexts is None or self.params.update_jacobian_sexts:
-            strengths, _, _, _ = self.parse_params_from_pos(pos)
-            jacobian_sexts = self.calc_jacobian_sexts(strengths, step)
+        # --- gains ---
+        fit_gains = (
+            self.params.fit_gain_bpms
+            or self.params.fit_coup_bpms
+            or self.params.fit_gain_corr
+        )
+
+        if fit_gains:
+            if self.jacobian_gains is not None:
+                jacobian_gains = self.jacobian_gains
 
-        if self.jacobian_gains is None or self.params.update_jacobian_gains:
-            jacobian_gains = self.calc_jacobian_gains(pos)
+            update_gains = self.params.update_jacobian_gains
+            if self.jacobian_gains is None or update_gains:
+                jacobian_gains = self.calc_jacobian_gains(pos)
 
-        gains = self.params.fit_gain_bpms or self.params.fit_coup_bpms
-        gains = gains or self.params.fit_gain_corr
+            jacobians.append(jacobian_gains)
 
-        if gains:
-            jacobian = _np.hstack((jacobian_sexts, jacobian_gains))
+        if jacobians:
+            jacobian = _np.hstack(jacobians)
         else:
-            jacobian = jacobian_sexts
+            jacobian = None
 
         self.jacobian = jacobian
 
@@ -540,12 +553,14 @@ def get_initial_pos(self):
 
     def parse_params_from_pos(self, pos):
         """."""
+        sexts = None
         gains_corrs = None
         gains_bpms = None
         coup_bpms = None
-
-        n = self.params.nr_sexts
-        sexts = pos[:n]
+        n = 0
+        if self.params.fit_sexts:
+            n = self.params.nr_sexts
+            sexts = pos[:n]
         if self.params.fit_gain_corr:
             gains_corrs = pos[n : n + self.params.nr_corrs]
             n += self.params.nr_corrs

From 19134e799dc74d9f97c1723e81861ad0d7147ab5 Mon Sep 17 00:00:00 2001
From: Fernando <fernandohds564@gmail.com>
Date: Thu, 23 Oct 2025 13:33:26 -0300
Subject: [PATCH 19/40] OPTANL.WIP: @vellosok75 and @fernandohds564 partial
 revistion of code.

---
 apsuite/optics_analysis/noeco.py      | 101 ++++++++++++--------------
 apsuite/optimization/base.py          |   8 +-
 apsuite/optimization/least_squares.py |  37 +++++-----
 3 files changed, 69 insertions(+), 77 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index da23c9d7..4208b7b1 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -4,9 +4,9 @@
 import numpy as _np
 import pyaccel as _pa
 from pyaccel.optics.miscellaneous import (
-    get_chromaticities,
-    get_mcf,
-    get_rf_frequency,
+    get_chromaticities as _get_chromaticities,
+    get_mcf as _get_mcf,
+    get_rf_frequency as _get_rf_frequency,
 )
 from pymodels import si as _si
 from scipy.linalg import block_diag
@@ -71,7 +71,6 @@ def __init__(self):
         self.update_jacobian_sexts = False
         self.update_jacobian_gains = True
 
-        self.nr_sexts = len(self.sextfams2fit)
         self.nr_bpms = 160
         self.nr_chs = 120
         self.nr_cvs = 160
@@ -81,6 +80,11 @@ def __init__(self):
         self.init_gain_bpms = _np.ones(2 * self.nr_bpms)  # Gx, Gy
         self.init_coup_bpms = _np.zeros(2 * self.nr_bpms)  # Cc, Cy
 
+    @property
+    def nr_sexts(self):
+        """."""
+        return len(self.sextfams2fit)
+
 
 class NOECOFit(LeastSquaresOptimize):
     """."""
@@ -100,7 +104,6 @@ def __init__(
         self._chromcorr = None
         self.jacobian_sexts = None
         self.jacobian_gains = None
-        self.jacobian = None
 
         super().__init__(
             params=params,
@@ -109,11 +112,8 @@ def __init__(
             isonline=isonline,
         )
 
-        if oeorm_goal is not None:
-            self.oeorm_goal = oeorm_goal
-
-        if bpms_noise is not None:
-            self.bpms_noise = bpms_noise
+        self.oeorm_goal = oeorm_goal
+        self.bpms_noise = bpms_noise
 
     @property
     def model(self):
@@ -126,6 +126,8 @@ def model(self, model):
         if not isinstance(model, _pa.accelerator.Accelerator):
             raise TypeError('model must be pyaccel.accelerator.Accelerator')
         self._model = model
+        self._model.radiation_on = 1
+        self._model.cavity_on = True
         self.famdata = _si.get_family_data(self._model)
         self.create_corr_objects()
 
@@ -137,6 +139,8 @@ def oeorm_goal(self):
     @oeorm_goal.setter
     def oeorm_goal(self, oeorm):
         """."""
+        if oeorm is None:
+            raise ValueError('oeorm_goal must not be None.')
         oeorm_ = self._select_matrix_blocks(oeorm)
         self._oeorm_goal = oeorm_
         self.merit_figure_goal = oeorm_.ravel()
@@ -149,6 +153,8 @@ def bpms_noise(self):
     @bpms_noise.setter
     def bpms_noise(self, bpms_noise):
         """."""
+        if bpms_noise is None:
+            return
         self._bpms_noise = bpms_noise
         self.params.errorbars = self._get_errorbars(bpms_noise)
 
@@ -158,14 +164,14 @@ def create_model(self):
 
     def create_corr_objects(self):
         """."""
-        if self.model is None:
+        if self._model is None:
             raise RuntimeError(
                 'self.model is None. Call create_model() or set model first'
             )
 
-        self._orbmat = OrbRespmat(model=self.model, acc='SI', use6dtrack=True)
-        self._tunecorr = TuneCorr(model=self.model, acc='SI')
-        self._chromcorr = ChromCorr(model=self.model, acc='SI')
+        self._orbmat = OrbRespmat(model=self._model, acc='SI', use6dtrack=True)
+        self._tunecorr = TuneCorr(model=self._model, acc='SI')
+        self._chromcorr = ChromCorr(model=self._model, acc='SI')
 
     def adjust_model(self, tunes=None, chroms=None):
         """."""
@@ -174,14 +180,14 @@ def adjust_model(self, tunes=None, chroms=None):
         if chroms is None:
             chroms = self.params.chromx, self.params.chromy
 
-        cav = self.model.cavity_on
-        rad = self.model.radiation_on
-        self.model.cavity_on = False
-        self.model.radiation_on = False
-        self._tunecorr.correct_parameters(tunes, model=self.model)
-        self._chromcorr.correct_parameters(chroms, model=self.model)
-        self.model.cavity_on = cav
-        self.model.radiation_on = rad
+        cav = self._model.cavity_on
+        rad = self._model.radiation_on
+        self._model.cavity_on = False
+        self._model.radiation_on = False
+        self._tunecorr.correct_parameters(tunes, model=self._model)
+        self._chromcorr.correct_parameters(chroms, model=self._model)
+        self._model.cavity_on = cav
+        self._model.radiation_on = rad
 
     def _select_matrix_blocks(self, oeorm):
         n, m = self.params.nr_bpms, self.params.nr_corrs + 1
@@ -199,19 +205,19 @@ def _select_matrix_blocks(self, oeorm):
             oeorm_[n:, -1] = dispy
         return oeorm_.reshape((2 * n, m))
 
-    def _set_energy_offset(self, energy_offset, model=None):
+    def _set_delta_energy_offset(self, energy_offset, model=None):
         """."""
         if model is None:
             model = self.model
-        freq0 = get_rf_frequency(model)
-        alpha = get_mcf(model)
+        freq0 = _get_rf_frequency(model)
+        alpha = _get_mcf(model)
         df = -alpha * freq0 * energy_offset
         self._set_delta_rf_frequency(df, model)
 
     def _set_delta_rf_frequency(self, delta_frequency, model=None):
         if model is None:
             model = self.model
-        freq0 = get_rf_frequency(model)
+        freq0 = _get_rf_frequency(model)
         # print(f'Changing RF freq. by {delta_frequency:.4e}')
         self._set_rf_frequency(freq0 + delta_frequency, model)
 
@@ -224,11 +230,11 @@ def _set_rf_frequency(self, frequency, model=None):
         model[cav_idx].frequency = frequency
 
     def _initialization(self):
-        if self.model is None:
+        if self._model is None:
             print('Cannot start optimization without machine model.')
             return False
 
-        if self.oeorm_goal is None:
+        if self._oeorm_goal is None:
             print('Cannot start optimization without oeorm_goal')
             return False
 
@@ -239,13 +245,6 @@ def _initialization(self):
                 print('No parameters to fit!')
                 return False
             self.params.initial_position = pos
-
-        self.params.limit_lower = _np.full(
-            self.params.initial_position.shape, -_np.inf
-        )
-        self.params.limit_upper = _np.full(
-            self.params.initial_position.shape, _np.inf
-        )
         return True
 
     def _get_errorbars(self, bpms_noise):
@@ -273,7 +272,6 @@ def get_oeorm(
         self,
         model=None,
         strengths=None,
-        idempotent=True,
         delta=1e-2,
         normalize_rf=True,
         ravel=False,
@@ -286,16 +284,18 @@ def get_oeorm(
             strengths0 = self.get_strengths()
             self.set_strengths(strengths)
 
-        rf_freq0 = get_rf_frequency(self.model)
-        self._set_energy_offset(energy_offset=delta, model=self.model)
+        rf_freq0 = _get_rf_frequency(model)
+        self._set_delta_energy_offset(energy_offset=delta, model=model)
+        # TODO: model not being used here!
         mat_pos = self._orbmat.get_respm(add_rfline=True)
 
-        self._set_rf_frequency(frequency=rf_freq0, model=self.model)
+        self._set_rf_frequency(frequency=rf_freq0, model=model)
 
-        self._set_energy_offset(energy_offset=-delta, model=self.model)
+        self._set_delta_energy_offset(energy_offset=-delta, model=model)
+        # TODO: model not being used here!
         mat_neg = self._orbmat.get_respm(add_rfline=True)
 
-        self._set_rf_frequency(frequency=rf_freq0, model=self.model)
+        self._set_rf_frequency(frequency=rf_freq0, model=model)
 
         oeorm = (mat_pos - mat_neg) / 2 / delta
 
@@ -307,38 +307,33 @@ def get_oeorm(
         if ravel:
             oeorm = _np.ravel(oeorm)
 
-        if strengths is not None and idempotent:
+        if strengths is not None:
             self.set_strengths(strengths0)
 
         return oeorm
 
-    def get_chroms(self, strengths=None, idempotent=True):
+    def get_chroms(self, strengths=None):
         """."""
         if strengths is not None:
             strengths0 = self.get_strengths()
             self.set_strengths(strengths)
 
-        chroms = get_chromaticities(self.model)
+        chroms = _get_chromaticities(self.model)
 
-        if strengths is not None and idempotent:
+        if strengths is not None:
             self.set_strengths(strengths0)
 
         return _np.array(chroms)
 
-    def calc_merit_figure(self, strengths=None):
+    def calc_merit_figure(self, pos):
         """."""
-        if strengths is None:
-            pos = self.get_default_pos()
-            strengths, *_ = self.parse_params_from_pos(pos)
-
-        oeorm = self.get_oeorm(
+        strengths, *_ = self.parse_params_from_pos(pos)
+        return self.get_oeorm(
             strengths=strengths,
-            idempotent=True,
             delta=self.params.denergy_oeorm_calc,
             normalize_rf=True,
             ravel=True,
         )
-        return oeorm
 
     def calc_residual(self, merit_figure_meas, merit_figure_goal=None):
         """."""
diff --git a/apsuite/optimization/base.py b/apsuite/optimization/base.py
index 3802d492..6198884f 100644
--- a/apsuite/optimization/base.py
+++ b/apsuite/optimization/base.py
@@ -27,8 +27,8 @@ class OptimizeParams(_Params):
     def __init__(self):
         """."""
         self._initial_position = _np.array([])
-        self._limit_upper = _np.array([])
-        self._limit_lower = _np.array([])
+        self._limit_upper = _np.array([_np.inf])
+        self._limit_lower = _np.array([-_np.inf])
         self.max_number_iters = 100
         self.max_number_evals = 1000
         self._boundary_policy = self.BoundaryPolicy.ToNaN
@@ -295,9 +295,7 @@ def _finalization(self):
         pass
 
     def _objective_func(self, pos):
-        self._num_objective_evals += 1
         pos = self.params.check_and_adjust_boundary(pos)
-        self.positions_evaluated.extend(_np.array(pos, ndmin=2))
         res = []
         for posi in _np.array(pos, ndmin=2):
             if self._stopevt.is_set():
@@ -307,6 +305,8 @@ def _objective_func(self, pos):
                 res.append(_np.nan)
             else:
                 res.append(self.objective_function(posi))
+            self.positions_evaluated.append(posi)
+            self._num_objective_evals += 1
         res = _np.array(res)
         # the objective function must be a (m, n)-array
         # for the m individuals values of the n objectives
diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 8758c107..22c363d9 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -1,5 +1,6 @@
 """."""
 
+import math as _math
 import numpy as _np
 
 from apsuite.optimization.base import Optimize, OptimizeParams
@@ -74,23 +75,21 @@ def calc_residual(self, merit_figure_meas, merit_figure_goal=None):
             res = res / self.params.errorbars
         return res
 
-    def calc_merit_figure(self):
+    def calc_merit_figure(self, pos):
         """."""
+        _ = pos
         raise NotImplementedError(
             'Problem-specific figure of merit needs to be implemented'
         )
 
-    def calc_jacobian(self, pos=None, step=1e-4):
+    def calc_jacobian(self, pos, step=1e-4):
         """."""
         jacobian_t = list()
-        if pos is None:
-            raise ValueError('Specify position to calculate Jacobian.')
         pos0 = pos.copy()
         for i in range(len(pos)):
             pos[i] += step / 2
             figm_pos = self.calc_merit_figure(pos)
-            pos[i] = pos0[i]
-            pos[i] -= step / 2
+            pos[i] -= step
             figm_neg = self.calc_merit_figure(pos)
             pos[i] = pos0[i]
 
@@ -159,24 +158,22 @@ def print_(*args, **kwargs):
                 _np.linalg.pinv(matrix, rcond=self.params.rcond) @ M.T @ res
             )  # LM/GN Normal equation
             # TODO: include processing of pseudo-inverse for singvals selection
-            pos -= lr * delta  # position update
-
-            if _np.any(_np.isnan(pos)):
-                print_('\tNaNs detected in pos. Aborting.')
+            if _np.any(_np.isnan(delta)):
+                print_('\tProblem in step calculation. Aborting.')
                 break
 
-            fails = 0
-            while fails <= patience:
+            pos -= lr * delta  # position update
+            for _ in range(patience):
                 try:
                     res = self._objective_func(pos)
-                    break
+                    if not _np.any(_np.isnan(res)):
+                        break
+                    print_('NaN in objective function. Back-tracking.')
                 except Exception:
                     print_('Merit figure evaluation failed. Back-tracking.')
-                    fails += 1
-                    lr /= lr_factor  # reduce step size
-                    lr = min(lr, lr_min)
-                    pos = pos_init  # restore pos
-                    pos -= lr * delta  # apply smaller step
+                lr /= lr_factor  # reduce step size
+                lr = max(lr, lr_min)
+                pos = pos_init - lr * delta  # apply smaller step
 
             chi2_old = self.history_chi2[-1]
             chi2_new = self.calc_chi2(res)
@@ -184,11 +181,11 @@ def print_(*args, **kwargs):
             #     chi2 += w * _np.std(res) * _np.linalg.norm(pos)
             print_(f'\tchi²: {chi2_new:.6g}')
 
-            delta_chi2 = chi2_old - chi2_new
-            if abs(delta_chi2) < atol or abs(delta_chi2) / chi2_old < rtol:
+            if _math.isclose(chi2_new, chi2_old, rel_tol=rtol, abs_tol=atol):
                 print_('\tConvergence tolerance reached. Exiting.')
                 break
 
+            delta_chi2 = chi2_old - chi2_new
             if damping_constant:
                 if delta_chi2 > 0:
                     damping_constant /= damping_factor

From 73d07bb87fb5bd7520051fa73f47ebef516de630 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Mon, 27 Oct 2025 09:59:40 -0300
Subject: [PATCH 20/40] OPT.LEASTSQUARES.ENH: make `calc_residual`  an explicit
 function of `pos`

- residual evaluation was an implicit function of position. The objective func would evaluate the merit figure at current pos then calculate the residual with `calc_pos`
- for NOECO, the residual evaluation depends not only on the position for calculating the figure of merit, but also for the application of gains to the goal figure of merit
- to avoid workarounds, better make the residual calculation an explicit function of `pos`
---
 apsuite/optimization/least_squares.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 22c363d9..5900a4df 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -55,9 +55,7 @@ def __init__(
 
     def objective_function(self, pos):
         """."""
-        merit_figure = self.calc_merit_figure(pos)
-        residual = self.calc_residual(merit_figure)
-        return residual
+        return self.calc_residual(pos)
 
     def calc_chi2(self, residual):
         """."""
@@ -65,8 +63,10 @@ def calc_chi2(self, residual):
         self.history_chi2.append(chi2)
         return chi2
 
-    def calc_residual(self, merit_figure_meas, merit_figure_goal=None):
+    def calc_residual(self, pos, merit_figure_goal=None):
         """."""
+        merit_figure_meas = self.calc_merit_figure(pos)
+
         if merit_figure_goal is None:
             merit_figure_goal = self.merit_figure_goal
 

From 674f014ac88a9318b67e5015ed61c0c13c2cae12 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Mon, 27 Oct 2025 10:17:42 -0300
Subject: [PATCH 21/40] OPTANLY.NOECO.ENH: improvements in functions default
 args & `calc_residual` w/ explicit `pos` dependence

- default `pos` arg: `calc_merit_figure`, `calc_residual`, `calc_jacobian` and `apply_gains` now accept `None` arg and use default pos, which is the initial pos (machine initial state and gains initial guesses)
- `calc_residual` with explicit `pos` dependence, in accordance with previous commit in `least_squares.py`
- automatic handling of possibly raveled oeorm now inside `apply_gains`
- `default_pos`: default pos is the initial pos (machine init strengths & gains initial guesses)
---
 apsuite/optics_analysis/noeco.py | 59 +++++++++++++++++---------------
 1 file changed, 31 insertions(+), 28 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 4208b7b1..34735955 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -325,8 +325,10 @@ def get_chroms(self, strengths=None):
 
         return _np.array(chroms)
 
-    def calc_merit_figure(self, pos):
+    def calc_merit_figure(self, pos=None):
         """."""
+        if pos is None:
+            pos = self.get_initial_pos()
         strengths, *_ = self.parse_params_from_pos(pos)
         return self.get_oeorm(
             strengths=strengths,
@@ -335,23 +337,22 @@ def calc_merit_figure(self, pos):
             ravel=True,
         )
 
-    def calc_residual(self, merit_figure_meas, merit_figure_goal=None):
+    def calc_residual(self, pos=None, merit_figure_goal=None):
         """."""
+        if pos is None:
+            pos = self.get_initial_pos()
+
         if merit_figure_goal is None:
             merit_figure_goal = self.oeorm_goal
 
-        if merit_figure_goal.ndim == 1:
-            merit_figure_goal = merit_figure_goal.reshape((
-                2 * self.params.nr_bpms,
-                self.params.nr_corrs + 1,
-            ))
+        merit_figure_goal = self.apply_gains(merit_figure_goal, pos).ravel()
 
-        merit_figure_goal = self.apply_gains(merit_figure_goal).ravel()
+        return super().calc_residual(pos, merit_figure_goal)
 
-        return super().calc_residual(merit_figure_meas, merit_figure_goal)
-
-    def calc_jacobian_sexts(self, strengths, step):
+    def calc_jacobian_sexts(self, strengths=None, step=1e-4):
         """."""
+        if strengths is None:
+            strengths, *_ = self.parse_params_from_pos(self.get_default_pos())
         jacobian = super().calc_jacobian(strengths, step)
         self.jacobian_sexts = jacobian
         return jacobian
@@ -361,7 +362,7 @@ def calc_jacobian_gains(self, pos=None):
         if pos is None:
             pos = self.get_default_pos()
 
-        oeorm_meas = self.oeorm_goal
+        oeorm = self.oeorm_goal
         jacobian = None
 
         # corr. gains Jacobian jth col = measured OEORM jth col & zeros
@@ -375,7 +376,7 @@ def calc_jacobian_gains(self, pos=None):
         if self.params.fit_gain_corr:
             # apply BPMs gains and calculate corrs gains jacobian
             oeorm_bpms_gains = self.apply_gains(
-                oeorm_meas, pos, bpms=True, corrs=False
+                oeorm, pos, bpms=True, corrs=False
             )
             jacobian_gains_corrs = block_diag(*oeorm_bpms_gains.T).T
             if self.params.errorbars is not None:
@@ -383,9 +384,7 @@ def calc_jacobian_gains(self, pos=None):
 
             jacobian = jacobian_gains_corrs[:, : self.params.nr_corrs]
 
-        oeorm_corr_gains = self.apply_gains(
-            oeorm_meas, pos, bpms=False, corrs=True
-        )
+        oeorm_corr_gains = self.apply_gains(oeorm, pos, bpms=False, corrs=True)
         # BPM gains Jacobian jth col = measured OEORM jth line, padded with
         # zeros
         # J_gains = [[oeorm_corr_gains[0, :].T, zeros, ...      zeros        ],
@@ -429,7 +428,7 @@ def calc_jacobian_gains(self, pos=None):
 
         return jacobian
 
-    def calc_jacobian(self, pos=None, step=0.0001):
+    def calc_jacobian(self, pos=None, step=1e-4):
         """."""
         if pos is None:
             pos = self.get_default_pos()
@@ -443,7 +442,7 @@ def calc_jacobian(self, pos=None, step=0.0001):
 
             update_sexts = self.params.update_jacobian_sexts
             if self.jacobian_sexts is None or update_sexts:
-                strengths, _, _, _ = self.parse_params_from_pos(pos)
+                strengths, *_ = self.parse_params_from_pos(pos)
                 jacobian_sexts = self.calc_jacobian_sexts(strengths, step)
 
             jacobians.append(jacobian_sexts)
@@ -493,10 +492,18 @@ def apply_gains(self, oeorm, pos=None, bpms=True, corrs=True):
         """."""
         if pos is None:
             pos = self.get_default_pos()
-        _, *ret = self.parse_params_from_pos(pos)
-        gains_corrs, gains_bpms, coup_bpms = ret
 
         oeorm_ = oeorm.copy()
+        ravel = False
+        if oeorm_.ndim == 1:
+            ravel = True
+            oeorm_ = oeorm_.reshape((
+                2 * self.params.nr_bpms,
+                self.params.nr_corrs + 1,
+            ))
+
+        _, *ret = self.parse_params_from_pos(pos)
+        gains_corrs, gains_bpms, coup_bpms = ret
 
         if corrs and gains_corrs is not None:
             oeorm_ = oeorm_ * _np.r_[gains_corrs, [1]]
@@ -505,7 +512,7 @@ def apply_gains(self, oeorm, pos=None, bpms=True, corrs=True):
             Gb = self.bpms_gains_matrix(gains_bpms, coup_bpms)
             oeorm_ = Gb @ oeorm_
 
-        return oeorm_
+        return oeorm_.ravel() if ravel else oeorm_
 
     def bpms_gains_matrix(self, gains, coups):
         """."""
@@ -523,13 +530,9 @@ def bpms_gains_matrix(self, gains, coups):
     def get_default_pos(self):
         """."""
         pos = (
-            self.positions_evaluated[-1]
-            if len(self.positions_evaluated)
-            else (
-                self.params.initial_position
-                if len(self.params.initial_position)
-                else self.get_initial_pos()
-            )
+            self.params.initial_position
+            if len(self.params.initial_position)
+            else self.get_initial_pos()
         )
         return pos
 

From 8b575bce4accefc1a673661aa5bb160f267c48a9 Mon Sep 17 00:00:00 2001
From: Fernando <fernandohds564@gmail.com>
Date: Mon, 27 Oct 2025 14:56:48 -0300
Subject: [PATCH 22/40] OPTANL.WIP: @vellosok75 and @fernandohds564 partial
 revistion of code.

---
 apsuite/optics_analysis/noeco.py | 31 ++++++-------------------------
 1 file changed, 6 insertions(+), 25 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 34735955..d00d0116 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -353,9 +353,7 @@ def calc_jacobian_sexts(self, strengths=None, step=1e-4):
         """."""
         if strengths is None:
             strengths, *_ = self.parse_params_from_pos(self.get_default_pos())
-        jacobian = super().calc_jacobian(strengths, step)
-        self.jacobian_sexts = jacobian
-        return jacobian
+        return super().calc_jacobian(strengths, step)
 
     def calc_jacobian_gains(self, pos=None):
         """."""
@@ -365,20 +363,14 @@ def calc_jacobian_gains(self, pos=None):
         oeorm = self.oeorm_goal
         jacobian = None
 
-        # corr. gains Jacobian jth col = measured OEORM jth col & zeros
-        # J_gains = [[oeorm_bpm_gains[:, 0].T, zeros, ...      zeros        ],
-        #            [zeros, oeorm_bpm_gains[:, 1].T, ...      zeros        ],
-        #            [zeros, zeros, oeorm_bpm_gains[:, 2].T    zeros        ],
-        #            [  .      .            .                     .         ],
-        #            [  .      .            .                     .         ],
-        #            [zeros,  zeros,       ...    oeorm_bpm_gains[:, 280].T]]
-        # this is implemented below
         if self.params.fit_gain_corr:
             # apply BPMs gains and calculate corrs gains jacobian
             oeorm_bpms_gains = self.apply_gains(
                 oeorm, pos, bpms=True, corrs=False
             )
-            jacobian_gains_corrs = block_diag(*oeorm_bpms_gains.T).T
+            jacobian_gains_corrs = _np.vstack([
+                block_diag(m) for m in oeorm_bpms_gains
+            ])
             if self.params.errorbars is not None:
                 jacobian_gains_corrs /= self.params.errorbars[:, None]
 
@@ -424,8 +416,6 @@ def calc_jacobian_gains(self, pos=None):
             else:
                 jacobian = _np.hstack((jacobian, jacobian_coup_bpms))
 
-        self.jacobian_gains = jacobian
-
         return jacobian
 
     def calc_jacobian(self, pos=None, step=1e-4):
@@ -437,9 +427,6 @@ def calc_jacobian(self, pos=None, step=1e-4):
 
         # --- sexts ---
         if self.params.fit_sexts:
-            if self.jacobian_sexts is not None:
-                jacobian_sexts = self.jacobian_sexts
-
             update_sexts = self.params.update_jacobian_sexts
             if self.jacobian_sexts is None or update_sexts:
                 strengths, *_ = self.parse_params_from_pos(pos)
@@ -455,21 +442,15 @@ def calc_jacobian(self, pos=None, step=1e-4):
         )
 
         if fit_gains:
-            if self.jacobian_gains is not None:
-                jacobian_gains = self.jacobian_gains
-
             update_gains = self.params.update_jacobian_gains
             if self.jacobian_gains is None or update_gains:
                 jacobian_gains = self.calc_jacobian_gains(pos)
 
             jacobians.append(jacobian_gains)
 
+        jacobian = None
         if jacobians:
             jacobian = _np.hstack(jacobians)
-        else:
-            jacobian = None
-
-        self.jacobian = jacobian
 
         return jacobian
 
@@ -506,7 +487,7 @@ def apply_gains(self, oeorm, pos=None, bpms=True, corrs=True):
         gains_corrs, gains_bpms, coup_bpms = ret
 
         if corrs and gains_corrs is not None:
-            oeorm_ = oeorm_ * _np.r_[gains_corrs, [1]]
+            oeorm_ = oeorm_ * _np.r_[gains_corrs, 1]
 
         if bpms and (gains_bpms is not None or coup_bpms is not None):
             Gb = self.bpms_gains_matrix(gains_bpms, coup_bpms)

From 768e680a77e07a40b0e15d29322f051ce0d42278 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Mon, 27 Oct 2025 16:11:37 -0300
Subject: [PATCH 23/40] OPTANLY.NOECO: fix corr. gains jacobian & restore
 temporary logic for constructing full jacobians

---
 apsuite/optics_analysis/noeco.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index d00d0116..cee4078e 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -369,7 +369,7 @@ def calc_jacobian_gains(self, pos=None):
                 oeorm, pos, bpms=True, corrs=False
             )
             jacobian_gains_corrs = _np.vstack([
-                block_diag(m) for m in oeorm_bpms_gains
+                _np.diag(m) for m in oeorm_bpms_gains
             ])
             if self.params.errorbars is not None:
                 jacobian_gains_corrs /= self.params.errorbars[:, None]
@@ -427,6 +427,9 @@ def calc_jacobian(self, pos=None, step=1e-4):
 
         # --- sexts ---
         if self.params.fit_sexts:
+            if self.jacobian_sexts is not None:
+                jacobian_sexts = self.jacobian_sexts
+
             update_sexts = self.params.update_jacobian_sexts
             if self.jacobian_sexts is None or update_sexts:
                 strengths, *_ = self.parse_params_from_pos(pos)
@@ -442,6 +445,9 @@ def calc_jacobian(self, pos=None, step=1e-4):
         )
 
         if fit_gains:
+            if self.jacobian_gains is not None:
+                jacobian_gains = self.jacobian_gains
+
             update_gains = self.params.update_jacobian_gains
             if self.jacobian_gains is None or update_gains:
                 jacobian_gains = self.calc_jacobian_gains(pos)

From 11f01a16ec45dedb70466b8582ae076cc71e3e2c Mon Sep 17 00:00:00 2001
From: Fernando <fernandohds564@gmail.com>
Date: Mon, 27 Oct 2025 17:19:45 -0300
Subject: [PATCH 24/40] OPTANL.WIP: @vellosok75 and @fernandohds564 partial
 revistion of code.

---
 apsuite/optics_analysis/noeco.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index cee4078e..88f54ade 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -416,7 +416,7 @@ def calc_jacobian_gains(self, pos=None):
             else:
                 jacobian = _np.hstack((jacobian, jacobian_coup_bpms))
 
-        return jacobian
+        return -jacobian
 
     def calc_jacobian(self, pos=None, step=1e-4):
         """."""

From e65a0051117c9ff28db3c6c9fcfa9e160ac2f3e0 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 30 Oct 2025 15:43:45 -0300
Subject: [PATCH 25/40] OPTANLY.NOECO.ENH: OEORM blocks selection

- enables selecting which OEORM blocks (diag, off-diag, dispersions) will be included for figure of merit calculations
- does so by keeping track of `oeorm_mask`, which applies to the OEORM or its ravelled vector and the jacobian columns for filtering entries of interest
---
 apsuite/optics_analysis/noeco.py | 140 ++++++++++++++++++++++++-------
 1 file changed, 111 insertions(+), 29 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 88f54ade..3a6db7e8 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -50,7 +50,6 @@ class NOECOParams(LeastSquaresParams):
 
     def __init__(self):
         """."""
-        super().__init__()
         self.denergy_oeorm_calc = 1e-2  # 1 %
         self.tunex = 49.16
         self.tuney = 14.22
@@ -63,10 +62,11 @@ def __init__(self):
         self.fit_coup_bpms = True
         self.fit_gain_corr = True
 
-        self.use_diag_blocks = True
-        self.use_hor_disp = True
-        self.use_ver_disp = False
-        self.use_offdiag_blocks = False
+        self._use_diag_blocks = None
+        self._use_hor_disp = None
+        self._use_ver_disp = None
+        self._use_offdiag_blocks = None
+        self.oeorm_mask = None
 
         self.update_jacobian_sexts = False
         self.update_jacobian_gains = True
@@ -79,12 +79,89 @@ def __init__(self):
         self.init_gain_corr = _np.ones(self.nr_corrs)
         self.init_gain_bpms = _np.ones(2 * self.nr_bpms)  # Gx, Gy
         self.init_coup_bpms = _np.zeros(2 * self.nr_bpms)  # Cc, Cy
+        super().__init__()
+
+        self.use_diag_blocks = True
+        self.use_hor_disp = True
+        self.use_offdiag_blocks = False
+        self.use_ver_disp = False
 
     @property
     def nr_sexts(self):
         """."""
         return len(self.sextfams2fit)
 
+    @property
+    def use_diag_blocks(self):
+        """."""
+        return self._use_diag_blocks
+
+    @use_diag_blocks.setter
+    def use_diag_blocks(self, value):
+        """."""
+        self._use_diag_blocks = value
+        self._update_oeorm_mask()
+
+    @property
+    def use_offdiag_blocks(self):
+        """."""
+        return self._use_offdiag_blocks
+
+    @use_offdiag_blocks.setter
+    def use_offdiag_blocks(self, value):
+        """."""
+        self._use_offdiag_blocks = value
+        if not value:
+            self.fit_coup_bpms = value  # don't fit coupling
+        self._update_oeorm_mask()
+
+    @property
+    def use_hor_disp(self):
+        """."""
+        return self._use_hor_disp
+
+    @use_hor_disp.setter
+    def use_hor_disp(self, value):
+        """."""
+        self._use_hor_disp = value
+        self._update_oeorm_mask()
+
+    @property
+    def use_ver_disp(self):
+        """."""
+        return self._use_ver_disp
+
+    @use_ver_disp.setter
+    def use_ver_disp(self, value):
+        """."""
+        self._use_ver_disp = value
+        self._update_oeorm_mask()
+
+    def _update_oeorm_mask(self):
+        """."""
+        self.oeorm_mask = self._get_oeorm_mask()
+
+    def _get_oeorm_mask(self):
+        """."""
+        n, m = self.nr_bpms, self.nr_corrs + 1
+        mask = _np.zeros((2 * n, m), dtype=bool)
+
+        if self.use_diag_blocks:
+            mask[:n, : self.nr_chs] = True
+            mask[n:, self.nr_chs : self.nr_corrs] = True
+
+        if self.use_offdiag_blocks:
+            mask[:n, self.nr_chs : self.nr_corrs] = True
+            mask[n:, : self.nr_chs] = True
+
+        if self.use_hor_disp:
+            mask[:n, -1] = True
+
+        if self.use_ver_disp:
+            mask[n:, -1] = True
+
+        return mask
+
 
 class NOECOFit(LeastSquaresOptimize):
     """."""
@@ -141,9 +218,17 @@ def oeorm_goal(self, oeorm):
         """."""
         if oeorm is None:
             raise ValueError('oeorm_goal must not be None.')
-        oeorm_ = self._select_matrix_blocks(oeorm)
-        self._oeorm_goal = oeorm_
-        self.merit_figure_goal = oeorm_.ravel()
+        self._oeorm_goal = oeorm
+
+    @property
+    def merit_figure_goal(self):
+        """."""
+        return self.oeorm_goal.ravel()[self.params.oeorm_mask.ravel()]
+
+    @merit_figure_goal.setter
+    def merit_figure_goal(self, merit_figure):
+        """."""
+        pass  # compatibility only
 
     @property
     def bpms_noise(self):
@@ -189,23 +274,7 @@ def adjust_model(self, tunes=None, chroms=None):
         self._model.cavity_on = cav
         self._model.radiation_on = rad
 
-    def _select_matrix_blocks(self, oeorm):
-        n, m = self.params.nr_bpms, self.params.nr_corrs + 1
-        xx, xy, dispx, yx, yy, dispy = self.get_oeorm_blocks(oeorm)
-        oeorm_ = _np.zeros((2 * n, m))
-        if self.params.use_diag_blocks:
-            oeorm_[:n, : self.params.nr_chs] = xx
-            oeorm_[n:, self.params.nr_chs : self.params.nr_corrs] = yy
-        if self.params.use_offdiag_blocks:
-            oeorm_[:n, self.params.nr_chs : self.params.nr_corrs] = xy
-            oeorm_[n:, : self.params.nr_chs] = yx
-        if self.params.use_hor_disp:
-            oeorm_[:n, -1] = dispx
-        if self.params.use_ver_disp:
-            oeorm_[n:, -1] = dispy
-        return oeorm_.reshape((2 * n, m))
-
-    def _set_delta_energy_offset(self, energy_offset, model=None):
+        def _set_delta_energy_offset(self, energy_offset, model=None):
         """."""
         if model is None:
             model = self.model
@@ -275,6 +344,7 @@ def get_oeorm(
         delta=1e-2,
         normalize_rf=True,
         ravel=False,
+        filter_blocks=False,
     ):
         """."""
         if model is None:
@@ -302,7 +372,8 @@ def get_oeorm(
         if normalize_rf:
             oeorm[:, -1] *= 1e6  # [um/Hz]
 
-        oeorm = self._select_matrix_blocks(oeorm)
+        if filter_blocks:
+            oeorm = oeorm[self.params.oeorm_mask]
 
         if ravel:
             oeorm = _np.ravel(oeorm)
@@ -335,6 +406,7 @@ def calc_merit_figure(self, pos=None):
             delta=self.params.denergy_oeorm_calc,
             normalize_rf=True,
             ravel=True,
+            filter_blocks=True,
         )
 
     def calc_residual(self, pos=None, merit_figure_goal=None):
@@ -345,12 +417,16 @@ def calc_residual(self, pos=None, merit_figure_goal=None):
         if merit_figure_goal is None:
             merit_figure_goal = self.oeorm_goal
 
-        merit_figure_goal = self.apply_gains(merit_figure_goal, pos).ravel()
+        merit_figure_goal = self.apply_gains(merit_figure_goal, pos)
+        merit_figure_goal = merit_figure_goal[self.params.oeorm_mask]
 
         return super().calc_residual(pos, merit_figure_goal)
 
     def calc_jacobian_sexts(self, strengths=None, step=1e-4):
         """."""
+        # TODO: bug when calculating sextupole jacobian
+        # with errorbars & block selections
+        # array shape issue
         if strengths is None:
             strengths, *_ = self.parse_params_from_pos(self.get_default_pos())
         return super().calc_jacobian(strengths, step)
@@ -374,7 +450,9 @@ def calc_jacobian_gains(self, pos=None):
             if self.params.errorbars is not None:
                 jacobian_gains_corrs /= self.params.errorbars[:, None]
 
-            jacobian = jacobian_gains_corrs[:, : self.params.nr_corrs]
+            jacobian = jacobian_gains_corrs[
+                self.params.oeorm_mask.ravel(), : self.params.nr_corrs
+            ]
 
         oeorm_corr_gains = self.apply_gains(oeorm, pos, bpms=False, corrs=True)
         # BPM gains Jacobian jth col = measured OEORM jth line, padded with
@@ -391,6 +469,10 @@ def calc_jacobian_gains(self, pos=None):
             if self.params.errorbars is not None:
                 jacobian_gains_bpms /= self.params.errorbars[:, None]
 
+            jacobian_gains_bpms = jacobian_gains_bpms[
+                self.params.oeorm_mask.ravel()
+            ]
+
             if jacobian is None:
                 jacobian = jacobian_gains_bpms
             else:
@@ -416,7 +498,7 @@ def calc_jacobian_gains(self, pos=None):
             else:
                 jacobian = _np.hstack((jacobian, jacobian_coup_bpms))
 
-        return -jacobian
+        return -jacobian  # oposite sign convention for numeric & analytic jacs
 
     def calc_jacobian(self, pos=None, step=1e-4):
         """."""

From efd7fb1180915fc7dd0fc414055105c489c60a04 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 30 Oct 2025 15:51:14 -0300
Subject: [PATCH 26/40] OPTANLY.NOECO.FEAT: new plotting features

- method for calculating chromatic optics
- matrix diffs visualization
- plot chromatic optics
- allows plot comparison of strenghts with respect to initial strength
---
 apsuite/optics_analysis/noeco.py | 162 ++++++++++++++++++++++++++++++-
 1 file changed, 158 insertions(+), 4 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 3a6db7e8..ae8a43cd 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -274,7 +274,7 @@ def adjust_model(self, tunes=None, chroms=None):
         self._model.cavity_on = cav
         self._model.radiation_on = rad
 
-        def _set_delta_energy_offset(self, energy_offset, model=None):
+    def _set_delta_energy_offset(self, energy_offset, model=None):
         """."""
         if model is None:
             model = self.model
@@ -396,6 +396,43 @@ def get_chroms(self, strengths=None):
 
         return _np.array(chroms)
 
+    def get_chrom_optics(self, strengths=None):
+        """."""
+        # NOTE: get this to pyaccel maybe?
+        if strengths is not None:
+            strengths0 = self.get_strengths()
+            self.set_strengths(strengths)
+
+        model = self.model
+
+        cav = model.cavity_on
+        rad = model.radiation_on
+        model.cavity_on = False
+        model.radiation_on = False
+        bpms_idcs = _pa.lattice.flatten(self.famdata['BPM']['index'])
+
+        delta = 1e-3
+        twissp, *_ = _pa.optics.calc_twiss(
+            model, indices=bpms_idcs, energy_offset=delta / 2
+        )
+        twissn, *_ = _pa.optics.calc_twiss(
+            model, indices=bpms_idcs, energy_offset=-delta / 2
+        )
+
+        model.cavity_on = cav
+        model.radiation_on = rad
+
+        betax = (twissp.betax - twissn.betax) / delta
+        betay = (twissp.betay - twissn.betay) / delta
+        mux = (twissp.mux - twissn.mux) / delta
+        muy = (twissp.muy - twissn.muy) / delta
+        etax = (twissp.etax - twissn.etax) / delta
+
+        if strengths is not None:
+            self.set_strengths(strengths0)
+
+        return betax, betay, mux, muy, etax
+
     def calc_merit_figure(self, pos=None):
         """."""
         if pos is None:
@@ -638,21 +675,138 @@ def parse_params_from_pos(self, pos):
             coup_bpms = pos[n : n + 2 * self.params.nr_bpms]
         return sexts, gains_corrs, gains_bpms, coup_bpms
 
-    def plot_strengths(self, strengths=None, fig=None, ax=None, label=None):
+    def plot_strengths(
+        self,
+        strengths=None,
+        fig=None,
+        ax=None,
+        label=None,
+        diff_from_init=False,
+        percentual_diff=True,
+    ):
         """."""
         if fig is None or ax is None:
             fig, ax = _mplt.subplots(figsize=(12, 4))
         if strengths is None:
             strengths = self.get_strengths()
 
-        ax.plot(strengths, 'o-', mfc='none', label=label)
+        strengths_ = strengths.copy()
+
+        if diff_from_init:
+            str0, *_ = self.parse_params_from_pos(self.params.initial_position)
+            strengths_ -= str0
+            if percentual_diff:
+                strengths_ /= str0 / 100
+
+        ax.plot(strengths_, 'o-', mfc='none', label=label)
         tick_label = self.params.sextfams2fit
         tick_pos = list(range(len(tick_label)))
         ax.set_xticks(tick_pos, labels=tick_label, rotation=45)
-        ax.set_ylabel(r'SL [m$^{-2}$]')
+        line_label = 'SL'
+        if diff_from_init:
+            line_label = "delta " + line_label
+        if diff_from_init and percentual_diff:
+            line_label = "percentual " + line_label + ' [\%]'
+        else:
+            line_label = line_label + r' [m$^{-2}$]'
+        ax.set_ylabel(line_label)
         ax.set_xlabel('chromatic sextupole families')
         if label:
             ax.legend()
         fig.tight_layout()
         fig.show()
         return fig, ax
+
+    def plot_matrix_diff(self, mat1=None, mat2=None, residual=None, title=''):
+        """."""
+        if mat1 is not None and mat2 is not None:
+            residual = _np.abs(mat1 - mat2)
+
+        if residual.ndim == 1:
+            residual.reshape(2 * self.params.nr_bpms, self.params.nr_corrs + 1)
+
+        x = _np.arange(residual.shape[1])  # colunas
+        y = _np.arange(residual.shape[0])  # linhas
+        X, Y = _np.meshgrid(x, y)
+
+        fig = _mplt.figure()
+        fig.suptitle(title)
+        ax = fig.add_subplot(111, projection='3d')
+        ax.plot_surface(X, Y, _np.abs(residual), cmap='viridis')
+
+        ax.set_xlabel('corrs. index')
+        ax.set_ylabel('BPMs index')
+        ax.set_zlabel('OEORM abs. diff [m/rad]')
+
+        _mplt.show()
+
+        return fig, ax
+
+    def plot_chrom_optics(
+        self,
+        betax=None,
+        betay=None,
+        mux=None,
+        muy=None,
+        etax=None,
+        fig=None,
+        axes=None,
+        title='',
+        diff_from_init=False,
+    ):
+        """."""
+        if fig is None:
+            fig, axes = _mplt.subplots(3, 1, sharex=True, figsize=(12, 10))
+
+        if (betax, betay, mux, muy, etax) == (None, None, None, None, None):
+            betax, betay, mux, muy, etax = self.get_chrom_optics()
+
+        betax_ = betax.copy()
+        betay_ = betay.copy()
+        mux_ = mux.copy()
+        muy_ = muy.copy()
+        etax_ = etax.copy()
+
+        if diff_from_init:
+            # plot deviation from initial strengths optics
+            str0, *_ = self.parse_params_from_pos(self.params.initial_position)
+            betax0, betay0, mux0, muy0, etax0 = self.get_chrom_optics(str0)
+            betax_ -= betax0
+            betay_ -= betay0
+            mux_ -= mux0
+            muy_ -= muy0
+            etax_ -= etax0
+
+        ax0 = axes[0]
+        ax0.set_title('chromatic beta')
+        ax0.set_ylabel(r'$\beta_x^{(1)}$ [m]', color='blue')
+        ax0.plot(betax_, color='blue', label='hor')
+        ax0.tick_params(axis='y', labelcolor='blue')
+
+        ax0tw = ax0.twinx()
+        ax0tw.set_ylabel(r'$\beta_y^{(1)}$ [m]', color='red')
+        ax0tw.plot(betay_, color='red', label=r'$\beta_y$')
+        ax0tw.tick_params(axis='y', labelcolor='red')
+
+        axes[1].set_title('2nd-order dispersion')
+        axes[1].set_ylabel(r'$\eta^{(2)}$ [cm]')
+        axes[1].plot(etax_ * 100, color='green')
+
+        ax1 = axes[2]
+        ax1.set_title('chromatic phase advance')
+        ax1.set_xlabel('BPMs indices')
+        ax1.set_ylabel(r'$\phi_x^{(1)}$ [rad]', color='blue')
+        ax1.plot(mux_, color='blue', label=r'$\beta_y$')
+        ax1.tick_params(axis='y', labelcolor='blue')
+
+        ax1tw = ax1.twinx()
+        ax1tw.set_ylabel(r'$\phi_y^{(1)}$ [rad]', color='red')
+        ax1tw.plot(muy_, color='red', label=r'$\mu_y$')
+        ax1tw.tick_params(axis='y', labelcolor='red')
+
+        if title:
+            fig.suptitle(title)
+        fig.tight_layout()
+        fig.show()
+
+        return fig, axes

From ad7fd1c99df8b414b4a7fb72ccac89b8301e46db Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Tue, 11 Nov 2025 11:28:48 -0300
Subject: [PATCH 27/40] OPTANLY.NOECO: include chromaticity fitting

- fitting to measured OEORMs is failing to reproduce vertical chromaticity
- including chromaticities to the fitting with adequate weights (200 for fitting w/o off-diag blocks) rendered more well-behaved solutions
---
 apsuite/optics_analysis/noeco.py | 53 +++++++++++++++++++++++++++-----
 1 file changed, 45 insertions(+), 8 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index ae8a43cd..2b99b727 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -54,7 +54,7 @@ def __init__(self):
         self.tunex = 49.16
         self.tuney = 14.22
         self.chromx = 3.11
-        self.chromy = 2.5
+        self.chromy = 2.66
 
         self.sextfams2fit = self.SEXT_FAMS_CHROM
         self.fit_sexts = True
@@ -79,12 +79,14 @@ def __init__(self):
         self.init_gain_corr = _np.ones(self.nr_corrs)
         self.init_gain_bpms = _np.ones(2 * self.nr_bpms)  # Gx, Gy
         self.init_coup_bpms = _np.zeros(2 * self.nr_bpms)  # Cc, Cy
+        self.chrom_weights = _np.array([1, 1])
         super().__init__()
 
         self.use_diag_blocks = True
-        self.use_hor_disp = True
         self.use_offdiag_blocks = False
+        self.use_hor_disp = True
         self.use_ver_disp = False
+        self.use_chroms = False
 
     @property
     def nr_sexts(self):
@@ -173,6 +175,7 @@ def __init__(
         params = NOECOParams()
 
         self._oeorm_goal = None
+        self._chroms_goal = _np.array([params.chromx, params.chromy])
         self._bpms_noise = None
         self._model = None
         self.famdata = None
@@ -220,10 +223,26 @@ def oeorm_goal(self, oeorm):
             raise ValueError('oeorm_goal must not be None.')
         self._oeorm_goal = oeorm
 
+    @property
+    def chroms_goal(self):
+        """."""
+        return self._chroms_goal
+
+    @chroms_goal.setter
+    def chroms_goal(self, chroms):
+        """."""
+        if chroms is None:
+            raise ValueError('chroms_goal must not be None.')
+        self._chroms_goal = chroms
+
     @property
     def merit_figure_goal(self):
         """."""
-        return self.oeorm_goal.ravel()[self.params.oeorm_mask.ravel()]
+        figgoal = self.oeorm_goal.ravel()[self.params.oeorm_mask.ravel()]
+        if self.params.use_chroms:
+            chroms = (self.params.chrom_weights * self.chroms_goal).ravel()
+            figgoal = _np.concatenate((figgoal, chroms))
+        return figgoal
 
     @merit_figure_goal.setter
     def merit_figure_goal(self, merit_figure):
@@ -438,13 +457,19 @@ def calc_merit_figure(self, pos=None):
         if pos is None:
             pos = self.get_initial_pos()
         strengths, *_ = self.parse_params_from_pos(pos)
-        return self.get_oeorm(
+        merit_fig = self.get_oeorm(
             strengths=strengths,
             delta=self.params.denergy_oeorm_calc,
             normalize_rf=True,
             ravel=True,
             filter_blocks=True,
         )
+        if self.params.use_chroms:
+            chroms = self.get_chroms(strengths)
+            chroms *= self.params.chrom_weights
+            merit_fig = _np.concatenate((merit_fig, chroms.ravel()))
+
+        return merit_fig
 
     def calc_residual(self, pos=None, merit_figure_goal=None):
         """."""
@@ -457,6 +482,12 @@ def calc_residual(self, pos=None, merit_figure_goal=None):
         merit_figure_goal = self.apply_gains(merit_figure_goal, pos)
         merit_figure_goal = merit_figure_goal[self.params.oeorm_mask]
 
+        if self.params.use_chroms:
+            merit_figure_goal = _np.concatenate((
+                merit_figure_goal,
+                (self.params.chrom_weights * self.chroms_goal).ravel(),
+            ))
+
         return super().calc_residual(pos, merit_figure_goal)
 
     def calc_jacobian_sexts(self, strengths=None, step=1e-4):
@@ -523,8 +554,8 @@ def calc_jacobian_gains(self, pos=None):
         if self.params.fit_coup_bpms:
             jacobian_coup_bpms = block_diag(
                 *_np.vstack((
-                    oeorm_corr_gains[self.params.nr_bpms :],
-                    oeorm_corr_gains[: self.params.nr_bpms],
+                    oeorm_corr_gains[self.params.nr_bpms :],  #TODO: coups
+                    oeorm_corr_gains[: self.params.nr_bpms],  # and no disps
                 ))
             ).T
             if self.params.errorbars is not None:
@@ -535,6 +566,12 @@ def calc_jacobian_gains(self, pos=None):
             else:
                 jacobian = _np.hstack((jacobian, jacobian_coup_bpms))
 
+        if self.params.use_chroms:
+            jacobian = _np.vstack((
+                jacobian,
+                _np.zeros((2, jacobian.shape[-1])),
+            ))
+
         return -jacobian  # oposite sign convention for numeric & analytic jacs
 
     def calc_jacobian(self, pos=None, step=1e-4):
@@ -704,9 +741,9 @@ def plot_strengths(
         ax.set_xticks(tick_pos, labels=tick_label, rotation=45)
         line_label = 'SL'
         if diff_from_init:
-            line_label = "delta " + line_label
+            line_label = 'delta ' + line_label
         if diff_from_init and percentual_diff:
-            line_label = "percentual " + line_label + ' [\%]'
+            line_label = 'percentual ' + line_label + ' [\%]'
         else:
             line_label = line_label + r' [m$^{-2}$]'
         ax.set_ylabel(line_label)

From 8b0a40c41bdcb4c74430fc260a5e9de1febeda38 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Wed, 19 Nov 2025 16:58:02 -0300
Subject: [PATCH 28/40] OPT.ANLY.ENH: include dipersion column weight & fix
 errorbar setting

- introduces  `disp_weight` for weighting the dispersion column
- experimental noise: now supports natrix with  noise for  individual OEORM entries, as well as a single BPM noise vector used for a given line
- fix errorbar: errorbars were not being conformed to the `oeorm_mask` object, so the selection of diag, off-diag, hor disp and vert disp noise components was broken
- handling of merit figure shape outside `apply_gains` function, with new method `reshape_merit_figure`
---
 apsuite/optics_analysis/noeco.py | 64 +++++++++++++++++++-------------
 1 file changed, 38 insertions(+), 26 deletions(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 2b99b727..139fde3a 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -80,6 +80,7 @@ def __init__(self):
         self.init_gain_bpms = _np.ones(2 * self.nr_bpms)  # Gx, Gy
         self.init_coup_bpms = _np.zeros(2 * self.nr_bpms)  # Cc, Cy
         self.chrom_weights = _np.array([1, 1])
+        self.disp_weight = 1
         super().__init__()
 
         self.use_diag_blocks = True
@@ -260,7 +261,10 @@ def bpms_noise(self, bpms_noise):
         if bpms_noise is None:
             return
         self._bpms_noise = bpms_noise
-        self.params.errorbars = self._get_errorbars(bpms_noise)
+        if bpms_noise.squeeze().ndim == 1:
+            self.params.errorbars = self._get_errorbars(bpms_noise)
+        else:
+            self.params.errorbars = bpms_noise[self.params.oeorm_mask].ravel()
 
     def create_model(self):
         """."""
@@ -340,7 +344,7 @@ def _get_errorbars(self, bpms_noise):
         errorbar = _np.ones(2 * n * m)
         for i, sigma in enumerate(bpms_noise, start=0):
             errorbar[i * m : (i + 1) * m] = _np.repeat(sigma, m)
-        return errorbar
+        return errorbar[self.params.oeorm_mask.ravel()]
 
     def get_strengths(self):
         """."""
@@ -363,7 +367,6 @@ def get_oeorm(
         delta=1e-2,
         normalize_rf=True,
         ravel=False,
-        filter_blocks=False,
     ):
         """."""
         if model is None:
@@ -391,9 +394,6 @@ def get_oeorm(
         if normalize_rf:
             oeorm[:, -1] *= 1e6  # [um/Hz]
 
-        if filter_blocks:
-            oeorm = oeorm[self.params.oeorm_mask]
-
         if ravel:
             oeorm = _np.ravel(oeorm)
 
@@ -461,9 +461,13 @@ def calc_merit_figure(self, pos=None):
             strengths=strengths,
             delta=self.params.denergy_oeorm_calc,
             normalize_rf=True,
-            ravel=True,
-            filter_blocks=True,
+            ravel=False,
         )
+
+        if self.params.disp_weight:
+            merit_fig[:, -1] *= self.params.disp_weight
+        merit_fig = merit_fig[self.params.oeorm_mask].ravel()
+
         if self.params.use_chroms:
             chroms = self.get_chroms(strengths)
             chroms *= self.params.chrom_weights
@@ -479,8 +483,15 @@ def calc_residual(self, pos=None, merit_figure_goal=None):
         if merit_figure_goal is None:
             merit_figure_goal = self.oeorm_goal
 
+        if merit_figure_goal.ndim == 1:
+            merit_figure_goal = self.reshape_merit_fig(merit_figure_goal)
+
         merit_figure_goal = self.apply_gains(merit_figure_goal, pos)
-        merit_figure_goal = merit_figure_goal[self.params.oeorm_mask]
+
+        if self.params.disp_weight:
+            merit_figure_goal[:, -1] *= self.params.disp_weight
+
+        merit_figure_goal = merit_figure_goal[self.params.oeorm_mask].ravel()
 
         if self.params.use_chroms:
             merit_figure_goal = _np.concatenate((
@@ -490,6 +501,14 @@ def calc_residual(self, pos=None, merit_figure_goal=None):
 
         return super().calc_residual(pos, merit_figure_goal)
 
+    def reshape_merit_fig(self, merit_fig):
+        """."""
+        merit_fig = merit_fig.reshape((
+            2 * self.params.nr_bpms,
+            self.params.nr_corrs + 1,
+        ))
+        return merit_fig
+
     def calc_jacobian_sexts(self, strengths=None, step=1e-4):
         """."""
         # TODO: bug when calculating sextupole jacobian
@@ -514,13 +533,11 @@ def calc_jacobian_gains(self, pos=None):
             )
             jacobian_gains_corrs = _np.vstack([
                 _np.diag(m) for m in oeorm_bpms_gains
-            ])
+            ])[self.params.oeorm_mask.ravel()]
+
             if self.params.errorbars is not None:
                 jacobian_gains_corrs /= self.params.errorbars[:, None]
-
-            jacobian = jacobian_gains_corrs[
-                self.params.oeorm_mask.ravel(), : self.params.nr_corrs
-            ]
+            jacobian = jacobian_gains_corrs[:, : self.params.nr_corrs]
 
         oeorm_corr_gains = self.apply_gains(oeorm, pos, bpms=False, corrs=True)
         # BPM gains Jacobian jth col = measured OEORM jth line, padded with
@@ -534,12 +551,11 @@ def calc_jacobian_gains(self, pos=None):
         # this is implemented below
         if self.params.fit_gain_bpms:
             jacobian_gains_bpms = block_diag(*oeorm_corr_gains).T
-            if self.params.errorbars is not None:
-                jacobian_gains_bpms /= self.params.errorbars[:, None]
-
             jacobian_gains_bpms = jacobian_gains_bpms[
                 self.params.oeorm_mask.ravel()
             ]
+            if self.params.errorbars is not None:
+                jacobian_gains_bpms /= self.params.errorbars[:, None]
 
             if jacobian is None:
                 jacobian = jacobian_gains_bpms
@@ -554,10 +570,13 @@ def calc_jacobian_gains(self, pos=None):
         if self.params.fit_coup_bpms:
             jacobian_coup_bpms = block_diag(
                 *_np.vstack((
-                    oeorm_corr_gains[self.params.nr_bpms :],  #TODO: coups
+                    oeorm_corr_gains[self.params.nr_bpms :],  # TODO: coups
                     oeorm_corr_gains[: self.params.nr_bpms],  # and no disps
                 ))
             ).T
+            jacobian_coup_bpms = jacobian_gains_corrs[
+                self.params.oeorm_mask.ravel()
+            ]
             if self.params.errorbars is not None:
                 jacobian_coup_bpms /= self.params.errorbars[:, None]
 
@@ -637,13 +656,6 @@ def apply_gains(self, oeorm, pos=None, bpms=True, corrs=True):
             pos = self.get_default_pos()
 
         oeorm_ = oeorm.copy()
-        ravel = False
-        if oeorm_.ndim == 1:
-            ravel = True
-            oeorm_ = oeorm_.reshape((
-                2 * self.params.nr_bpms,
-                self.params.nr_corrs + 1,
-            ))
 
         _, *ret = self.parse_params_from_pos(pos)
         gains_corrs, gains_bpms, coup_bpms = ret
@@ -655,7 +667,7 @@ def apply_gains(self, oeorm, pos=None, bpms=True, corrs=True):
             Gb = self.bpms_gains_matrix(gains_bpms, coup_bpms)
             oeorm_ = Gb @ oeorm_
 
-        return oeorm_.ravel() if ravel else oeorm_
+        return oeorm_
 
     def bpms_gains_matrix(self, gains, coups):
         """."""

From 93d6600555607c783dbde9f73b4245e01d482f6d Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 15 Jan 2026 16:51:22 -0300
Subject: [PATCH 29/40] OPT.LEASTSQUARES.ENH: include tentative steps to LM
 loop

- these changes introduce tentative steps to the LM loop
- prior to them, steps were always accepted, whether they improved or worsened the chi2.  Changes in the damping constan to control steps were applied only to the subsequent step
- now the steps are tested prior to acceptance, and are accepted only if improve the objective
- this change removes the need for back-tracking when the objective function evaluation fails
-  steps control now concentrates in the damping constant of the LM loop, with no need for a learning_rate constant
---
 apsuite/optimization/least_squares.py | 106 +++++++++++---------------
 1 file changed, 44 insertions(+), 62 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 5900a4df..08394774 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -20,14 +20,11 @@ def __init__(self):
         self.rcond = 1e-15
         self.damping_constant = 1.0
         self.damping_factor = 10.0
+        self.max_damping_constant = 1e10
         self.ridge_constant = 0.0
         self.jacobian = None
         self.jacobian_update_rate = 0
         self.verbose = True
-        self.learning_rate = 1.0
-        self.min_learning_rate = 1e-3
-        self.backtracking_factor = 2
-        self.patience = 5
         self.errorbars = None
 
 
@@ -110,86 +107,71 @@ def print_(*args, **kwargs):
         niter = self.params.max_number_iters
         atol = self.params.abs_tol_convergence
         rtol = self.params.rel_tol_convergence
-        jacobian = self.jacobian
-        damping_constant = self.params.damping_constant
+        rcond = self.params.rcond
+
+        damping = self.params.damping_constant
         damping_factor = self.params.damping_factor
-        ridge_constant = self.params.ridge_constant
+        damping_max = self.params.max_damping_constant
+
+        ridge = self.params.ridge_constant
         jacobian_update_rate = self.params.jacobian_update_rate
-        lr = self.params.learning_rate
-        lr_min = self.params.min_learning_rate
-        lr_factor = self.params.backtracking_factor
-        patience = self.params.patience
+        jacobian = self.jacobian
 
-        pos0 = self.params.initial_position
-        pos = pos0.copy()
+        pos = self.params.initial_position.copy()
 
         res = self._objective_func(pos)
         chi2 = self.calc_chi2(res)
         M = self.calc_jacobian(pos) if jacobian is None else jacobian
-
-        # if c:
-        #     chi2 += w * _np.std(res0) * _np.linalg.norm(pos0)
+        MTM = M.T @ M
+        ridge_reg = ridge * _np.eye(MTM.shape[0])
 
         print_(f'initial chi²: {chi2:.6g}')
 
-        MTM = M.T @ M
-        ridge_reg = ridge_constant * _np.eye(
-            MTM.shape[0]
-        )  # Ridge regularization
-
         for it in range(niter):
             print_(f'iteration {it:03d}')
 
-            pos_init = pos.copy()
-
             if jacobian_update_rate and it:
                 if not it % jacobian_update_rate:
                     M = self.calc_jacobian(pos)
                     MTM = M.T @ M
 
-            lm_reg = _np.diag(
-                damping_constant * _np.diag(MTM)
-            )  # Levenberg-Macquardt regularization
-
+            lm_reg = _np.diag(damping * _np.diag(MTM))
             matrix = MTM + ridge_reg + lm_reg
 
-            res = self.objfuncs_evaluated[-1]  # last evaluated residual
-            delta = (
-                _np.linalg.pinv(matrix, rcond=self.params.rcond) @ M.T @ res
-            )  # LM/GN Normal equation
-            # TODO: include processing of pseudo-inverse for singvals selection
+            res = self.objfuncs_evaluated[-1]
+            delta = _np.linalg.pinv(matrix, rcond=rcond) @ (M.T @ res)
+            # TODO: chi2 for Ridge
+
             if _np.any(_np.isnan(delta)):
-                print_('\tProblem in step calculation. Aborting.')
+                print_('\tInvalid step direction. Aborting.')
                 break
 
-            pos -= lr * delta  # position update
-            for _ in range(patience):
-                try:
-                    res = self._objective_func(pos)
-                    if not _np.any(_np.isnan(res)):
-                        break
-                    print_('NaN in objective function. Back-tracking.')
-                except Exception:
-                    print_('Merit figure evaluation failed. Back-tracking.')
-                lr /= lr_factor  # reduce step size
-                lr = max(lr, lr_min)
-                pos = pos_init - lr * delta  # apply smaller step
-
-            chi2_old = self.history_chi2[-1]
-            chi2_new = self.calc_chi2(res)
-            # if c:
-            #     chi2 += w * _np.std(res) * _np.linalg.norm(pos)
-            print_(f'\tchi²: {chi2_new:.6g}')
-
-            if _math.isclose(chi2_new, chi2_old, rel_tol=rtol, abs_tol=atol):
-                print_('\tConvergence tolerance reached. Exiting.')
+            pos_trial = pos - delta
+            chi2_old = chi2
+
+            try:
+                res_trial = self._objective_func(pos_trial)
+                if _np.any(_np.isnan(res_trial)):
+                    raise ValueError
+                chi2_trial = self.calc_chi2(res_trial)
+                success = chi2_trial < chi2_old
+            except Exception:
+                success = False
+
+            if success:
+                pos = pos_trial
+                res = res_trial
+                chi2 = chi2_trial
+                damping /= damping_factor
+                print_(f'\tchi²: {chi2:.6g} (accepted)')
+            else:
+                damping *= damping_factor
+                print_('\tchi² increased. Step rejected.')
+
+            if damping > damping_max:
+                print_('\tDamping exceeded maximum value. Aborting.')
                 break
 
-            delta_chi2 = chi2_old - chi2_new
-            if damping_constant:
-                if delta_chi2 > 0:
-                    damping_constant /= damping_factor
-                    print_('\tImproved fit. Decreasing LM damping_constant.')
-                else:
-                    damping_constant *= damping_factor
-                    print_('\tWorsened fit. Increasing LM damping_constant.')
+            if _math.isclose(chi2, chi2_old, rel_tol=rtol, abs_tol=atol):
+                print_('\tConvergence tolerance reached.')
+                break

From 02dacde26ffe971880591431379f5e778997e6bd Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Tue, 3 Mar 2026 10:11:14 -0300
Subject: [PATCH 30/40] OPT.LSTSQR.ENH: use squared error instead of MSQE as
 chi2

---
 apsuite/optimization/least_squares.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 08394774..5ca83c5b 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -56,7 +56,7 @@ def objective_function(self, pos):
 
     def calc_chi2(self, residual):
         """."""
-        chi2 = _np.mean(residual**2)
+        chi2 = _np.sum(residual**2)
         self.history_chi2.append(chi2)
         return chi2
 

From 0e5d267802226d1362bf23e1e1c1563e00c4b0d8 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 5 Mar 2026 09:26:18 -0300
Subject: [PATCH 31/40] OPT.BASE.ENH: float type formatting in __str__ method

- properties `max_number_iters` or `max_number_evals` may be set to `np.inf`
- in this case, the integer formatting method breaks
- `np.inf` requires float format
---
 apsuite/optimization/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/apsuite/optimization/base.py b/apsuite/optimization/base.py
index 6198884f..e6d75f78 100644
--- a/apsuite/optimization/base.py
+++ b/apsuite/optimization/base.py
@@ -84,8 +84,8 @@ def limit_lower(self, arr):
     def __str__(self):
         """."""
         stg = ''
-        stg += self._TMPD('max_number_iters', self.max_number_iters, '')
-        stg += self._TMPD('max_number_evals', self.max_number_evals, '')
+        stg += self._TMPF('max_number_iters', self.max_number_iters, '')
+        stg += self._TMPF('max_number_evals', self.max_number_evals, '')
         stg += self._TMPS(
             'boundary_policy',
             self.BoundaryPolicy._fields[self.boundary_policy], '')

From d34c8afc9914305b6e0419d1da115cda99779d9c Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 5 Mar 2026 10:07:28 -0300
Subject: [PATCH 32/40] OPTANLY.NOECO.ENH: include default limits to
 initialization method

- `_initialization` method now also specifies conssitent upper/lower limits for the optimization params
---
 apsuite/optics_analysis/noeco.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 139fde3a..7ed63bb8 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -324,10 +324,11 @@ def _set_rf_frequency(self, frequency, model=None):
     def _initialization(self):
         if self._model is None:
             print('Cannot start optimization without machine model.')
+            print('Set model to `model` property or call `create_model()`.')
             return False
 
         if self._oeorm_goal is None:
-            print('Cannot start optimization without oeorm_goal')
+            print('Cannot start optimization without `oeorm_goal`')
             return False
 
         if not len(self.params.initial_position):
@@ -337,6 +338,19 @@ def _initialization(self):
                 print('No parameters to fit!')
                 return False
             self.params.initial_position = pos
+
+        if not self.params.are_positions_consistent():
+            print(
+                'Initial position and positions limits are not consistent  '
+                + 'or posistions limits have not been specified by user.'
+            )
+            print('Using default limits for all parameters: (-inf, inf).')
+            self.params.limit_lower = _np.full_like(
+                self.params.initial_position, -_np.inf
+            )
+            self.params.limit_upper = _np.full_like(
+                self.params.initial_position, _np.inf
+            )
         return True
 
     def _get_errorbars(self, bpms_noise):

From 7dcace72e0d189f7d2e6a46a763c463b48a4296c Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 5 Mar 2026 10:55:34 -0300
Subject: [PATCH 33/40] OPT.LSTSQ.DOC: add `__str__`  method

---
 apsuite/optimization/least_squares.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 5ca83c5b..1b2dbdeb 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -8,6 +8,7 @@
 
 class LeastSquaresParams(OptimizeParams):
     """."""
+    _TMPE = '{:30s}: {:10.3e} {:s}\n'.format
 
     def __init__(self):
         """."""
@@ -27,6 +28,27 @@ def __init__(self):
         self.verbose = True
         self.errorbars = None
 
+    def __str__(self):
+        """."""
+        stg = super().__str__()
+        stg += '\n'
+        stg += '\nLeastSquaresParams:\n'
+        stg += '\n'
+        stg += self._TMPE('abs_tol_convergence', self.abs_tol_convergence, '')
+        stg += self._TMPE('rel_tol_convergence', self.rel_tol_convergence, '')
+        stg += self._TMPE('rcond', self.rcond, '')
+        stg += self._TMPF('damping_constant', self.damping_constant, '')
+        stg += self._TMPF('damping_factor', self.damping_factor, '')
+        stg += self._TMPE(
+            'max_damping_constant', self.max_damping_constant, ''
+        )
+        stg += self._TMPF('ridge_constant', self.ridge_constant, '')
+        stg += self._TMPD(
+            'jacobian_update_rate', self.jacobian_update_rate, ''
+        )
+        stg += self._TMPS('verbose', str(self.verbose), '')
+        return stg
+
 
 class LeastSquaresOptimize(Optimize):
     """Least-squares optimization via Levenberg–Marquardt style loop."""

From 4544858306e0013b93d2f9b1ada0fa296570ae2c Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 5 Mar 2026 11:22:24 -0300
Subject: [PATCH 34/40] OPT.BASE.ENH: force right alignment for strings in
 __str__ method

- other formats are always right-aligned
- however, strings not necessarily.
- for instance, `str(bool)` would always lose alignment
---
 apsuite/optimization/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/apsuite/optimization/base.py b/apsuite/optimization/base.py
index e6d75f78..b2fea0b0 100644
--- a/apsuite/optimization/base.py
+++ b/apsuite/optimization/base.py
@@ -20,7 +20,7 @@ class OptimizeParams(_Params):
 
     _TMPD = '{:30s}: {:10d} {:s}\n'.format
     _TMPF = '{:30s}: {:10.3f} {:s}\n'.format
-    _TMPS = '{:30s}: {:10s} {:s}\n'.format
+    _TMPS = '{:30s}: {:>10s} {:s}\n'.format
 
     BoundaryPolicy = _get_namedtuple('BoundaryPolicy', ('ToBoundary', 'ToNaN'))
 

From d0762d2cb0238b544cdf4383ac7688ed1496d9ad Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Thu, 5 Mar 2026 11:23:16 -0300
Subject: [PATCH 35/40] OPTANLY.NOECO.DOC: add `__str__` method

- adds str method for NOECO params
---
 apsuite/optics_analysis/noeco.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/apsuite/optics_analysis/noeco.py b/apsuite/optics_analysis/noeco.py
index 7ed63bb8..64a3c284 100644
--- a/apsuite/optics_analysis/noeco.py
+++ b/apsuite/optics_analysis/noeco.py
@@ -89,6 +89,34 @@ def __init__(self):
         self.use_ver_disp = False
         self.use_chroms = False
 
+    def __str__(self):
+        """."""
+        stg = super().__str__()
+        stg += '\n'
+        stg += '\nNOECO Params\n'
+        stg += '\n'
+        stg += self._TMPE('denergy_oeorm_calc', self.denergy_oeorm_calc, '')
+        stg += self._TMPF('tunex', self.tunex, '')
+        stg += self._TMPF('tuney', self.tuney, '')
+        stg += self._TMPF('chromx', self.chromx, '')
+        stg += self._TMPF('chromy', self.chromy, '')
+        stg += self._TMPS('fit_sexts', str(self.fit_sexts), '')
+        stg += self._TMPS('fit_gain_bpms', str(self.fit_gain_bpms), '')
+        stg += self._TMPS('fit_coup_bpms', str(self.fit_coup_bpms), '')
+        stg += self._TMPS('fit_gain_corr', str(self.fit_gain_corr), '')
+        stg += self._TMPS('use_diag_blocks', str(self.use_diag_blocks), '')
+        stg += self._TMPS(
+            'use_offdiag_blocks', str(self.use_offdiag_blocks), ''
+        )
+        stg += self._TMPS('use_hor_disp', str(self.use_hor_disp), '')
+        stg += self._TMPS('use_ver_disp', str(self.use_ver_disp), '')
+        stg += self._TMPS('use_chroms', str(self.use_chroms), '')
+        stg += '{:30s}: [{:.3e}, {:.3e}] {:s}\n'.format(
+            'chrom_weights', self.chrom_weights[0], self.chrom_weights[1], ''
+        )
+        stg += self._TMPF('disp_weight', self.disp_weight, '')
+        return stg
+
     @property
     def nr_sexts(self):
         """."""

From ebb23582b8f1175c0d3b1b36abccefc90af25c27 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Fri, 6 Mar 2026 09:41:37 -0300
Subject: [PATCH 36/40] WIP (OPT.LEASTSQRS): include Tikhonovo/Ridge
 regularization

- change `ridge_constant`  to `tikhonov_reg_constants`
- allowing anisotropic constraining of particular singular modes
- changing of the pseudo-inversion method: from pinv of the normal matrix to the calculation of the SVD and filtering/regularization of singular values
- also requires including the singular modes penalization to the chi2
- not yet thoroughly tested
---
 apsuite/optimization/least_squares.py | 65 ++++++++++++++++++++-------
 1 file changed, 50 insertions(+), 15 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 1b2dbdeb..3fffedbb 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -8,6 +8,7 @@
 
 class LeastSquaresParams(OptimizeParams):
     """."""
+
     _TMPE = '{:30s}: {:10.3e} {:s}\n'.format
 
     def __init__(self):
@@ -22,7 +23,7 @@ def __init__(self):
         self.damping_constant = 1.0
         self.damping_factor = 10.0
         self.max_damping_constant = 1e10
-        self.ridge_constant = 0.0
+        self.tikhonov_reg_constants = 0.0
         self.jacobian = None
         self.jacobian_update_rate = 0
         self.verbose = True
@@ -42,7 +43,9 @@ def __str__(self):
         stg += self._TMPE(
             'max_damping_constant', self.max_damping_constant, ''
         )
-        stg += self._TMPF('ridge_constant', self.ridge_constant, '')
+        # stg += self._TMPF(
+        #     'tikhonov_reg_constants', self.tikhonov_reg_constants, ''
+        # )
         stg += self._TMPD(
             'jacobian_update_rate', self.jacobian_update_rate, ''
         )
@@ -82,6 +85,12 @@ def calc_chi2(self, residual):
         self.history_chi2.append(chi2)
         return chi2
 
+    def calc_chi2_tikhonov(self, residual, vtmat, tikhonov, delta_pos):
+        """."""
+        chi2 = self.calc_chi2(residual)
+        chi2 += _np.linalg.norm(vtmat @ (tikhonov * delta_pos))**2
+        return chi2
+
     def calc_residual(self, pos, merit_figure_goal=None):
         """."""
         merit_figure_meas = self.calc_merit_figure(pos)
@@ -135,34 +144,55 @@ def print_(*args, **kwargs):
         damping_factor = self.params.damping_factor
         damping_max = self.params.max_damping_constant
 
-        ridge = self.params.ridge_constant
+        tikhonov = _np.asarray(self.params.tikhonov_reg_constants)
         jacobian_update_rate = self.params.jacobian_update_rate
         jacobian = self.jacobian
 
         pos = self.params.initial_position.copy()
+        self.normal_eq_matrix = []
 
         res = self._objective_func(pos)
         chi2 = self.calc_chi2(res)
-        M = self.calc_jacobian(pos) if jacobian is None else jacobian
-        MTM = M.T @ M
-        ridge_reg = ridge * _np.eye(MTM.shape[0])
-
         print_(f'initial chi²: {chi2:.6g}')
 
+        if jacobian is None:
+            jacobian = self.calc_jacobian(pos)
+            self.jacobian = jacobian
+        U, S, Vt = _np.linalg.svd(jacobian, full_matrices=False)
+
+        if tikhonov.ndim == 0:
+            tikhonov = _np.full(S.size, tikhonov)
+        elif tikhonov.ndim == 1:
+            if tikhonov.size != S.size:
+                raise ValueError('tikhonov_reg_constants size mismatch')
+
         for it in range(niter):
             print_(f'iteration {it:03d}')
 
             if jacobian_update_rate and it:
                 if not it % jacobian_update_rate:
-                    M = self.calc_jacobian(pos)
-                    MTM = M.T @ M
-
-            lm_reg = _np.diag(damping * _np.diag(MTM))
-            matrix = MTM + ridge_reg + lm_reg
+                    jacobian = self.calc_jacobian(pos)
+                    self.jacobian = jacobian
+                    U, S, Vt = _np.linalg.svd(jacobian, full_matrices=False)
 
             res = self.objfuncs_evaluated[-1]
-            delta = _np.linalg.pinv(matrix, rcond=rcond) @ (M.T @ res)
-            # TODO: chi2 for Ridge
+
+            matrix = jacobian.T @ jacobian
+            matrix += damping * _np.diag(_np.diag(matrix))
+            if _np.any(tikhonov):
+                matrix += _np.diag(tikhonov**2)
+            self.normal_eq_matrix.append(matrix)
+
+            if rcond is not None:
+                mask = S < rcond * S[0]  # not sure mask should compare S
+            else:
+                mask = _np.zeros_like(S, dtype=bool)
+
+            den = S**2 * (1 + damping) + tikhonov**2
+            S_inv = S / den
+            S_inv[mask] = 0
+
+            delta = (Vt.T  @ (S_inv * (U.T @ res)))
 
             if _np.any(_np.isnan(delta)):
                 print_('\tInvalid step direction. Aborting.')
@@ -175,7 +205,12 @@ def print_(*args, **kwargs):
                 res_trial = self._objective_func(pos_trial)
                 if _np.any(_np.isnan(res_trial)):
                     raise ValueError
-                chi2_trial = self.calc_chi2(res_trial)
+                if _np.any(tikhonov):
+                    chi2_trial = self.calc_chi2_tikhonov(
+                        res_trial, Vt.T, tikhonov, delta
+                    )
+                else:
+                    chi2_trial = self.calc_chi2(res_trial)
                 success = chi2_trial < chi2_old
             except Exception:
                 success = False

From a1a1f4a9899ecb28a7d8e5ac610750899b8787b3 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Tue, 10 Mar 2026 16:42:00 -0300
Subject: [PATCH 37/40] BUG (OPT.LEASTSQRS): avoid division by zero when LM
 damping param is zero

---
 apsuite/optimization/least_squares.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 3fffedbb..90f4655c 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -219,10 +219,10 @@ def print_(*args, **kwargs):
                 pos = pos_trial
                 res = res_trial
                 chi2 = chi2_trial
-                damping /= damping_factor
+                damping /= damping_factor if damping_factor else 1
                 print_(f'\tchi²: {chi2:.6g} (accepted)')
             else:
-                damping *= damping_factor
+                damping *= damping_factor if damping_factor else 1
                 print_('\tchi² increased. Step rejected.')
 
             if damping > damping_max:

From 284d2d73f6660eb65b28fed64fcd1246158e96f8 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Fri, 13 Mar 2026 16:49:42 -0300
Subject: [PATCH 38/40] ENH (LSTSQRS):  handle tikhonov constants size

---
 apsuite/optimization/least_squares.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 90f4655c..4f52705a 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -88,7 +88,7 @@ def calc_chi2(self, residual):
     def calc_chi2_tikhonov(self, residual, vtmat, tikhonov, delta_pos):
         """."""
         chi2 = self.calc_chi2(residual)
-        chi2 += _np.linalg.norm(vtmat @ (tikhonov * delta_pos))**2
+        chi2 += _np.linalg.norm(vtmat @ (tikhonov * delta_pos)) ** 2
         return chi2
 
     def calc_residual(self, pos, merit_figure_goal=None):
@@ -162,9 +162,23 @@ def print_(*args, **kwargs):
 
         if tikhonov.ndim == 0:
             tikhonov = _np.full(S.size, tikhonov)
+
         elif tikhonov.ndim == 1:
-            if tikhonov.size != S.size:
-                raise ValueError('tikhonov_reg_constants size mismatch')
+            if tikhonov.size > S.size:
+                print_(
+                    'Warning: more than necessary tikhonov_reg_constants ' +
+                    'were specified. Clipping the null-space associated ' +
+                    'coefficients.'
+                )
+            elif tikhonov.size < S.size:
+                print_(
+                    'Warning: less than necessary tikhonov_reg_constants ' +
+                    'were specified. Padding zeros to  the last singular ' +
+                    'modes. They will will be unconstrained.'
+                )
+            tikhonov = _np.zeros(S.size)[: min(S.size, tikhonov.size)] = (
+                tikhonov[: S.size]
+            )
 
         for it in range(niter):
             print_(f'iteration {it:03d}')
@@ -192,7 +206,7 @@ def print_(*args, **kwargs):
             S_inv = S / den
             S_inv[mask] = 0
 
-            delta = (Vt.T  @ (S_inv * (U.T @ res)))
+            delta = Vt.T @ (S_inv * (U.T @ res))
 
             if _np.any(_np.isnan(delta)):
                 print_('\tInvalid step direction. Aborting.')

From 1ab61af306b3a7c0166b2ce36e526fe5769a96d3 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Tue, 31 Mar 2026 11:42:48 -0300
Subject: [PATCH 39/40] ENH (OPT.LSTSQRS): tikhonov regulariszation in
 parameter space

- preivous commits implemented anisotropic tikhonov-like regularization in singular space
- this means penalizations occurred for particular singular modes
- also, the least squares problem inversion was incorrect (particularly the levenberg-macquardt part)
- this commit changes from singular space to regular parameter space tikhonov regularization and fixes the problem inversion
---
 apsuite/optimization/least_squares.py | 55 ++++++++++++++-------------
 1 file changed, 28 insertions(+), 27 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 4f52705a..8ce80558 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -85,10 +85,10 @@ def calc_chi2(self, residual):
         self.history_chi2.append(chi2)
         return chi2
 
-    def calc_chi2_tikhonov(self, residual, vtmat, tikhonov, delta_pos):
+    def calc_chi2_tikhonov(self, residual, tikhonov, delta_pos):
         """."""
         chi2 = self.calc_chi2(residual)
-        chi2 += _np.linalg.norm(vtmat @ (tikhonov * delta_pos)) ** 2
+        chi2 += _np.linalg.norm(tikhonov * delta_pos) ** 2
         return chi2
 
     def calc_residual(self, pos, merit_figure_goal=None):
@@ -149,7 +149,7 @@ def print_(*args, **kwargs):
         jacobian = self.jacobian
 
         pos = self.params.initial_position.copy()
-        self.normal_eq_matrix = []
+        self.normal_matrices = []  # Save history of normal mats
 
         res = self._objective_func(pos)
         chi2 = self.calc_chi2(res)
@@ -161,24 +161,23 @@ def print_(*args, **kwargs):
         U, S, Vt = _np.linalg.svd(jacobian, full_matrices=False)
 
         if tikhonov.ndim == 0:
-            tikhonov = _np.full(S.size, tikhonov)
+            tikhonov = _np.full(jacobian.shape[1], tikhonov)
 
         elif tikhonov.ndim == 1:
-            if tikhonov.size > S.size:
+            if tikhonov.size > jacobian.shape[1]:
                 print_(
-                    'Warning: more than necessary tikhonov_reg_constants ' +
-                    'were specified. Clipping the null-space associated ' +
-                    'coefficients.'
+                    'Warning: more than necessary tikhonov_reg_constants '
+                    + 'were specified. Ignoring the exceeding constants.'
                 )
-            elif tikhonov.size < S.size:
+                tikhonov = tikhonov[: jacobian.shape[1]]
+            elif tikhonov.size < jacobian.shape[1]:
                 print_(
-                    'Warning: less than necessary tikhonov_reg_constants ' +
-                    'were specified. Padding zeros to  the last singular ' +
-                    'modes. They will will be unconstrained.'
+                    'Warning: less than necessary tikhonov_reg_constants '
+                    + 'were specified. Padding zeros to match dimensions.'
                 )
-            tikhonov = _np.zeros(S.size)[: min(S.size, tikhonov.size)] = (
-                tikhonov[: S.size]
-            )
+            tikhonov_ = _np.zeros(jacobian.shape[1])
+            tikhonov_[: tikhonov.size] = tikhonov
+            tikhonov = tikhonov_
 
         for it in range(niter):
             print_(f'iteration {it:03d}')
@@ -187,32 +186,34 @@ def print_(*args, **kwargs):
                 if not it % jacobian_update_rate:
                     jacobian = self.calc_jacobian(pos)
                     self.jacobian = jacobian
-                    U, S, Vt = _np.linalg.svd(jacobian, full_matrices=False)
 
             res = self.objfuncs_evaluated[-1]
 
-            matrix = jacobian.T @ jacobian
-            matrix += damping * _np.diag(_np.diag(matrix))
+            matrix = jacobian.T @ jacobian  # Gauss-Newton
+            matrix += damping * _np.diag(_np.diag(matrix))  # Levenberg-Macq.
+
             if _np.any(tikhonov):
-                matrix += _np.diag(tikhonov**2)
-            self.normal_eq_matrix.append(matrix)
+                matrix += _np.diag(tikhonov**2)  # Tikhonov regularization
+
+            self.normal_matrices.append(matrix)  # Save history of normal mats
+
+            U, S, Vt = _np.linalg.svd(matrix, full_matrices=False)
 
             if rcond is not None:
-                mask = S < rcond * S[0]  # not sure mask should compare S
+                mask = S < rcond * S[0]
             else:
                 mask = _np.zeros_like(S, dtype=bool)
 
-            den = S**2 * (1 + damping) + tikhonov**2
-            S_inv = S / den
+            S_inv = 1 / S
             S_inv[mask] = 0
 
-            delta = Vt.T @ (S_inv * (U.T @ res))
+            delta = -Vt.T @ (S_inv * (U.T @ (jacobian.T @ res)))
 
             if _np.any(_np.isnan(delta)):
                 print_('\tInvalid step direction. Aborting.')
                 break
 
-            pos_trial = pos - delta
+            pos_trial = pos + delta
             chi2_old = chi2
 
             try:
@@ -221,7 +222,7 @@ def print_(*args, **kwargs):
                     raise ValueError
                 if _np.any(tikhonov):
                     chi2_trial = self.calc_chi2_tikhonov(
-                        res_trial, Vt.T, tikhonov, delta
+                        res_trial, tikhonov, delta
                     )
                 else:
                     chi2_trial = self.calc_chi2(res_trial)
@@ -244,5 +245,5 @@ def print_(*args, **kwargs):
                 break
 
             if _math.isclose(chi2, chi2_old, rel_tol=rtol, abs_tol=atol):
-                print_('\tConvergence tolerance reached.')
+                print_('\tConvergence tolerance reached. Exiting')
                 break

From 20e8e46036d3a209efdb191f83c2f77843793ba1 Mon Sep 17 00:00:00 2001
From: Matheus <vellosok75@gmail.com>
Date: Tue, 7 Apr 2026 16:34:45 -0300
Subject: [PATCH 40/40] ENH (LSTSQRS) : add suport for additional
 regularizations

- regularization along specified directions is included by specifying a matrix `reg_mat`
- the constraints consists of linear combinations of knobs
- each row corresponds to a constraint
---
 apsuite/optimization/least_squares.py | 27 ++++++++++++++++++++++-----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/apsuite/optimization/least_squares.py b/apsuite/optimization/least_squares.py
index 8ce80558..a8afc785 100644
--- a/apsuite/optimization/least_squares.py
+++ b/apsuite/optimization/least_squares.py
@@ -24,6 +24,7 @@ def __init__(self):
         self.damping_factor = 10.0
         self.max_damping_constant = 1e10
         self.tikhonov_reg_constants = 0.0
+        self.reg_mat = None
         self.jacobian = None
         self.jacobian_update_rate = 0
         self.verbose = True
@@ -145,9 +146,16 @@ def print_(*args, **kwargs):
         damping_max = self.params.max_damping_constant
 
         tikhonov = _np.asarray(self.params.tikhonov_reg_constants)
+        reg_mat = self.params.reg_mat
+
         jacobian_update_rate = self.params.jacobian_update_rate
         jacobian = self.jacobian
 
+        if _np.any(tikhonov):
+            self.chi2_tikhonov_history = []
+        if reg_mat is not None:
+            self.chi2_diff_history = []
+
         pos = self.params.initial_position.copy()
         self.normal_matrices = []  # Save history of normal mats
 
@@ -195,6 +203,9 @@ def print_(*args, **kwargs):
             if _np.any(tikhonov):
                 matrix += _np.diag(tikhonov**2)  # Tikhonov regularization
 
+            if reg_mat is not None:
+                matrix += reg_mat.T @ reg_mat  # Difference regularization
+
             self.normal_matrices.append(matrix)  # Save history of normal mats
 
             U, S, Vt = _np.linalg.svd(matrix, full_matrices=False)
@@ -220,12 +231,18 @@ def print_(*args, **kwargs):
                 res_trial = self._objective_func(pos_trial)
                 if _np.any(_np.isnan(res_trial)):
                     raise ValueError
+                chi2_trial = self.calc_chi2(res_trial)
+
                 if _np.any(tikhonov):
-                    chi2_trial = self.calc_chi2_tikhonov(
-                        res_trial, tikhonov, delta
-                    )
-                else:
-                    chi2_trial = self.calc_chi2(res_trial)
+                    chi2_tikho = _np.linalg.norm(tikhonov * delta) ** 2
+                    self.chi2_tikhonov_history.append(chi2_tikho)
+                    chi2_trial += chi2_tikho
+
+                if reg_mat is not None:
+                    chi2_diff = _np.linalg.norm(_np.dot(reg_mat, delta)) ** 2
+                    self.chi2_diff_history.append(chi2_diff)
+                    chi2_trial += chi2_diff
+
                 success = chi2_trial < chi2_old
             except Exception:
                 success = False