mimicpy 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

mimicpy/force_matching/opt_ff.py

@@ -0,0 +1,2114 @@
+import numpy as np
+import logging
+from ..utils.errors import MiMiCPyError
+from ..utils.constants import au_to_nm, kb_au2gmx, au_kjm, nm_to_au, kb_gmx2au, kjm_au, kb_au2g96, kb_g962au
+from .bonded_forces import *
+from .nonbonded_forces import *
+from ..scripts.fm_input import FMInput
+from scipy.optimize import least_squares
+import MDAnalysis as mda 
+import multiprocessing as mp
+import random
+
+
+section_atom_counts = {
+        "bonds": 2,
+        "angles": 3,
+        "dihedrals": 4,
+        "pairs": 2,
+    }
+
+class L2Regularizer:
+    """
+    L2 regularization for force field parameters.
+    
+    Implements the harmonic penalty function:
+    Θ(p) = α * Σ_i (p_i - p_i^0)^2 / (2 * γ_i^2)
+    
+    where:
+    - p_i are the current parameters
+    - p_i^0 are the initial/reference parameters  
+    - γ_i are the prior widths (controls how much parameters can deviate)
+    - α is the regularization strength
+    """
+    
+    def __init__(self, initial_params, prior_widths=None, alpha=0.1, param_types=None):
+        """
+        Initialize L2 regularizer.
+        
+        Args:
+            initial_params (np.ndarray): Initial parameter values (p^0)
+            prior_widths (np.ndarray, optional): Prior widths (γ). If None, auto-generated
+            alpha (float): Regularization strength (α)
+            param_types (list, optional): List of parameter type strings for type-specific widths
+        """
+        self.initial_params = np.array(initial_params)
+        self.alpha = alpha
+        self.param_types = param_types
+        
+        if prior_widths is None:
+            # Auto-generate prior widths based on parameter classes
+            self.prior_widths = self._auto_generate_prior_widths(initial_params)
+        else:
+            self.prior_widths = np.array(prior_widths)
+            
+        # Ensure all arrays have same length
+        assert len(self.initial_params) == len(self.prior_widths), \
+            "Initial parameters and prior widths must have same length"
+    
+    def _auto_generate_prior_widths(self, params):
+        """
+        Auto-generate prior widths based on parameter classes.
+        This is a simplified version - in practice you'd want to group by parameter type.
+        """
+        # For now, use a simple approach: 10% of parameter value or 0.1 if parameter is 0
+        widths = np.abs(params) * 0.05
+        widths[widths == 0] = 0.05  # Default width for zero parameters
+        return widths
+    
+    def set_prior_widths_by_type(self, type_widths):
+        """
+        Set prior widths based on parameter types.
+        
+        Args:
+            type_widths (dict): Dictionary mapping parameter types to widths
+                Example: {'bond_length': 0.1, 'bond_force': 50.0, 'angle_value': 0.1, ...}
+        """
+        if self.param_types is None:
+            raise ValueError("Parameter types must be set to use type-specific prior widths")
+        
+        # Add safety check for array bounds
+        if len(self.param_types) != len(self.prior_widths):
+            raise MiMiCPyError(f"Parameter types length ({len(self.param_types)}) does not match prior widths length ({len(self.prior_widths)})")
+        
+        for i, param_type in enumerate(self.param_types):
+            if param_type in type_widths:
+                self.prior_widths[i] = type_widths[param_type]
+    
+    def set_prior_widths_from_fm_input(self, fm_input):
+        """
+        Set prior widths based on FMInput parameters.
+        
+        Args:
+            fm_input: FMInput object containing regularization parameters
+        """
+        type_widths = {}
+        
+        # Map FMInput parameters to type widths
+        if fm_input.regularization_bond_length_width is not None:
+            type_widths['bond_length'] = fm_input.regularization_bond_length_width
+        if fm_input.regularization_bond_force_width is not None:
+            type_widths['bond_force'] = fm_input.regularization_bond_force_width
+        if fm_input.regularization_angle_value_width is not None:
+            type_widths['angle_value'] = fm_input.regularization_angle_value_width
+        if fm_input.regularization_angle_force_width is not None:
+            type_widths['angle_force'] = fm_input.regularization_angle_force_width
+        if fm_input.regularization_dihedral_force_width is not None:
+            type_widths['dihedral_force'] = fm_input.regularization_dihedral_force_width
+        
+        if type_widths:
+            self.set_prior_widths_by_type(type_widths)
+    
+    def compute_regularization_term(self, current_params):
+        """
+        Compute the L2 regularization term.
+        
+        Args:
+            current_params (np.ndarray): Current parameter values
+            
+        Returns:
+            float: Regularization penalty
+        """
+        diff = current_params - self.initial_params
+        penalty = np.sum(diff**2 / (2 * self.prior_widths**2))
+        return self.alpha * penalty
+    
+    def compute_regularization_gradient(self, current_params):
+        """
+        Compute the gradient of the L2 regularization term.
+        
+        Args:
+            current_params (np.ndarray): Current parameter values
+            
+        Returns:
+            np.ndarray: Gradient of regularization term
+        """
+        diff = current_params - self.initial_params
+        gradient = diff / self.prior_widths**2
+        return self.alpha * gradient
+    
+    def compute_regularization_residuals(self, current_params):
+        """
+        Compute regularization residuals for augmented residual approach.
+        
+        This method returns residuals that can be appended to the force matching
+        residuals for use with scipy.least_squares Levenberg-Marquardt method.
+        
+        Args:
+            current_params (np.ndarray): Current parameter values
+            
+        Returns:
+            np.ndarray: Regularization residuals
+        """
+        diff = current_params - self.initial_params
+        # Scale the residuals by sqrt(alpha/2) and prior widths to match the regularization term
+        # The regularization term is: α * Σ(p_i - p_i^0)² / (2 * γ_i²)
+        # So the residuals should be: √(α/2) * (p_i - p_i^0) / γ_i
+        residuals = np.sqrt(self.alpha / 2) * diff / self.prior_widths
+        return residuals
+
+class AdaptiveL2Regularizer:
+    """
+    Adaptive L2 regularization that applies different regularization strengths
+    based on the energy hierarchy of interactions.
+    
+    Bonds (highest energy) get stronger regularization to prevent overfitting,
+    while dihedrals (lowest energy) get weaker regularization to allow more flexibility.
+    """
+    
+    def __init__(self, initial_params, bond2params, bonds, angles, dihedrals, 
+                 base_alpha=0.1, energy_hierarchy_scale=10.0, param_types=None):
+        """
+        Initialize adaptive regularizer.
+        
+        Args:
+            initial_params (np.ndarray): Initial parameter values
+            bond2params (dict): Mapping of interaction indices to parameter indices
+            bonds, angles, dihedrals (list): Lists of interactions
+            base_alpha (float): Base regularization strength
+            energy_hierarchy_scale (float): Scale factor for energy hierarchy
+            param_types (list, optional): List of parameter type strings
+        """
+        self.initial_params = np.array(initial_params)
+        self.base_alpha = base_alpha
+        self.energy_hierarchy_scale = energy_hierarchy_scale
+        self.param_types = param_types
+        
+        # Create parameter masks for each interaction type
+        self.bond_params_mask = np.zeros(len(initial_params), dtype=bool)
+        self.angle_params_mask = np.zeros(len(initial_params), dtype=bool)
+        self.dihedral_params_mask = np.zeros(len(initial_params), dtype=bool)
+        
+        # Identify which parameters belong to which interaction types
+        for bond in bonds:
+            if bond.get('optimize', False):
+                param_indices = bond2params.get(bond['index'], [])
+                for idx in param_indices:
+                    if idx is not None:
+                        self.bond_params_mask[idx] = True
+        
+        for angle in angles:
+            if angle.get('optimize', False):
+                param_indices = bond2params.get(angle['index'], [])
+                for idx in param_indices:
+                    if idx is not None:
+                        self.angle_params_mask[idx] = True
+        
+        for dihedral in dihedrals:
+            if dihedral.get('optimize', False):
+                param_indices = bond2params.get(dihedral['index'], [])
+                for idx in param_indices:
+                    if idx is not None:
+                        self.dihedral_params_mask[idx] = True
+        
+        # Generate adaptive prior widths based on energy hierarchy
+        self.prior_widths = self._generate_adaptive_prior_widths(initial_params)
+        
+        # Generate adaptive alpha values for each parameter
+        self.adaptive_alphas = self._generate_adaptive_alphas()
+    
+    def _generate_adaptive_prior_widths(self, params):
+        """Generate prior widths that respect energy hierarchy."""
+        widths = np.abs(params) * 0.1
+        widths[widths == 0] = 0.1  # Default width for zero parameters
+        
+        # Apply energy hierarchy scaling
+        # Bonds: tighter regularization (smaller widths)
+        widths[self.bond_params_mask] *= 0.1
+        
+        # Angles: medium regularization
+        widths[self.angle_params_mask] *= 0.5
+        
+        # Dihedrals: looser regularization (larger widths)
+        widths[self.dihedral_params_mask] *= 2.0
+        
+        return widths
+    
+    def _generate_adaptive_alphas(self):
+        """Generate adaptive alpha values based on energy hierarchy."""
+        alphas = np.full(len(self.initial_params), self.base_alpha)
+        
+        # Bonds: stronger regularization (higher alpha)
+        alphas[self.bond_params_mask] *= self.energy_hierarchy_scale
+        
+        # Angles: medium regularization
+        alphas[self.angle_params_mask] *= np.sqrt(self.energy_hierarchy_scale)
+        
+        # Dihedrals: weaker regularization (lower alpha)
+        alphas[self.dihedral_params_mask] *= 1.0 / self.energy_hierarchy_scale
+        
+        return alphas
+    
+    def compute_regularization_term(self, current_params):
+        """Compute adaptive regularization term."""
+        diff = current_params - self.initial_params
+        penalty = np.sum(self.adaptive_alphas * diff**2 / (2 * self.prior_widths**2))
+        return penalty
+    
+    def compute_regularization_gradient(self, current_params):
+        """Compute gradient of adaptive regularization term."""
+        diff = current_params - self.initial_params
+        gradient = self.adaptive_alphas * diff / self.prior_widths**2
+        return gradient
+    
+    def compute_regularization_residuals(self, current_params):
+        """Compute adaptive regularization residuals."""
+        diff = current_params - self.initial_params
+        # Scale residuals by sqrt(adaptive_alpha/2) and prior widths
+        residuals = np.sqrt(self.adaptive_alphas / 2) * diff / self.prior_widths
+        return residuals
+
+class ParameterOptimizer:
+    def __init__(self, eq_mapping):
+        self.eq_mapping = eq_mapping
+        self.ff_optimize = []
+        self.bond2params = {}
+        self.param_types = []  # Track parameter types as they're added
+        self.counter = 0
+        self._selected_interactions = {
+            'bonds': set(),
+            'angles': set(),
+            'dihedrals': set()
+        }
+        # Store interactions for equivalence checking
+        self._bonds = []
+        self._angles = []
+        self._dihedrals = []
+
+    def select_interactions(self, interaction_type, indices=None, all=False):
+        """
+        Select specific interactions for optimization.
+        
+        Args:
+            interaction_type (str): Type of interaction ('bonds', 'angles', or 'dihedrals')
+            indices (list, optional): List of interaction indices to select. If None, selects all.
+            all (bool, optional): If True, selects all interactions of the given type.
+        """
+        if interaction_type not in self._selected_interactions:
+            raise ValueError(f"Invalid interaction type: {interaction_type}")
+            
+        if all:
+            self._selected_interactions[interaction_type] = set()
+        elif indices is not None:
+            self._selected_interactions[interaction_type] = set(indices)
+        else:
+            self._selected_interactions[interaction_type] = set()
+
+    def is_selected(self, interaction_type, index):
+        """Check if an interaction is selected for optimization."""
+        if not self._selected_interactions[interaction_type]:
+            return True  # If no specific selections, optimize all
+        return index in self._selected_interactions[interaction_type]
+
+    def optimize_bond(self, bond, optimize_length=True, optimize_force_constant=True, hydrogen_indices=None, exclude_hydrogen=False):
+      
+        if exclude_hydrogen and hydrogen_indices is not None and any(atom in hydrogen_indices for atom in bond['atoms']):
+            bond['optimize'] = False
+            self.bond2params[bond['index']] = [None, None]
+            return
+        if not self.is_selected('bonds', bond['index']):
+            bond['optimize'] = False
+            self.bond2params[bond['index']] = [None, None]
+            return
+        if bond['parameters'][1] == 0.0:
+            idx = [None, None]
+            self.bond2params[bond['index']] = idx
+            bond['optimize'] = False
+            return
+        bond['optimize'] = True
+        params_to_optimize = []
+        param_types = []
+        # Always maintain the order: [length, force_constant]
+        if optimize_length:
+            params_to_optimize.append(bond['parameters'][0])
+            param_types.append('bond_length')
+        else:
+            params_to_optimize.append(None)
+            param_types.append(None)
+        if optimize_force_constant:
+            params_to_optimize.append(bond['parameters'][1])
+            param_types.append('bond_force')
+        else:
+            params_to_optimize.append(None)
+            param_types.append(None)
+        self._add_parameters(bond, params_to_optimize, param_types)
+        self._bonds.append(bond)
+
+    def optimize_angle(self, angle, optimize_angle=True, optimize_force_constant=True, hydrogen_indices=None, exclude_hydrogen=False):
+        if exclude_hydrogen and hydrogen_indices is not None and any(atom in hydrogen_indices for atom in angle['atoms']):
+            angle['optimize'] = False
+            self.bond2params[angle['index']] = [None, None]
+            return
+        if not self.is_selected('angles', angle['index']):
+            angle['optimize'] = False
+            self.bond2params[angle['index']] = [None, None]
+            return
+        if angle['parameters'][1] == 0.0:
+            idx = [None, None]
+            self.bond2params[angle['index']] = idx
+            angle['optimize'] = False
+            return
+        angle['optimize'] = True
+        params_to_optimize = []
+        param_types = []
+        # Always maintain the order: [angle_value, force_constant]
+        if optimize_angle:
+            params_to_optimize.append(angle['parameters'][0])
+            param_types.append('angle_value')
+        else:
+            params_to_optimize.append(None)
+            param_types.append(None)
+        if optimize_force_constant:
+            params_to_optimize.append(angle['parameters'][1])
+            param_types.append('angle_force')
+        else:
+            params_to_optimize.append(None)
+            param_types.append(None)
+        self._add_parameters(angle, params_to_optimize, param_types)
+        self._angles.append(angle)
+
+    def optimize_dihedral(self, dihedral, optimize_force_constant=True, hydrogen_indices=None, exclude_hydrogen=False):
+        if exclude_hydrogen and hydrogen_indices is not None and any(atom in hydrogen_indices for atom in dihedral['atoms']):
+            dihedral['optimize'] = False
+            if dihedral['function'] in [1, 4, 9]:
+                self.bond2params[dihedral['index']] = [None, None, None]
+            elif dihedral['function'] == 2:
+                self.bond2params[dihedral['index']] = [None, None]
+            elif dihedral['function'] == 3:
+                self.bond2params[dihedral['index']] = [None] * 6
+            return
+        if not self.is_selected('dihedrals', dihedral['index']):
+            dihedral['optimize'] = False
+            if dihedral['function'] in [1, 4, 9]:
+                self.bond2params[dihedral['index']] = [None, None, None]
+            elif dihedral['function'] == 2:
+                self.bond2params[dihedral['index']] = [None, None]
+            elif dihedral['function'] == 3:
+                self.bond2params[dihedral['index']] = [None] * 6
+            return
+        dihedral['optimize'] = True
+        params_to_optimize = []
+        param_types = []
+        if dihedral['function'] in [1, 4, 9]:  # Format 1
+            if dihedral['parameters'][1] == 0.0:  # Check force constant
+                self.bond2params[dihedral['index']] = [None, None, None]
+                dihedral['optimize'] = False
+                return
+            params_to_optimize = [None, None, None]  # [phi0, cp, mult]
+            param_types = [None, None, None]
+            if optimize_force_constant:
+                params_to_optimize[1] = dihedral['parameters'][1]  # cp
+                param_types[1] = 'dihedral_force'
+        elif dihedral['function'] == 2:  # Format 2
+            if dihedral['parameters'][1] == 0.0:  # Check force constant
+                self.bond2params[dihedral['index']] = [None, None]
+                dihedral['optimize'] = False
+                return
+            params_to_optimize = [None, None]  # [param1, param2]
+            param_types = [None, None]
+            if optimize_force_constant:
+                params_to_optimize[1] = dihedral['parameters'][1]  # param2
+                param_types[1] = 'dihedral_force'
+        elif dihedral['function'] == 3:  # Format 3 (Ryckaert-Bellemans)
+            if all(p == 0.0 for p in dihedral['parameters']):
+                self.bond2params[dihedral['index']] = [None] * 6
+                dihedral['optimize'] = False
+                return
+            params_to_optimize = [None] * 6  # [C0, C1, C2, C3, C4, C5]
+            param_types = [None] * 6
+            if optimize_force_constant:
+                for i, param in enumerate(dihedral['parameters']):
+                    if param != 0.0:
+                        params_to_optimize[i] = param
+                        param_types[i] = 'dihedral_force'
+        self._add_parameters(dihedral, params_to_optimize, param_types)
+        self._dihedrals.append(dihedral)
+
+    def _add_parameters(self, interaction, params_to_optimize, param_types_to_add=None):
+        # Check if equivalent interaction already exists
+        for existing_idx, existing_params in self.bond2params.items():
+            if self._is_equivalent(interaction, existing_idx):
+                self.bond2params[interaction['index']] = existing_params
+                return
+
+        # Determine the expected structure based on interaction type
+        n_atoms = len(interaction['atoms'])
+        if n_atoms == 2:  # bond
+            expected_length = 2
+        elif n_atoms == 3:  # angle
+            expected_length = 2
+        elif n_atoms == 4:  # dihedral
+            if interaction['function'] in [1, 4, 9]:
+                expected_length = 3
+            elif interaction['function'] == 2:
+                expected_length = 2
+            elif interaction['function'] == 3:
+                expected_length = 6
+            else:
+                expected_length = 3  # default
+        else:
+            expected_length = 2  # default
+
+        # Initialize param_indices with None values
+        param_indices = [None] * expected_length
+        
+        # Add parameters that are being optimized
+        for i, param in enumerate(params_to_optimize):
+            if param is not None:
+                self.ff_optimize.append(param)
+                param_indices[i] = len(self.ff_optimize) - 1
+                # Add parameter type if provided
+                if param_types_to_add and i < len(param_types_to_add) and param_types_to_add[i] is not None:
+                    self.param_types.append(param_types_to_add[i])
+            # If param is None, param_indices[i] remains None
+        
+        self.bond2params[interaction['index']] = param_indices
+
+    def _is_equivalent(self, interaction1, interaction2_idx):
+        """Check if two interactions are equivalent based on their atoms and function type.
+        
+        Args:
+            interaction1 (dict): First interaction to compare
+            interaction2_idx (int): Index of second interaction in bond2params
+            
+        Returns:
+            bool: True if interactions are equivalent, False otherwise
+        """
+        # Get the second interaction from the stored interactions
+        interaction2 = None
+        for interaction_type in ['bonds', 'angles', 'dihedrals']:
+            for interaction in getattr(self, f'_{interaction_type}', []):
+                if interaction['index'] == interaction2_idx:
+                    interaction2 = interaction
+                    break
+            if interaction2:
+                break
+                
+        if not interaction2:
+            return False
+            
+        # Check if function types match
+        if interaction1['function'] != interaction2['function']:
+            return False
+            
+        # Determine interaction type based on number of atoms
+        n_atoms1 = len(interaction1['atoms'])
+        n_atoms2 = len(interaction2['atoms'])
+        
+        if n_atoms1 != n_atoms2:
+            return False
+            
+        # Check equivalence based on number of atoms
+        if n_atoms1 == 2:  # bonds
+            return check_bond_equivalence(interaction1['atoms'], interaction2['atoms'], self.eq_mapping)
+        elif n_atoms1 == 3:  # angles
+            return check_angle_equivalence(interaction1['atoms'], interaction2['atoms'], self.eq_mapping)
+        elif n_atoms1 == 4:  # dihedrals
+            return check_dihedral_equivalence(interaction1['atoms'], interaction2['atoms'], self.eq_mapping)
+        return False
+
+    def get_optimized_parameters(self):
+        return np.array(self.ff_optimize), self.bond2params, self.param_types
+
+
+
+def get_configurations_optff(fmdata, tpr_file, trr_file, begin, end, step,
+                             qm_region):
+    """
+    Get configurations for force field optimization.
+    
+    Args:
+        fmdata (FMDataset): FMDataset object containing QM data (assumed to use CPMD IDs).
+        tpr_file (str): GROMACS TPR file.
+        trr_file (str): GROMACS TRR file.
+        begin (int): Starting frame.
+        end (int): Ending frame.
+        step (int): Step size.
+        qm_region (QMRegion): QMRegion object to get GROMACS to CPMD mapping.
+        
+    Returns:
+        list: List of configuration dictionaries with QM data reordered to match
+              GROMACS topology order, including MM atoms bonded to QM atoms.
+    """
+    u = mda.Universe(tpr_file, trr_file)
+    configurations = []
+    
+    # Get GROMACS to CPMD mapping (CPMD IDs are 1-based)
+    gmx_to_cpmd_map = qm_region.gmx_to_cpmd_map
+    qm_atoms_gmx_indices = qm_region.qm_atoms.index
+    
+    # Get the list of CPMD IDs from FMdata for the first configuration (assuming it's constant)
+    # These are assumed to be 1-based CPMD IDs.
+    cpmd_ids_fmdata_order = fmdata.get_configuration_properties(begin, 'id', 'qm')
+    
+
+    # Create mapping from FMdata's CPMD-based order to the desired GROMACS topology order
+    # fm_to_top_order[i] will be the index in cpmd_ids_fmdata_order (and thus in fm_coords/fm_forces)
+    # that corresponds to the i-th GROMACS atom in qm_atoms_gmx_indices.
+    fm_to_top_order = np.zeros(len(qm_atoms_gmx_indices), dtype=int)
+    for i, gmx_idx in enumerate(qm_atoms_gmx_indices):
+        target_cpmd_id = gmx_to_cpmd_map.get(gmx_idx)
+        if target_cpmd_id is None:
+            raise ValueError(f"GROMACS atom {gmx_idx} not found in gmx_to_cpmd_map.")
+            
+        fm_idx_arr = np.where(cpmd_ids_fmdata_order == target_cpmd_id)[0]
+        if len(fm_idx_arr) > 0:
+            fm_to_top_order[i] = fm_idx_arr[0]
+        else:
+            raise ValueError(
+                f"CPMD ID {target_cpmd_id} (for GMX atom {gmx_idx}) "
+                f"not found in FMdata's list of CPMD IDs: {cpmd_ids_fmdata_order}"
+            )
+    if not qm_region.boundary_atoms.empty:
+        gmx_ids_mm_order = fmdata.get_configuration_properties(begin, 'id', 'mm')
+        # Create mapping from GROMACS IDs to indices in MM coordinates
+        gmx_to_mm_idx = {gmx_id: idx for idx, gmx_id in enumerate(gmx_ids_mm_order)}
+        for idx in range(begin, end, step):
+            config = dict()
+            
+            # Get QM data from FMdata
+            fm_coords = fmdata.get_configuration_properties(idx, 'coordinate', 'qm')
+            fm_forces = fmdata.get_configuration_properties(idx, 'force', 'qm')
+           
+            # Reorder coordinates and forces to match GROMACS topology order
+            config['qm_coordinates'] = fm_coords[fm_to_top_order]
+            config['qm_forces'] = fm_forces[fm_to_top_order]
+            
+            # Get MM coordinates from FMData
+            fm_mm_coords = fmdata.get_configuration_properties(idx, 'coordinate', 'mm')
+            
+            # Get extended QM atoms DataFrame
+            extended_df = qm_region.extended_qm_atoms
+            
+            # Step 1: Get QM atom coordinates (already have them in config['qm_coordinates'])
+            qm_coords = config['qm_coordinates']
+            
+            # Step 2: Get MM atom coordinates
+            mm_atoms = extended_df[extended_df['is_qm'] == 0]  # MM atoms
+            mm_coords = np.zeros((len(mm_atoms), 3))
+            
+            for mm_idx, (gmx_idx, _) in enumerate(mm_atoms.iterrows()):
+                if gmx_idx in gmx_to_mm_idx:
+                    fm_mm_idx = gmx_to_mm_idx[gmx_idx]
+                    mm_coords[mm_idx] = fm_mm_coords[fm_mm_idx]
+                else:
+                    logging.warning(f'Could not find MM atom {gmx_idx} in FMData MM coordinates')
+            
+            # Step 3: Concatenate QM and MM coordinates to create extended coordinates
+            extended_coords = np.vstack([qm_coords, mm_coords])
+            
+            # Get GROMACS forces for extended QM atom set
+            forces, positions = get_qm_gmx_forces(u, idx, qm_atoms_gmx_indices)
+            
+            config['qm_gmx_forces'] = forces
+            config['qm_coordinates'] = extended_coords
+            
+            configurations.append(config)
+    else:
+        for idx in range(begin, end, step):
+            config = dict()
+            
+            # Get QM data from FMdata
+            fm_coords = fmdata.get_configuration_properties(idx, 'coordinate', 'qm')
+            fm_forces = fmdata.get_configuration_properties(idx, 'force', 'qm')
+            
+            # Reorder coordinates and forces to match GROMACS topology order
+            config['qm_coordinates'] = fm_coords[fm_to_top_order]
+            config['qm_forces'] = fm_forces[fm_to_top_order]
+            
+            # Get GROMACS forces and positions (already in GROMACS topology order via qm_atoms_gmx_indices)
+            forces, positions = get_qm_gmx_forces(u, idx, qm_atoms_gmx_indices)
+            config['qm_gmx_forces'] = forces
+            config['qm_gmx_coordinates'] = positions
+            
+            configurations.append(config)
+    return configurations
+
+
+def compute_residue(ff_optimize,
+                     qmmm_forces,
+                    gmx_force, qm_coordinates, 
+                    bonds,
+                    angles,
+                    dihedrals,
+                    bond2params):
+    
+
+    # Compute bonded forces for the extended set of atoms
+    # Pass qm_atoms_count to get only QM atom forces
+    qm_bonded_forces = compute_bonded_forces(ff_optimize, qm_coordinates,
+                                          bonds, angles, dihedrals, bond2params,
+                                          qm_atoms_count=len(qmmm_forces))
+    
+    reference_bonded_force = qmmm_forces - gmx_force
+    residue = qm_bonded_forces - reference_bonded_force
+
+    return residue, reference_bonded_force
+
+
+def _process_single_configuration(args):
+    """
+    Process a single configuration for parallel computation.
+    
+    Args:
+        args: Tuple containing (ff_optimize, config, bonds, angles, dihedrals, bond2params)
+        
+    Returns:
+        tuple: (residual, ref_bonded_force)
+    """
+    ff_optimize, config, bonds, angles, dihedrals, bond2params = args
+    
+    qmmm_forces = config['qm_forces']
+    residual, ref_bonded = compute_residue(ff_optimize,
+                                          qmmm_forces,
+                                          config["qm_gmx_forces"],
+                                          config["qm_coordinates"], 
+                                          bonds,
+                                          angles,
+                                          dihedrals,
+                                          bond2params)
+    
+    return residual, ref_bonded
+
+
+def compute_ff_obj(ff_optimize, configurations, 
+                   bonds, angles, dihedrals, bond2params, 
+                   n_processes=None):
+    """
+    Compute force field objective function with optional parallelization.
+    
+    Args:
+        ff_optimize (numpy.ndarray): Force field parameters to optimize
+        configurations (list): List of configuration dictionaries
+        bonds, angles, dihedrals (list): Interaction lists
+        bond2params (dict): Parameter mapping
+        n_processes (int, optional): Number of processes to use. If None, uses serial processing
+        
+    Returns:
+        tuple: (total_residual, ref_bonded_forces)
+    """
+    if type(configurations) != list:
+        configurations = [configurations]
+    
+    # Use parallel processing if requested and beneficial
+    if n_processes is not None and n_processes > 1 and len(configurations) > 1:
+        return _compute_ff_obj_parallel(ff_optimize, configurations, 
+                                       bonds, angles, dihedrals, bond2params, n_processes)
+    else:
+        return _compute_ff_obj_serial(ff_optimize, configurations, 
+                                     bonds, angles, dihedrals, bond2params)
+
+
+def _compute_ff_obj_serial(ff_optimize, configurations, 
+                          bonds, angles, dihedrals, bond2params):
+    """Serial version of compute_ff_obj (original implementation)."""
+    total_residual = []
+    ref_bonded_forces = []
+    for config in configurations:
+        qmmm_forces = config['qm_forces']
+
+        residual, ref_bonded  = compute_residue(ff_optimize,
+                                                qmmm_forces,
+                                                config["qm_gmx_forces"],
+                                                config["qm_coordinates"], 
+                                                bonds,
+                                                angles,
+                                                dihedrals,
+                                                bond2params)
+                    
+        ref_bonded_forces.append(ref_bonded)
+        total_residual.append(residual)
+    return np.vstack(total_residual), np.vstack(ref_bonded_forces)
+
+
+def _compute_ff_obj_parallel(ff_optimize, configurations, 
+                            bonds, angles, dihedrals, bond2params, n_processes):
+    """Parallel version of compute_ff_obj using multiprocessing."""
+    # Prepare arguments for parallel processing
+    args_list = [(ff_optimize, config, bonds, angles, dihedrals, bond2params) 
+                 for config in configurations]
+    
+    # Use multiprocessing pool
+    with mp.Pool(processes=n_processes) as pool:
+        results = pool.map(_process_single_configuration, args_list)
+    
+    # Unpack results
+    total_residual = [result[0] for result in results]
+    ref_bonded_forces = [result[1] for result in results]
+    
+    return np.vstack(total_residual), np.vstack(ref_bonded_forces)
+
+
+def check_bond_equivalence(bond1, bond2, eq_mapping):
+    eq = False
+    a1 = eq_mapping.get(bond1[0])
+    b1 = eq_mapping.get(bond1[1])
+    a2 = eq_mapping.get(bond2[0])
+    b2 = eq_mapping.get(bond2[1])
+
+    if a1 is not None and b1 is not None and a2 is not None and b2 is not None:
+        if (a1 == a2 and b1 == b2) or (a1 == b2 and b1 == a2):
+            eq = True
+    return eq
+
+
+def check_angle_equivalence(bond1, bond2, eq_mapping):
+    eq = False
+    a1 = eq_mapping.get(bond1[0])
+    b1 = eq_mapping.get(bond1[1])
+    c1 = eq_mapping.get(bond1[2])
+    a2 = eq_mapping.get(bond2[0])
+    b2 = eq_mapping.get(bond2[1])
+    c2 = eq_mapping.get(bond2[2])
+
+    if b1 == b2:
+        if (a1 == a2 and c1 == c2) or (a1 == c2 and c1 == a2):
+            eq = True
+    return eq
+
+
+def check_dihedral_equivalence(bond1, bond2, eq_mapping):
+    eq = False
+    a1 = eq_mapping.get(bond1[0])
+    b1 = eq_mapping.get(bond1[1])
+    c1 = eq_mapping.get(bond1[2])
+    d1 = eq_mapping.get(bond1[3])
+    a2 = eq_mapping.get(bond2[0])
+    b2 = eq_mapping.get(bond2[1])
+    c2 = eq_mapping.get(bond2[2])
+    d2 = eq_mapping.get(bond2[3])
+
+    if ((a1 == a2 and b1 == b2 and c1 == c2 and d1 == d2) or 
+        (a1 == d2 and b1 == c2 and c1 == b2 and d1 == a2)):
+        eq = True
+    
+    return eq
+
+
+def get_optimize_ff_parameters(qm_region, eq_mapping, fm_input=None):
+    """
+    Get optimized force field parameters with fine-grained control over which interactions to optimize.
+    
+    Args:
+        qm_region (QMRegion): QMRegion object containing QM interactions
+        eq_mapping (dict): Mapping of equivalent atoms
+        fm_input (FMInput, optional): FMInput object containing all optimization and regularization parameters
+        
+    Returns:
+        tuple: (ff_optimize, bond2params, regularizer) where ff_optimize is array of parameters to optimize,
+               bond2params maps interaction indices to parameter indices, and regularizer is the L2Regularizer object
+    """
+    # Get interactions from QMRegion
+    if qm_region.qm_interactions is None:
+        qm_interactions = qm_region.extract_qm_interactions()
+    else:
+        qm_interactions = qm_region.qm_interactions
+    bonds = qm_interactions['bonds']
+    angles = qm_interactions['angles']
+    dihedrals = qm_interactions['dihedrals']
+    
+    # atoms indices are 0-based in the QMRegion
+    hydrogen_indices = set(qm_region.qm_atoms[qm_region.qm_atoms['element'].str.startswith('H')].index - 1)
+    optimizer = ParameterOptimizer(eq_mapping)
+    
+    # Optional: skip optimization for solvent molecules
+    skip_solvent_optimization = getattr(fm_input, 'skip_solvent_optimization', True)
+    
+    # Process bonds
+    for bond in bonds:
+        if skip_solvent_optimization and qm_region.is_solvent_interaction(bond):
+            bond['optimize'] = False
+            optimizer.bond2params[bond['index']] = [None, None]
+            continue
+        # Skip optimization if bond involves boundary atoms (already marked in extract_qm_interactions)
+        if bond.get('involves_boundary', False):
+            bond['optimize'] = False
+            optimizer.bond2params[bond['index']] = [None, None]
+            continue
+            
+        optimizer.optimize_bond(
+            bond,
+            fm_input.optimize_bond_length,
+            fm_input.optimize_bond_force,
+            hydrogen_indices=hydrogen_indices,
+            exclude_hydrogen=fm_input.exclude_hydrogen_bonds
+        )
+    
+    # Process angles
+    for angle in angles:
+        if skip_solvent_optimization and qm_region.is_solvent_interaction(angle):
+            angle['optimize'] = False
+            optimizer.bond2params[angle['index']] = [None, None]
+            continue
+        # Skip optimization if angle involves boundary atoms (already marked in extract_qm_interactions)
+        if angle.get('involves_boundary', False):
+            angle['optimize'] = False
+            optimizer.bond2params[angle['index']] = [None, None]
+            continue
+            
+        optimizer.optimize_angle(
+            angle,
+            fm_input.optimize_angle_value,
+            fm_input.optimize_angle_force,
+            hydrogen_indices=hydrogen_indices,
+            exclude_hydrogen=fm_input.exclude_hydrogen_angles
+        )
+    
+    # Process dihedrals
+    for dihedral in dihedrals:
+        if skip_solvent_optimization and qm_region.is_solvent_interaction(dihedral):
+            dihedral['optimize'] = False
+            if dihedral['function'] in [1, 4, 9]:
+                optimizer.bond2params[dihedral['index']] = [None, None, None]
+            elif dihedral['function'] == 2:
+                optimizer.bond2params[dihedral['index']] = [None, None]
+            elif dihedral['function'] == 3:
+                optimizer.bond2params[dihedral['index']] = [None] * 6
+            continue
+        # Skip optimization if dihedral involves boundary atoms (already marked in extract_qm_interactions)
+        if dihedral.get('involves_boundary', False):
+            dihedral['optimize'] = False
+            if dihedral['function'] in [1, 4, 9]:
+                optimizer.bond2params[dihedral['index']] = [None, None, None]
+            elif dihedral['function'] == 2:
+                optimizer.bond2params[dihedral['index']] = [None, None]
+            elif dihedral['function'] == 3:
+                optimizer.bond2params[dihedral['index']] = [None] * 6
+            continue
+            
+        optimizer.optimize_dihedral(
+            dihedral,
+            fm_input.optimize_dihedral_force,
+            hydrogen_indices=hydrogen_indices,
+            exclude_hydrogen=fm_input.exclude_hydrogen_dihedrals
+        )
+
+    ff_optimize, bond2params, param_types = optimizer.get_optimized_parameters()
+    
+    # Setup regularization if parameters are being optimized
+    if len(ff_optimize) > 0:
+        # Parameter types are now generated during optimization in the ParameterOptimizer
+        # This ensures consistency between parameter optimization and type assignment
+        logging.debug(f"Generated {len(param_types)} parameter types for {len(ff_optimize)} parameters")
+        if len(param_types) <= 20:  # Only log if not too many
+            logging.debug(f"Parameter types: {param_types}")
+        else:
+            logging.debug(f"Parameter types (first 10): {param_types[:10]}...")
+        
+        # Log parameter type distribution
+        type_counts = {}
+        for pt in param_types:
+            type_counts[pt] = type_counts.get(pt, 0) + 1
+        logging.debug(f"Parameter type distribution: {type_counts}")
+        
+        if fm_input.regularization:
+        # Create regularizer with parameter types
+            regularizer = L2Regularizer(ff_optimize, None, fm_input.regularization_alpha, param_types)
+        
+        # Set type-specific prior widths from FMInput
+            regularizer.set_prior_widths_from_fm_input(fm_input)
+        else:
+            regularizer = None
+
+    else:
+        # No parameters to optimize      
+        raise MiMiCPyError("No parameters to optimize")  
+    
+    return ff_optimize, bond2params, regularizer
+
+
+def compute_obj_lm(ff_optimize,
+                   configurations, 
+                   bonds, angles, dihedrals, bond2params,
+                   *args):
+    
+    # Extract parallelization settings from args
+    n_processes = None
+    for arg in args:
+        if isinstance(arg, dict) and 'n_processes' in arg:
+            n_processes = arg.get('n_processes')
+            break
+    
+    residue, ref_bonded = compute_ff_obj(ff_optimize, configurations,
+                           bonds, angles, dihedrals, bond2params, n_processes)
+    
+    # Extract filename and regularization from args
+    filename = args[0] if len(args) > 0 else "output.txt"
+    regularization = None
+    
+    # Look for regularization object in args
+    for arg in args:
+        if hasattr(arg, 'compute_regularization_residuals'):
+            regularization = arg
+            break
+    
+    # For Levenberg-Marquardt, we need to augment the residual vector with regularization
+    if regularization is not None:
+        # Compute regularization residuals
+        reg_residuals = regularization.compute_regularization_residuals(ff_optimize)
+        # Augment the residual vector
+        augmented_residuals = np.concatenate([residue.ravel(), reg_residuals])
+    else:
+        augmented_residuals = residue.ravel()
+    
+    # Compute objective value for logging
+    sum_residual = np.sum(residue**2)
+    if regularization is not None:
+        reg_term = regularization.compute_regularization_term(ff_optimize)
+        sum_residual += reg_term
+    
+    sdf = np.sqrt(sum_residual/ np.sum(ref_bonded**2))  
+    abs_val = np.sqrt(sum_residual/ np.size(residue))
+
+    with open(filename, '+a' ) as file: 
+        file.write(f"{sum_residual}    {sdf}      {abs_val}     ") 
+        file.write('       '.join([str(x) for x in ff_optimize]))
+        file.write('\n')
+    
+    return augmented_residuals
+
+
+def jacobian_ff_obj(ff_optimize,
+                   configurations, 
+                   bonds, angles, dihedrals,
+                   bond2params, *args):
+
+    if type(configurations) != list:
+        configurations = [configurations]
+
+    # Extract parallelization parameters from args if available
+    n_processes = None
+    
+    # Look for parallelization settings in args
+    for arg in args:
+        if isinstance(arg, dict) and 'n_processes' in arg:
+            n_processes = arg.get('n_processes')
+            break
+    
+    # Use parallel processing if requested and beneficial
+    if n_processes is not None and n_processes > 1 and len(configurations) > 1:
+        total_jacobian = _compute_jacobian_parallel(ff_optimize, configurations, 
+                                                  bonds, angles, dihedrals, bond2params, n_processes)
+    else:
+        total_jacobian = _compute_jacobian_serial(ff_optimize, configurations, 
+                                                bonds, angles, dihedrals, bond2params)
+    
+    jac = np.vstack(total_jacobian)
+    
+    # Add regularization Jacobian if provided
+    regularization = None
+    if len(args) > 0:
+        # Look for regularization object in args
+        for arg in args:
+            if hasattr(arg, 'compute_regularization_residuals'):
+                regularization = arg
+                break
+    
+    if regularization is not None:
+        # For L2 regularization, the Jacobian is a diagonal matrix
+        n_params = len(ff_optimize)
+        reg_jacobian = np.zeros((n_params, n_params))
+        np.fill_diagonal(reg_jacobian, np.sqrt(regularization.alpha / 2) / regularization.prior_widths)
+        
+        # Augment the Jacobian matrix
+        augmented_jacobian = np.vstack([jac, reg_jacobian])
+        return augmented_jacobian
+    
+    return jac
+
+
+def _compute_jacobian_serial(ff_optimize, configurations, 
+                            bonds, angles, dihedrals, bond2params):
+    """Serial version of Jacobian computation (original implementation)."""
+    total_jacobian = []
+    for config in configurations:
+        # Get QM atoms count from the forces array length
+        qm_atoms_count = len(config['qm_forces'])
+        jac = jacobian_ff(ff_optimize,
+                           config["qm_coordinates"], 
+                           bond2params,
+                           bonds, angles, dihedrals,
+                           qm_atoms_count=qm_atoms_count)
+        total_jacobian.append(jac)
+    return total_jacobian
+
+
+def _compute_jacobian_parallel(ff_optimize, configurations, 
+                              bonds, angles, dihedrals, bond2params, n_processes):
+    """Parallel version of Jacobian computation using multiprocessing."""
+    # Prepare arguments for parallel processing
+    args_list = [(ff_optimize, config, bonds, angles, dihedrals, bond2params) 
+                 for config in configurations]
+    
+    with mp.Pool(processes=n_processes) as pool:
+        results = pool.map(_process_single_jacobian, args_list)
+    
+    return results
+
+
+def optimization_ff(qm_region, ff_optimize, configurations, bond2params, filename, regularizer=None, 
+                   n_processes=None, validation_fraction=0.1, random_seed=None):
+    """
+    Optimize force field parameters using least squares minimization with optional L2 regularization.
+    
+    Args:
+        qm_region (QMRegion): QMRegion object containing QM interactions
+        ff_optimize (numpy.ndarray): Initial force field parameters to optimize
+        configurations (list): List of configuration dictionaries
+        bond2params (dict): Mapping of interaction indices to parameter indices
+        filename (str): Output file to write optimization results
+        regularizer (L2Regularizer, optional): L2 regularization object
+        n_processes (int, optional): Number of processes to use for parallelization
+        validation_fraction (float): Fraction of configurations to use for validation (default: 0.1)
+        random_seed (int, optional): Random seed for reproducible train/val split
+        
+    Returns:
+        scipy.optimize.OptimizeResult: Result of the optimization
+    """
+    # Get interactions from QMRegion
+    qm_interactions = qm_region.qm_interactions
+    bonds = qm_interactions['bonds']
+    angles = qm_interactions['angles']
+    dihedrals = qm_interactions['dihedrals']
+    
+    # Split configurations into training and validation sets
+    train_configurations, val_configurations = split_configurations_train_val(
+        configurations, validation_fraction, random_seed
+    )
+    
+    # Create parallelization settings
+    parallel_settings = {'n_processes': n_processes}
+    
+    lower_bound = np.zeros(ff_optimize.shape)
+    upper_bound = np.inf*np.ones(ff_optimize.shape)
+    bounds = (lower_bound, upper_bound)
+    
+    # Use Jacobian if regularization is provided for better convergence
+    # if regularizer is not None:
+    res = least_squares(fun=compute_obj_lm,
+                    x0=ff_optimize,
+                    jac=jacobian_ff_obj,
+                    method='lm',
+                    args=(train_configurations,
+                        bonds, angles, dihedrals, bond2params, filename, regularizer, parallel_settings),
+                    verbose=1)
+    
+    # Compute validation SDB
+    if val_configurations:
+        val_sdf, val_cost, val_residuals, val_ref_bonded = compute_validation_sdf(
+            res.x, val_configurations, bonds, angles, dihedrals, bond2params, n_processes
+        )
+        
+        # Log validation results
+        with open(filename, 'a') as f:
+            f.write(f"\n# VALIDATION RESULTS\n")
+            f.write(f"# Validation SDF: {val_sdf:.6f}\n")
+            f.write(f"# Validation cost: {val_cost:.6f}\n")
+            f.write(f"# Validation configurations: {len(val_configurations)}\n\n")
+        
+        # Store validation results in the optimization result object
+        res.validation_sdf = val_sdf
+        res.validation_cost = val_cost
+        res.validation_residuals = val_residuals
+        res.validation_ref_bonded = val_ref_bonded
+    
+    return res
+
+
+def hierarchical_optimization_ff(qm_region, ff_optimize, configurations, bond2params, filename, 
+                                regularizer=None, fm_input=None, n_processes=None, validation_fraction=0.1, random_seed=None):
+    """
+    Optimize force field parameters using hierarchical approach:
+    1. Optimize bonds first (highest energy contribution)
+    2. Fix bond parameters, optimize angles 
+    3. Fix bond and angle parameters, optimize dihedrals
+    
+    This approach respects the energy hierarchy and should produce more physically realistic parameters.
+    
+    Args:
+        qm_region (QMRegion): QMRegion object containing QM interactions
+        ff_optimize (numpy.ndarray): Initial force field parameters to optimize
+        configurations (list): List of configuration dictionaries
+        bond2params (dict): Mapping of interaction indices to parameter indices
+        filename (str): Output file to write optimization results
+        regularizer (L2Regularizer, optional): L2 regularization object
+        fm_input (FMInput, optional): FMInput object containing optimization controls
+        n_processes (int, optional): Number of processes for parallel computation
+        validation_fraction (float): Fraction of configurations to use for validation (default: 0.1)
+        random_seed (int, optional): Random seed for reproducible train/val split
+        
+    Returns:
+        tuple: (optimized_parameters, optimization_results) where optimized_parameters is the final
+               parameter array and optimization_results is a list of optimization results for each stage
+    """
+    # Get interactions from QMRegion
+    qm_interactions = qm_region.qm_interactions
+    bonds = qm_interactions['bonds']
+    angles = qm_interactions['angles']
+    dihedrals = qm_interactions['dihedrals']
+    
+    # Split configurations into training and validation sets
+    train_configurations, val_configurations = split_configurations_train_val(
+        configurations, validation_fraction, random_seed
+    )
+    
+    # Create parameter masks for each interaction type
+    bond_params_mask = np.zeros(len(ff_optimize), dtype=bool)
+    angle_params_mask = np.zeros(len(ff_optimize), dtype=bool)
+    dihedral_params_mask = np.zeros(len(ff_optimize), dtype=bool)
+    
+    # Validate that we have parameters to optimize
+    if len(ff_optimize) == 0:
+        logging.warning("Warning: No parameters to optimize")
+        return ff_optimize, []
+    
+    # Identify which parameters belong to which interaction types
+    for bond in bonds:
+        if bond.get('optimize', False):
+            param_indices = bond2params.get(bond['index'], [])
+            for idx in param_indices:
+                if idx is not None:
+                    bond_params_mask[idx] = True
+    
+    for angle in angles:
+        if angle.get('optimize', False):
+            param_indices = bond2params.get(angle['index'], [])
+            for idx in param_indices:
+                if idx is not None:
+                    angle_params_mask[idx] = True
+    
+    for dihedral in dihedrals:
+        if dihedral.get('optimize', False):
+            param_indices = bond2params.get(dihedral['index'], [])
+            for idx in param_indices:
+                if idx is not None:
+                    dihedral_params_mask[idx] = True
+    
+    # Debug information
+    logging.info(f"Total parameters: {len(ff_optimize)}")
+    logging.info(f"Bond parameters: {np.sum(bond_params_mask)}")
+    logging.info(f"Angle parameters: {np.sum(angle_params_mask)}")
+    logging.info(f"Dihedral parameters: {np.sum(dihedral_params_mask)}")
+    
+    # Validate regularizer dimensions
+    if regularizer is not None:
+        if len(regularizer.initial_params) != len(ff_optimize):
+            logging.warning(f"Warning: Regularizer parameter count ({len(regularizer.initial_params)}) doesn't match optimization parameters ({len(ff_optimize)})")
+        if len(regularizer.prior_widths) != len(ff_optimize):
+            logging.warning(f"Warning: Regularizer prior widths count ({len(regularizer.prior_widths)}) doesn't match optimization parameters ({len(ff_optimize)})")
+        if regularizer.param_types and len(regularizer.param_types) != len(ff_optimize):
+            logging.warning(f"Warning: Regularizer parameter types count ({len(regularizer.param_types)}) doesn't match optimization parameters ({len(ff_optimize)})")
+        else:
+            logging.info(f"Parameter types: {len(regularizer.param_types) if regularizer.param_types else 0}")
+    else:
+        logging.info("No regularizer defined")
+    
+    current_params = ff_optimize.copy()
+    optimization_results = []
+    
+    # Stage 1: Optimize bonds only
+    if np.any(bond_params_mask):
+        logging.info("Stage 1: Optimizing bond parameters...")
+        bond_params = current_params.copy()
+        
+        # Create bond-only regularizer if needed
+        bond_regularizer = None
+        if regularizer is not None:
+            bond_param_types = _create_param_types_for_mask(regularizer.param_types, bond_params_mask)
+            
+            bond_regularizer = L2Regularizer(
+                regularizer.initial_params[bond_params_mask],
+                regularizer.prior_widths[bond_params_mask],
+                regularizer.alpha,
+                bond_param_types
+            )
+        
+        # Optimize only bond parameters
+        bond_result = _optimize_parameter_subset(
+            current_params, bond_params_mask, train_configurations, 
+            bonds, angles, dihedrals, bond2params, filename, 
+            bond_regularizer, "bonds", n_processes
+        )
+        
+        # Update current parameters with optimized bond parameters
+        current_params[bond_params_mask] = bond_result.x
+        optimization_results.append(("bonds", bond_result))
+        
+        logging.info(f"Bond optimization completed. Final objective: {bond_result.cost}")
+    
+    # Stage 2: Fix bonds, optimize angles
+    if np.any(angle_params_mask):
+        logging.info("Stage 2: Optimizing angle parameters (bonds fixed)...")
+        
+        # Create angle-only regularizer if needed
+        angle_regularizer = None
+        if regularizer is not None:
+            angle_param_types = _create_param_types_for_mask(regularizer.param_types, angle_params_mask)
+            
+            angle_regularizer = L2Regularizer(
+                regularizer.initial_params[angle_params_mask],
+                regularizer.prior_widths[angle_params_mask],
+                regularizer.alpha,
+                angle_param_types
+            )
+        
+        # Optimize only angle parameters
+        angle_result = _optimize_parameter_subset(
+            current_params, angle_params_mask, train_configurations, 
+            bonds, angles, dihedrals, bond2params, filename, 
+            angle_regularizer, "angles", n_processes
+        )
+        
+        # Update current parameters with optimized angle parameters
+        current_params[angle_params_mask] = angle_result.x
+        optimization_results.append(("angles", angle_result))
+        
+        logging.info(f"Angle optimization completed. Final objective: {angle_result.cost}")
+    
+    # Stage 3: Fix bonds and angles, optimize dihedrals
+    if np.any(dihedral_params_mask):
+        logging.info("Stage 3: Optimizing dihedral parameters (bonds and angles fixed)...")
+        
+        # Create dihedral-only regularizer if needed
+        dihedral_regularizer = None
+        if regularizer is not None:
+            dihedral_param_types = _create_param_types_for_mask(regularizer.param_types, dihedral_params_mask)
+            
+            dihedral_regularizer = L2Regularizer(
+                regularizer.initial_params[dihedral_params_mask],
+                regularizer.prior_widths[dihedral_params_mask],
+                regularizer.alpha,
+                dihedral_param_types
+            )
+        
+        # Optimize only dihedral parameters
+        dihedral_result = _optimize_parameter_subset(
+            current_params, dihedral_params_mask, train_configurations, 
+            bonds, angles, dihedrals, bond2params, filename, 
+            dihedral_regularizer, "dihedrals", n_processes
+        )
+        
+        # Update current parameters with optimized dihedral parameters
+        current_params[dihedral_params_mask] = dihedral_result.x
+        optimization_results.append(("dihedrals", dihedral_result))
+        
+        logging.info(f"Dihedral optimization completed. Final objective: {dihedral_result.cost}")
+    
+    # Compute validation SDB on final optimized parameters
+    if val_configurations:
+        val_sdf, val_cost, val_residuals, val_ref_bonded = compute_validation_sdf(
+            current_params, val_configurations, bonds, angles, dihedrals, bond2params, n_processes
+        )
+        
+        # Log validation results
+        with open(filename, 'a') as f:
+            f.write(f"\n# HIERARCHICAL OPTIMIZATION VALIDATION RESULTS\n")
+            f.write(f"# Final Validation SDF: {val_sdf:.6f}\n")
+            f.write(f"# Final Validation cost: {val_cost:.6f}\n")
+            f.write(f"# Validation configurations: {len(val_configurations)}\n")
+            f.write(f"# Training configurations: {len(train_configurations)}\n\n")
+        
+        logging.info(f"Final validation SDF: {val_sdf:.6f}")
+        logging.info(f"Final validation cost: {val_cost:.6f}")
+    
+    return current_params, optimization_results
+
+
+def _optimize_parameter_subset(current_params, param_mask, configurations, 
+                              bonds, angles, dihedrals, bond2params, filename, 
+                              regularizer, stage_name, n_processes=None):
+    """
+    Optimize a subset of parameters while keeping others fixed.
+    
+    Args:
+        current_params (np.ndarray): Current parameter values
+        param_mask (np.ndarray): Boolean mask indicating which parameters to optimize
+        configurations (list): List of configuration dictionaries
+        bonds, angles, dihedrals: Interaction lists
+        bond2params (dict): Parameter mapping
+        filename (str): Output filename
+        regularizer (L2Regularizer): Regularizer for this subset
+        stage_name (str): Name of optimization stage for logging
+        
+    Returns:
+        scipy.optimize.OptimizeResult: Optimization result
+    """
+    # Create objective function that only optimizes the specified parameters
+    def subset_objective(subset_params):
+        # Create full parameter array with fixed values
+        full_params = current_params.copy()
+        full_params[param_mask] = subset_params
+        
+        # Compute objective with full parameter array
+        residue, _ = compute_ff_obj(full_params, configurations, 
+                                   bonds, angles, dihedrals, bond2params, n_processes)
+        
+        obj_value = np.sum(residue**2)
+        
+        # Add regularization term if provided
+        if regularizer is not None:
+            reg_term = regularizer.compute_regularization_term(subset_params)
+            obj_value += reg_term
+        
+        return obj_value
+    
+    def subset_residuals(subset_params):
+        # Create full parameter array with fixed values
+        full_params = current_params.copy()
+        full_params[param_mask] = subset_params
+        
+        # Compute residuals with full parameter array
+        residue, ref_bonded = compute_ff_obj(full_params, configurations, 
+                                   bonds, angles, dihedrals, bond2params, n_processes)
+        
+        # Compute objective value for logging
+        sum_residual = np.sum(residue**2)
+        if regularizer is not None:
+            reg_term = regularizer.compute_regularization_term(subset_params)
+            sum_residual += reg_term
+        
+        sdf = np.sqrt(sum_residual / np.sum(ref_bonded**2))  
+        abs_val = np.sqrt(sum_residual / np.size(residue))
+
+        with open(filename, '+a') as file: 
+            file.write(f"{sum_residual}    {sdf}      {abs_val}     ") 
+            file.write('       '.join([str(x) for x in full_params]))
+            file.write('\n')
+        
+        # Augment with regularization residuals if provided
+        if regularizer is not None:
+            reg_residuals = regularizer.compute_regularization_residuals(subset_params)
+            residue = np.concatenate([residue.ravel(), reg_residuals])
+        
+        return residue
+    
+    def subset_jacobian(subset_params):
+        # Create full parameter array with fixed values
+        full_params = current_params.copy()
+        full_params[param_mask] = subset_params
+        
+        # Compute full Jacobian
+        full_jac = jacobian_ff_obj(full_params, configurations, 
+                                  bonds, angles, dihedrals, bond2params, {'n_processes': n_processes})
+        
+        # Extract only the columns corresponding to parameters being optimized
+        subset_jac = full_jac[:, param_mask]
+        
+        # Add regularization Jacobian if provided
+        if regularizer is not None:
+            n_params = np.sum(param_mask)
+            reg_jacobian = np.zeros((n_params, n_params))
+            np.fill_diagonal(reg_jacobian, np.sqrt(regularizer.alpha / 2) / regularizer.prior_widths)
+            subset_jac = np.vstack([subset_jac, reg_jacobian])
+        
+        return subset_jac
+    
+    # Initial parameters for this subset
+    x0 = current_params[param_mask]
+    
+    # Bounds for this subset (only positive values)
+    lower_bound = np.zeros(x0.shape)
+    upper_bound = np.inf * np.ones(x0.shape)
+    bounds = (lower_bound, upper_bound)
+    
+    # Optimize using least squares
+    result = least_squares(
+        fun=subset_residuals,
+        x0=x0,
+        jac=subset_jacobian,
+        method='lm',
+        verbose=1
+    )
+    
+    # Log results
+    with open(filename, 'a') as f:
+        f.write(f"\n# {stage_name.upper()} OPTIMIZATION RESULTS\n")
+        f.write(f"# Final cost: {result.cost}\n")
+        f.write(f"# Number of iterations: {result.nfev}\n")
+        f.write(f"# Success: {result.success}\n")
+        f.write(f"# Optimized parameters: {result.x}\n\n")
+    
+    return result
+
+
+def energy_weighted_optimization_ff(qm_region, ff_optimize, configurations, bond2params, filename, 
+                                   regularizer=None, fm_input=None, n_processes=None, validation_fraction=0.1, random_seed=None):
+    """
+    Optimize force field parameters using energy-weighted objective function.
+    
+    This approach weights the contribution of different interaction types based on their
+    typical energy scales: bonds (highest), angles (medium), dihedrals (lowest).
+    
+    Args:
+        qm_region (QMRegion): QMRegion object containing QM interactions
+        ff_optimize (numpy.ndarray): Initial force field parameters to optimize
+        configurations (list): List of configuration dictionaries
+        bond2params (dict): Mapping of interaction indices to parameter indices
+        filename (str): Output file to write optimization results
+        regularizer (L2Regularizer, optional): L2 regularization object
+        fm_input (FMInput, optional): FMInput object containing optimization controls
+        n_processes (int, optional): Number of processes for parallel computation
+        validation_fraction (float): Fraction of configurations to use for validation (default: 0.1)
+        random_seed (int, optional): Random seed for reproducible train/val split
+        
+    Returns:
+        scipy.optimize.OptimizeResult: Result of the optimization
+    """
+    # Get interactions from QMRegion
+    qm_interactions = qm_region.qm_interactions
+    bonds = qm_interactions['bonds']
+    angles = qm_interactions['angles']
+    dihedrals = qm_interactions['dihedrals']
+    
+    # Split configurations into training and validation sets
+    train_configurations, val_configurations = split_configurations_train_val(
+        configurations, validation_fraction, random_seed
+    )
+    
+    # Get energy weights from fm_input or use defaults
+    if fm_input is not None:
+        bond_weight = getattr(fm_input, 'bond_energy_weight', 1.0)
+        angle_weight = getattr(fm_input, 'angle_energy_weight', 0.1)
+        dihedral_weight = getattr(fm_input, 'dihedral_energy_weight', 0.01)
+    else:
+        bond_weight = 1.0
+        angle_weight = 0.1
+        dihedral_weight = 0.01
+    
+    # Create interaction type masks
+    bond_mask = np.zeros(len(ff_optimize), dtype=bool)
+    angle_mask = np.zeros(len(ff_optimize), dtype=bool)
+    dihedral_mask = np.zeros(len(ff_optimize), dtype=bool)
+    
+    # Identify which parameters belong to which interaction types
+    for bond in bonds:
+        if bond.get('optimize', False):
+            param_indices = bond2params.get(bond['index'], [])
+            for idx in param_indices:
+                if idx is not None:
+                    bond_mask[idx] = True
+    
+    for angle in angles:
+        if angle.get('optimize', False):
+            param_indices = bond2params.get(angle['index'], [])
+            for idx in param_indices:
+                if idx is not None:
+                    angle_mask[idx] = True
+    
+    for dihedral in dihedrals:
+        if dihedral.get('optimize', False):
+            param_indices = bond2params.get(dihedral['index'], [])
+            for idx in param_indices:
+                if idx is not None:
+                    dihedral_mask[idx] = True
+    
+    def energy_weighted_residuals(params):
+        """Compute residuals with energy-based weighting."""
+        residue, ref_bonded = compute_ff_obj(params, train_configurations, 
+                                   bonds, angles, dihedrals, bond2params, n_processes)
+        
+        # Apply energy weights to different interaction types
+        weighted_residue = residue.copy()
+        
+        # Weight bond contributions
+        if np.any(bond_mask):
+            weighted_residue *= bond_weight
+        
+        # Weight angle contributions (this is approximate - we weight all residues)
+        # In practice, you might want to identify which residues come from which interactions
+        if np.any(angle_mask):
+            # For simplicity, we apply angle weight to all residues
+            # A more sophisticated approach would track which residues come from which interactions
+            weighted_residue *= angle_weight
+        
+        # Weight dihedral contributions
+        if np.any(dihedral_mask):
+            weighted_residue *= dihedral_weight
+        
+        # Compute objective value for logging (using unweighted residuals for consistency)
+        sum_residual = np.sum(residue**2)
+        if regularizer is not None:
+            reg_term = regularizer.compute_regularization_term(params)
+            sum_residual += reg_term
+        
+        sdf = np.sqrt(sum_residual / np.sum(ref_bonded**2))  
+        abs_val = np.sqrt(sum_residual / np.size(residue))
+
+        with open(filename, '+a') as file: 
+            file.write(f"{sum_residual}    {sdf}      {abs_val}     ") 
+            file.write('       '.join([str(x) for x in params]))
+            file.write('\n')
+        
+        # Augment with regularization residuals if provided
+        if regularizer is not None:
+            reg_residuals = regularizer.compute_regularization_residuals(params)
+            weighted_residue = np.concatenate([weighted_residue.ravel(), reg_residuals])
+        
+        return weighted_residue
+    
+    def energy_weighted_jacobian(params):
+        """Compute Jacobian with energy-based weighting."""
+        jac = jacobian_ff_obj(params, train_configurations, 
+                              bonds, angles, dihedrals, bond2params, {'n_processes': n_processes})
+        
+        # Apply energy weights to Jacobian rows
+        weighted_jac = jac.copy()
+        
+        # Weight bond contributions
+        if np.any(bond_mask):
+            weighted_jac *= bond_weight
+        
+        # Weight angle contributions
+        if np.any(angle_mask):
+            weighted_jac *= angle_weight
+        
+        # Weight dihedral contributions
+        if np.any(dihedral_mask):
+            weighted_jac *= dihedral_weight
+        
+        # Add regularization Jacobian if provided
+        if regularizer is not None:
+            n_params = len(params)
+            reg_jacobian = np.zeros((n_params, n_params))
+            np.fill_diagonal(reg_jacobian, np.sqrt(regularizer.alpha / 2) / regularizer.prior_widths)
+            weighted_jac = np.vstack([weighted_jac, reg_jacobian])
+        
+        return weighted_jac
+    
+    # Bounds for parameters (only positive values)
+    lower_bound = np.zeros(ff_optimize.shape)
+    upper_bound = np.inf * np.ones(ff_optimize.shape)
+    bounds = (lower_bound, upper_bound)
+    
+    # Optimize using least squares with energy weighting
+    result = least_squares(
+        fun=energy_weighted_residuals,
+        x0=ff_optimize,
+        jac=energy_weighted_jacobian,
+        method='lm',
+        verbose=1
+    )
+    
+    # Compute validation SDB
+    if val_configurations:
+        val_sdf, val_cost, val_residuals, val_ref_bonded = compute_validation_sdf(
+            result.x, val_configurations, bonds, angles, dihedrals, bond2params, n_processes
+        )
+        
+        # Store validation results in the optimization result object
+        result.validation_sdf = val_sdf
+        result.validation_cost = val_cost
+        result.validation_residuals = val_residuals
+        result.validation_ref_bonded = val_ref_bonded
+    
+    # Log results
+    with open(filename, 'a') as f:
+        f.write(f"\n# ENERGY-WEIGHTED OPTIMIZATION RESULTS\n")
+        f.write(f"# Bond weight: {bond_weight}\n")
+        f.write(f"# Angle weight: {angle_weight}\n")
+        f.write(f"# Dihedral weight: {dihedral_weight}\n")
+        f.write(f"# Final cost: {result.cost}\n")
+        f.write(f"# Number of iterations: {result.nfev}\n")
+        f.write(f"# Success: {result.success}\n")
+        f.write(f"# Optimized parameters: {result.x}\n")
+        if val_configurations:
+            f.write(f"# Validation SDF: {val_sdf:.6f}\n")
+            f.write(f"# Validation cost: {val_cost:.6f}\n")
+            f.write(f"# Validation configurations: {len(val_configurations)}\n")
+            f.write(f"# Training configurations: {len(train_configurations)}\n")
+        f.write("\n")
+    
+    return result
+
+def adaptive_regularization_optimization_ff(qm_region, ff_optimize, configurations, bond2params, filename, 
+                                          fm_input=None, n_processes=None, validation_fraction=0.1, random_seed=None):
+    """
+    Optimize force field parameters using adaptive regularization based on energy hierarchy.
+    
+    This approach applies different regularization strengths to different interaction types:
+    - Bonds: Strong regularization (prevent overfitting of high-energy terms)
+    - Angles: Medium regularization
+    - Dihedrals: Weak regularization (allow flexibility for low-energy terms)
+    
+    Args:
+        qm_region (QMRegion): QMRegion object containing QM interactions
+        ff_optimize (numpy.ndarray): Initial force field parameters to optimize
+        configurations (list): List of configuration dictionaries
+        bond2params (dict): Mapping of interaction indices to parameter indices
+        filename (str): Output file to write optimization results
+        fm_input (FMInput, optional): FMInput object containing optimization controls
+        n_processes (int, optional): Number of processes for parallel computation
+        validation_fraction (float): Fraction of configurations to use for validation (default: 0.1)
+        random_seed (int, optional): Random seed for reproducible train/val split
+        
+    Returns:
+        scipy.optimize.OptimizeResult: Result of the optimization
+    """
+    # Get interactions from QMRegion
+    qm_interactions = qm_region.qm_interactions
+    bonds = qm_interactions['bonds']
+    angles = qm_interactions['angles']
+    dihedrals = qm_interactions['dihedrals']
+    
+    # Split configurations into training and validation sets
+    train_configurations, val_configurations = split_configurations_train_val(
+        configurations, validation_fraction, random_seed
+    )
+    
+    # Get adaptive regularization parameters from fm_input or use defaults
+    if fm_input is not None:
+        base_alpha = getattr(fm_input, 'adaptive_base_alpha', 0.1)
+        energy_hierarchy_scale = getattr(fm_input, 'energy_hierarchy_scale', 10.0)
+    else:
+        base_alpha = 0.1
+        energy_hierarchy_scale = 10.0
+    
+    # Create adaptive regularizer
+    adaptive_regularizer = AdaptiveL2Regularizer(
+        ff_optimize, bond2params, bonds, angles, dihedrals,
+        base_alpha=base_alpha,
+        energy_hierarchy_scale=energy_hierarchy_scale
+    )
+    
+    def adaptive_residuals(params):
+        """Compute residuals with adaptive regularization."""
+        residue, ref_bonded = compute_ff_obj(params, train_configurations, 
+                                   bonds, angles, dihedrals, bond2params, n_processes)
+        
+        # Compute objective value for logging
+        sum_residual = np.sum(residue**2)
+        reg_term = adaptive_regularizer.compute_regularization_term(params)
+        sum_residual += reg_term
+        
+        sdf = np.sqrt(sum_residual / np.sum(ref_bonded**2))  
+        abs_val = np.sqrt(sum_residual / np.size(residue))
+
+        with open(filename, '+a') as file: 
+            file.write(f"{sum_residual}    {sdf}      {abs_val}     ") 
+            file.write('       '.join([str(x) for x in params]))
+            file.write('\n')
+        
+        # Augment with adaptive regularization residuals
+        reg_residuals = adaptive_regularizer.compute_regularization_residuals(params)
+        augmented_residuals = np.concatenate([residue.ravel(), reg_residuals])
+        
+        return augmented_residuals
+    
+    def adaptive_jacobian(params):
+        """Compute Jacobian with adaptive regularization."""
+        jac = jacobian_ff_obj(params, train_configurations, 
+                              bonds, angles, dihedrals, bond2params, {'n_processes': n_processes})
+        
+        # Add adaptive regularization Jacobian
+        n_params = len(params)
+        reg_jacobian = np.zeros((n_params, n_params))
+        np.fill_diagonal(reg_jacobian, np.sqrt(adaptive_regularizer.adaptive_alphas / 2) / adaptive_regularizer.prior_widths)
+        augmented_jacobian = np.vstack([jac, reg_jacobian])
+        
+        return augmented_jacobian
+    
+    # Bounds for parameters (only positive values)
+    lower_bound = np.zeros(ff_optimize.shape)
+    upper_bound = np.inf * np.ones(ff_optimize.shape)
+    bounds = (lower_bound, upper_bound)
+    
+    # Optimize using least squares with adaptive regularization
+    result = least_squares(
+        fun=adaptive_residuals,
+        x0=ff_optimize,
+        jac=adaptive_jacobian,
+        method='lm',
+        verbose=1
+    )
+    
+    # Compute validation SDB
+    if val_configurations:
+        val_sdf, val_cost, val_residuals, val_ref_bonded = compute_validation_sdf(
+            result.x, val_configurations, bonds, angles, dihedrals, bond2params, n_processes
+        )
+        
+        # Store validation results in the optimization result object
+        result.validation_sdf = val_sdf
+        result.validation_cost = val_cost
+        result.validation_residuals = val_residuals
+        result.validation_ref_bonded = val_ref_bonded
+    
+    # Log results
+    with open(filename, 'a') as f:
+        f.write(f"\n# ADAPTIVE REGULARIZATION OPTIMIZATION RESULTS\n")
+        f.write(f"# Base alpha: {base_alpha}\n")
+        f.write(f"# Energy hierarchy scale: {energy_hierarchy_scale}\n")
+        f.write(f"# Bond parameters: {np.sum(adaptive_regularizer.bond_params_mask)}\n")
+        f.write(f"# Angle parameters: {np.sum(adaptive_regularizer.angle_params_mask)}\n")
+        f.write(f"# Dihedral parameters: {np.sum(adaptive_regularizer.dihedral_params_mask)}\n")
+        f.write(f"# Final cost: {result.cost}\n")
+        f.write(f"# Number of iterations: {result.nfev}\n")
+        f.write(f"# Success: {result.success}\n")
+        f.write(f"# Optimized parameters: {result.x}\n")
+        if val_configurations:
+            f.write(f"# Validation SDF: {val_sdf:.6f}\n")
+            f.write(f"# Validation cost: {val_cost:.6f}\n")
+            f.write(f"# Validation configurations: {len(val_configurations)}\n")
+            f.write(f"# Training configurations: {len(train_configurations)}\n")
+        f.write("\n")
+    
+    return result
+
+def unified_optimization_ff(qm_region, ff_optimize, configurations, bond2params, filename, 
+                           regularizer=None, fm_input=None, n_processes=None, validation_fraction=0.1, random_seed=None):
+    """
+    Unified force field optimization function that supports multiple hierarchical approaches.
+    
+    This function automatically selects the appropriate optimization method based on the
+    FMInput parameters and provides a consistent interface for all optimization strategies.
+    
+    Args:
+        qm_region (QMRegion): QMRegion object containing QM interactions
+        ff_optimize (numpy.ndarray): Initial force field parameters to optimize
+        configurations (list): List of configuration dictionaries
+        bond2params (dict): Mapping of interaction indices to parameter indices
+        filename (str): Output file to write optimization results
+        regularizer (L2Regularizer, optional): L2 regularization object
+        fm_input (FMInput, optional): FMInput object containing optimization controls
+        n_processes (int, optional): Number of processes for parallel computation
+        validation_fraction (float): Fraction of configurations to use for validation (default: 0.1)
+        random_seed (int, optional): Random seed for reproducible train/val split
+        
+    Returns:
+        tuple: (optimized_parameters, optimization_result) where optimized_parameters is the final
+               parameter array and optimization_result is the optimization result object
+    """
+    if fm_input is None:
+        fm_input = FMInput()
+    
+    # Ensure optimization_method is set
+    if not hasattr(fm_input, 'optimization_method') or fm_input.optimization_method is None:
+        fm_input.optimization_method = 'hierarchical'
+    
+    optimization_method = fm_input.optimization_method
+    
+    print(f"Using optimization method: {optimization_method}")
+    
+    if optimization_method == 'hierarchical':
+        # Sequential hierarchical optimization
+        optimized_params, optimization_results = hierarchical_optimization_ff(
+            qm_region, ff_optimize, configurations, bond2params, filename, 
+            regularizer, fm_input, n_processes, validation_fraction, random_seed
+        )
+        # Return the last optimization result for compatibility
+        if optimization_results:
+            return optimized_params, optimization_results[-1][1]
+        else:
+            return optimized_params, None
+            
+    elif optimization_method == 'energy_weighted':
+        # Energy-weighted simultaneous optimization
+        result = energy_weighted_optimization_ff(
+            qm_region, ff_optimize, configurations, bond2params, filename, 
+            regularizer, fm_input, n_processes, validation_fraction, random_seed
+        )
+        return result.x, result
+        
+    elif optimization_method == 'adaptive':
+        # Adaptive regularization optimization
+        result = adaptive_regularization_optimization_ff(
+            qm_region, ff_optimize, configurations, bond2params, filename, 
+            fm_input, n_processes, validation_fraction, random_seed
+        )
+        return result.x, result
+        
+    elif optimization_method == 'simultaneous':
+        # Traditional simultaneous optimization
+        result = optimization_ff(
+            qm_region, ff_optimize, configurations, bond2params, filename, 
+            regularizer, n_processes, validation_fraction, random_seed
+        )
+        return result.x, result
+        
+    else:
+        raise ValueError(f"Unknown optimization method: {optimization_method}. "
+                        f"Supported methods: 'hierarchical', 'energy_weighted', 'adaptive', 'simultaneous'")
+
+
+def _create_param_types_for_mask(param_types, param_mask):
+    """
+    Create a parameter types list for a specific parameter mask.
+    
+    Args:
+        param_types (list): Original parameter types list
+        param_mask (np.ndarray): Boolean mask indicating which parameters to include
+        
+    Returns:
+        list: Parameter types for the masked parameters
+    """
+    if param_types is None:
+        return None
+    
+    # Create a list of parameter types for the masked parameters only
+    masked_param_types = []
+    for i, is_included in enumerate(param_mask):
+        if is_included and i < len(param_types):
+            masked_param_types.append(param_types[i])
+    
+    return masked_param_types
+
+
+def _copy_fm_input(fm_input, new_method=None):
+    """
+    Create a copy of FMInput object with optional method override.
+    
+    Args:
+        fm_input: FMInput object to copy
+        new_method: Optional new optimization method to set
+        
+    Returns:
+        FMInput: Copy of the input object
+    """
+    if fm_input is None:
+        return FMInput()
+    
+    # Try to use the copy method if available
+    if hasattr(fm_input, 'copy'):
+        try:
+            copy_input = fm_input.copy()
+        except:
+            # Fallback: create new object with same parameters
+            copy_input = FMInput()
+            # Copy all attributes manually
+            for attr in dir(fm_input):
+                if not attr.startswith('_') and not callable(getattr(fm_input, attr)):
+                    try:
+                        setattr(copy_input, attr, getattr(fm_input, attr))
+                    except:
+                        pass
+    else:
+        # Fallback: create new object with same parameters
+        copy_input = FMInput()
+        # Copy all attributes manually
+        for attr in dir(fm_input):
+            if not attr.startswith('_') and not callable(getattr(fm_input, attr)):
+                try:
+                    setattr(copy_input, attr, getattr(fm_input, attr))
+                except:
+                    pass
+    
+    # Set new method if specified
+    if new_method is not None:
+        copy_input.optimization_method = new_method
+    
+    return copy_input
+
+
+def compare_optimization_methods(qm_region, ff_optimize, configurations, bond2params, filename, 
+                               regularizer=None, fm_input=None, n_processes=None, validation_fraction=0.1, random_seed=None):
+    """
+    Compare different optimization methods and report their performance.
+    
+    This function runs all available optimization methods and compares their results
+    to help users choose the best approach for their system.
+    
+    Args:
+        qm_region (QMRegion): QMRegion object containing QM interactions
+        ff_optimize (numpy.ndarray): Initial force field parameters to optimize
+        configurations (list): List of configuration dictionaries
+        bond2params (dict): Mapping of interaction indices to parameter indices
+        filename (str): Output file to write optimization results
+        regularizer (L2Regularizer, optional): L2 regularization object
+        fm_input (FMInput, optional): FMInput object containing optimization controls
+        n_processes (int, optional): Number of processes for parallel computation
+        validation_fraction (float): Fraction of configurations to use for validation (default: 0.1)
+        random_seed (int, optional): Random seed for reproducible train/val split
+        
+    Returns:
+        dict: Dictionary containing results for each optimization method
+    """
+    if fm_input is None:
+        fm_input = FMInput()
+    
+    methods = ['hierarchical', 'energy_weighted', 'adaptive', 'simultaneous']
+    results = {}
+    
+    logging.info("Comparing optimization methods...")
+    
+    for method in methods:
+        logging.info(f"\nTesting method: {method}")
+        
+        # Create a copy of fm_input with the current method
+        test_fm_input = _copy_fm_input(fm_input, method)
+        
+        try:
+            # Run optimization with current method
+            optimized_params, opt_result = unified_optimization_ff(
+                qm_region, ff_optimize.copy(), configurations, bond2params, 
+                f"{filename}_{method}", regularizer, test_fm_input, n_processes, validation_fraction, random_seed
+            )
+            
+            # Get validation SDF from optimization result if available
+            if hasattr(opt_result, 'validation_sdf') and opt_result.validation_sdf is not None:
+                val_sdf = opt_result.validation_sdf
+                val_cost = opt_result.validation_cost
+            else:
+                # Compute validation SDF manually if not available
+                train_configs, val_configs = split_configurations_train_val(
+                    configurations, validation_fraction, random_seed
+                )
+                val_sdf, val_cost, _, _ = compute_validation_sdf(
+                    optimized_params, val_configs, 
+                    qm_region.qm_interactions['bonds'],
+                    qm_region.qm_interactions['angles'],
+                    qm_region.qm_interactions['dihedrals'],
+                    bond2params, n_processes
+                )
+            
+            # Compute training SDF for comparison
+            train_configs, _ = split_configurations_train_val(
+                configurations, validation_fraction, random_seed
+            )
+            train_residue, train_ref_bonded = compute_ff_obj(optimized_params, train_configs, 
+                                                           qm_region.qm_interactions['bonds'],
+                                                           qm_region.qm_interactions['angles'],
+                                                           qm_region.qm_interactions['dihedrals'],
+                                                           bond2params)
+            
+            train_cost = np.sum(train_residue**2)
+            train_sdf = np.sqrt(train_cost / np.sum(train_ref_bonded**2))
+            
+            results[method] = {
+                'optimized_params': optimized_params,
+                'train_cost': train_cost,
+                'train_sdf': train_sdf,
+                'val_cost': val_cost,
+                'val_sdf': val_sdf,
+                'success': opt_result.success if opt_result else True,
+                'nfev': opt_result.nfev if opt_result else 0
+            }
+            
+            logging.info(f"  Training cost: {train_cost:.6f}")
+            logging.info(f"  Training SDF: {train_sdf:.6f}")
+            logging.info(f"  Validation cost: {val_cost:.6f}")
+            logging.info(f"  Validation SDF: {val_sdf:.6f}")
+            logging.info(f"  Success: {results[method]['success']}")
+            
+        except Exception as e:
+            logging.error(f"  Error in method {method}: {str(e)}")
+            results[method] = {'error': str(e)}
+    
+    # Write comparison summary
+    with open(f"{filename}_comparison", 'w') as f:
+        f.write("# OPTIMIZATION METHOD COMPARISON\n")
+        f.write("# Method\tTrain Cost\tTrain SDF\tVal Cost\tVal SDF\tSuccess\tFunction Evaluations\n")
+        
+        for method in methods:
+            if method in results and 'error' not in results[method]:
+                f.write(f"{method}\t{results[method]['train_cost']:.6f}\t"
+                       f"{results[method]['train_sdf']:.6f}\t"
+                       f"{results[method]['val_cost']:.6f}\t"
+                       f"{results[method]['val_sdf']:.6f}\t"
+                       f"{results[method]['success']}\t"
+                       f"{results[method]['nfev']}\n")
+            else:
+                f.write(f"{method}\tERROR\tERROR\tERROR\tERROR\tFalse\t0\n")
+    
+    return results
+
+
+def _process_single_jacobian(args):
+    """
+    Process a single configuration for Jacobian computation in parallel.
+    
+    Args:
+        args: Tuple containing (ff_optimize, config, bonds, angles, dihedrals, bond2params)
+        
+    Returns:
+        numpy.ndarray: Jacobian matrix for this configuration
+    """
+    ff_optimize, config, bonds, angles, dihedrals, bond2params = args
+    
+    # Get QM atoms count from the forces array length
+    qm_atoms_count = len(config['qm_forces'])
+    jac = jacobian_ff(ff_optimize,
+                      config["qm_coordinates"], 
+                      bond2params,
+                      bonds, angles, dihedrals,
+                      qm_atoms_count=qm_atoms_count)
+    
+    return jac
+
+
+def split_configurations_train_val(configurations, validation_fraction=0.1, random_seed=None):
+    """
+    Split configurations into training and validation sets.
+    
+    Args:
+        configurations (list): List of configuration dictionaries
+        validation_fraction (float): Fraction of configurations to use for validation (default: 0.1)
+        random_seed (int, optional): Random seed for reproducible splits
+        
+    Returns:
+        tuple: (train_configurations, val_configurations)
+    """
+    if random_seed is not None:
+        random.seed(random_seed)
+        np.random.seed(random_seed)
+    
+    n_configs = len(configurations)
+    n_val = max(1, int(n_configs * validation_fraction))
+    n_train = n_configs - n_val
+    
+    # Create indices and shuffle them
+    indices = list(range(n_configs))
+    random.shuffle(indices)
+    
+    # Split indices
+    train_indices = indices[:n_train]
+    val_indices = indices[n_train:]
+    
+    # Create configuration lists
+    train_configurations = [configurations[i] for i in train_indices]
+    val_configurations = [configurations[i] for i in val_indices]
+    
+    logging.info(f"Split {n_configs} configurations into {n_train} training and {n_val} validation")
+    
+    return train_configurations, val_configurations
+
+
+def compute_validation_sdf(ff_optimize, val_configurations, bonds, angles, dihedrals, bond2params, n_processes=None):
+    """
+    Compute SDF on validation set.
+    
+    Args:
+        ff_optimize (numpy.ndarray): Optimized force field parameters
+        val_configurations (list): Validation configuration list
+        bonds, angles, dihedrals (list): Interaction lists
+        bond2params (dict): Parameter mapping
+        n_processes (int, optional): Number of processes for parallel computation
+        
+    Returns:
+        tuple: (val_sdf, val_cost, val_residuals, val_ref_bonded)
+    """
+    if not val_configurations:
+        logging.warning("No validation configurations provided")
+        return None, None, None, None
+    
+    # Compute residuals on validation set
+    val_residuals, val_ref_bonded = compute_ff_obj(ff_optimize, val_configurations,
+                                                   bonds, angles, dihedrals, bond2params, n_processes)
+    
+    # Compute validation metrics
+    val_cost = np.sum(val_residuals**2)
+    val_sdf = np.sqrt(val_cost / np.sum(val_ref_bonded**2))
+    
+    logging.info(f"Validation SDF: {val_sdf:.6f}")
+    logging.info(f"Validation cost: {val_cost:.6f}")
+    
+    return val_sdf, val_cost, val_residuals, val_ref_bonded