Repository: MannLabs/alphadia (language: python)

Issue:
New dynamic lock with RT proportion
Take a proportion of the total gradient length given an initial RT tolerance less than 1. 


Discussion / hints:

Fix patch:
diff --git a/alphadia/constants/default.yaml b/alphadia/constants/default.yaml
index cd19ccc8..df2cd3ab 100644
--- a/alphadia/constants/default.yaml
+++ b/alphadia/constants/default.yaml
@@ -98,9 +98,6 @@ search_advanced:
 
 calibration:
 
-  # minimum number of steps taken during the optimization lock (during which the elution groups used for optimization are extracted)
-  optimization_lock_min_steps: 0
-
   # Number of precursors searched and scored per batch
   batch_size: 8000
 
@@ -110,7 +107,7 @@ calibration:
   # the maximum number of steps that a given optimizer is permitted to take
   max_steps: 20
 
-  # the maximum number of steps that a given optimizer is permitted to take
+  # the minimum number of steps that a given optimizer must take before it can be said to have converged
   min_steps: 2
 
   # TODO: remove this parameter
diff --git a/alphadia/exceptions.py b/alphadia/exceptions.py
index 8cd35be6..94efbc7e 100644
--- a/alphadia/exceptions.py
+++ b/alphadia/exceptions.py
@@ -37,19 +37,19 @@ class NoPsmFoundError(BusinessError):
     _msg = "No psm files accumulated, can't continue"
 
 
-class NoRecalibrationTargetError(BusinessError):
-    """Raise when no recalibration target is found."""
+class NoOptimizationLockTargetError(BusinessError):
+    """Raise when the optimization lock target is not found."""
 
-    _error_code = "NO_RECALIBRATION_TARGET"
+    _error_code = "NO_OPTIMIZATION_LOCK_TARGET"
 
-    _msg = "Searched all data without finding recalibration target"
+    _msg = "Searched all data without finding optimization lock target"
 
-    _detail_msg = """Search for raw file failed as not enough precursors were found for calibration.
+    _detail_msg = """Search for raw file failed as not enough precursors were found for calibration and optimization.
                  This can have the following reasons:
                    1. The sample was empty and therefore no precursors were found.
                    2. The sample contains only very few precursors.
                       For small libraries, try to set recalibration_target to a lower value.
-                      For large libraries, try to reduce the library size and reduce the calibration MS1 and MS2 tolerance.
+                      For large libraries, try to reduce the library size and reduce the initial MS1 and MS2 tolerance.
                    3. There was a fundamental issue with search parameters."""
 
 
diff --git a/alphadia/outputtransform.py b/alphadia/outputtransform.py
index 024a2786..f075a545 100644
--- a/alphadia/outputtransform.py
+++ b/alphadia/outputtransform.py
@@ -297,6 +297,7 @@ class SearchPlanOutput:
     PSM_INPUT = "psm"
     PRECURSOR_OUTPUT = "precursors"
     STAT_OUTPUT = "stat"
+    INTERNAL_OUTPUT = "internal"
     PG_OUTPUT = "protein_groups"
     LIBRARY_OUTPUT = "speclib.mbr"
     TRANSFER_OUTPUT = "speclib.transfer"
@@ -365,6 +366,7 @@ def build(
             folder_list, save=False, base_spec_lib=base_spec_lib
         )
         _ = self.build_stat_df(folder_list, psm_df=psm_df, save=True)
+        _ = self.build_internal_df(folder_list, save=True)
         _ = self.build_lfq_tables(folder_list, psm_df=psm_df, save=True)
         _ = self.build_library(
             base_spec_lib,
@@ -678,6 +680,50 @@ def build_stat_df(
 
         return stat_df
 
+    def build_internal_df(
+        self,
+        folder_list: list[str],
+        save: bool = True,
+    ):
+        """Build internal data table from a list of seach outputs
+
+        Parameters
+        ----------
+
+        folder_list: List[str]
+            List of folders containing the search outputs
+
+        save: bool
+            Save the precursor table to disk
+
+        Returns
+        -------
+
+        stat_df: pd.DataFrame
+            Precursor table
+        """
+        logger.progress("Building internal statistics")
+
+        internal_df_list = []
+        for folder in folder_list:
+            internal_df_list.append(
+                _build_run_internal_df(
+                    folder,
+                )
+            )
+
+        internal_df = pd.concat(internal_df_list)
+
+        if save:
+            logger.info("Writing internal output to disk")
+            write_df(
+                internal_df,
+                os.path.join(self.output_folder, self.INTERNAL_OUTPUT),
+                file_format="tsv",
+            )
+
+        return internal_df
+
     def build_lfq_tables(
         self,
         folder_list: list[str],
@@ -850,9 +896,6 @@ def _build_run_stat_df(
     optimization_manager_path = os.path.join(
         folder, peptidecentric.PeptideCentricWorkflow.OPTIMIZATION_MANAGER_PATH
     )
-    timing_manager_path = os.path.join(
-        folder, peptidecentric.PeptideCentricWorkflow.TIMING_MANAGER_PATH
-    )
 
     if channels is None:
         channels = [0]
@@ -894,20 +937,47 @@ def _build_run_stat_df(
             base_dict["rt_error"] = np.nan
             base_dict["mobility_error"] = np.nan
 
-        if os.path.exists(timing_manager_path):
-            timing_manager = manager.TimingManager(path=timing_manager_path)
+        out_df.append(base_dict)
 
-            base_dict["optimization_duration"] = timing_manager.optimization["duration"]
-            base_dict["extraction_duration"] = timing_manager.extraction["duration"]
+    return pd.DataFrame(out_df)
 
-        else:
-            logger.warning(f"Error reading timing manager for {raw_name}")
-            base_dict["optimization_duration"] = np.nan
-            base_dict["extraction_duration"] = np.nan
 
-        out_df.append(base_dict)
+def _build_run_internal_df(
+    folder_path: str,
+):
+    """Build stat dataframe for a single run.
 
-    return pd.DataFrame(out_df)
+    Parameters
+    ----------
+
+    folder_path: str
+        Path (from the base directory of the output_folder attribute of the Plan class) to the directory containing the raw file and the managers
+
+
+    Returns
+    -------
+    pd.DataFrame
+        Dataframe containing the statistics
+
+    """
+    timing_manager_path = os.path.join(
+        folder_path, peptidecentric.PeptideCentricWorkflow.TIMING_MANAGER_PATH
+    )
+    raw_name = os.path.basename(folder_path)
+
+    internal_dict = {
+        "run": raw_name,
+    }
+
+    if os.path.exists(timing_manager_path):
+        timing_manager = manager.TimingManager(path=timing_manager_path)
+        for key in timing_manager.timings:
+            internal_dict[f"duration_{key}"] = [timing_manager.timings[key]["duration"]]
+
+    else:
+        logger.warning(f"Error reading timing manager for {raw_name}")
+
+    return pd.DataFrame(internal_dict)
 
 
 def perform_protein_fdr(psm_df):
diff --git a/alphadia/planning.py b/alphadia/planning.py
index 57481f4a..f093db22 100644
--- a/alphadia/planning.py
+++ b/alphadia/planning.py
@@ -321,9 +321,16 @@ def run(
                     logger.info(f"No existing quantification found for {raw_name}")
 
                 workflow.load(dia_path, speclib)
-                workflow.calibration()
 
+                workflow.timing_manager.set_start_time("optimization")
+                workflow.search_parameter_optimization()
+                workflow.timing_manager.set_end_time("optimization")
+
+                workflow.timing_manager.set_start_time("extraction")
                 psm_df, frag_df = workflow.extraction()
+                workflow.timing_manager.set_end_time("extraction")
+                workflow.timing_manager.save()
+
                 psm_df = psm_df[psm_df["qval"] <= self.config["fdr"]["fdr"]]
 
                 if self.config["multiplexing"]["enabled"]:
diff --git a/alphadia/workflow/base.py b/alphadia/workflow/base.py
index ba81995c..c3e5f067 100644
--- a/alphadia/workflow/base.py
+++ b/alphadia/workflow/base.py
@@ -102,9 +102,9 @@ def load(
             self._calibration_manager.disable_mobility_calibration()
 
         # initialize the optimization manager
-
         self._optimization_manager = manager.OptimizationManager(
             self.config,
+            gradient_length=self.dia_data.rt_values.max(),
             path=os.path.join(self.path, self.OPTIMIZATION_MANAGER_PATH),
             load_from_file=self.config["general"]["reuse_calibration"],
             figure_path=os.path.join(self.path, self.FIGURE_PATH),
diff --git a/alphadia/workflow/manager.py b/alphadia/workflow/manager.py
index fdc1e196..6dc2e3e8 100644
--- a/alphadia/workflow/manager.py
+++ b/alphadia/workflow/manager.py
@@ -455,6 +455,7 @@ class OptimizationManager(BaseManager):
     def __init__(
         self,
         config: None | dict = None,
+        gradient_length: None | float = None,
         path: None | str = None,
         load_from_file: bool = True,
         **kwargs,
@@ -464,10 +465,15 @@ def __init__(
         self.reporter.log_event("initializing", {"name": f"{self.__class__.__name__}"})
 
         if not self.is_loaded_from_file:
+            rt_error = (
+                config["search_initial"]["initial_rt_tolerance"]
+                if config["search_initial"]["initial_rt_tolerance"] > 1
+                else config["search_initial"]["initial_rt_tolerance"] * gradient_length
+            )
             initial_parameters = {
                 "ms1_error": config["search_initial"]["initial_ms1_tolerance"],
                 "ms2_error": config["search_initial"]["initial_ms2_tolerance"],
-                "rt_error": config["search_initial"]["initial_rt_tolerance"],
+                "rt_error": rt_error,
                 "mobility_error": config["search_initial"][
                     "initial_mobility_tolerance"
                 ],
@@ -487,7 +493,6 @@ def fit(self, update_dict):
         """Update the parameters dict with the values in update_dict."""
         self.__dict__.update(update_dict)
         self.is_fitted = True
-        self.save()
 
     def predict(self):
         """Return the parameters dict."""
@@ -761,34 +766,32 @@ def __init__(
         load_from_file: bool = True,
         **kwargs,
     ):
+        """Contains and updates timing information for the portions of the workflow."""
         super().__init__(path=path, load_from_file=load_from_file, **kwargs)
         self.reporter.log_string(f"Initializing {self.__class__.__name__}")
         self.reporter.log_event("initializing", {"name": f"{self.__class__.__name__}"})
-
         if not self.is_loaded_from_file:
-            self.__dict__.update(
-                {
-                    "optimization": {
-                        "start": None,
-                        "end": None,
-                        "duration": None,
-                    },
-                    "extraction": {
-                        "start": None,
-                        "end": None,
-                        "duration": None,
-                    },
-                }
-            )
+            self.timings = {}
 
-    def start(self, workflow_stage: str):
-        self.__dict__.update({workflow_stage: {"start": pd.Timestamp.now()}})
-        self.save()
+    def set_start_time(self, workflow_stage: str):
+        """Stores the start time of the given stage of the workflow in the timings attribute. Also saves the timing manager to disk.
+
+        Parameters
+        ----------
+        workflow_stage : str
+            The name under which the timing will be stored in the timings dict
+        """
+        self.timings.update({workflow_stage: {"start": pd.Timestamp.now()}})
+
+    def set_end_time(self, workflow_stage: str):
+        """Stores the end time of the given stage of the workflow in the timings attribute and calculates the duration. Also saves the timing manager to disk.
+        Parameters
+        ----------
+        workflow_stage : str
+            The name under which the timing will be stored in the timings dict
 
-    def end(self, workflow_stage: str):
-        self.__dict__[workflow_stage]["end"] = pd.Timestamp.now()
-        self.__dict__[workflow_stage]["duration"] = (
-            self.__dict__[workflow_stage]["end"]
-            - self.__dict__[workflow_stage]["start"]
+        """
+        self.timings[workflow_stage]["end"] = pd.Timestamp.now()
+        self.timings[workflow_stage]["duration"] = (
+            self.timings[workflow_stage]["end"] - self.timings[workflow_stage]["start"]
         ).total_seconds() / 60
-        self.save()
diff --git a/alphadia/workflow/optimization.py b/alphadia/workflow/optimization.py
index 16e2e2b6..a629a89d 100644
--- a/alphadia/workflow/optimization.py
+++ b/alphadia/workflow/optimization.py
@@ -2,12 +2,15 @@
 from abc import ABC, abstractmethod
 
 import matplotlib.pyplot as plt
+import numpy as np
 
 # alpha family imports
 # third party imports
 import pandas as pd
 import seaborn as sns
 
+from alphadia.exceptions import NoOptimizationLockTargetError
+
 # alphadia imports
 from alphadia.workflow import reporting
 
@@ -30,9 +33,9 @@ def __init__(
             The reporter object used to log information about the optimization process. If None, a new LogBackend object is created.
 
         """
-        self.optimal_parameter = None
         self.workflow = workflow
         self.reporter = reporting.LogBackend() if reporter is None else reporter
+        self.num_prev_optimizations = 0
 
     @abstractmethod
     def step(self, precursors_df: pd.DataFrame, fragments_df: pd.DataFrame):
@@ -81,12 +84,12 @@ def __init__(
         self.history_df = pd.DataFrame()
         self.workflow.optimization_manager.fit({self.parameter_name: initial_parameter})
         self.has_converged = False
+        self.num_prev_optimizations = 0
 
     def step(
         self,
         precursors_df: pd.DataFrame,
         fragments_df: pd.DataFrame,
-        current_step: int = -1,
     ):
         """See base class. The feature is used to track the progres of the optimization (stored in .feature) and determine whether it has converged."""
         if self.has_converged:
@@ -96,6 +99,7 @@ def step(
             )
             return
 
+        self.num_prev_optimizations += 1
         new_row = pd.DataFrame(
             [
                 {
@@ -113,7 +117,7 @@ def step(
             ]
         )
         self.history_df = pd.concat([self.history_df, new_row], ignore_index=True)
-        just_converged = self._check_convergence(current_step)
+        just_converged = self._check_convergence()
 
         if just_converged:
             self.has_converged = True
@@ -217,7 +221,7 @@ def _propose_new_parameter(self, df):
         """
         pass
 
-    def _check_convergence(self, current_step: int = -1):
+    def _check_convergence(self):
         """Optimization should stop if continued narrowing of the parameter is not improving the feature value.
         This function checks if the previous rounds of optimization have led to a meaningful improvement in the feature value.
         If so, it continues optimization and appends the proposed new parameter to the list of parameters. If not, it stops optimization and sets the optimal parameter attribute.
@@ -226,15 +230,12 @@ def _check_convergence(self, current_step: int = -1):
         -----
             Because the check for an increase in feature value requires two previous rounds, the function will also initialize for another round of optimization if there have been fewer than 3 rounds.
 
-        Parameters
-        ----------
-        current_step: int
-            The current step in the optimization process. By default it is set to -1, which prevents the optimizer from converging unless min_steps has been set to 0.
 
         """
 
         min_steps_reached = (
-            current_step >= self.workflow.config["calibration"]["min_steps"] - 1
+            self.num_prev_optimizations
+            >= self.workflow.config["calibration"]["min_steps"]
         )
         return (
             min_steps_reached
@@ -292,7 +293,7 @@ def __init__(
         self.target_parameter = target_parameter
         self.has_converged = False
 
-    def _check_convergence(self, proposed_parameter: float, current_step: int = -1):
+    def _check_convergence(self, proposed_parameter: float):
         """The optimization has converged if the proposed parameter is equal to or less than the target parameter and the a sufficient number of steps has been taken.
 
         Parameters
@@ -300,9 +301,6 @@ def _check_convergence(self, proposed_parameter: float, current_step: int = -1):
         proposed_parameter: float
             The proposed parameter for the next round of optimization.
 
-        current_step: int
-            The current step in the optimization process. By default it is set to -1, which prevents the optimizer from converging unless min_steps has been set to 0.
-
         Returns
         -------
         bool
@@ -310,11 +308,11 @@ def _check_convergence(self, proposed_parameter: float, current_step: int = -1):
 
 
         """
-
-        return (
-            proposed_parameter <= self.target_parameter
-            and current_step >= self.workflow.config["calibration"]["min_steps"] - 1
+        min_steps_reached = (
+            self.num_prev_optimizations
+            >= self.workflow.config["calibration"]["min_steps"]
         )
+        return proposed_parameter <= self.target_parameter and min_steps_reached
 
     def _propose_new_parameter(self, df: pd.DataFrame):
         """See base class. The update rule is
@@ -334,7 +332,6 @@ def step(
         self,
         precursors_df: pd.DataFrame,
         fragments_df: pd.DataFrame,
-        current_step: int = -1,
     ):
         """See base class."""
         if self.has_converged:
@@ -343,11 +340,11 @@ def step(
                 verbosity="progress",
             )
             return
-
+        self.num_prev_optimizations += 1
         new_parameter = self._propose_new_parameter(
             precursors_df if self.estimator_group_name == "precursor" else fragments_df
         )
-        just_converged = self._check_convergence(new_parameter, current_step)
+        just_converged = self._check_convergence(new_parameter)
         self.workflow.optimization_manager.fit({self.parameter_name: new_parameter})
         self.workflow.optimization_manager.fit(
             {"classifier_version": self.workflow.fdr_manager.current_version}
@@ -382,7 +379,7 @@ def __init__(
         self.parameter_name = "rt_error"
         self.estimator_group_name = "precursor"
         self.estimator_name = "rt"
-        self.feature_name = "precursor_count"
+        self.feature_name = "precursor_proportion_detected"
         super().__init__(initial_parameter, workflow, reporter)
 
     def _propose_new_parameter(self, df: pd.DataFrame):
@@ -405,7 +402,7 @@ def _propose_new_parameter(self, df: pd.DataFrame):
     def _get_feature_value(
         self, precursors_df: pd.DataFrame, fragments_df: pd.DataFrame
     ):
-        return len(precursors_df)
+        return len(precursors_df) / self.workflow.optlock.total_precursors
 
 
 class AutomaticMS2Optimizer(AutomaticOptimizer):
@@ -419,7 +416,7 @@ def __init__(
         self.parameter_name = "ms2_error"
         self.estimator_group_name = "fragment"
         self.estimator_name = "mz"
-        self.feature_name = "precursor_count"
+        self.feature_name = "precursor_proportion_detected"
         super().__init__(initial_parameter, workflow, reporter)
 
     def _propose_new_parameter(self, df: pd.DataFrame):
@@ -442,7 +439,7 @@ def _propose_new_parameter(self, df: pd.DataFrame):
     def _get_feature_value(
         self, precursors_df: pd.DataFrame, fragments_df: pd.DataFrame
     ):
-        return len(precursors_df)
+        return len(precursors_df) / self.workflow.optlock.total_precursors
 
 
 class AutomaticMS1Optimizer(AutomaticOptimizer):
@@ -494,7 +491,7 @@ def __init__(
         self.parameter_name = "mobility_error"
         self.estimator_group_name = "precursor"
         self.estimator_name = "mobility"
-        self.feature_name = "precursor_count"
+        self.feature_name = "precursor_proportion_detected"
         super().__init__(initial_parameter, workflow, reporter)
 
     def _propose_new_parameter(self, df: pd.DataFrame):
@@ -518,7 +515,7 @@ def _propose_new_parameter(self, df: pd.DataFrame):
     def _get_feature_value(
         self, precursors_df: pd.DataFrame, fragments_df: pd.DataFrame
     ):
-        return len(precursors_df)
+        return len(precursors_df) / self.workflow.optlock.total_precursors
 
 
 class TargetedRTOptimizer(TargetedOptimizer):
@@ -579,3 +576,164 @@ def __init__(
         self.estimator_group_name = "precursor"
         self.estimator_name = "mobility"
         super().__init__(initial_parameter, target_parameter, workflow, reporter)
+
+
+class OptimizationLock:
+    def __init__(self, library, config: dict):
+        """Sets and updates the optimization lock, which is the data used for calibration and optimization of the search parameters.
+
+        Parameters
+        ----------
+        library: alphabase.spectral_library.flat.SpecLibFlat
+            The library object from the PeptideCentricWorkflow object, which includes the precursor and fragment library dataframes.
+
+        config: dict
+            The configuration dictionary from the PeptideCentricWorkflow object.
+        """
+        self._library = library
+        self._config = config
+
+        self.first_time_to_reach_target = True
+        self.has_target_num_precursors = False
+
+        self.elution_group_order = library._precursor_df["elution_group_idx"].unique()
+        rng = np.random.default_rng(seed=772)
+        rng.shuffle(self.elution_group_order)
+
+        self.target_count = self._config["calibration"]["optimization_lock_target"]
+
+        self.batch_idx = 0
+        self.set_batch_plan()
+
+        eg_idxes = self.elution_group_order[self.start_idx : self.stop_idx]
+        self.batch_precursor_df = self._library._precursor_df[
+            self._library._precursor_df["elution_group_idx"].isin(eg_idxes)
+        ]
+
+        self.feature_dfs = []
+        self.fragment_dfs = []
+
+    def _get_exponential_batches(self, step):
+        """Get the number of batches for a given step
+        This plan has the shape:
+        1, 2, 4, 8, 16, 32, 64, ...
+        """
+        return int(2**step)
+
+    def set_batch_plan(self):
+        """Gets an exponential batch plan based on the batch_size value in the config."""
+        n_eg = self._library._precursor_df["elution_group_idx"].nunique()
+
+        plan = []
+
+        batch_size = self._config["calibration"]["batch_size"]
+        step = 0
+        start_idx = 0
+
+        while start_idx < n_eg:
+            n_batches = self._get_exponential_batches(step)
+            stop_idx = min(start_idx + n_batches * batch_size, n_eg)
+            plan.append((start_idx, stop_idx))
+            step += 1
+            start_idx = stop_idx
+
+        self.batch_plan = plan
+
+    def update_with_extraction(self, feature_df, fragment_df):
+        """Extract features and fragments from the current batch of the optimization lock.
+
+        Parameters
+        ----------
+        feature_df: pd.DataFrame
+            The feature dataframe for the current batch of the optimization lock (from workflow.extract_batch).
+
+        fragment_df: pd.DataFrame
+            The fragment dataframe for the current batch of the optimization lock (from workflow.extract_batch).
+        """
+
+        self.feature_dfs += [feature_df]
+        self.fragment_dfs += [fragment_df]
+
+        self.total_precursors = len(self.features_df)
+
+    def update_with_fdr(self, precursor_df):
+        """Calculates the number of precursors at 1% FDR for the current optimization lock and determines if it is sufficient to perform calibration and optimization.
+
+        Parameters
+        ----------
+        precursor_df: pd.DataFrame
+            The precursor dataframe for the current batch of the optimization lock (from workflow.perform_fdr).
+        """
+
+        self.count = len(precursor_df[precursor_df["qval"] < 0.01])
+
+        self.has_target_num_precursors = self.count >= self.target_count
+
+    def increase_batch_idx(self):
+        """If the optimization lock does not contain enough precursors at 1% FDR, the optimization lock proceeds to include the next step in the batch plan in the library attribute.
+        This is done by incrementing self.batch_idx.
+        """
+        self.batch_idx += 1
+
+    def decrease_batch_idx(self):
+        """If the optimization lock contains enough precursors at 1% FDR, checks if enough precursors can be obtained using a smaller library and updates the library attribute accordingly.
+        If not, the same library is used as before.
+        This is done by taking the smallest step in the batch plan which gives more precursors than the target number of precursors.
+        """
+
+        batch_plan_diff = np.array(
+            [
+                stop_at_given_idx - self.stop_idx * self.target_count / self.count
+                for _, stop_at_given_idx in self.batch_plan
+            ]
+        )  # Calculate the difference between the number of precursors expected at the given idx and the target number of precursors for each idx in the batch plan.
+        smallest_value = np.min(
+            batch_plan_diff[batch_plan_diff > 0]
+        )  # Take the smallest positive difference (i.e. the smallest idx that is expected to yield more than the target number of precursors).
+        self.batch_idx = np.where(batch_plan_diff == smallest_value)[0][
+            0
+        ]  # Set the batch idx to the index of the smallest positive difference.
+
+    def update(self):
+        """Updates the library to use for the next round of optimization, either adjusting it upwards or downwards depending on whether the target has been reached.
+        If the target has been reached, the feature and fragment dataframes are reset (since decrease will )"""
+        if not self.has_target_num_precursors:
+            self.increase_batch_idx()
+            eg_idxes = self.elution_group_order[
+                self.start_idx : self.stop_idx
+            ]  # Take only additional elution groups in the new step of the batch plan. The target has not been reached and hence recalibration and optimization have not been performed, so previous extractions can be re-used.
+
+        else:
+            self.first_time_to_reach_target = False
+            self.decrease_batch_idx()
+            eg_idxes = self.elution_group_order[
+                : self.stop_idx
+            ]  # Take all elution groups up the the stop index. The target has been reached and hence the recalibration and optimization have been performed, so everything must be extracted afresh.
+            self.feature_dfs = []
+            self.fragment_dfs = []
+
+        # Sets the batch to be used in the next round of optimization. Note that this batch will only include additional precursors if the target was not reached in this step (as the precursors from this step are stored in the feature_dfs attribute).
+        self.batch_precursor_df = self._library._precursor_df[
+            self._library._precursor_df["elution_group_idx"].isin(eg_idxes)
+        ]
+
+    @property
+    def features_df(self):
+        return pd.concat(self.feature_dfs)
+
+    @property
+    def fragments_df(self):
+        return pd.concat(self.fragment_dfs)
+
+    @property
+    def start_idx(self):
+        if self.has_target_num_precursors:
+            return 0
+        elif self.batch_idx >= len(self.batch_plan):
+            raise NoOptimizationLockTargetError()
+        else:
+            return self.batch_plan[self.batch_idx][0]
+
+    @property
+    def stop_idx(self):
+        return self.batch_plan[self.batch_idx][1]
diff --git a/alphadia/workflow/peptidecentric.py b/alphadia/workflow/peptidecentric.py
index 9f3dcfcb..efe144ee 100644
--- a/alphadia/workflow/peptidecentric.py
+++ b/alphadia/workflow/peptidecentric.py
@@ -256,156 +256,6 @@ def norm_to_rt(
         else:
             raise ValueError(f"Unknown norm_rt_mode {mode}")
 
-    def get_exponential_batches(self, step):
-        """Get the number of batches for a given step
-        This plan has the shape:
-        1, 2, 4, 8, 16, 32, 64, ...
-        """
-        return int(2**step)
-
-    def get_batch_plan(self):
-        """Gets an exponential batch plan based on the batch_size value in the config."""
-        n_eg = self.spectral_library._precursor_df["elution_group_idx"].nunique()
-
-        plan = []
-
-        batch_size = self.config["calibration"]["batch_size"]
-        step = 0
-        start_index = 0
-
-        while start_index < n_eg:
-            n_batches = self.get_exponential_batches(step)
-            stop_index = min(start_index + n_batches * batch_size, n_eg)
-            plan.append((start_index, stop_index))
-            step += 1
-            start_index = stop_index
-
-        return plan
-
-    def get_optimization_lock(self):
-        """Search parameter optimization (i.e. refinement of tolerances for RT, MS2, etc.) is performed on a subset of the elution groups in the spectral library.
-        This subset is termed the optimization lock.
-        The number of elution groups which must be searched to get a sufficiently large number for robust calibration varies depending the library used and the data.
-        This function searches an increasing number of elution groups until a sufficient number of precursors are identified at 1% FDR and a sufficient number of steps have been taken.
-        The values deemed sufficient are specified in by "optimization_lock_target" and "optmization_lock_min_steps" in the config.
-
-        Returns
-        -------
-        optimization_lock_eg_idxes : np.ndarray
-            The elution group (eg) indices (in .spectral_library._precursor_df) of the precursors which will be used for calibration.
-
-        precursor_df : pd.DataFrame
-            Dataframe of all precursors accumulated during the optimization lock, including q-values from FDR correction.
-
-        fragments_df : pd.DataFrame
-            Dataframe of all fragments accumulated during the optimization lock, including q-values from FDR correction.
-
-        """
-
-        self.elution_group_order = self.spectral_library.precursor_df[
-            "elution_group_idx"
-        ].unique()
-        np.random.shuffle(self.elution_group_order)
-
-        batch_plan = self.get_batch_plan()
-
-        features = []
-        fragments = []
-        for current_step, (start_index, stop_index) in enumerate(batch_plan):
-            self.reporter.log_string(
-                f"=== Step {current_step}, extracting elution groups {start_index} to {stop_index} ===",
-                verbosity="progress",
-            )
-
-            eg_idxes = self.elution_group_order[start_index:stop_index]
-            batch_precursor_df = self.spectral_library._precursor_df[
-                self.spectral_library._precursor_df["elution_group_idx"].isin(eg_idxes)
-            ]
-
-            feature_df, fragment_df = self.extract_batch(batch_precursor_df)
-            features += [feature_df]
-            fragments += [fragment_df]
-            features_df = pd.concat(features)
-            fragments_df = pd.concat(fragments)
-
-            self.reporter.log_string(
-                f"=== Step {current_step}, extracted {len(feature_df)} precursors and {len(fragment_df)} fragments ===",
-                verbosity="progress",
-            )
-
-            precursor_df = self.fdr_correction(
-                features_df, fragments_df, self.optimization_manager.classifier_version
-            )
-
-            self.reporter.log_string(
-                f"=== FDR correction performed with classifier version {self.optimization_manager.classifier_version} ===",
-                verbosity="info",
-            )
-
-            num_precursors_at_01FDR = len(precursor_df[precursor_df["qval"] < 0.01])
-
-            self.reporter.log_string(
-                f"=== Checking if minimum number of precursors for optimization found yet; minimum number is {self.config['calibration']['optimization_lock_target']} ===",
-                verbosity="progress",
-            )
-
-            self.log_precursor_df(precursor_df)
-
-            self.reporter.log_string(
-                f"=== Classifier has been trained for {self.fdr_manager.current_version + 1} iteration(s); minimum number is {self.config['calibration']['optimization_lock_min_steps']} ===",
-                verbosity="progress",
-            )
-
-            if (
-                num_precursors_at_01FDR
-                > self.config["calibration"]["optimization_lock_target"]
-                and current_step
-                >= self.config["calibration"]["optimization_lock_min_steps"] - 1
-            ):
-                final_stop_index = stop_index  # final_stop_index is the number of elution groups that will be included in the calibration data
-                break
-
-        optimization_lock_eg_idxes = self.elution_group_order[:final_stop_index]
-
-        optimization_lock_library_precursor_df = self.spectral_library._precursor_df[
-            self.spectral_library._precursor_df["elution_group_idx"].isin(
-                optimization_lock_eg_idxes
-            )
-        ]
-
-        optimization_lock_fragment_idxes = np.concatenate(
-            [
-                np.arange(row["flat_frag_start_idx"], row["flat_frag_stop_idx"])
-                for _, row in optimization_lock_library_precursor_df.iterrows()
-            ]
-        )
-
-        # Extract the fragments for the optimization lock and reset the indices to a consecutive range of positive integers. This simplifies future access based on position
-        optimization_lock_library_fragment_df = self.spectral_library._fragment_df.iloc[
-            optimization_lock_fragment_idxes
-        ].reset_index(drop=True)
-
-        # Change the fragment indices in the precursor_df to match the fragment indices in the optimization lock fragment_df instead of the full spectral library.
-        num_frags = (
-            optimization_lock_library_precursor_df["flat_frag_stop_idx"]
-            - optimization_lock_library_precursor_df["flat_frag_start_idx"]
-        )
-        optimization_lock_library_precursor_df["flat_frag_stop_idx"] = (
-            num_frags.cumsum()
-        )
-        optimization_lock_library_precursor_df["flat_frag_start_idx"] = (
-            optimization_lock_library_precursor_df[
-                "flat_frag_stop_idx"
-            ].shift(fill_value=0)
-        )
-
-        return (
-            optimization_lock_library_precursor_df,
-            optimization_lock_library_fragment_df,
-            precursor_df,
-            fragments_df,
-        )
-
     def get_ordered_optimizers(self):
         """Select appropriate optimizers. Targeted optimization is used if a valid target value (i.e. a number greater than 0) is specified in the config;
         if a value less than or equal to 0 is supplied, automatic optimization is used.
@@ -419,42 +269,41 @@ def get_ordered_optimizers(self):
             List of lists of optimizers
 
         """
-        config_search_initial = self.config["search_initial"]
         config_search = self.config["search"]
 
         if config_search["target_ms2_tolerance"] > 0:
             ms2_optimizer = optimization.TargetedMS2Optimizer(
-                config_search_initial["initial_ms2_tolerance"],
+                self.optimization_manager.ms2_error,
                 config_search["target_ms2_tolerance"],
                 self,
             )
         else:
             ms2_optimizer = optimization.AutomaticMS2Optimizer(
-                config_search_initial["initial_ms2_tolerance"],
+                self.optimization_manager.ms2_error,
                 self,
             )
 
         if config_search["target_rt_tolerance"] > 0:
             rt_optimizer = optimization.TargetedRTOptimizer(
-                config_search_initial["initial_rt_tolerance"],
+                self.optimization_manager.rt_error,
                 config_search["target_rt_tolerance"],
                 self,
             )
         else:
             rt_optimizer = optimization.AutomaticRTOptimizer(
-                config_search_initial["initial_rt_tolerance"],
+                self.optimization_manager.rt_error,
                 self,
             )
         if self.dia_data.has_ms1:
             if config_search["target_ms1_tolerance"] > 0:
                 ms1_optimizer = optimization.TargetedMS1Optimizer(
-                    config_search_initial["initial_ms1_tolerance"],
+                    self.optimization_manager.ms1_error,
                     config_search["target_ms1_tolerance"],
                     self,
                 )
             else:
                 ms1_optimizer = optimization.AutomaticMS1Optimizer(
-                    config_search_initial["initial_ms1_tolerance"],
+                    self.optimization_manager.ms1_error,
                     self,
                 )
         else:
@@ -462,16 +311,14 @@ def get_ordered_optimizers(self):
         if self.dia_data.has_mobility:
             if config_search["target_mobility_tolerance"] > 0:
                 mobility_optimizer = optimization.TargetedMobilityOptimizer(
-                    config_search_initial["initial_mobility_tolerance"],
+                    self.optimization_manager.mobility_error,
                     config_search["target_mobility_tolerance"],
                     self,
                 )
             else:
                 mobility_optimizer = optimization.AutomaticMobilityOptimizer(
-                    config_search_initial["initial_mobility_tolerance"],
-                    self.calibration_manager,
-                    self.optimization_manager,
-                    self.fdr_manager,
+                    self.optimization_manager.mobility_error,
+                    self,
                 )
         else:
             mobility_optimizer = None
@@ -505,11 +352,10 @@ def get_ordered_optimizers(self):
 
         return ordered_optimizers
 
-    def first_recalibration_and_optimization(
+    def first_recalibration(
         self,
         precursor_df: pd.DataFrame,
         fragments_df: pd.DataFrame,
-        ordered_optimizers: list,
     ):
         """Performs the first recalibration and optimization step.
 
@@ -521,25 +367,18 @@ def first_recalibration_and_optimization(
         fragments_df : pd.DataFrame
             Fragment dataframe from optimization lock
 
-        ordered_optimizers : list
-            List of lists of optimizers in correct order
         """
+
+        self.optimization_manager.fit(
+            {"classifier_version": self.fdr_manager.current_version}
+        )
         precursor_df_filtered, fragments_df_filtered = self.filter_dfs(
             precursor_df, fragments_df
         )
 
         self.recalibration(precursor_df_filtered, fragments_df_filtered)
 
-        self.reporter.log_string(
-            "=== Performing initial optimization on extracted data. ===",
-            verbosity="info",
-        )
-
-        for optimizers in ordered_optimizers:
-            for optimizer in optimizers:
-                optimizer.step(precursor_df_filtered, fragments_df_filtered)
-
-    def calibration(self):
+    def search_parameter_optimization(self):
         """Performs optimization of the search parameters. This occurs in two stages:
         1) Optimization lock: the data are searched to acquire a locked set of precursors which is used for search parameter optimization. The classifier is also trained during this stage.
         2) Optimization loop: the search parameters are optimized iteratively using the locked set of precursors.
@@ -557,40 +396,22 @@ def calibration(self):
                 verbosity="progress",
             )
             return
-        self.timing_manager.start("optimization")
+
         # Get the order of optimization
         ordered_optimizers = self.get_ordered_optimizers()
 
         self.reporter.log_string(
-            "Starting initial classifier training and precursor identification.",
+            "Starting initial search for precursors.",
             verbosity="progress",
         )
 
-        # Get the optimization lock
-        (
-            self.optimization_lock_library_precursor_df,
-            self.optimization_lock_library_fragment_df,
-            precursor_df,
-            fragments_df,
-        ) = self.get_optimization_lock()
-
-        self.optimization_manager.fit(
-            {"classifier_version": self.fdr_manager.current_version}
-        )
-
-        self.reporter.log_string(
-            "Required number of precursors found and required number of training iterations performed. Starting search parameter optimization.",
-            verbosity="progress",
-        )
-
-        # Perform a first recalibration on the optimization lock.
-        self.first_recalibration_and_optimization(
-            precursor_df, fragments_df, ordered_optimizers
-        )
+        self.optlock = optimization.OptimizationLock(self.spectral_library, self.config)
 
         # Start of optimization/recalibration loop
         for optimizers in ordered_optimizers:
-            for current_step in range(self.config["calibration"]["max_steps"]):
+            for current_step in range(
+                self.config["calibration"]["max_steps"]
+            ):  # Note current_step here refers to a different step than the attribute of the same name in the optimizer -- this should be rectified
                 if np.all([optimizer.has_converged for optimizer in optimizers]):
                     self.reporter.log_string(
                         f"Optimization finished for {', '.join([optimizer.parameter_name for optimizer in optimizers])}.",
@@ -601,23 +422,29 @@ def calibration(self):
                         optimizer.plot()
 
                     break
+                self.reporter.log_string(
+                    f"=== Step {current_step}, extracting elution groups {self.optlock.start_idx} to {self.optlock.stop_idx} ===",
+                    verbosity="progress",
+                )
 
-                features_df, fragments_df = self.extract_batch(
-                    self.optimization_lock_library_precursor_df,
-                    self.optimization_lock_library_fragment_df,
+                feature_df, fragment_df = self.extract_batch(
+                    self.optlock.batch_precursor_df
                 )
+                self.optlock.update_with_extraction(feature_df, fragment_df)
 
                 self.reporter.log_string(
-                    f"=== Step {current_step}, extracted {len(features_df)} precursors and {len(fragments_df)} fragments ===",
+                    f"=== Step {current_step}, extracted {len(self.optlock.features_df)} precursors and {len(self.optlock.fragments_df)} fragments ===",
                     verbosity="progress",
                 )
 
                 precursor_df = self.fdr_correction(
-                    features_df,
-                    fragments_df,
+                    self.optlock.features_df,
+                    self.optlock.fragments_df,
                     self.optimization_manager.classifier_version,
                 )
 
+                self.optlock.update_with_fdr(precursor_df)
+
                 self.reporter.log_string(
                     f"=== FDR correction performed with classifier version {self.optimization_manager.classifier_version} ===",
                     verbosity="info",
@@ -625,30 +452,43 @@ def calibration(self):
 
                 self.log_precursor_df(precursor_df)
 
-                precursor_df_filtered, fragments_df_filtered = self.filter_dfs(
-                    precursor_df, fragments_df
-                )
+                if self.optlock.has_target_num_precursors:
+                    if self.optlock.first_time_to_reach_target:  # Recalibrates and updates classifier, but does not optimize the first time the target is reached. Optimization is more stable when done with calibrated values.
+                        self.first_recalibration(
+                            precursor_df, self.optlock.fragments_df
+                        )
+                        self.optlock.update()
+                        self.reporter.log_string(
+                            "Required number of precursors found. Starting search parameter optimization.",
+                            verbosity="progress",
+                        )
+                        continue
+
+                    precursor_df_filtered, fragments_df_filtered = self.filter_dfs(
+                        precursor_df, self.optlock.fragments_df
+                    )
 
-                self.recalibration(precursor_df_filtered, fragments_df_filtered)
+                    self.recalibration(precursor_df_filtered, fragments_df_filtered)
 
-                self.reporter.log_string(
-                    "=== checking if optimization conditions were reached ===",
-                    verbosity="info",
-                )
+                    self.reporter.log_string(
+                        "=== checking if optimization conditions were reached ===",
+                        verbosity="info",
+                    )
 
-                for optimizer in optimizers:
-                    optimizer.step(
-                        precursor_df_filtered, fragments_df_filtered, current_step
+                    for optimizer in optimizers:
+                        optimizer.step(precursor_df_filtered, fragments_df_filtered)
+
+                    self.reporter.log_string(
+                        "==============================================",
+                        verbosity="info",
                     )
 
-                self.reporter.log_string(
-                    "==============================================", verbosity="info"
-                )
+                    self.reporter.log_string(
+                        f"=== Optimization has been performed {optimizer.num_prev_optimizations} times; minimum number is {self.config['calibration']['min_steps']} ===",
+                        verbosity="progress",
+                    )
 
-                self.reporter.log_string(
-                    f"=== Optimization has been performed for {current_step + 1} step(s); minimum number is {self.config['calibration']['min_steps']} ===",
-                    verbosity="progress",
-                )
+                self.optlock.update()
 
             else:
                 self.reporter.log_string(
@@ -673,7 +513,7 @@ def calibration(self):
             "==============================================", verbosity="progress"
         )
 
-        self.timing_manager.end("optimization")
+        self.save_managers()
 
     def filter_dfs(self, precursor_df, fragments_df):
         """Filters precursor and fragment dataframes to extract the most reliable examples for calibration.
@@ -706,22 +546,18 @@ def filter_dfs(self, precursor_df, fragments_df):
             fragments_df["precursor_idx"].isin(precursor_df_filtered["precursor_idx"])
         ]
 
-        min_fragments = 500
-        max_fragments = 5000
-        min_correlation = 0.7
         fragments_df_filtered = fragments_df_filtered.sort_values(
-            by=["correlation"], ascending=False
+            by="correlation", ascending=False
         )
-        stop_rank = min(
-            max(
-                np.searchsorted(
-                    fragments_df_filtered["correlation"].values, min_correlation
-                ),
-                min_fragments,
-            ),
-            max_fragments,
-        )
-        fragments_df_filtered = fragments_df_filtered.iloc[:stop_rank]
+        # Determine the number of fragments to keep
+        min_fragments, max_fragments = 500, 5000
+        min_correlation = 0.7
+
+        high_corr_count = (fragments_df_filtered["correlation"] > min_correlation).sum()
+        stop_rank = min(max(high_corr_count, min_fragments), max_fragments)
+
+        # Select top fragments
+        fragments_df_filtered = fragments_df_filtered.head(stop_rank)
 
         self.reporter.log_string(
             f"fragments_df_filtered: {len(fragments_df_filtered)}", verbosity="info"
@@ -758,13 +594,11 @@ def recalibration(self, precursor_df_filtered, fragments_df_filtered):
         )
 
         self.calibration_manager.predict(
-            self.optimization_lock_library_precursor_df,
+            self.spectral_library._precursor_df,
             "precursor",
         )
 
-        self.calibration_manager.predict(
-            self.optimization_lock_library_fragment_df, "fragment"
-        )
+        self.calibration_manager.predict(self.spectral_library._fragment_df, "fragment")
 
         self.optimization_manager.fit(
             {
@@ -914,8 +748,15 @@ def extract_batch(
 
         return features_df, fragments_df
 
+    def save_managers(self):
+        """Saves the calibration, optimization and FDR managers to disk so that they can be reused if needed.
+        Note the timing manager is not saved at this point as it is saved with every call to it.
+        The FDR manager is not saved because it is not used in subsequent parts of the workflow.
+        """
+        self.calibration_manager.save()
+        self.optimization_manager.save()  # this replaces the .save() call when the optimization manager is fitted, since there seems little point in saving an intermediate optimization manager.
+
     def extraction(self):
-        self.timing_manager.start("extraction")
         self.optimization_manager.fit(
             {
                 "num_candidates": self.config["search"]["target_num_candidates"],
@@ -960,8 +801,6 @@ def extraction(self):
 
         self.log_precursor_df(precursor_df)
 
-        self.timing_manager.end("extraction")
-
         return precursor_df, fragments_df
 
     def log_precursor_df(self, precursor_df):
diff --git a/nbs/debug/debug_lvl1.ipynb b/nbs/debug/debug_lvl1.ipynb
index 73e18f1b..db0282cc 100644
--- a/nbs/debug/debug_lvl1.ipynb
+++ b/nbs/debug/debug_lvl1.ipynb
@@ -112,7 +112,7 @@
     "    plan.config,\n",
     ")\n",
     "workflow.load(dia_path, speclib)\n",
-    "workflow.calibration()"
+    "workflow.search_parameter_optimization()"
    ]
   },
   {
@@ -221,7 +221,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.7"
+   "version": "3.11.9"
   },
   "orig_nbformat": 4
  },


Test patch:
diff --git a/tests/unit_tests/test_outputtransform.py b/tests/unit_tests/test_outputtransform.py
index 6129a6c4..0fd5187e 100644
--- a/tests/unit_tests/test_outputtransform.py
+++ b/tests/unit_tests/test_outputtransform.py
@@ -77,25 +77,27 @@ def test_output_transform():
 
         optimization_manager = manager.OptimizationManager(
             config,
-            os.path.join(
+            path=os.path.join(
                 raw_folder,
                 peptidecentric.PeptideCentricWorkflow.OPTIMIZATION_MANAGER_PATH,
             ),
         )
         optimization_manager.fit({"ms2_error": 6})
-
+        optimization_manager.save()
         timing_manager = manager.TimingManager(
-            os.path.join(
+            path=os.path.join(
                 raw_folder, peptidecentric.PeptideCentricWorkflow.TIMING_MANAGER_PATH
             )
         )
 
-        timing_manager.start("extraction")
-        timing_manager.end("extraction")
+        timing_manager.set_start_time("extraction")
+        timing_manager.set_end_time("extraction")
+        timing_manager.save()
 
     output = outputtransform.SearchPlanOutput(config, temp_folder)
     _ = output.build_precursor_table(raw_folders, save=True)
     _ = output.build_stat_df(raw_folders, save=True)
+    _ = output.build_internal_df(raw_folders, save=True)
     _ = output.build_lfq_tables(raw_folders, save=True)
 
     # validate psm_df output
@@ -128,7 +130,7 @@ def test_output_transform():
     assert len(stat_df) == 3
     assert stat_df["ms2_error"][0] == 6
     assert stat_df["rt_error"][0] == 200
-    assert isinstance(stat_df["extraction_duration"][0], float)
+
     assert all(
         [
             col in stat_df.columns
@@ -136,12 +138,14 @@ def test_output_transform():
                 "run",
                 "precursors",
                 "proteins",
-                "optimization_duration",
-                "extraction_duration",
             ]
         ]
     )
 
+    internal_df = pd.read_csv(
+        os.path.join(temp_folder, f"{output.INTERNAL_OUTPUT}.tsv"), sep="\t"
+    )
+    assert isinstance(internal_df["duration_extraction"][0], float)
     # validate protein_df output
     protein_df = pd.read_parquet(os.path.join(temp_folder, "pg.matrix.parquet"))
     assert all([col in protein_df.columns for col in ["run_0", "run_1", "run_2"]])
diff --git a/tests/unit_tests/test_workflow.py b/tests/unit_tests/test_workflow.py
index 6c675344..1c8ea430 100644
--- a/tests/unit_tests/test_workflow.py
+++ b/tests/unit_tests/test_workflow.py
@@ -1,6 +1,7 @@
 import os
 import shutil
 import tempfile
+from copy import deepcopy
 from pathlib import Path
 
 import numpy as np
@@ -232,6 +233,21 @@ def test_optimization_manager():
     assert optimization_manager.is_fitted is False
 
 
+def test_optimization_manager_rt_proportion():
+    TEST_OPTIMIZATION_CONFIG_PROPORTION = deepcopy(TEST_OPTIMIZATION_CONFIG)
+    TEST_OPTIMIZATION_CONFIG_PROPORTION["search_initial"]["initial_rt_tolerance"] = 0.5
+    optimization_manager = manager.OptimizationManager(
+        TEST_OPTIMIZATION_CONFIG_PROPORTION, 1200
+    )
+
+    assert optimization_manager.fwhm_rt == 5
+    assert optimization_manager.fwhm_mobility == 0.01
+    optimization_manager.rt_error = 600
+
+    assert optimization_manager.is_loaded_from_file is False
+    assert optimization_manager.is_fitted is False
+
+
 def test_optimization_manager_save_load():
     temp_path = os.path.join(tempfile.tempdir, "optimization_manager.pkl")
 
@@ -271,6 +287,7 @@ def test_optimization_manager_fit():
     assert optimization_manager.fwhm_cycles == 10
     assert optimization_manager.fwhm_mobility == 0.02
 
+    optimization_manager.save()
     os.remove(temp_path)
 
 
@@ -430,6 +447,11 @@ def create_workflow_instance():
 
     workflow.init_fdr_manager()
 
+    class MockOptlock:
+        total_precursors = 2000
+
+    workflow.optlock = MockOptlock()
+
     return workflow
 
 
@@ -450,7 +472,7 @@ def test_automatic_ms2_optimizer():
     assert ms2_optimizer.parameter_name == "ms2_error"
 
     workflow.fdr_manager._current_version += 1
-    ms2_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=0)
+    ms2_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     assert len(ms2_optimizer.history_df) == 1
 
@@ -462,7 +484,7 @@ def test_automatic_ms2_optimizer():
 
     assert workflow.optimization_manager.classifier_version == -1
 
-    ms2_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=1)
+    ms2_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     calibration_test_df1 = pd.concat(
         [calibration_test_df1, pd.DataFrame(calibration_test_df1.loc[0]).T],
@@ -472,16 +494,17 @@ def test_automatic_ms2_optimizer():
 
     assert workflow.optimization_manager.classifier_version == -1
 
-    ms2_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=2)
+    ms2_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     assert ms2_optimizer.has_converged is True
     assert (
-        ms2_optimizer.history_df.precursor_count == pd.Series([1000, 1001, 1002])
+        ms2_optimizer.history_df.precursor_proportion_detected
+        == pd.Series([1000 / 2000, 1001 / 2000, 1002 / 2000])
     ).all()
     assert (
         workflow.optimization_manager.ms2_error
         == ms2_optimizer.history_df.parameter[
-            ms2_optimizer.history_df.precursor_count.idxmax()
+            ms2_optimizer.history_df.precursor_proportion_detected.idxmax()
         ]
     )
     assert workflow.optimization_manager.classifier_version == 2
@@ -504,7 +527,7 @@ def test_automatic_rt_optimizer():
     assert rt_optimizer.parameter_name == "rt_error"
 
     workflow.fdr_manager._current_version += 1
-    rt_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=0)
+    rt_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     assert len(rt_optimizer.history_df) == 1
 
@@ -516,7 +539,7 @@ def test_automatic_rt_optimizer():
 
     assert workflow.optimization_manager.classifier_version == -1
 
-    rt_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=1)
+    rt_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     calibration_test_df1 = pd.concat(
         [calibration_test_df1, pd.DataFrame(calibration_test_df1.loc[0]).T],
@@ -526,16 +549,17 @@ def test_automatic_rt_optimizer():
 
     assert workflow.optimization_manager.classifier_version == -1
 
-    rt_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=2)
+    rt_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     assert rt_optimizer.has_converged is True
     assert (
-        rt_optimizer.history_df.precursor_count == pd.Series([1000, 1001, 1002])
+        rt_optimizer.history_df.precursor_proportion_detected
+        == pd.Series([1000 / 2000, 1001 / 2000, 1002 / 2000])
     ).all()
     assert (
         workflow.optimization_manager.rt_error
         == rt_optimizer.history_df.parameter[
-            rt_optimizer.history_df.precursor_count.idxmax()
+            rt_optimizer.history_df.precursor_proportion_detected.idxmax()
         ]
     )
     assert workflow.optimization_manager.classifier_version == 2
@@ -558,7 +582,7 @@ def test_automatic_ms1_optimizer():
     assert ms1_optimizer.parameter_name == "ms1_error"
 
     workflow.fdr_manager._current_version += 1
-    ms1_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=0)
+    ms1_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     assert len(ms1_optimizer.history_df) == 1
 
@@ -570,7 +594,7 @@ def test_automatic_ms1_optimizer():
 
     assert workflow.optimization_manager.classifier_version == -1
 
-    ms1_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=1)
+    ms1_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     calibration_test_df1 = pd.concat(
         [calibration_test_df1, pd.DataFrame(calibration_test_df1.loc[0]).T],
@@ -580,7 +604,7 @@ def test_automatic_ms1_optimizer():
 
     assert workflow.optimization_manager.classifier_version == -1
 
-    ms1_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=2)
+    ms1_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     assert ms1_optimizer.has_converged is True
     assert (
@@ -609,7 +633,7 @@ def test_automatic_mobility_optimizer():
     assert mobility_optimizer.parameter_name == "mobility_error"
 
     workflow.fdr_manager._current_version += 1
-    mobility_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=0)
+    mobility_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     assert len(mobility_optimizer.history_df) == 1
 
@@ -620,7 +644,7 @@ def test_automatic_mobility_optimizer():
     workflow.fdr_manager._current_version += 1
 
     assert workflow.optimization_manager.classifier_version == -1
-    mobility_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=1)
+    mobility_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     calibration_test_df1 = pd.concat(
         [calibration_test_df1, pd.DataFrame(calibration_test_df1.loc[0]).T],
@@ -630,16 +654,17 @@ def test_automatic_mobility_optimizer():
 
     assert workflow.optimization_manager.classifier_version == -1
 
-    mobility_optimizer.step(calibration_test_df1, calibration_test_df2, current_step=2)
+    mobility_optimizer.step(calibration_test_df1, calibration_test_df2)
 
     assert mobility_optimizer.has_converged is True
     assert (
-        mobility_optimizer.history_df.precursor_count == pd.Series([1000, 1001, 1002])
+        mobility_optimizer.history_df.precursor_proportion_detected
+        == pd.Series([1000 / 2000, 1001 / 2000, 1002 / 2000])
     ).all()
     assert (
         workflow.optimization_manager.mobility_error
         == mobility_optimizer.history_df.parameter[
-            mobility_optimizer.history_df.precursor_count.idxmax()
+            mobility_optimizer.history_df.precursor_proportion_detected.idxmax()
         ]
     )
     assert workflow.optimization_manager.classifier_version == 2
@@ -665,9 +690,7 @@ def test_targeted_ms2_optimizer():
         assert optimizer.has_converged is False
 
         workflow.fdr_manager._current_version += 1
-        optimizer.step(
-            calibration_test_df1, calibration_test_df2, current_step=current_step
-        )
+        optimizer.step(calibration_test_df1, calibration_test_df2)
 
         assert workflow.optimization_manager.classifier_version == current_step
 
@@ -695,9 +718,7 @@ def test_targeted_rt_optimizer():
         assert optimizer.has_converged is False
 
         workflow.fdr_manager._current_version += 1
-        optimizer.step(
-            calibration_test_df1, calibration_test_df2, current_step=current_step
-        )
+        optimizer.step(calibration_test_df1, calibration_test_df2)
 
         assert workflow.optimization_manager.classifier_version == current_step
 
@@ -725,9 +746,7 @@ def test_targeted_ms1_optimizer():
         assert optimizer.has_converged is False
 
         workflow.fdr_manager._current_version += 1
-        optimizer.step(
-            calibration_test_df1, calibration_test_df2, current_step=current_step
-        )
+        optimizer.step(calibration_test_df1, calibration_test_df2)
 
         assert workflow.optimization_manager.classifier_version == current_step
 
@@ -755,12 +774,99 @@ def test_targeted_mobility_optimizer():
         assert optimizer.has_converged is False
 
         workflow.fdr_manager._current_version += 1
-        optimizer.step(
-            calibration_test_df1, calibration_test_df2, current_step=current_step
-        )
+        optimizer.step(calibration_test_df1, calibration_test_df2)
 
         assert workflow.optimization_manager.classifier_version == current_step
 
     assert optimizer.has_converged is True
 
     assert workflow.optimization_manager.mobility_error == optimizer.target_parameter
+
+
+class TEST_LIBRARY:
+    def __init__(self):
+        idxes = np.arange(0, 100000)
+        self._precursor_df = pd.DataFrame({"elution_group_idx": idxes})
+        self._fragments_df = pd.DataFrame(
+            {
+                "precursor_idx": np.concatenate(
+                    [np.full(10, precursor_idx) for precursor_idx in idxes]
+                )
+            }
+        )
+
+
+TEST_OPTLOCK_CONFIG = {
+    "calibration": {
+        "batch_size": 8000,
+        "optimization_lock_target": 200,
+    }
+}
+
+
+def test_optlock():
+    library = TEST_LIBRARY()
+    optlock = optimization.OptimizationLock(library, TEST_OPTLOCK_CONFIG)
+
+    assert optlock.start_idx == optlock.batch_plan[0][0]
+
+    feature_df = pd.DataFrame({"col": np.arange(0, 1000)})
+    fragment_df = pd.DataFrame({"col": np.arange(0, 10000)})
+
+    optlock.update_with_extraction(feature_df, fragment_df)
+
+    assert optlock.total_precursors == 1000
+    precursor_df = pd.DataFrame(
+        {"qval": np.concatenate([np.full(100, 0.005), np.full(1000, 0.05)])}
+    )
+    optlock.update_with_fdr(precursor_df)
+
+    assert optlock.has_target_num_precursors is False
+    assert optlock.first_time_to_reach_target is True
+    optlock.update()
+
+    assert optlock.start_idx == optlock.batch_plan[1][0]
+
+    feature_df = pd.DataFrame({"col": np.arange(0, 1000)})
+    fragment_df = pd.DataFrame({"col": np.arange(0, 10000)})
+
+    optlock.update_with_extraction(feature_df, fragment_df)
+
+    assert optlock.total_precursors == 2000
+
+    precursor_df = pd.DataFrame(
+        {"qval": np.concatenate([np.full(200, 0.005), np.full(1000, 0.05)])}
+    )
+
+    optlock.update_with_fdr(precursor_df)
+
+    assert optlock.has_target_num_precursors is True
+    assert optlock.first_time_to_reach_target is True
+    optlock.update()
+
+    assert optlock.first_time_to_reach_target is False
+    assert optlock.start_idx == 0
+
+
+def test_optlock_batch_idx():
+    library = TEST_LIBRARY()
+    optlock = optimization.OptimizationLock(library, TEST_OPTLOCK_CONFIG)
+
+    optlock.batch_plan = [[0, 100], [100, 2000], [2000, 8000]]
+
+    assert optlock.start_idx == 0
+
+    optlock.increase_batch_idx()
+    assert optlock.start_idx == 100
+
+    optlock.increase_batch_idx()
+    assert optlock.start_idx == 2000
+
+    precursor_df = pd.DataFrame({"qval": np.full(4500, 0.005)})
+
+    optlock.update_with_fdr(precursor_df)
+
+    optlock.decrease_batch_idx()
+
+    assert optlock.start_idx == 0
+    assert optlock.stop_idx == 2000